From 00e12cebb4f574c578431db920cf7533cc7835f9 Mon Sep 17 00:00:00 2001 From: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:20:12 +0800 Subject: [PATCH] [Fix]: Defer DSA MLA CP KV gather for fp8 trtllm prefill in PD mode (#29161) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> --- .../srt/layers/attention/dsa_backend.py | 43 ++++++++++++++++++- .../attention_forward_methods/forward_mla.py | 19 +++++++- 2 files changed, 59 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 1502b6a2b..341c3a2dc 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -41,7 +41,9 @@ from sglang.srt.layers.attention.dsa.utils import ( compute_dsa_seqlens, dsa_cp_round_robin_split_data, dsa_cp_round_robin_split_q_seqs, + dsa_use_prefill_cp, is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_in_seq_split, pad_dsa_cache_seqlens, ) from sglang.srt.layers.attention.utils import ( @@ -49,6 +51,10 @@ from sglang.srt.layers.attention.utils import ( mla_quantize_and_rope_for_fp8, seqlens_expand_triton, ) +from sglang.srt.layers.utils.cp_utils import ( + cp_all_gather_rerange_output, + cp_split_and_rebuild_position, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import ( get_bool_env_var, @@ -74,6 +80,24 @@ if TYPE_CHECKING: from sglang.srt.speculative.spec_info import SpecInput +def _all_gather_dsa_trtllm_fp8_kv( + forward_batch: ForwardBatch, + k: torch.Tensor, + k_rope: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + kv_lora_rank = k.shape[-1] + qk_rope_head_dim = k_rope.shape[-1] + kv_dtype = k.dtype + kv = torch.cat((k, k_rope), dim=-1).view(torch.uint8) + kv = cp_all_gather_rerange_output( + kv, + get_parallel().attn_cp_size, + forward_batch, + torch.cuda.current_stream(), + ).view(kv_dtype) + return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1) + + _is_hip = is_hip() if _is_hip: @@ -2488,17 +2512,25 @@ class DeepseekSparseAttnBackend( cos_sin_cache is not None ), "For FP8 path cos_sin_cache should not be None." + rope_positions = forward_batch.positions + if dsa_use_prefill_cp(forward_batch): + rope_positions = cp_split_and_rebuild_position( + forward_batch, rope_positions + ) + q, k, k_rope = mla_quantize_and_rope_for_fp8( q, q_rope, k.squeeze(1), k_rope.squeeze(1), - forward_batch.positions, + rope_positions, cos_sin_cache, is_neox, self.kv_lora_rank, self.qk_rope_head_dim, ) + if save_kv_cache and dsa_use_prefill_cp(forward_batch): + k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope) merge_query = False # Save KV cache if requested @@ -2561,6 +2593,15 @@ class DeepseekSparseAttnBackend( block_tables = page_table_1.unsqueeze(1) seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens + if ( + dsa_use_prefill_cp(forward_batch) + and is_dsa_prefill_cp_in_seq_split() + and forward_batch.attn_cp_metadata is not None + ): + cp_meta = forward_batch.attn_cp_metadata + seq_chunks = list(torch.split(seq_lens, cp_meta.split_list, dim=0)) + seq_lens = torch.cat([seq_chunks[i] for i in cp_meta.zigzag_index], dim=0) + out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( query=q, kv_cache=kv, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 3de047b36..cffc957de 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -164,6 +164,14 @@ if _use_aiter_gfx95: from sglang.srt.layers.rocm_linear_utils import fused_qk_rope_cat_and_cache_mla +def _should_defer_dsa_cp_kv_gather( + *, + dsa_prefill_cp: bool, + fuse_rope_for_trtllm_mla: bool, +) -> bool: + return dsa_prefill_cp and fuse_rope_for_trtllm_mla + + class DeepseekMLAForwardMixin: def init_mla_forward(self: DeepseekV2AttentionMLA): self.flashinfer_mla_disable_ragged = ( @@ -522,18 +530,25 @@ class DeepseekMLAForwardMixin: elif is_kv_b_lora_active(self): q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out) + fuse_rope_for_trtllm_mla = self._fuse_rope_for_trtllm_mla(forward_batch) skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused() skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla() if ( self.rotary_emb is not None - and (not self._fuse_rope_for_trtllm_mla(forward_batch)) + and (not fuse_rope_for_trtllm_mla) and (not skip_rope_for_dsa_tilelang_fused) and (not skip_rope_for_aiter_fused_mla) and (not _use_aiter or not _is_gfx95_supported or self.use_dsa) ): q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) - if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch): + dsa_prefill_cp = dsa_use_prefill_cp(forward_batch) + mla_prefill_cp = mla_use_prefill_cp(forward_batch) + defer_kv_gather_until_after_rope = _should_defer_dsa_cp_kv_gather( + dsa_prefill_cp=dsa_prefill_cp, + fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla, + ) + if (dsa_prefill_cp or mla_prefill_cp) and not defer_kv_gather_until_after_rope: # support allgather+rerrange k_nope, k_pe = self.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe