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

[FlashInfer v0.6.10] [RL] [DSv32] [GLM-5] Add --dsa-topk-backend and integrate FlashInfer and pytorch topk (#22851)

Browse Source
This commit is contained in:
Ziang Li
2026-05-25 13:08:03 -07:00
committed by GitHub
parent b13d3d18c6
commit 2b9dd9c8b3
9 changed files with 706 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/server_arguments.md
+1
View File
@@ -273,6 +273,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` |
| `--dsa-prefill-backend` | Choose the DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention). `--nsa-prefill-backend` is a deprecated alias. | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter`, `trtllm` |
| `--dsa-decode-backend` | Choose the DSA backend for the decode stage when running DeepSeek DSA-style attention. Overrides `--attention-backend` for decoding. `--nsa-decode-backend` is a deprecated alias. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`, `trtllm` |
| `--dsa-topk-backend` | Choose the DSA indexer top-k backend. The `torch` backend currently requires `SGLANG_DSA_FUSE_TOPK=false`. | `sgl-kernel` | `sgl-kernel`, `torch`, `flashinfer` |
| `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only).| `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `flashinfer_cutlass`, `flashinfer_deepgemm`, `cutlass`, `triton`, `aiter` |
| `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'flashinfer_cutlass' (default), 'auto' (auto-selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). All backends are from FlashInfer; when FlashInfer is unavailable, sgl-kernel CUTLASS is used as an automatic fallback.| `flashinfer_cutlass` | `auto`, `flashinfer_cudnn`, `flashinfer_cutlass`, `flashinfer_trtllm` |
| `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) |
docs/references/environment_variables.md
+2
View File
@@ -95,6 +95,8 @@ SGLang supports various environment variables that can be used to configure its
| Environment Variable | Description | Default Value |
| --- | --- | --- |
| `SGLANG_DSA_FUSE_TOPK` | Fuse the operation of picking topk logits and picking topk indices from page table (`SGLANG_NSA_FUSE_TOPK` is a deprecated alias) | `true` |
| `SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC` | Use deterministic FlashInfer topk kernels when `--dsa-topk-backend=flashinfer` | `false` |
| `SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK` | Tie-break mode for FlashInfer DSA topk when `--dsa-topk-backend=flashinfer`: unset disables explicit tie-breaking, `small` prefers the smaller candidate index for equal scores, and `large` prefers the larger candidate index for equal scores. Setting this variable makes FlashInfer use deterministic topk. | `unset` |
| `SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA` | Precompute metadata that can be shared among different draft steps when MTP is enabled (`SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA` is a deprecated alias) | `true` |
| `SGLANG_USE_FUSED_METADATA_COPY` | Control whether to use fused metadata copy kernel for cuda graph replay | `true` |
| `SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` | When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2) (`SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` is a deprecated alias) | `2048` |
docs_new/docs/advanced_features/server_arguments.mdx
+6
View File
@@ -1206,6 +1206,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`fa3`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-topk-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the DSA indexer top-k backend. The `torch` backend currently requires `SGLANG_DSA_FUSE_TOPK=false`.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`sgl-kernel`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>sgl-kernel</code>, <code>torch</code>, <code>flashinfer</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--fp8-gemm-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only).</td>
docs_new/docs/references/environment_variables.mdx
+10
View File
@@ -416,6 +416,16 @@ SGLang supports various environment variables that can be used to configure its
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Fuse the operation of picking topk logits and picking topk indices from page table. <code>SGLANG_NSA_FUSE_TOPK</code> is a deprecated alias.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use deterministic FlashInfer topk kernels when <code>--dsa-topk-backend=flashinfer</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Tie-break mode for FlashInfer DSA topk when <code>--dsa-topk-backend=flashinfer</code>: unset disables explicit tie-breaking, <code>small</code> prefers the smaller candidate index for equal scores, and <code>large</code> prefers the larger candidate index for equal scores. Setting this variable makes FlashInfer use deterministic topk.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>unset</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Precompute metadata that can be shared among different draft steps when MTP is enabled. <code>SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA</code> is a deprecated alias.</td>
python/sglang/srt/environ.py
+2
View File
@@ -465,6 +465,8 @@ class Envs:
# DSA Backend (canonical names; fall back to SGLANG_NSA_* with deprecation warning)
SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK")
SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC = EnvBool(False)
SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK = EnvStr(None)
SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBoolWithAlias(
True, deprecated_name="SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA"
)
python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
+271
View File
@@ -0,0 +1,271 @@
from __future__ import annotations
from enum import Enum, IntEnum, auto
from typing import Callable, Dict, List, Optional, Tuple
import torch
from sglang.srt.environ import envs
_FLASHINFER_TIE_BREAK_VALUES = {
"small": 1,
"large": 2,
}
class TopkTransformMethod(IntEnum):
# Transform topk indices to indices to the page table (page_size = 1)
PAGED = auto()
# Transform topk indices to indices to ragged kv (non-paged)
RAGGED = auto()
class DSATopKBackend(Enum):
SGL_KERNEL = "sgl-kernel"
TORCH = "torch"
FLASHINFER = "flashinfer"
def is_sgl_kernel(self) -> bool:
return self == DSATopKBackend.SGL_KERNEL
def is_torch(self) -> bool:
return self == DSATopKBackend.TORCH
def is_flashinfer(self) -> bool:
return self == DSATopKBackend.FLASHINFER
def topk_func(
self,
score: torch.Tensor,
lengths: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if self.is_sgl_kernel():
from sgl_kernel import fast_topk_v2
return fast_topk_v2(score, lengths, topk, row_starts=row_starts)
if self.is_torch():
return _topk_unfused(
score,
lengths,
topk,
row_starts=row_starts,
topk_op=torch.topk,
topk_op_kwargs={"dim": -1},
)
if self.is_flashinfer():
import flashinfer
return _topk_unfused(
score,
lengths,
topk,
row_starts=row_starts,
topk_op=flashinfer.top_k,
topk_op_kwargs={
"sorted": False,
"deterministic": envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
"tie_break": _flashinfer_tie_break_value(),
"dsa_graph_safe": True,
},
)
raise RuntimeError(f"Unsupported {self = }.")
def topk_transform(
self,
logits: torch.Tensor,
lengths: torch.Tensor,
topk: int,
topk_transform_method: TopkTransformMethod,
attn_metadata,
cu_seqlens_q_topk: Optional[torch.Tensor] = None,
topk_indices_offset: Optional[torch.Tensor] = None,
row_starts: Optional[torch.Tensor] = None,
batch_idx_list: Optional[List[int]] = None,
force_unfused_topk: bool = False,
) -> torch.Tensor:
if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk:
return self.topk_func(logits, lengths, topk, row_starts=row_starts)
if self.is_sgl_kernel():
from sgl_kernel import (
fast_topk_transform_fused,
fast_topk_transform_ragged_fused,
)
if topk_transform_method == TopkTransformMethod.PAGED:
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,
)
if topk_transform_method == TopkTransformMethod.RAGGED:
if topk_indices_offset is None:
raise RuntimeError(
"RAGGED topk_transform requires topk_indices_offset; "
"expected extend-without-speculative metadata."
)
return fast_topk_transform_ragged_fused(
score=logits,
lengths=lengths,
topk_indices_offset=topk_indices_offset,
topk=topk,
row_starts=row_starts,
)
raise RuntimeError(f"Unsupported {topk_transform_method = }.")
if self.is_flashinfer():
import flashinfer
if topk_transform_method == TopkTransformMethod.PAGED:
row_to_batch, local_row_starts = _build_flashinfer_paged_args(
attn_metadata=attn_metadata,
row_starts=row_starts,
cu_seqlens_q_topk=cu_seqlens_q_topk,
batch_idx_list=batch_idx_list,
device=logits.device,
num_rows=logits.shape[0],
)
return flashinfer.top_k_page_table_transform(
logits.contiguous(),
attn_metadata.page_table_1.contiguous(),
lengths.contiguous(),
topk,
row_to_batch=row_to_batch,
deterministic=envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
tie_break=_flashinfer_tie_break_value(),
dsa_graph_safe=True,
row_starts=local_row_starts,
)
if topk_transform_method == TopkTransformMethod.RAGGED:
if topk_indices_offset is None:
raise RuntimeError(
"RAGGED topk_transform requires topk_indices_offset; "
"expected extend-without-speculative metadata."
)
return flashinfer.top_k_ragged_transform(
logits.contiguous(),
topk_indices_offset.contiguous(),
lengths.contiguous(),
topk,
deterministic=envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
tie_break=_flashinfer_tie_break_value(),
dsa_graph_safe=True,
row_starts=row_starts,
)
raise RuntimeError(f"Unsupported {topk_transform_method = }.")
raise RuntimeError(f"Unsupported {self = } for SGLANG_DSA_FUSE_TOPK.")
def _topk_unfused(
score: torch.Tensor,
lengths: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
topk_op: Callable[..., Tuple[torch.Tensor, torch.Tensor]] = torch.topk,
topk_op_kwargs: Optional[Dict[str, object]] = None,
) -> torch.Tensor:
batch_size, max_score_len = score.shape
topk_indices = score.new_full((batch_size, topk), -1, dtype=torch.int32)
if batch_size == 0 or topk == 0 or max_score_len == 0:
return topk_indices
if row_starts is None:
row_starts = torch.zeros_like(lengths, dtype=torch.int32, device=score.device)
else:
row_starts = row_starts.to(dtype=torch.int32, device=score.device)
lengths = lengths.to(dtype=torch.int32, device=score.device)
col_indices = torch.arange(max_score_len, dtype=torch.int32, device=score.device)
col_indices = col_indices.unsqueeze(0)
row_starts_unsqueezed = row_starts.unsqueeze(1)
row_ends_unsqueezed = (row_starts + lengths).unsqueeze(1)
valid_mask = (col_indices >= row_starts_unsqueezed) & (
col_indices < row_ends_unsqueezed
)
masked_logits = score.masked_fill(~valid_mask, float("-inf"))
valid_topk = min(topk, max_score_len)
topk_kwargs = topk_op_kwargs or {}
topk_scores, topk_col_indices = topk_op(masked_logits, valid_topk, **topk_kwargs)
topk_local_indices = topk_col_indices.to(torch.int32) - row_starts_unsqueezed
topk_local_indices = topk_local_indices.masked_fill(
topk_scores == float("-inf"), -1
)
topk_indices[:, :valid_topk] = topk_local_indices
return topk_indices
def _build_flashinfer_paged_args(
attn_metadata,
row_starts: Optional[torch.Tensor],
cu_seqlens_q_topk: Optional[torch.Tensor],
batch_idx_list: Optional[List[int]],
device: torch.device,
num_rows: int,
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
row_to_batch = (
torch.as_tensor(batch_idx_list, dtype=torch.int32, device=device)
if batch_idx_list is not None
else None
)
if (
row_to_batch is not None
and cu_seqlens_q_topk is not None
and row_to_batch.shape[0] != num_rows
):
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)
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.
# Avoid dynamic tensor construction in this branch to keep CUDA graph capture safe.
num_batches = cu_seqlens_q_topk.shape[0] - 1
if not (row_starts is None and num_rows == num_batches):
q_lens = (cu_seqlens_q_topk[1:] - cu_seqlens_q_topk[:-1]).to(
dtype=torch.int32, device=device
)
row_to_batch = torch.repeat_interleave(
torch.arange(q_lens.shape[0], dtype=torch.int32, device=device),
q_lens,
)
if row_starts is not None and row_to_batch is None:
raise RuntimeError(
"PAGED topk_transform with row_starts requires cu_seqlens_q metadata."
)
local_row_starts = row_starts
if local_row_starts is not None and row_to_batch is not None:
local_row_starts = (
local_row_starts - attn_metadata.cu_seqlens_k[:-1][row_to_batch]
)
return row_to_batch, local_row_starts
def _flashinfer_tie_break_value() -> int:
mode = envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.get()
if mode is None:
return 0
mode = mode.lower()
if mode not in _FLASHINFER_TIE_BREAK_VALUES:
raise RuntimeError(
"SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK must be one of "
f"{tuple(_FLASHINFER_TIE_BREAK_VALUES)} or unset, got {mode!r}."
)
return _FLASHINFER_TIE_BREAK_VALUES[mode]
python/sglang/srt/layers/attention/dsa_backend.py
+46 -53
View File
@@ -2,8 +2,15 @@ from __future__ import annotations
import logging
from dataclasses import dataclass
from enum import IntEnum, auto
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias
from typing import (
TYPE_CHECKING,
Dict,
List,
Literal,
Optional,
Tuple,
TypeAlias,
)
import torch
@@ -19,6 +26,10 @@ from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import (
compute_cu_seqlens,
)
from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.attention.dsa.dsa_topk_backend import (
DSATopKBackend,
TopkTransformMethod,
)
from sglang.srt.layers.attention.dsa.quant_k_cache import quantize_k_cache
from sglang.srt.layers.attention.dsa.transform_index import (
transform_index_page_table_decode,
@@ -161,13 +172,6 @@ class DSAMetadata:
token_to_batch_idx: Optional[torch.Tensor] = None
class TopkTransformMethod(IntEnum):
# Transform topk indices to indices to the page table (page_size = 1)
PAGED = auto()
# Transform topk indices to indices to ragged kv (non-paged)
RAGGED = auto()
@torch.compile
def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor:
return torch.cat(tensors, dim=dim)
@@ -193,6 +197,7 @@ def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor:
class DSAIndexerMetadata(BaseIndexerMetadata):
attn_metadata: DSAMetadata
topk_transform_method: TopkTransformMethod
topk_backend: DSATopKBackend = DSATopKBackend.SGL_KERNEL
paged_mqa_schedule_metadata: Optional[torch.Tensor] = None
force_unfused_topk: bool = False
@@ -231,17 +236,11 @@ class DSAIndexerMetadata(BaseIndexerMetadata):
logits: torch.Tensor,
topk: int,
ks: Optional[torch.Tensor] = None,
cu_seqlens_q: torch.Tensor = None,
ke_offset: torch.Tensor = None,
batch_idx_list: List[int] = None,
cu_seqlens_q: Optional[torch.Tensor] = None,
ke_offset: Optional[torch.Tensor] = None,
batch_idx_list: Optional[List[int]] = None,
topk_indices_offset_override: Optional[torch.Tensor] = None,
) -> torch.Tensor:
from sgl_kernel import (
fast_topk_transform_fused,
fast_topk_transform_ragged_fused,
fast_topk_v2,
)
if topk_indices_offset_override is not None:
cu_topk_indices_offset = topk_indices_offset_override
cu_seqlens_q_topk = None
@@ -259,38 +258,18 @@ class DSAIndexerMetadata(BaseIndexerMetadata):
seq_lens_topk = ke_offset
else:
seq_lens_topk = self.get_seqlens_expanded()
if batch_idx_list is not None:
page_table_size_1 = self.attn_metadata.page_table_1[batch_idx_list]
else:
page_table_size_1 = self.attn_metadata.page_table_1
if not envs.SGLANG_DSA_FUSE_TOPK.get() or self.force_unfused_topk:
return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks)
elif self.topk_transform_method == TopkTransformMethod.PAGED:
# NOTE(dark): if fused, we return a transformed page table directly
return fast_topk_transform_fused(
score=logits,
lengths=seq_lens_topk,
page_table_size_1=page_table_size_1,
cu_seqlens_q=cu_seqlens_q_topk,
topk=topk,
row_starts=ks,
)
elif self.topk_transform_method == TopkTransformMethod.RAGGED:
if cu_topk_indices_offset is None:
raise RuntimeError(
"RAGGED topk_transform requires topk_indices_offset; "
"expected extend-without-speculative metadata."
)
return fast_topk_transform_ragged_fused(
score=logits,
lengths=seq_lens_topk,
topk_indices_offset=cu_topk_indices_offset,
topk=topk,
row_starts=ks,
)
else:
assert False, f"Unsupported {self.topk_transform_method = }"
return self.topk_backend.topk_transform(
logits=logits,
lengths=seq_lens_topk,
topk=topk,
topk_transform_method=self.topk_transform_method,
attn_metadata=self.attn_metadata,
cu_seqlens_q_topk=cu_seqlens_q_topk,
topk_indices_offset=cu_topk_indices_offset,
row_starts=ks,
batch_idx_list=batch_idx_list,
force_unfused_topk=self.force_unfused_topk,
)
_DSA_IMPL_T: TypeAlias = Literal[
@@ -343,6 +322,9 @@ class DeepseekSparseAttnBackend(
model_runner.server_args.dsa_prefill_backend
)
self.dsa_decode_impl: _DSA_IMPL_T = model_runner.server_args.dsa_decode_backend
self.dsa_topk_backend: DSATopKBackend = DSATopKBackend(
model_runner.server_args.dsa_topk_backend
)
if self.num_q_heads <= 64:
self.flashmla_kv_num_q_heads = 64
elif self.num_q_heads <= 128:
@@ -398,6 +380,16 @@ class DeepseekSparseAttnBackend(
else:
self.workspace_buffer = None
def _get_fused_topk_page_table(self, topk_indices: torch.Tensor) -> torch.Tensor:
if (
self.dsa_topk_backend.is_sgl_kernel()
or self.dsa_topk_backend.is_flashinfer()
):
return topk_indices
raise RuntimeError(
f"Unsupported {self.dsa_topk_backend = } for SGLANG_DSA_FUSE_TOPK."
)
def get_device_int32_arange(self, l: int) -> torch.Tensor:
if l > len(self._arange_buf):
next_pow_of_2 = 1 << (l - 1).bit_length()
@@ -1431,7 +1423,7 @@ class DeepseekSparseAttnBackend(
forward_batch.forward_mode
)
if envs.SGLANG_DSA_FUSE_TOPK.get():
page_table_1 = topk_indices
page_table_1 = self._get_fused_topk_page_table(topk_indices)
else:
if topk_transform_method == TopkTransformMethod.RAGGED:
topk_indices_offset = metadata.topk_indices_offset
@@ -1619,7 +1611,7 @@ class DeepseekSparseAttnBackend(
layer.layer_id,
)
elif envs.SGLANG_DSA_FUSE_TOPK.get():
page_table_1 = topk_indices
page_table_1 = self._get_fused_topk_page_table(topk_indices)
else:
page_table_1 = transform_index_page_table_decode(
page_table=metadata.page_table_1,
@@ -2126,7 +2118,7 @@ class DeepseekSparseAttnBackend(
topk_indices = self._pad_topk_indices(topk_indices, q.shape[0])
if envs.SGLANG_DSA_FUSE_TOPK.get():
page_table_1 = topk_indices
page_table_1 = self._get_fused_topk_page_table(topk_indices)
elif is_prefill:
page_table_1 = transform_index_page_table_prefill(
page_table=metadata.page_table_1,
@@ -2287,6 +2279,7 @@ class DeepseekSparseAttnBackend(
topk_transform_method=self.get_topk_transform_method(
forward_batch.forward_mode
),
topk_backend=self.dsa_topk_backend,
paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata,
force_unfused_topk=force_unfused,
)
python/sglang/srt/server_args.py
+12
View File
@@ -266,6 +266,8 @@ DSA_CHOICES = [
]
NSA_CHOICES = DSA_CHOICES # deprecated alias
DSA_TOPK_BACKEND_CHOICES = ["sgl-kernel", "torch", "flashinfer"]
MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"]
MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"]
@@ -557,6 +559,7 @@ class ServerArgs:
dsa_decode_backend: Optional[str] = (
None # auto-detect based on hardware/kv_cache_dtype
)
dsa_topk_backend: str = "sgl-kernel"
disable_flashinfer_autotune: bool = False
mamba_backend: str = "triton"
@@ -5494,6 +5497,15 @@ class ServerArgs:
choices=DSA_CHOICES,
help="[Deprecated] Use --dsa-decode-backend instead.",
)
parser.add_argument(
"--dsa-topk-backend",
dest="dsa_topk_backend",
default=ServerArgs.dsa_topk_backend,
type=str,
choices=DSA_TOPK_BACKEND_CHOICES,
help="DSA indexer top-k backend. Options: 'sgl-kernel', 'torch', 'flashinfer'. "
"The 'torch' backend currently requires SGLANG_DSA_FUSE_TOPK=false.",
)
parser.add_argument(
"--fp8-gemm-backend",
type=str,
test/registered/kernels/test_dsa_indexer.py
+356 -1
View File
@@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch
import torch
from sglang.srt.environ import envs
from sglang.srt.layers import dp_attention as _dp_attn
from sglang.test.ci.ci_register import register_cuda_ci
@@ -16,7 +17,15 @@ from sglang.srt.layers.attention.dsa.dsa_indexer import (
Indexer,
rotate_activation,
)
from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
from sglang.srt.layers.attention.dsa.dsa_topk_backend import (
DSATopKBackend,
TopkTransformMethod,
)
from sglang.srt.layers.attention.dsa_backend import (
DeepseekSparseAttnBackend,
DSAIndexerMetadata,
DSAMetadata,
)
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.linear import LinearBase
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
@@ -250,6 +259,7 @@ class MockModelRunner:
"enable_deterministic_inference": False,
"dsa_prefill_backend": "flashmla_sparse",
"dsa_decode_backend": "fa3",
"dsa_topk_backend": "sgl-kernel",
},
)()
self.hisparse_coordinator = None
@@ -423,6 +433,271 @@ class TestDSAIndexer(CustomTestCase):
"Output should have padding or exact topk size",
)
def _make_tie_free_logits(
self, batch_size: int, max_score_len: int
) -> torch.Tensor:
perm = torch.argsort(
torch.randn(
batch_size, max_score_len, dtype=torch.float32, device=self.device
),
dim=-1,
)
return torch.gather(
torch.arange(max_score_len, device=self.device, dtype=torch.float32)
.unsqueeze(0)
.expand(batch_size, -1),
dim=1,
index=perm,
)
def _run_unfused_topk_backend_validity_test(
self,
batch_size: int,
max_score_len: int,
topk: int,
topk_backend: DSATopKBackend,
with_row_starts: bool,
):
logits = self._make_tie_free_logits(batch_size, max_score_len)
if with_row_starts:
row_starts = torch.randint(
0,
max_score_len - 1,
(batch_size,),
dtype=torch.int32,
device=self.device,
)
max_lengths = max_score_len - row_starts
random_lengths = torch.randint(
0,
max_score_len - 1,
(batch_size,),
dtype=torch.int32,
device=self.device,
)
seq_lens_expanded = torch.minimum(max_lengths, random_lengths)
else:
row_starts = None
seq_lens_expanded = torch.randint(
0,
max_score_len - 1,
(batch_size,),
dtype=torch.int32,
device=self.device,
)
seq_lens_expanded = seq_lens_expanded.to(dtype=torch.int32, device=self.device)
max_seq_len_k = int(seq_lens_expanded.max().item())
cu_seqlens_q = torch.arange(
batch_size + 1, dtype=torch.int32, device=self.device
)
dsa_cu_seqlens_k = torch.zeros(
batch_size + 1, dtype=torch.int32, device=self.device
)
dsa_cu_seqlens_k[1:] = torch.cumsum(seq_lens_expanded, dim=0)
page_table_1 = (
torch.arange(max_seq_len_k, dtype=torch.int32, device=self.device)
.unsqueeze(0)
.expand(batch_size, -1)
.contiguous()
)
metadata = DSAIndexerMetadata(
attn_metadata=DSAMetadata(
page_size=1,
cache_seqlens_int32=seq_lens_expanded.clone(),
max_seq_len_q=1,
max_seq_len_k=max_seq_len_k,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_q.clone(),
page_table_1=page_table_1,
real_page_table=page_table_1,
dsa_cache_seqlens_int32=seq_lens_expanded.clone(),
dsa_cu_seqlens_q=cu_seqlens_q.clone(),
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
dsa_extend_seq_lens_list=seq_lens_expanded.cpu().tolist(),
dsa_seqlens_expanded=seq_lens_expanded,
),
topk_transform_method=TopkTransformMethod.PAGED,
topk_backend=topk_backend,
)
with envs.SGLANG_DSA_FUSE_TOPK.override(False):
topk_test = metadata.topk_transform(logits, topk, ks=row_starts)
self.assertEqual(topk_test.shape, (batch_size, topk))
self.assertEqual(topk_test.dtype, torch.int32)
expected_valid = torch.minimum(
seq_lens_expanded,
torch.full_like(seq_lens_expanded, topk),
)
actual_valid = (topk_test >= 0).sum(dim=-1).to(torch.int32)
self.assertTrue(torch.equal(actual_valid, expected_valid))
starts = (
row_starts.to(torch.int32)
if row_starts is not None
else torch.zeros(
(topk_test.shape[0],), dtype=torch.int32, device=topk_test.device
)
)
for row in range(topk_test.shape[0]):
test_row = topk_test[row]
valid_test = test_row[test_row >= 0]
expected_k = int(expected_valid[row].item())
self.assertEqual(valid_test.numel(), expected_k)
if expected_k == 0:
continue
start = int(starts[row].item())
row_len = int(seq_lens_expanded[row].item())
self.assertTrue(torch.all((valid_test >= 0) & (valid_test < row_len)))
self.assertEqual(torch.unique(valid_test).numel(), valid_test.numel())
row_scores = logits[row, start : start + row_len]
ref_topk = torch.topk(row_scores, expected_k, dim=-1, sorted=False).indices
self.assertTrue(
torch.equal(
torch.sort(valid_test.to(torch.int32)).values,
torch.sort(ref_topk.to(torch.int32)).values,
)
)
def _run_fused_topk_backend_equivalence_test(
self,
batch_size: int,
max_score_len: int,
topk: int,
topk_transform_method: TopkTransformMethod,
with_row_starts: bool,
query_lens: Optional[List[int]] = None,
):
num_rows = sum(query_lens) if query_lens is not None else batch_size
logits = self._make_tie_free_logits(num_rows, max_score_len)
if with_row_starts:
row_starts = torch.randint(
0,
max_score_len - 1,
(num_rows,),
dtype=torch.int32,
device=self.device,
)
max_lengths = max_score_len - row_starts
random_lengths = torch.randint(
1,
max_score_len,
(num_rows,),
dtype=torch.int32,
device=self.device,
)
seq_lens_expanded = torch.minimum(max_lengths, random_lengths)
else:
row_starts = None
seq_lens_expanded = torch.randint(
1,
max_score_len,
(num_rows,),
dtype=torch.int32,
device=self.device,
)
topk_indices_offset = (
torch.arange(num_rows, dtype=torch.int32, device=self.device)
* max_score_len
)
if query_lens is None:
cu_seqlens_q = torch.arange(
batch_size + 1, dtype=torch.int32, device=self.device
)
q_lens = None
batch_idx_list = None
else:
q_lens = torch.tensor(query_lens, dtype=torch.int32, device=self.device)
cu_seqlens_q = torch.zeros(
batch_size + 1, dtype=torch.int32, device=self.device
)
cu_seqlens_q[1:] = torch.cumsum(q_lens, dim=0)
batch_idx_list = list(range(batch_size))
cu_seqlens_k = torch.zeros(
batch_size + 1, dtype=torch.int32, device=self.device
)
dsa_cu_seqlens_k = torch.zeros(
num_rows + 1, dtype=torch.int32, device=self.device
)
dsa_cu_seqlens_k[1:] = torch.cumsum(seq_lens_expanded, dim=0)
page_table_1 = (
(
torch.arange(max_score_len, dtype=torch.int32, device=self.device)
.unsqueeze(0)
.expand(batch_size, -1)
)
+ (
torch.arange(
batch_size, dtype=torch.int32, device=self.device
).unsqueeze(1)
* max_score_len
)
).contiguous()
attn_metadata = DSAMetadata(
page_size=1,
cache_seqlens_int32=seq_lens_expanded.clone(),
max_seq_len_q=1,
max_seq_len_k=max_score_len,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
page_table_1=page_table_1,
real_page_table=page_table_1,
dsa_cache_seqlens_int32=seq_lens_expanded.clone(),
dsa_cu_seqlens_q=cu_seqlens_q.clone(),
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
dsa_extend_seq_lens_list=seq_lens_expanded.cpu().tolist(),
dsa_seqlens_expanded=seq_lens_expanded,
topk_indices_offset=(
topk_indices_offset
if topk_transform_method == TopkTransformMethod.RAGGED
else None
),
)
metadata_sgl = DSAIndexerMetadata(
attn_metadata=attn_metadata,
topk_transform_method=topk_transform_method,
topk_backend=DSATopKBackend.SGL_KERNEL,
)
metadata_flashinfer = DSAIndexerMetadata(
attn_metadata=attn_metadata,
topk_transform_method=topk_transform_method,
topk_backend=DSATopKBackend.FLASHINFER,
)
with envs.SGLANG_DSA_FUSE_TOPK.override(True):
out_sgl = metadata_sgl.topk_transform(
logits,
topk,
ks=row_starts,
cu_seqlens_q=q_lens,
batch_idx_list=batch_idx_list,
)
out_flashinfer = metadata_flashinfer.topk_transform(
logits,
topk,
ks=row_starts,
cu_seqlens_q=q_lens,
batch_idx_list=batch_idx_list,
)
self.assertEqual(out_sgl.shape, out_flashinfer.shape)
self.assertEqual(out_sgl.dtype, out_flashinfer.dtype)
self.assertEqual(out_sgl.dtype, torch.int32)
self.assertTrue(
torch.equal(
torch.sort(out_sgl, dim=-1).values,
torch.sort(out_flashinfer, dim=-1).values,
)
)
@patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm")
def test_indexer_basic_creation(self, mock_deep_gemm):
"""Test basic indexer creation and initialization."""
@@ -632,6 +907,86 @@ class TestDSAIndexer(CustomTestCase):
topk_indices = metadata.topk_transform(logits, topk)
self.assertEqual(topk_indices.shape, (batch_size, topk))
def test_topk_unfused_backends_valid_selection(self):
batch_size = 8
max_score_len = 16 * 1024
topk = 2048
for topk_backend in [
DSATopKBackend.SGL_KERNEL,
DSATopKBackend.TORCH,
DSATopKBackend.FLASHINFER,
]:
tie_break_values = (
[None, "small", "large"]
if topk_backend == DSATopKBackend.FLASHINFER
else [None]
)
for tie_break in tie_break_values:
for with_row_starts in [False, True]:
with self.subTest(
topk_backend=topk_backend.value,
tie_break=tie_break,
with_row_starts=with_row_starts,
):
with envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override(
tie_break
):
self._run_unfused_topk_backend_validity_test(
batch_size,
max_score_len,
topk,
topk_backend=topk_backend,
with_row_starts=with_row_starts,
)
def test_topk_fused_backends_equivalence(self):
batch_size = 8
max_score_len = 16 * 1024
topk = 2048
for tie_break in [None, "small", "large"]:
for topk_transform_method in [
TopkTransformMethod.PAGED,
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,
with_row_starts=with_row_starts,
):
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=topk_transform_method,
with_row_starts=with_row_starts,
)
with self.subTest(
tie_break=tie_break,
topk_transform_method=TopkTransformMethod.PAGED.name,
with_row_starts=False,
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=False,
query_lens=[1, 2, 3, 1, 2, 1, 3, 2],
)
# TODO: enable this test after indexer accuracy aligned
# @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm")
# def test_indexer_with_different_topk(self, mock_deep_gemm):