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 5, 14:45 UTC
[ROCm] Fix aiter persistent mode mla with q/o nhead<16 for kimi-k2.5 tp8 (#38615)
Signed-off-by: wufann <36477220+wufann@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
wufann committed
Apr 3, 2026 at 18:54 UTC
1b117cb0ac51b84dd7aed364e6e802dea5147ca6
1 file changed
+4
-3
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+4
-3
@@ -129,9 +129,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
129
130
from aiter import dtypes, get_mla_metadata_info_v1
131
132
- self._num_attention_heads = vllm_config.model_config.get_num_attention_heads(
133
- vllm_config.parallel_config
134
- )
132
+ # For num_attention_heads < 16 (e.g. kimi-k2.5 head=8 with TP8),
133
+ # make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent
134
+ # with the actual tensor shape passed to mla_decode_fwd.
135
+ self._num_attention_heads = max(16, self.num_heads)
136
q_dtype = self.decode_attn_out_dtype
137
kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
138
if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):