diff --git a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py index f9207b4fe..137718c99 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 0da94696f..9dee901dd 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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 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 b3d666f8c..e140a117d 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 @@ -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: diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index cefc57621..8c7f5731d 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/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 diff --git a/test/registered/kernels/test_dcp_lse_combine.py b/test/registered/kernels/test_dcp_lse_combine.py index 9d26a6428..0249b8a1e 100644 --- a/test/registered/kernels/test_dcp_lse_combine.py +++ b/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