[DCP] Drop two per-layer launches from the MLA target-verify path (#34240)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/srt/layers/attention/cutedsl_mla_backend.py

@@ -19,7 +19,6 @@ path is stable (see the TODO in tokenspeed_mla_backend.py).
 from __future__ import annotations
 
 import logging
-import math
 from typing import TYPE_CHECKING, Optional
 
 import torch
@@ -56,19 +55,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-# The flashinfer cute-dsl MLA decode kernel returns a natural-log (base-e) LSE,
-# whereas sglang's DCP cross-rank merge (forward_mla: dcp_a2a_lse_reduce /
-# cp_lse_ag_out_rs_mla) assumes the FlashInfer-MLA/FlashMLA base-2 convention
-# (is_lse_base_on_e=False). Multiplying a natural-log LSE by log2(e) rebases it
-# to base-2 (the softmax output is base-invariant; only the LSE value changes).
-# CONFIRMED base-e (not base-2), so this rebase is required, not optional:
-# the flashinfer-dcp-backport public-API unit test asserts the public
-# trtllm_batch_decode_with_kv_cache_mla LSE against a torch.logsumexp
-# (natural-log) reference at atol=1e-2 and passes (a base-2 LSE would be
-# off by 1/ln2 ~= 44%). GPU job 467640:
-# tests/attention/test_cute_dsl_mla_dcp*.py 27/27 + 17/17 pass.
-_LSE_BASE2_FROM_NATURAL_LOG = math.log2(math.e)
-
 
 class CuteDslMLABackend(TRTLLMMLABackend):
     """flashinfer cute-dsl MLA decode backend with decode context parallelism."""
@@ -299,7 +285,7 @@ class CuteDslMLABackend(TRTLLMMLABackend):
         Without DCP (``cp_world <= 1``) this defers to the base cute-dsl path.
         With DCP, ``seq_lens`` are this rank's cyclic-local KV lengths and
         ``causal_seqs`` the global per-request KV lengths; the kernel returns a
-        rank-local ``(out, lse)`` (LSE rebased to base-2 for the sglang merge).
+        rank-local ``(out, lse)``, the LSE in natural log.
         """
         if cp_world <= 1:
             return super()._run_decode_kernel(
@@ -336,7 +322,7 @@ class CuteDslMLABackend(TRTLLMMLABackend):
             ),
             return_lse=True,  # DCP requires the rank-local LSE for the merge
         )
-        return raw_out, lse * _LSE_BASE2_FROM_NATURAL_LOG
+        return raw_out, lse
 
     def forward_decode(
         self,

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -273,6 +273,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         self.disable_chunked_prefix_cache = get_schedule().disable_chunked_prefix_cache
 
         self.num_draft_tokens = get_spec().speculative_num_draft_tokens
+        self.dense_q_indptr_verify = (
+            self.q_indptr_decode * self.num_draft_tokens
+            if self.num_draft_tokens
+            else None
+        )
         self._verify_mask = None
         # Tree-mask scratch is fetched from the target backend only.
         self.is_draft_runner = model_runner.is_draft_worker
@@ -821,6 +826,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             k_scale = 1.0
         return q_scale * k_scale * layer.scaling
 
+    def _dense_q_indptr(self, bs: int, draft_token_num: int) -> torch.Tensor:
+        """Query indptr for a dense [bs, draft_token_num] verify batch."""
+        if draft_token_num == self.num_draft_tokens:
+            return self.dense_q_indptr_verify[: bs + 1]
+        return self.q_indptr_decode[: bs + 1] * draft_token_num
+
     def _run_decode_kernel(
         self,
         query: torch.Tensor,
@@ -1407,18 +1418,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                     layer.v_head_dim,
                 )
                 lse = lse.view(bs * draft_token_num, layer.tp_q_head_num)
-                dense_q_indptr = torch.arange(
-                    0,
-                    (bs + 1) * draft_token_num,
-                    draft_token_num,
-                    dtype=torch.int32,
-                    device=q.device,
-                )
                 fixup_zero_kv_rows(
                     output,
                     lse,
                     metadata.seq_lens_k,
-                    dense_q_indptr,
+                    self._dense_q_indptr(bs, draft_token_num),
                     draft_token_num,
                 )
                 return output.flatten(1), lse

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -103,9 +103,7 @@ def is_dcp_mla_decode_phase(forward_batch: ForwardBatch) -> bool:
 
 
 def is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
-    # FlashMLA exposes natural-log softmax LSE. FlashInfer MLA and the other
-    # currently supported MLA DCP decode backends expose base-2 LSE.
-    return attention_backend == "flashmla"
+    return attention_backend in {"flashmla", "cutedsl_mla"}
 
 
 if _is_cuda: