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]