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
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")