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

fix(dsa): correct packed FlashInfer top-k and backend selection semantics (#32490)

Browse Source 
Co-authored-by: Parth Chadha <parth@humansand.ai>
This commit is contained in:
Ziang Li
2026-07-30 20:20:16 -07:00
committed by GitHub
co-authored by Parth Chadha
parent c039e1a7ee
commit 0aefba7283
3 changed files with 132 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
+37 -9
View File
@@ -34,6 +34,9 @@ class DSATopKBackend(Enum):
def is_flashinfer(self) -> bool:
return self == DSATopKBackend.FLASHINFER
def should_use_topk_v2(self) -> bool:
return self.is_sgl_kernel() and envs.SGLANG_OPT_USE_TOPK_V2.get()
def topk_func(
self,
score: torch.Tensor,
@@ -88,18 +91,19 @@ class DSATopKBackend(Enum):
if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk:
return self.topk_func(logits, lengths, topk, row_starts=row_starts)
# Decode-shaped PAGED top-k (plain decode AND spec verify / draft-extend,
# whose expanded rows match the same shape) routes to the DeepSeek-V4 top-k
# v2 JIT kernel, which fuses top-k selection and the page-table transform in
# one launch and consumes the indexer's own page_size>=1 table directly, so
# no page_size=1 table is materialized. Shared by DeepSeek-V3.2 and GLM DSA.
# Decode-shaped PAGED top-k for the SGL backend (plain decode AND spec
# verify / draft-extend, whose expanded rows match the same shape) routes
# to the DeepSeek-V4 top-k v2 JIT kernel. It fuses top-k selection and the
# page-table transform in one launch and consumes the indexer's own
# page_size>=1 table directly, so no page_size=1 table is materialized.
# Shared by DeepSeek-V3.2 and GLM DSA.
# This is a deterministic dispatch on the work shape, not a best-effort
# attempt: the fused-decode CUDA graph drops the page_size=1 table for
# exactly this case (see dsa_drop_wide_page_table), so once the shape
# matches we commit to v2 and never silently fall back to the legacy
# page_size=1 path from here.
if (
envs.SGLANG_OPT_USE_TOPK_V2.get()
self.should_use_topk_v2()
and topk_transform_method == TopkTransformMethod.PAGED
and row_starts is None
and batch_idx_list is None
@@ -154,6 +158,26 @@ class DSATopKBackend(Enum):
import flashinfer
if topk_transform_method == TopkTransformMethod.PAGED:
if row_starts is not None:
# Packed PAGED extend uses batch-global score offsets with
# request-local page tables. FlashInfer applies row_starts
# to both, so reuse the SGL transform.
from sgl_kernel import fast_topk_transform_fused
page_table_size_1 = (
attn_metadata.page_table_1[batch_idx_list]
if batch_idx_list is not None
else attn_metadata.page_table_1
)
return fast_topk_transform_fused(
score=logits,
lengths=lengths,
page_table_size_1=page_table_size_1,
cu_seqlens_q=cu_seqlens_q_topk,
topk=topk,
row_starts=row_starts,
)
row_to_batch, local_row_starts = _build_flashinfer_paged_args(
attn_metadata=attn_metadata,
row_starts=row_starts,
@@ -268,7 +292,6 @@ def _topk_transform_v2_paged(
``seqlens_expand_kernel``); 0 takes the trivial all-(-1) output path.
"""
from sglang.kernels.ops.attention.dsv4.topk import topk_transform_512_v2
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
num_rows = logits.shape[0]
@@ -297,7 +320,7 @@ def _topk_transform_v2_paged(
plan is not None and plan.shape[0] == num_rows + 1
), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)"
page_size = get_token_to_kv_pool().page_size
page_size = attn_metadata.page_size
out = logits.new_full((num_rows, topk), -1, dtype=torch.int32)
topk_transform_512_v2(logits, lengths_i32, page_table, out, page_size, plan)
return out
@@ -317,6 +340,8 @@ def _build_flashinfer_paged_args(
else None
)
# Both dynamic mappings contain one entry per logit row. Supplying the known
# size avoids synchronizing CUDA to infer the sum of the repeat counts.
if (
row_to_batch is not None
and cu_seqlens_q_topk is not None
@@ -325,7 +350,9 @@ def _build_flashinfer_paged_args(
q_lens = (cu_seqlens_q_topk[1:] - cu_seqlens_q_topk[:-1]).to(
dtype=torch.int32, device=device
)
row_to_batch = torch.repeat_interleave(row_to_batch, q_lens)
row_to_batch = torch.repeat_interleave(
row_to_batch, q_lens, output_size=num_rows
)
if row_to_batch is None and cu_seqlens_q_topk is not None:
# Decode-like case (one query row per batch) does not need an explicit mapping.
@@ -338,6 +365,7 @@ def _build_flashinfer_paged_args(
row_to_batch = torch.repeat_interleave(
torch.arange(q_lens.shape[0], dtype=torch.int32, device=device),
q_lens,
output_size=num_rows,
)
if row_starts is not None and row_to_batch is None:
python/sglang/srt/layers/attention/dsa_backend.py
+13 -11
View File
@@ -229,7 +229,7 @@ class DSAMetadata:
paged_mqa_ctx_lens_2d: Optional[torch.Tensor] = None
# Precomputed once per forward batch and reused across layers: the
# DeepSeek-V4 top-k v2 plan (cluster-threshold metadata) for the folded
# decode top-k transform. None unless SGLANG_OPT_USE_TOPK_V2 and decode.
# decode top-k transform. None unless the SGL top-k v2 path is enabled.
topk_v2_plan: Optional[torch.Tensor] = None
# The sum of sequence lengths for key, prefill only
seq_lens_sum: Optional[int] = None
@@ -329,6 +329,8 @@ class DSAIndexerMetadata(BaseIndexerMetadata):
cu_topk_indices_offset = torch.repeat_interleave(
cu_seqlens_q_topk[:-1],
cu_seqlens_q,
# Avoid reading sum(cu_seqlens_q) back to the host.
output_size=logits.shape[0],
)
else:
cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q
@@ -778,8 +780,8 @@ class DeepseekSparseAttnBackend(
# that dispatches to `_topk_transform_v2_paged` -- decode AND MTP
# target-verify / draft-extend, whose expanded row count is exactly what v2
# sees -- otherwise the helper's plan-present assertion fires. None only
# when the fold is disabled; such metadata is never dispatched to v2.
if not envs.SGLANG_OPT_USE_TOPK_V2.get():
# when the SGL v2 path is disabled; such metadata is never dispatched to v2.
if not self.dsa_topk_backend.should_use_topk_v2():
return None
from sglang.kernels.ops.attention.dsv4.topk import plan_topk_v2
@@ -1250,12 +1252,12 @@ class DeepseekSparseAttnBackend(
# page_size=1 table. This MUST match the exact condition under which
# `DSATopKBackend.topk_transform` dispatches decode PAGED to
# `_topk_transform_v2_paged` -- otherwise the legacy transform would read a
# dropped (None) table. Hence: fused top-k AND v2 enabled AND index_topk in
# the kernel's supported range, on CUDA with page_size>1. Excludes HIP (its
# indexer reads page_table_1), hisparse (needs page_size=1 loc translation),
# and spec decoding (MTP precompute fast-path + target-verify/draft-extend
# still consume the wide table). Computed once from stable config; the graph
# is captured once per process.
# dropped (None) table. Hence: SGL top-k backend AND fused top-k AND v2
# enabled AND index_topk in the kernel's supported range, on CUDA with
# page_size>1. Excludes HIP (its indexer reads page_table_1), hisparse
# (needs page_size=1 loc translation), and spec decoding (MTP precompute
# fast-path + target-verify/draft-extend still consume the wide table).
# Computed once from stable config; the graph is captured once per process.
self.dsa_drop_wide_page_table = (
is_cuda()
and not _is_hip
@@ -1263,9 +1265,9 @@ class DeepseekSparseAttnBackend(
and self.hisparse_coordinator is None
and not self.speculative_num_draft_tokens
and self.use_fused_topk
and envs.SGLANG_OPT_USE_TOPK_V2.get()
and self.dsa_topk_backend.should_use_topk_v2()
and self.dsa_index_topk is not None
and self.dsa_index_topk <= 2048
and 0 < self.dsa_index_topk <= 2048
)
max_ctx_len = self.req_to_token.shape[1]
test/registered/kernels/ops/attention/test_dsa_indexer.py
+82 -8
View File
@@ -679,7 +679,14 @@ class TestDSAIndexer(CustomTestCase):
topk_backend=DSATopKBackend.FLASHINFER,
)
with envs.SGLANG_DSA_FUSE_TOPK.override(True):
repeat_interleave = torch.repeat_interleave
with (
envs.SGLANG_DSA_FUSE_TOPK.override(True),
patch(
"sglang.srt.layers.attention.dsa_backend.torch.repeat_interleave",
wraps=repeat_interleave,
) as mock_repeat_interleave,
):
out_sgl = metadata_sgl.topk_transform(
logits,
topk,
@@ -695,6 +702,15 @@ class TestDSAIndexer(CustomTestCase):
batch_idx_list=batch_idx_list,
)
if query_lens is not None:
self.assertTrue(mock_repeat_interleave.call_args_list)
self.assertTrue(
all(
call.kwargs.get("output_size") == num_rows
for call in mock_repeat_interleave.call_args_list
)
)
self.assertEqual(out_sgl.shape, out_flashinfer.shape)
self.assertEqual(out_sgl.dtype, out_flashinfer.dtype)
self.assertEqual(out_sgl.dtype, torch.int32)
@@ -957,13 +973,6 @@ class TestDSAIndexer(CustomTestCase):
TopkTransformMethod.RAGGED,
]:
for with_row_starts in [False, True]:
if (
topk_transform_method == TopkTransformMethod.PAGED
and with_row_starts
):
# The synthetic paged fixture uses the decode-like row mapping.
# Ragged fused and unfused cases cover shifted row windows.
continue
with self.subTest(
tie_break=tie_break,
topk_transform_method=topk_transform_method.name,
@@ -994,6 +1003,71 @@ class TestDSAIndexer(CustomTestCase):
with_row_starts=False,
query_lens=[1, 2, 3, 1, 2, 1, 3, 2],
)
with self.subTest(
tie_break=tie_break,
topk_transform_method=TopkTransformMethod.PAGED.name,
with_row_starts=True,
query_lens="multi",
):
with envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override(tie_break):
self._run_fused_topk_backend_equivalence_test(
batch_size=batch_size,
max_score_len=max_score_len,
topk=topk,
topk_transform_method=TopkTransformMethod.PAGED,
with_row_starts=True,
query_lens=[1, 2, 3, 1, 2, 1, 3, 2],
)
def test_topk_v2_respects_topk_backend(self):
seq_lens = torch.tensor([2048, 4096], dtype=torch.int32, device=self.device)
expected_plan = torch.empty(3, dtype=torch.int32, device=self.device)
for topk_backend, should_use_topk_v2 in [
(DSATopKBackend.SGL_KERNEL, True),
(DSATopKBackend.FLASHINFER, False),
]:
with self.subTest(topk_backend=topk_backend.value):
backend = object.__new__(DeepseekSparseAttnBackend)
backend.device = self.device
backend.real_page_size = 64
backend.hisparse_coordinator = None
backend.speculative_num_draft_tokens = 0
backend.use_fused_topk = True
backend.dsa_topk_backend = topk_backend
backend.dsa_index_topk = 2048
backend.dsa_decode_impl = "fa3"
backend.req_to_token = torch.empty(
2, 4096, dtype=torch.int32, device=self.device
)
with (
envs.SGLANG_OPT_USE_TOPK_V2.override(True),
patch(
"sglang.kernels.ops.attention.dsv4.topk.plan_topk_v2",
return_value=expected_plan,
) as mock_plan_topk_v2,
):
self.assertEqual(
topk_backend.should_use_topk_v2(), should_use_topk_v2
)
actual_plan = backend._build_topk_v2_plan(seq_lens)
backend.init_cuda_graph_state(max_bs=2, max_num_tokens=2)
if should_use_topk_v2:
self.assertIs(actual_plan, expected_plan)
mock_plan_topk_v2.assert_called_once_with(seq_lens)
else:
self.assertIsNone(actual_plan)
mock_plan_topk_v2.assert_not_called()
self.assertEqual(
backend.dsa_drop_wide_page_table,
should_use_topk_v2,
)
self.assertEqual(
backend.decode_cuda_graph_metadata["page_table"] is None,
should_use_topk_v2,
)
# TODO: enable this test after indexer accuracy aligned
# @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm")