biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[ROCm][MLA] Support block-size > 1 for AITER MLA backend (#27224)

Browse Source 
Signed-off-by: ganyi <ygan@amd.com>
Co-authored-by: wuhuikx <hattie.wu@amd.com>
This commit is contained in:
Pleaplusone
2025-11-05 23:43:02 +08:00
committed by GitHub
parent 80c9275348
commit 6cae1e5332
3 changed files with 34 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
vllm/platforms/rocm.py
View File

@@ -252,16 +252,9 @@ class RocmPlatform(Platform):
f"does not support block size {block_size}."
)
if selected_backend == _Backend.ROCM_AITER_MLA:
if block_size == 1:
logger.info("Using AITER MLA backend.")
return (
"vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501
)
raise ValueError(
f" The selected backend, {selected_backend.name},"
f"does not support block size {block_size}."
"(currently only supports block size 1)"
)
logger.info("Using AITER MLA backend.")
return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501
raise ValueError(
f" The selected backend, {selected_backend.name},"
f"is not MLA type while requested for MLA backend."