[misc] Revert #38346, #33426, #39061 and #39219

python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py

@@ -12,16 +12,15 @@
 
 Scope (v1): chain speculation only (``speculative_eagle_topk == 1``, i.e.
 ``retrieve_next_token is None``). The tree path keeps the unfused reference
-kernels. Requires ``T >= kernel_width - 1``.
-
-State: conv_state and the SSM state are read-only. Verify is speculative, and
-the commit scatter advances them from the selected intermediate window.
-
-Numerics: aligned with the unfused pair. The conv output is rounded to the
-activation dtype (bf16) before entering the recurrence — exactly what the
-unfused path does through its intermediate tensor — and all expressions mirror
-the reference kernels line by line. Reduction order still splits differently
-where many V heads share one Q/K head, worth ~1 ulp on the output.
+kernels. Requires ``T >= kernel_width - 1`` (the rolled conv state is then
+exactly the last ``kernel_width - 1`` input tokens, matching the reference
+kernel's store).
+
+Numerics: deliberately bit-aligned with the unfused pair. The conv output is
+rounded to the activation dtype (bf16) before entering the recurrence —
+exactly what the unfused path does through its intermediate tensor — and all
+expressions mirror the reference kernels line by line, with the same
+num_warps so reduction order matches.
 """
 
 from typing import Optional
@@ -319,8 +318,19 @@ def fused_kda_conv_gating_verify_kernel(
                 )
                 tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
 
-    # No conv-state writeback: every V tile reads the same Q/K history, so a
-    # tile in a later wave would read what i_v == 0 had overwritten.
+    # Rolled conv state after consuming T >= W-1 tokens is exactly the last
+    # W-1 input tokens — which are the current window registers. The verify
+    # pass never writes the ssm state back (rollback happens at commit).
+    if is_qk_owner:
+        tl.store(cs_base + q_ch + 0 * stride_cs_tok, q_c0, mask=mask_k)
+        tl.store(cs_base + q_ch + 1 * stride_cs_tok, q_c1, mask=mask_k)
+        tl.store(cs_base + q_ch + 2 * stride_cs_tok, q_c2, mask=mask_k)
+        tl.store(cs_base + k_ch + 0 * stride_cs_tok, k_c0, mask=mask_k)
+        tl.store(cs_base + k_ch + 1 * stride_cs_tok, k_c1, mask=mask_k)
+        tl.store(cs_base + k_ch + 2 * stride_cs_tok, k_c2, mask=mask_k)
+    tl.store(cs_base + v_ch + 0 * stride_cs_tok, v_c0, mask=mask_v)
+    tl.store(cs_base + v_ch + 1 * stride_cs_tok, v_c1, mask=mask_v)
+    tl.store(cs_base + v_ch + 2 * stride_cs_tok, v_c2, mask=mask_v)
 
 
 def fused_kda_conv_gating_verify(
@@ -348,11 +358,14 @@ def fused_kda_conv_gating_verify(
     softplus_beta: float = 1.0,
     softplus_threshold: float = 20.0,
     use_qk_l2norm_in_kernel: bool = True,
-    # num_warps=4 is ~1.3x faster than the unfused pair in-graph; 1 restores the
-    # reference reduction order but is ~2.4x slower, for numerics debugging only.
-    # The fp32 intermediate-ssm rollback cache carries the reduction-order delta
-    # furthest: ~6e-8 at T=4 standard gate (the production MTP shape), ~2e-3 at
-    # T=8 safe gate. conv_state is not comparable to the reference at all.
+    # num_warps=4 is ~1.3x faster than the unfused pair in-graph; the output,
+    # conv_state and conv-window caches stay bit-identical to the reference.
+    # Only the fp32 intermediate-ssm rollback cache differs: the tl.sum
+    # reduction-order delta (~1 ulp/step) compounds through the delta-rule
+    # recurrence — measured ~6e-8 at T=4 standard gate (the production MTP
+    # shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe gate. num_warps=1
+    # reproduces the reference reduction order exactly (all buffers
+    # bit-identical) but is ~2.4x slower in-graph — numerics debugging only.
     num_warps: int = 4,
 ) -> torch.Tensor:
     """Chain-verify fast path. Returns ``o`` of shape [1, seq_len, HV, V],

python/sglang/srt/layers/attention/dsa/utils.py

@@ -115,14 +115,13 @@ def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool:
 
 
 def is_dsa_enable_prefill_cp():
-    if get_parallel().attn_cp_size <= 1:
-        return False
-
     if is_hip() or is_npu() or is_musa():
         return False
 
     # Generic prefill CP derives activation from the runtime topology and model
     # architecture.
+    if get_parallel().attn_cp_size <= 1:
+        return False
     from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4
 
     hf_config = process_model_config().hf_config

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

@@ -50,6 +50,11 @@
 
 from sgl_kernel import merge_state_v2
 
+from sglang.kernels.ops.attention.flash_attention import (
+    flash_attn_varlen_func,
+    flash_attn_with_kvcache,
+)
+
 
 def _should_disable_scheduler_metadata_precompute() -> bool:
     return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention)
