A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 23:00 UTC
[ROCm][Perf] Enable `sparse_mla`'s cudagraph on ROCm platform (#35719)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Mar 11, 2026 at 00:14 UTC
82f3f30e266e24b26c46916a8c9daaea7d5e32bd
2 files changed
+3
-4
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+3
-1
@@ -151,7 +151,9 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
151
class ROCMAiterMLASparseMetadataBuilder(
152
AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
153
):
154
- _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
154
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
155
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
156
+ )
157
158
def __init__(
159
self,
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
-3
@@ -327,9 +327,6 @@ def rocm_fp8_paged_mqa_logits(
327
aiter_paged_mqa_logits_module = None
328
if rocm_aiter_ops.is_enabled():
329
aiter_paged_mqa_logits_module = paged_mqa_logits_module()
330
- # FIXME(ganyi): Temporarily disable the aiter path until nightly docker
331
- # update aiter to the fix PR.
332
- aiter_paged_mqa_logits_module = None
330
331
if aiter_paged_mqa_logits_module is not None:
332
deepgemm_fp8_paged_mqa_logits_stage1 = (