[ROCm][K3] Dequantize the fp8 decode query for MLA backends without quant-query support - TRITON_MLA (#51860)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Hongxia Yang committed
Aug 12, 2026 at 12:40 UTC
b745d08de14fddf0cbb843c686b5d335f7d658cc
1 file changed
+7
-5
vllm/models/kimi_k3/nvidia/mla.py
+7
-5
@@ -664,14 +664,10 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
664
cos_sin_cache=cos_sin_cache,
665
)
666
if is_quantized_kv_cache(self.kv_cache_dtype):
667
- assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined]
668
- "Kimi-K3 fp8 KV cache decode requires a backend that accepts an "
669
- "fp8 (quantized) query input."
670
- )
667
cache = self.kv_cache
668
if cache.dtype != torch.float8_e4m3fn:
669
cache = cache.view(torch.float8_e4m3fn)
674
- return fused_mla_decode_q_concat_kv_cache_insert(
670
+ mqa_q = fused_mla_decode_q_concat_kv_cache_insert(
671
ql_nope,
672
q_pe,
673
kv_c_normed,
@@ -683,6 +679,12 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
679
positions=positions,
680
cos_sin_cache=cos_sin_cache,
681
)
682
+ if not self.impl.supports_quant_query_input: # type: ignore[attr-defined]
683
+ # Backend dequantizes fp8 KV on load and takes a bf16 query
684
+ # (e.g. TRITON_MLA, the DSpark draft); undo the query
685
+ # quantization.
686
+ mqa_q = (mqa_q.to(torch.float32) * self._q_scale).to(ql_nope.dtype)
687
+ return mqa_q
688
return fused_mla_decode_q_concat_kv_cache_insert(
689
ql_nope,
690
q_pe,