@@ -1278,13 +1283,7 @@ def forward_extend(
         aux_tensors=None,
         rel_bias=None,
         rel_bias_event=None,
-        # Returns (output, lse) with lse in [total_q, num_heads].
-        return_lse: bool = False,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        lse_out = None
-        # Bound in __init__ so a subclass can substitute a different FA4 build.
-        flash_attn_with_kvcache = self.flash_attn_with_kvcache
-        flash_attn_varlen_func = self.flash_attn_varlen_func
+    ):
         if score_mod is not None and self.fa_impl_ver != 4:
             raise RuntimeError("score_mod is only supported by the FA4 backend.")
         cp_active = is_cp_active(forward_batch)
@@ -1564,7 +1563,7 @@ def _fa_cp_attn(
                     causal=False if use_cascade_attn else causal,
                     window_size=window_size,
                     softcap=layer.logit_cap,
-                    return_softmax_lse=use_cascade_attn or return_lse,
+                    return_softmax_lse=use_cascade_attn,
                     num_splits=self.num_splits,
                     out=_fa_out,
                     ver=self.fa_impl_ver,
@@ -1624,8 +1623,6 @@ def _fa_cp_attn(
                     o_expand,
                     softmax_lse_expand.T.contiguous(),
                 )
-            elif return_lse:
-                o, lse_out, *_ = result
             else:
                 o = result
         else:
@@ -1826,12 +1823,7 @@ def _mla_cp_attn(
                     else:
                         o = result
 
-        o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
-        if return_lse:
-            assert lse_out is not None
-            # The varlen kernel emits LSE head-major [num_heads, total_q].
-            return o, lse_out.transpose(0, 1).contiguous()
-        return o
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
     def forward_decode(
         self,
@@ -1852,13 +1844,7 @@ def forward_decode(
         aux_tensors=None,
         rel_bias=None,
         rel_bias_event=None,
-        # Returns (output, lse) with lse in [total_q, num_heads].
-        return_lse: bool = False,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        lse_out = None
-        # Bound in __init__ so a subclass can substitute a different FA4 build.
-        flash_attn_with_kvcache = self.flash_attn_with_kvcache
-        flash_attn_varlen_func = self.flash_attn_varlen_func
+    ) -> torch.Tensor:
         if score_mod is not None and self.fa_impl_ver != 4:
             raise RuntimeError("score_mod is only supported by the FA4 backend.")
         if k is not None:
@@ -2059,7 +2045,7 @@ def forward_decode(
                     causal=False if use_cascade_attn else causal,
                     window_size=window_size,
                     softcap=layer.logit_cap,
-                    return_softmax_lse=use_cascade_attn or return_lse,
+                    return_softmax_lse=use_cascade_attn,
                     num_splits=(
                         self.decode_num_splits
                         if not is_swa_layer
@@ -2104,8 +2090,6 @@ def forward_decode(
                         o_expand,
                         softmax_lse_expand.T.contiguous(),
                     )
-                elif return_lse:
-                    o, lse_out, *_ = result
                 else:
                     o = result
         else:
@@ -2184,12 +2168,7 @@ def forward_decode(
             else:
                 o = result
 
-        o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
-        if return_lse:
-            assert lse_out is not None
-            # The varlen kernel emits LSE head-major [num_heads, total_q].
-            return o, lse_out.transpose(0, 1).contiguous()
-        return o
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         """Initialize CUDA graph state for the attention backend.

python/sglang/srt/layers/attention/qsa/qsa_indexer.py

@@ -321,7 +321,6 @@ def update_key_state_and_compress(
                 self.compress_ratio, device=member_rows.device, dtype=torch.long
             )
             source_keys = token_k
-            group_locs = group_locs.clamp_max(source_keys.shape[0] - 1)
             source_rope = metadata.extend_rope_matrix
             if source_rope is None:
                 source_rope = build_rope_position_matrix(

python/sglang/srt/utils/common.py

@@ -221,17 +221,13 @@ def is_cpu() -> bool:
     return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_supported
 
 
-try:
-    import torchada  # noqa: F401
-except ImportError:
-    _IS_MUSA = False
-else:
-    _IS_MUSA = hasattr(torch.version, "musa") and torch.version.musa is not None
-
-
 @lru_cache(maxsize=1)
 def is_musa() -> bool:
-    return _IS_MUSA
+    try:
+        import torchada  # noqa: F401
+    except ImportError:
+        return False
+    return hasattr(torch.version, "musa") and torch.version.musa is not None
 
 
 @lru_cache(maxsize=1)

test/registered/attention/unittests/dense/test_fa4.py

@@ -3,13 +3,9 @@
 import torch
 
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
-from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
-    DENSE_ATOL,
-    DENSE_RTOL,
     DenseAttentionCase,
-    build_dense_attention_fixture,
     make_dense_cases,
     run_dense_attention_case,
 )
@@ -427,102 +423,6 @@ def test_runner_mode_spec_verify_cuda_graph_cases(self):
                     hidden_size=self.HIDDEN_SIZE,
                 )
 
-    RETURN_LSE_CASES = (
-        DenseAttentionCase(
-            name="return_lse_mha_extend",
-            backend="fa4",
-            forward_mode=ForwardMode.EXTEND,
-            num_heads=4,
-            num_kv_heads=4,
-            page_size=1,
-            prefix_lens=(2, 4),
-            extend_lens=(3, 1),
-        ),
-        DenseAttentionCase(
-            name="return_lse_gqa_decode",
-            backend="fa4",
-            forward_mode=ForwardMode.DECODE,
-            num_heads=8,
-            num_kv_heads=2,
-            page_size=1,
-            prefix_lens=(5, 9),
-        ),
-    )
-
-    def _reference_out_and_lse(self, fixture):
-        """Causal fp32 reference: output and per-query LSE."""
-        case, module = fixture.case, fixture.reference_module
-        dim = module.head_dim
-        rep = case.num_heads // case.num_kv_heads
-        q, k, v = module.project_qkv(fixture.input_hidden)
-        q = q.view(-1, case.num_heads, dim).float()
-        k = k.view(-1, case.num_kv_heads, dim).float()
-        v = v.view(-1, case.num_kv_heads, dim).float()
-
-        outs, lses, seen = [], [], 0
-        for req, prefix in enumerate(fixture.prefix_hidden):
-            _, prefix_k, prefix_v = module.project_qkv(prefix)
-            n = case.input_lens[req]
-            keys = torch.cat(
-                [prefix_k.view(-1, case.num_kv_heads, dim).float(), k[seen : seen + n]]
-            ).repeat_interleave(rep, dim=1)
-            values = torch.cat(
-                [prefix_v.view(-1, case.num_kv_heads, dim).float(), v[seen : seen + n]]
-            ).repeat_interleave(rep, dim=1)
-            for offset in range(n):
-                end = case.prefix_lens[req] + offset + 1
-                scores = (
-                    torch.einsum("hd,khd->hk", q[seen + offset], keys[:end])
-                    * module.scaling
-                )
-                probs = torch.softmax(scores, dim=-1)
-                outs.append(torch.einsum("hk,khd->hd", probs, values[:end]).reshape(-1))
-                lses.append(torch.logsumexp(scores, dim=-1))
-            seen += n
-        return torch.stack(outs), torch.stack(lses)
-
-    def test_return_lse(self):
-        """Calls the backend directly: return_lse is a backend-level contract,
-        and the RadixAttention dispatcher's custom-op schema cannot carry it.
-        """
-        for case in self.RETURN_LSE_CASES:
-            with self.subTest(case=case.name):
-                fixture = build_dense_attention_fixture(
-                    self, case, head_dim=self.HEAD_DIM, hidden_size=self.HIDDEN_SIZE
-                )
-                module = fixture.actual_module
-                forward = (
-                    fixture.backend.forward_decode
-                    if case.forward_mode.is_decode()
-                    else fixture.backend.forward_extend
-                )
-                with (
-                    torch.no_grad(),
-                    forward_context(ForwardContext(attn_backend=fixture.backend)),
-                ):
-                    fixture.backend.init_forward_metadata(fixture.forward_batch)
-                    q, k, v = module.project_qkv(fixture.input_hidden)
-                    result = forward(
-                        q, k, v, module.attn, fixture.forward_batch, return_lse=True
-                    )
-
-                self.assertIsInstance(result, tuple)
-                out, lse = result
-                self.assertEqual(
-                    tuple(lse.shape), (case.num_input_tokens, case.num_heads)
-                )
-
-                expected_out, expected_lse = self._reference_out_and_lse(fixture)
-                torch.testing.assert_close(
-                    lse.float(), expected_lse, atol=DENSE_ATOL, rtol=DENSE_RTOL
-                )
-                torch.testing.assert_close(
-                    out.float().reshape(expected_out.shape),
-                    expected_out,
-                    atol=DENSE_ATOL,
-                    rtol=DENSE_RTOL,
-                )
-
 
 if __name__ == "__main__":
     unittest.main()

test/registered/cp/test_cp_strategy_unit.py

@@ -121,24 +121,6 @@ def test_hip_dsa_cp_is_disabled(self):
         ):
             self.assertFalse(is_dsa_enable_prefill_cp())
 
-    def test_disabled_dsa_cp_skips_platform_probes(self):
-        parallel = SimpleNamespace(attn_cp_size=1)
-
-        with (
-            patch(
-                "sglang.srt.layers.attention.dsa.utils.get_parallel",
-                return_value=parallel,
-            ),
-            patch("sglang.srt.layers.attention.dsa.utils.is_hip") as mock_is_hip,
-            patch("sglang.srt.layers.attention.dsa.utils.is_npu") as mock_is_npu,
-            patch("sglang.srt.layers.attention.dsa.utils.is_musa") as mock_is_musa,
-        ):
-            self.assertFalse(is_dsa_enable_prefill_cp())
-
-        mock_is_hip.assert_not_called()
-        mock_is_npu.assert_not_called()
-        mock_is_musa.assert_not_called()
-
 
 class TestPrefillCPBCGReplay(CustomTestCase):
     def tearDown(self):