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, 09:30 UTC
[ROCm][Bugfix] fix cache block size mismatch for aiter unified attention (#37606)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Mar 19, 2026 at 20:00 UTC
4ca3fa6bb4633fed1196292f764ce8cf13f647b5
2 files changed
+7
-24
vllm/platforms/rocm.py
-24
@@ -665,7 +665,6 @@ class RocmPlatform(Platform):
665
def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
666
from vllm.config.compilation import CUDAGraphMode
667
668
- cache_config = vllm_config.cache_config
668
compilation_config = vllm_config.compilation_config
669
parallel_config = vllm_config.parallel_config
670
@@ -687,32 +686,9 @@ class RocmPlatform(Platform):
686
)
687
compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
688
690
- if cache_config and not cache_config.user_specified_block_size:
691
- if (
692
- envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION and envs.VLLM_ROCM_USE_AITER
693
- # NOTE: This block has been deprecated
694
- # or get_env_variable_attn_backend()
695
- # == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN
696
- # TODO: monitor https://github.com/vllm-project/vllm/pull/30396
697
- # to see how we can transition to the new way of selecting
698
- # attention backends
699
- ):
700
- cache_config.block_size = 64
701
- logger.warning(
702
- "[ROCM_AITER_UNIFIED_ATTN]: Setting kv cache block size to 64."
703
- )
704
- else:
705
- cache_config.block_size = 16
706
-
689
if parallel_config.worker_cls == "auto":
690
parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
691
710
- @classmethod
711
- def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
712
- # TODO: ROCm still sets block_size in check_and_update_config.
713
- # Move that logic here so block_size is chosen by the backend.
714
- pass
715
-
692
@classmethod
693
def verify_model_arch(cls, model_arch: str) -> None:
694
if model_arch in _ROCM_UNSUPPORTED_MODELS:
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+7
@@ -29,6 +29,13 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
29
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
30
return [MultipleOf(16)]
31
32
+ @classmethod
33
+ def get_preferred_block_size(cls, default_block_size: int) -> int:
34
+ logger.warning_once(
35
+ "[ROCM_AITER_UNIFIED_ATTN]: Setting kv cache block size to 64."
36
+ )
37
+ return 64
38
+
39
@classmethod
40
def supports_block_size(cls, block_size: int | None) -> bool:
41
if block_size is None: