[ROCM][DSV32][Perf][MTP] Enable UNIFORM_BATCH CG mode in rocm_aiter_mla_sparse (#45149)
Signed-off-by: Teemu Virolainen <teemu.virolainen@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Teemu Virolainen committed
Jul 9, 2026 at 22:35 UTC
b0dec2a11b91711ac1893aa18491a77b0f443644
1 file changed
+4
-3
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+4
-3
@@ -347,9 +347,7 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
347
class ROCMAiterMLASparseMetadataBuilder(
348
AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
349
):
350
- _cudagraph_support: ClassVar[AttentionCGSupport] = (
351
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
352
- )
350
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
351
352
def __init__(
353
self,
@@ -365,6 +363,9 @@ class ROCMAiterMLASparseMetadataBuilder(
363
self.device = device
364
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
365
366
+ self.vllm_config = vllm_config
367
+ self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
368
+
369
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
370
self.mla_dims = get_mla_dims(self.model_config)
371
self.topk_tokens = vllm_config.model_config.hf_config.index_topk