From 9b0ab5dd53583ee25066b8d937e8bd3235daa6e2 Mon Sep 17 00:00:00 2001 From: James E T Smith Date: Fri, 14 Aug 2026 13:38:14 -0400 Subject: [PATCH] [ROCm][AMD] Enable preshuffled sparse indexing for 16-token blocks (#51216) Signed-off-by: jamesETsmith Co-authored-by: Andreas Karatzas --- tests/v1/worker/test_gpu_model_runner.py | 14 ++++++++++++++ vllm/v1/attention/backends/mla/indexer.py | 2 +- .../backends/mla/rocm_aiter_mla_sparse.py | 2 +- 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 06e3fe4fa7c..8bbe15bb2a4 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -35,6 +35,10 @@ from vllm.utils.mem_constants import GiB_bytes from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backend import MultipleOf +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend +from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput @@ -291,6 +295,16 @@ def test_select_common_block_size_uses_largest_shared_int(): assert selected_size == 64 +def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch): + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + + selected_size = select_common_block_size( + 16, + [DeepseekV32IndexerBackend, ROCMAiterMLASparseBackend], + ) + assert selected_size == 16 + + def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids( dist_init, ): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2b3e400de3f..717cfa621ca 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -141,7 +141,7 @@ class DeepseekV32IndexerBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [1, 64] if current_platform.is_rocm() else [64] + return [1, MultipleOf(16)] if current_platform.is_rocm() else [64] @classmethod def get_supported_head_sizes(cls) -> list[int]: diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index d5db42cb7f9..3f10d610534 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -275,7 +275,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [1, 64] + return [1, MultipleOf(16)] @staticmethod def get_name() -> str: