[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
+    return attention_backend in {"flashmla", "cutedsl_mla"}
-    # currently supported MLA DCP decode backends expose base-2 LSE.
-    return attention_backend == "flashmla"
 
 
 if _is_cuda:

test/registered/dcp/test_dcp_layout_unit.py

@@ -168,6 +168,36 @@ class TestGetDcpLens(CustomTestCase):
         self.assertTrue(torch.equal(kernel_k[0], k[:, 0:1]))
         self.assertTrue(torch.equal(out, q))
 
+    def test_dense_q_indptr_matches_the_arange_it_replaces(self):
+        from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend
+
+        max_bs = 16
+        for num_draft_tokens in (1, 2, 8):
+            backend = object.__new__(TRTLLMMLABackend)
+            backend.q_indptr_decode = torch.arange(0, max_bs + 1, dtype=torch.int32)
+            backend.num_draft_tokens = num_draft_tokens
+            backend.dense_q_indptr_verify = backend.q_indptr_decode * num_draft_tokens
+            # Equal hits the precomputed buffer, +1 hits the fallback.
+            for draft_token_num in (num_draft_tokens, num_draft_tokens + 1):
+                for bs in (1, 3, max_bs):
+                    with self.subTest(
+                        num_draft_tokens=num_draft_tokens,
+                        draft_token_num=draft_token_num,
+                        bs=bs,
+                    ):
+                        got = backend._dense_q_indptr(bs, draft_token_num)
+                        expected = torch.arange(
+                            0,
+                            (bs + 1) * draft_token_num,
+                            draft_token_num,
+                            dtype=torch.int32,
+                        )
+                        self.assertEqual(got.dtype, torch.int32)
+                        self.assertTrue(
+                            torch.equal(got, expected),
+                            f"{got.tolist()} != {expected.tolist()}",
+                        )
+
     def test_paged_allocator_exposes_dcp_virtual_capacity(self):
         real_kv_size = 1024
         dcp_size = 4

test/registered/kernels/test_dcp_lse_combine.py

@@ -265,13 +265,17 @@ class TestCPUReference(CustomTestCase):
 
         self.assertFalse(torch.allclose(result_e, result_2, atol=1e-3))
 
-    def test_flashmla_selects_natural_log_lse(self):
+    def test_natural_log_lse_backends(self):
         from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
             is_mla_dcp_lse_base_on_e,
         )
 
         self.assertTrue(is_mla_dcp_lse_base_on_e("flashmla"))
+        self.assertTrue(is_mla_dcp_lse_base_on_e("cutedsl_mla"))
         self.assertFalse(is_mla_dcp_lse_base_on_e("flashinfer_mla"))
+        self.assertFalse(is_mla_dcp_lse_base_on_e("tokenspeed_mla"))
+        self.assertFalse(is_mla_dcp_lse_base_on_e("trtllm_mla"))
+        self.assertFalse(is_mla_dcp_lse_base_on_e(None))
 
     def test_nan_lse_handled(self):
         from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu