iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Bugfix] Route GLM-5.3-Flash MTP through ragged sparse MLA (#55239)

Browse Source 
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>
This commit is contained in:
James E T Smith
2026-09-10 20:13:14 -05:00
committed by GitHub
co-authored by Cursor mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 84030bbe3d
commit 828f4f19b4
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_rocm_glm5next_sparse.py
+4 -1
View File
@@ -76,7 +76,9 @@ def test_fit_kpool_indices_rejects_narrow_input():
("auto", 512, 0, 2, 2, 1, True),
("fp8", 512, 1, 0, 0, 32, False),
("auto", 576, 1, 0, 0, 32, False),
("auto", 512, 0, 2, 4, 2, False),
("auto", 512, 0, 2, 4, 2, True),
("auto", 512, 0, 2, 12, 6, True),
("auto", 512, 0, 0, 0, 0, False),
],
)
def test_rocm_sparse_triton_route(
@@ -88,6 +90,7 @@ def test_rocm_sparse_triton_route(
max_query_len,
expected,
):
"""Validate Triton routing for prefill, decode, and MTP verification."""
assert (
_use_rocm_sparse_triton(
kv_cache_dtype=kv_cache_dtype,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+7 -4
View File
@@ -53,13 +53,16 @@ def _use_rocm_sparse_triton(
num_decode_tokens: int,
max_query_len: int,
) -> bool:
"""Select the rope-free BF16 path not supported by AITER sparse MLA."""
plain_decode = num_decode_tokens == num_decodes
"""Select the rope-free BF16 path not supported by AITER sparse MLA.
The ragged Triton kernel indexes metadata per query token, so multi-token
speculative verification rows have the same capability requirements as
plain decode rows.
"""
return (
not kv_cache_dtype.startswith("fp8")
and head_size == kv_lora_rank
and plain_decode
and (num_prefills > 0 or (num_decodes > 0 and max_query_len == 1))
and (num_prefills > 0 or num_decodes > 0)
)