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 7, 02:45 UTC
[ROCm] Support MLA with nhead<16 and FP8 KV cache for TP=8 (Kimi K2.5/Linear) (#35850)
Signed-off-by: Li <chuali@amd.com>
Chuan (Richard) Li committed
Mar 6, 2026 at 12:24 UTC
c188749bcdaa2c72cc3c8a4a28e722af2abc4bb8
1 file changed
+19
-3
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+19
-3
@@ -221,11 +221,17 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
221
kv_sharing_target_layer_name,
222
**mla_args,
223
)
224
- assert num_heads == 16 or num_heads == 128, (
225
- f"Aiter MLA only supports 16 or 128 number of heads.\n"
224
+ _valid_heads = num_heads in (4, 8) or (
225
+ num_heads % 16 == 0 and 16 <= num_heads <= 128
226
+ )
227
+ assert _valid_heads, (
228
+ f"Aiter MLA supports num_heads of 4, 8, or multiples of 16 "
229
+ f"in [16, 128].\n"
230
f"Provided {num_heads} number of heads.\n"
231
"Try adjusting tensor_parallel_size value."
232
)
233
+ self._needs_head_repeat = num_heads < 16
234
+ self._head_repeat_factor = 16 // num_heads if num_heads < 16 else 1
235
unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
236
if any(unsupported_features):
237
raise NotImplementedError(
@@ -267,9 +273,16 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
273
274
assert isinstance(q, torch.Tensor)
275
B = q.shape[0]
276
+
277
+ if self._needs_head_repeat:
278
+ q = q.repeat_interleave(self._head_repeat_factor, dim=1)
279
+ kernel_num_heads = 16
280
+ else:
281
+ kernel_num_heads = self.num_heads
282
+
283
o = torch.zeros(
284
B,
272
- self.num_heads,
285
+ kernel_num_heads,
286
self.kv_lora_rank,
287
dtype=attn_metadata.decode.attn_out_dtype,
288
device=q.device,
@@ -291,4 +304,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
304
kv_scale=layer._k_scale,
305
)
306
307
+ if self._needs_head_repeat:
308
+ o = o[:, :: self._head_repeat_factor, :]
309
+
310
return o, None