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

[Spec][DSA] Add --speculative-dsa-topk-backend (#36313)

Browse Source
This commit is contained in:
Ziang Li
2026-08-25 23:35:03 -07:00
committed by GitHub
parent 0c84eaba7f
commit 3c9febc68b
9 changed files with 51 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/advanced_features/server_arguments.mdx
+7 -1
View File
@@ -1458,7 +1458,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
</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.05)"}}>Choose the DSA indexer top-k backend for the target model. 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>
@@ -1607,6 +1607,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.05)"}}>Attention backend for speculative decoding drafting.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Same as attention backend options</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--speculative-dsa-topk-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the DSA indexer top-k backend for speculative draft workers. 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)"}}>`--speculative-moe-runner-backend`</td>
docs/docs/advanced_features/speculative_decoding.mdx
+13
View File
@@ -202,6 +202,11 @@ To enable EAGLE speculative decoding the following parameters are relevant:
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Override attention backend for the draft model.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>None</code> (same as target)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--speculative-dsa-topk-backend</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Select the DSA indexer top-k backend for speculative draft workers independently of <code>--dsa-topk-backend</code>. Options are <code>sgl-kernel</code>, <code>torch</code>, and <code>flashinfer</code>; <code>torch</code> requires <code>SGLANG_DSA_FUSE_TOPK=false</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>sgl-kernel</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--speculative-draft-model-quantization</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Quantization method for the draft model. Use <code>"unquant"</code> to force no quantization even when the target model is quantized.</td>
@@ -409,6 +414,8 @@ python3 -m sglang.launch_server \
--log-level warning
```
For DSA-based MTP, draft workers default to the `sgl-kernel` top-k backend. Use `--speculative-dsa-topk-backend` to override the draft independently of `--dsa-topk-backend` for the target model.
**Send a request:**
```python Example
@@ -826,6 +833,12 @@ Below is a comprehensive list of all speculative decoding parameters available i
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>None</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Override attention backend for the draft model</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--speculative-dsa-topk-backend</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>str</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>sgl-kernel</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DSA indexer top-k backend for speculative draft workers, independent of <code>--dsa-topk-backend</code> (<code>sgl-kernel</code>, <code>torch</code>, or <code>flashinfer</code>)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--speculative-moe-runner-backend</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>str</code></td>
docs/docs/references/environment_variables.mdx
+2 -2
View File
@@ -523,12 +523,12 @@ SGLang supports various environment variables that can be used to configure its
</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.05)"}}>Use deterministic FlashInfer topk kernels when either <code>--dsa-topk-backend=flashinfer</code> or <code>--speculative-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.05)"}}>Tie-break mode for FlashInfer DSA topk when either <code>--dsa-topk-backend=flashinfer</code> or <code>--speculative-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>
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+1 -3
View File
@@ -560,9 +560,7 @@ class DeepseekV4AttnBackend(
self.enable_deepseek_v4_fp4_indexer: bool = (
model_runner.server_args.enable_deepseek_v4_fp4_indexer
)
self.dsa_topk_backend: DSATopKBackend = DSATopKBackend(
model_runner.server_args.dsa_topk_backend
)
self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner)
self.dsv4_prefill_backend: str = getattr(
model_runner.server_args, "dsv4_prefill_backend", "auto"
)
python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
+16 -1
View File
@@ -1,11 +1,15 @@
from __future__ import annotations
from enum import Enum, IntEnum, auto
from typing import Callable, Dict, List, Optional, Tuple
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
import torch
from sglang.srt.environ import envs
from sglang.srt.runtime_context import get_exec, get_spec
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
_FLASHINFER_TIE_BREAK_VALUES = {
"small": 1,
@@ -25,6 +29,17 @@ class DSATopKBackend(Enum):
TORCH = "torch"
FLASHINFER = "flashinfer"
@classmethod
def resolve(cls, model_runner: ModelRunner) -> DSATopKBackend:
"""Resolve the DSA top-k backend for one model runner.
``--dsa-topk-backend`` selects the target backend, while
``--speculative-dsa-topk-backend`` independently selects the draft.
"""
if model_runner.is_draft_worker:
return cls(get_spec().speculative_dsa_topk_backend)
return cls(get_exec().kernel.dsa_topk_backend)
def is_sgl_kernel(self) -> bool:
return self == DSATopKBackend.SGL_KERNEL
python/sglang/srt/layers/attention/dsa_backend.py
+1 -3
View File
@@ -340,9 +340,7 @@ class DeepseekSparseAttnBackend(
self.supports_mha_one_shot: bool = True
self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend
self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend
self.dsa_topk_backend: DSATopKBackend = DSATopKBackend(
model_runner.server_args.dsa_topk_backend
)
self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner)
if self.num_q_heads <= 64:
self.flashmla_kv_num_q_heads = 64
elif self.num_q_heads <= 128:
python/sglang/srt/server_args.py
+9 -1
View File
@@ -1897,7 +1897,7 @@ class ServerArgs:
dsa_topk_backend: A[
str,
Arg(
help="DSA indexer top-k backend. Options: 'sgl-kernel', 'torch', 'flashinfer'. The 'torch' backend currently requires SGLANG_DSA_FUSE_TOPK=false.",
help="DSA indexer top-k backend for the target model. Options: 'sgl-kernel', 'torch', 'flashinfer'. The 'torch' backend currently requires SGLANG_DSA_FUSE_TOPK=false.",
choices=DSA_TOPK_BACKEND_CHOICES,
),
NS("exec.kernel"),
@@ -2264,6 +2264,14 @@ class ServerArgs:
),
NS("spec"),
] = None
speculative_dsa_topk_backend: A[
str,
Arg(
help="DSA indexer top-k backend for speculative draft workers. Options: 'sgl-kernel', 'torch', 'flashinfer'. The 'torch' backend currently requires SGLANG_DSA_FUSE_TOPK=false.",
choices=DSA_TOPK_BACKEND_CHOICES,
),
NS("spec"),
] = "sgl-kernel"
speculative_draft_kv_cache_dtype: A[
Optional[str],
Arg(
python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py
+1
View File
@@ -298,6 +298,7 @@ class DSAMockModelRunner(ModelRunner):
self.prefill_attention_backend_str = case.backend
self.decode_attention_backend_str = case.backend
self.draft_attention_backend = None
self.is_draft_worker = False
# For TARGET_VERIFY / DRAFT_EXTEND, the DSA backend uses
# `self.speculative_num_draft_tokens` to size `seqlens_expanded`
# (`dsa_backend.py:482-486,510-515`). When zero, deep_gemm's
test/registered/kernels/ops/attention/test_dsa_indexer.py
+1
View File
@@ -180,6 +180,7 @@ class MockModelRunner:
self.config = {**DEFAULT_CONFIG, **(config or {})}
self.dtype = self.config["dtype"]
self.kv_cache_dtype = self.config["kv_cache_dtype"]
self.is_draft_worker = False
self.is_hybrid_swa = False
# Model configuration