[ROCm][Bugfix] Route GLM-5.3-Flash MTP through ragged sparse MLA (#55239)
Signed-off-by: jamesETsmith <james.smith9113@gmail.com> Signed-off-by: Cursor <cursoragent@cursor.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
James E T Smith committed
Sep 10, 2026 at 21:13 UTC
828f4f19b4d8ea6a97a047409a90b563d166002f
2 files changed
+11
-5
tests/v1/attention/test_rocm_glm5next_sparse.py
+4
-1
@@ -76,7 +76,9 @@ def test_fit_kpool_indices_rejects_narrow_input():
76
("auto", 512, 0, 2, 2, 1, True),
77
("fp8", 512, 1, 0, 0, 32, False),
78
("auto", 576, 1, 0, 0, 32, False),
79
- ("auto", 512, 0, 2, 4, 2, False),
79
+ ("auto", 512, 0, 2, 4, 2, True),
80
+ ("auto", 512, 0, 2, 12, 6, True),
81
+ ("auto", 512, 0, 0, 0, 0, False),
82
],
83
)
84
def test_rocm_sparse_triton_route(
@@ -88,6 +90,7 @@ def test_rocm_sparse_triton_route(
90
max_query_len,
91
expected,
92
):
93
+ """Validate Triton routing for prefill, decode, and MTP verification."""
94
assert (
95
_use_rocm_sparse_triton(
96
kv_cache_dtype=kv_cache_dtype,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+7
-4
@@ -53,13 +53,16 @@ def _use_rocm_sparse_triton(
53
num_decode_tokens: int,
54
max_query_len: int,
55
) -> bool:
56
- """Select the rope-free BF16 path not supported by AITER sparse MLA."""
57
- plain_decode = num_decode_tokens == num_decodes
56
+ """Select the rope-free BF16 path not supported by AITER sparse MLA.
57
+
58
+ The ragged Triton kernel indexes metadata per query token, so multi-token
59
+ speculative verification rows have the same capability requirements as
60
+ plain decode rows.
61
+ """
62
return (
63
not kv_cache_dtype.startswith("fp8")
64
and head_size == kv_lora_rank
61
- and plain_decode
62
- and (num_prefills > 0 or (num_decodes > 0 and max_query_len == 1))
65
+ and (num_prefills > 0 or num_decodes > 0)
66
)
67
68