From c72f0756d28286fa3f42e97791514cc6322b8078 Mon Sep 17 00:00:00 2001 From: Fan Yin <1106310035@qq.com> Date: Mon, 1 Dec 2025 05:38:19 +0800 Subject: [PATCH] Fix: fix flashmla fp8 kv cache acc error (#13841) Co-authored-by: ybyang --- .../srt/layers/attention/flashmla_backend.py | 77 ++++++++++++++++--- 1 file changed, 65 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 47644c59e..c3d017583 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -14,6 +14,7 @@ from sgl_kernel.flash_mla import flash_mla_with_kvcache, get_mla_metadata from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode if TYPE_CHECKING: @@ -75,6 +76,11 @@ class FlashMLABackend(FlashInferMLAAttnBackend): self.data_type = model_runner.kv_cache_dtype self.q_data_type = model_runner.dtype self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim + # Check if KV cache is FP8 (supports both e4m3 and e5m2) + self.is_fp8_kvcache = self.data_type in { + torch.float8_e4m3fn, + torch.float8_e5m2, + } self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens @@ -104,6 +110,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): forward_batch.seq_lens.to(torch.int32), self.num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) self.forward_metadata = FlashMLADecodeMetadata( mla_metadata, @@ -134,6 +141,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): seq_lens.to(torch.int32), self.num_draft_tokens * self.num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) # Use FlashMLADecodeMetadata which has the attributes forward_extend expects @@ -168,6 +176,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): ), self.num_draft_tokens * self.num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) else: self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata( @@ -176,6 +185,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): ), self.num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) self.cuda_graph_kv_indices = cuda_graph_kv_indices @@ -206,6 +216,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): seq_lens.to(torch.int32), num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) @@ -231,6 +242,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): seq_lens.to(torch.int32), self.num_draft_tokens * self.num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) @@ -281,6 +293,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): seq_lens.to(torch.int32), num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) @@ -306,6 +319,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): seq_lens.to(torch.int32), self.num_draft_tokens * self.num_q_heads, 1, + is_fp8_kvcache=self.is_fp8_kvcache, ) self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) @@ -353,8 +367,28 @@ class FlashMLABackend(FlashInferMLAAttnBackend): k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) - if self.data_type == torch.float8_e4m3fn: - reshape_q_fp8 = reshape_q.to(torch.float8_e4m3fn) + if self.is_fp8_kvcache: + # For FP8 KV cache, Q needs to be converted to FP8 for FlashMLA kernel + # In SGLang, we use layer.k_scale for both q and k scales + if layer.k_scale is not None: + q_scale = layer.k_scale + descale_q = layer.k_scale.reshape(1) + descale_k = layer.k_scale.reshape(1) + else: + # Fallback to 1.0 if k_scale is not initialized + q_scale = torch.ones((1,), dtype=torch.float32, device=reshape_q.device) + descale_q = torch.ones( + (1,), dtype=torch.float32, device=reshape_q.device + ) + descale_k = torch.ones( + (1,), dtype=torch.float32, device=reshape_q.device + ) + + # Reshape to 2D for scaled_fp8_quant (which requires 2D input) + q_shape = reshape_q.shape + reshape_q_2d = reshape_q.reshape(-1, q_shape[-1]) + reshape_q_fp8_2d, _ = scaled_fp8_quant(reshape_q_2d, q_scale) + reshape_q_fp8 = reshape_q_fp8_2d.reshape(q_shape) o, _ = flash_mla_with_kvcache( q=reshape_q_fp8, k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), @@ -365,8 +399,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend): num_splits=self.forward_metadata.num_splits, softmax_scale=layer.scaling, causal=True, - descale_q=torch.ones((1), dtype=torch.float32, device=reshape_q.device), - descale_k=torch.ones((1), dtype=torch.float32, device=reshape_q.device), + descale_q=descale_q, + descale_k=descale_k, ) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) @@ -412,8 +446,31 @@ class FlashMLABackend(FlashInferMLAAttnBackend): k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) - if self.data_type == torch.float8_e4m3fn: - reshape_q_fp8 = reshape_q.to(torch.float8_e4m3fn) + if self.is_fp8_kvcache: + # For FP8 KV cache, Q needs to be converted to FP8 for FlashMLA kernel + # In SGLang, we use layer.k_scale for both q and k scales + if layer.k_scale is not None: + q_scale = layer.k_scale + descale_q = layer.k_scale.reshape(1) + descale_k = layer.k_scale.reshape(1) + else: + # Fallback to 1.0 if k_scale is not initialized + q_scale = torch.ones( + (1,), dtype=torch.float32, device=reshape_q.device + ) + descale_q = torch.ones( + (1,), dtype=torch.float32, device=reshape_q.device + ) + descale_k = torch.ones( + (1,), dtype=torch.float32, device=reshape_q.device + ) + + # Quantize Q using scaled_fp8_quant (matching vLLM's approach) + # Reshape to 2D for scaled_fp8_quant (which requires 2D input) + q_shape = reshape_q.shape + reshape_q_2d = reshape_q.reshape(-1, q_shape[-1]) + reshape_q_fp8_2d, _ = scaled_fp8_quant(reshape_q_2d, q_scale) + reshape_q_fp8 = reshape_q_fp8_2d.reshape(q_shape) o, _ = flash_mla_with_kvcache( q=reshape_q_fp8, k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), @@ -425,12 +482,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend): num_splits=self.forward_metadata.num_splits, softmax_scale=layer.scaling, causal=True, - descale_q=torch.ones( - (1), dtype=torch.float32, device=reshape_q.device - ), - descale_k=torch.ones( - (1), dtype=torch.float32, device=reshape_q.device - ), + descale_q=descale_q, + descale_k=descale_k, ) else: o, _ = flash_mla_with_kvcache(