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

[ROCm][Perf] Allow MTP lens > 1 in Sparse MLA (#36681)

Browse Source 
Signed-off-by: Teemu Virolainen <teemu.virolainen@amd.com>
This commit is contained in:
tvirolai-amd
2026-03-11 16:43:03 +02:00
committed by GitHub
parent f3163bba67
commit a9e532afe2
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/v1/spec_decode/eagle.py
View File

@@ -214,11 +214,15 @@ class SpecDecodeBaseProposer:
# Determine allowed attention backends once during initialization. # Determine allowed attention backends once during initialization.
self.allowed_attn_types: tuple | None = None self.allowed_attn_types: tuple | None = None
if current_platform.is_rocm(): if current_platform.is_rocm():
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
ROCMAiterMLASparseMetadata,
)
from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata
rocm_types = [ rocm_types = [
TritonAttentionMetadata, TritonAttentionMetadata,
RocmAttentionMetadata, RocmAttentionMetadata,
ROCMAiterMLASparseMetadata,
] ]
# ROCM_AITER_FA is an optional backend # ROCM_AITER_FA is an optional backend
# We check is_enabled() here to avoid importing the backend module during # We check is_enabled() here to avoid importing the backend module during