diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 654f83b452..8f70574142 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -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,