minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[misc] Revert #38346, #33426, #39061 and #39219 (#39405)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-14 03:04:07 -07:00
committed by GitHub
parent 5aa9b8fb3e
commit 66c7bc838e
9 changed files with 49 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py
+29 -16
View File
@@ -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],
python/sglang/srt/layers/attention/dsa/utils.py
+2 -3
View File
@@ -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
+11 -32
View File
@@ -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.
python/sglang/srt/layers/attention/qsa/qsa_indexer.py
-1
View File
@@ -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(
python/sglang/srt/utils/common.py
+5 -9
View File
@@ -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
-100
View File
@@ -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()
test/registered/cp/test_cp_strategy_unit.py
-18
View File
@@ -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):
test/registered/kernels/test_fused_kda_conv_recurrent_verify.py
+2 -24
View File
@@ -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__]))
test/registered/unit/utils/test_common.py
-15
View File
@@ -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`