iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][K3] Dequantize the fp8 decode query for MLA backends without quant-query support - TRITON_MLA (#51860)

Browse Source 
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Hongxia Yang
2026-08-12 12:40:41 -04:00
committed by GitHub
co-authored by Claude Fable 5
parent 324f452f64
commit b745d08de1
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/kimi_k3/nvidia/mla.py
+7 -5
View File
@@ -664,14 +664,10 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
cos_sin_cache=cos_sin_cache,
)
if is_quantized_kv_cache(self.kv_cache_dtype):
assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined]
"Kimi-K3 fp8 KV cache decode requires a backend that accepts an "
"fp8 (quantized) query input."
)
cache = self.kv_cache
if cache.dtype != torch.float8_e4m3fn:
cache = cache.view(torch.float8_e4m3fn)
return fused_mla_decode_q_concat_kv_cache_insert(
mqa_q = fused_mla_decode_q_concat_kv_cache_insert(
ql_nope,
q_pe,
kv_c_normed,
@@ -683,6 +679,12 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
positions=positions,
cos_sin_cache=cos_sin_cache,
)
if not self.impl.supports_quant_query_input: # type: ignore[attr-defined]
# Backend dequantizes fp8 KV on load and takes a bf16 query
# (e.g. TRITON_MLA, the DSpark draft); undo the query
# quantization.
mqa_q = (mqa_q.to(torch.float32) * self._q_scale).to(ql_nope.dtype)
return mqa_q
return fused_mla_decode_q_concat_kv_cache_insert(
ql_nope,
q_pe,