[Model Bash][DSR1] Add selective dynamic shape marking for CustomOp (#34900)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-02-22 04:28:01 +04:00
parent a4047d4ea9
commit 74d90b1ce4
2 changed files with 46 additions and 7 deletions
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -950,7 +950,10 @@ def dynamic_per_batched_tensor_quant(
 logger = init_logger(__name__)


-@CustomOp.register("mla_decode_concat_quant_fp8")
+@CustomOp.register(
+    "mla_decode_concat_quant_fp8",
+    dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},
+)
 class _DecodeConcatQuantFP8(QuantFP8):
    """
    QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before