[ROCm][AMD] Enable preshuffled sparse indexing for 16-token blocks (#51216)
Signed-off-by: jamesETsmith <james.smith9113@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
James E T Smith committed
Aug 14, 2026 at 13:38 UTC
9b0ab5dd53583ee25066b8d937e8bd3235daa6e2
3 files changed
+16
-2
tests/v1/worker/test_gpu_model_runner.py
+14
@@ -35,6 +35,10 @@ from vllm.utils.mem_constants import GiB_bytes
35
from vllm.utils.system_utils import update_environment_variables
36
from vllm.utils.torch_utils import set_random_seed
37
from vllm.v1.attention.backend import MultipleOf
38
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
39
+from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
40
+ ROCMAiterMLASparseBackend,
41
+)
42
from vllm.v1.attention.backends.registry import AttentionBackendEnum
43
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
44
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
@@ -291,6 +295,16 @@ def test_select_common_block_size_uses_largest_shared_int():
295
assert selected_size == 64
296
297
298
+def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch):
299
+ monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
300
+
301
+ selected_size = select_common_block_size(
302
+ 16,
303
+ [DeepseekV32IndexerBackend, ROCMAiterMLASparseBackend],
304
+ )
305
+ assert selected_size == 16
306
+
307
+
308
def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids(
309
dist_init,
310
):
vllm/v1/attention/backends/mla/indexer.py
+1
-1
@@ -141,7 +141,7 @@ class DeepseekV32IndexerBackend(AttentionBackend):
141
142
@staticmethod
143
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
144
- return [1, 64] if current_platform.is_rocm() else [64]
144
+ return [1, MultipleOf(16)] if current_platform.is_rocm() else [64]
145
146
@classmethod
147
def get_supported_head_sizes(cls) -> list[int]:
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+1
-1
@@ -275,7 +275,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
275
276
@staticmethod
277
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
278
- return [1, 64]
278
+ return [1, MultipleOf(16)]
279
280
@staticmethod
281
def get_name() -> str: