diff --git a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py index 2d4a56fde..dfb5bc0ef 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py +++ b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py @@ -12,16 +12,15 @@ two transpose copies the unfused path needs to feed the conv kernel. 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``. +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). -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. +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], diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index f685334ab..db99b5b20 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 45a6572cf..e36841e44 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -50,6 +50,11 @@ if TYPE_CHECKING: 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 @@ class FlashAttentionBackend(AttentionBackend): 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 @@ class FlashAttentionBackend(AttentionBackend): 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 @@ class FlashAttentionBackend(AttentionBackend): o_expand, softmax_lse_expand.T.contiguous(), ) - elif return_lse: - o, lse_out, *_ = result else: o = result else: @@ -1826,12 +1823,7 @@ class FlashAttentionBackend(AttentionBackend): 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 @@ class FlashAttentionBackend(AttentionBackend): 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 @@ class FlashAttentionBackend(AttentionBackend): 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 @@ class FlashAttentionBackend(AttentionBackend): o_expand, softmax_lse_expand.T.contiguous(), ) - elif return_lse: - o, lse_out, *_ = result else: o = result else: @@ -2184,12 +2168,7 @@ class FlashAttentionBackend(AttentionBackend): 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. diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py index a15585835..4de4c6a8d 100644 --- a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py +++ b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py @@ -321,7 +321,6 @@ class QSAIndexer(MultiPlatformOp): 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( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index cb6efd8da..62ad5159f 100644 --- a/python/sglang/srt/utils/common.py +++ b/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) diff --git a/test/registered/attention/unittests/dense/test_fa4.py b/test/registered/attention/unittests/dense/test_fa4.py index c150248b7..2adbce5d6 100644 --- a/test/registered/attention/unittests/dense/test_fa4.py +++ b/test/registered/attention/unittests/dense/test_fa4.py @@ -3,13 +3,9 @@ import unittest 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 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase): 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() diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 23b652bd6..1bc6a537d 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -121,24 +121,6 @@ class TestCPStrategyUnit(CustomTestCase): ): 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): diff --git a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py index 3645cc1db..f67fe4634 100644 --- a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py +++ b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py @@ -166,12 +166,12 @@ def _compare_case(case, num_warps): idx_vals = inp["idx_vals"] valid_rows = [i for i, slot in enumerate(idx_vals) if slot >= 0] + touched_slots = [slot for slot in idx_vals if slot >= 0] o_ref_v = o_ref.reshape(B, T, HV, V)[valid_rows] o_fus_v = o_fus.reshape(B, T, HV, V)[valid_rows] assert torch.equal(o_ref_v, o_fus_v) - # conv_state is read-only in verify; the commit scatter advances it. - assert torch.equal(inp["conv_pool"], conv_fus) + assert torch.equal(conv_ref[touched_slots], conv_fus[touched_slots]) assert torch.equal(win_ref[valid_rows], win_fus[valid_rows]) torch.testing.assert_close( ic_ref[valid_rows], ic_fus[valid_rows], atol=4e-3, rtol=0 @@ -183,27 +183,5 @@ def test_matches_unfused_reference(case): _compare_case(case, num_warps=4) -def test_output_does_not_depend_on_cta_scheduling(): - """The verify output must not change with how the CTAs happen to be - scheduled. H=1 with HV=16 shares one Q/K history across 16 V tiles.""" - if torch.cuda.get_device_capability()[0] < 9: - pytest.skip("green contexts need SM90 or newer") - from flashinfer.green_ctx import split_device_green_ctx_by_sm_count - - case = (1, 6, 1, 16, 128, 128, 4, False, None, False, 1) - B, T, H, HV, K, V, W, has_bias, lower_bound, neg_slot, seed = case - inp = _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed) - full = _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps=4)[0] - - streams, _ = split_device_green_ctx_by_sm_count(torch.device("cuda:0"), [8]) - # The green stream is non-blocking, so it must be told to wait for the - # inputs produced above; synchronize() afterwards only waits on the consumer. - streams[0].wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(streams[0]): - squeezed = _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps=4)[0] - streams[0].synchronize() - assert torch.equal(full, squeezed) - - if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/unit/utils/test_common.py b/test/registered/unit/utils/test_common.py index 4d772b9a0..f93a65ecd 100644 --- a/test/registered/unit/utils/test_common.py +++ b/test/registered/unit/utils/test_common.py @@ -7,7 +7,6 @@ from sglang.srt.utils.common import ( flatten_arrays_to_int64_tensor, get_device_sm_nvidia_smi, get_nvidia_driver_version_str, - is_musa, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -16,20 +15,6 @@ register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd") -class TestMusaDetection(CustomTestCase): - def test_is_musa_is_torch_compile_safe(self): - is_musa.cache_clear() - - @torch.compile(backend="eager", fullgraph=True) - def add_platform_offset(value): - return value + 1 if is_musa() else value - 1 - - value = torch.zeros(1) - actual = add_platform_offset(value) - expected = torch.ones(1) if is_musa() else -torch.ones(1) - torch.testing.assert_close(actual, expected) - - @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") class TestFlattenArraysToInt64Tensor(CustomTestCase): """`flatten_arrays_to_int64_tensor` is invoked by `prepare_for_extend`