From 3c9febc68b23c5e5e8db4ac20b7fb1e5df44e0bb Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Tue, 25 Aug 2026 23:35:03 -0700 Subject: [PATCH] [Spec][DSA] Add --speculative-dsa-topk-backend (#36313) --- .../docs/advanced_features/server_arguments.mdx | 8 +++++++- .../advanced_features/speculative_decoding.mdx | 13 +++++++++++++ docs/docs/references/environment_variables.mdx | 4 ++-- .../srt/layers/attention/deepseek_v4_backend.py | 4 +--- .../layers/attention/dsa/dsa_topk_backend.py | 17 ++++++++++++++++- .../sglang/srt/layers/attention/dsa_backend.py | 4 +--- python/sglang/srt/server_args.py | 10 +++++++++- .../attention_methods/dsa_attention.py | 1 + .../kernels/ops/attention/test_dsa_indexer.py | 1 + 9 files changed, 51 insertions(+), 11 deletions(-) diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index ac4291753..9ecfff75a 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1458,7 +1458,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--dsa-topk-backend` - Choose the DSA indexer top-k backend. The `torch` backend currently requires `SGLANG_DSA_FUSE_TOPK=false`. + Choose the DSA indexer top-k backend for the target model. The `torch` backend currently requires `SGLANG_DSA_FUSE_TOPK=false`. `sgl-kernel` sgl-kernel, torch, flashinfer @@ -1607,6 +1607,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s Attention backend for speculative decoding drafting. `None` Same as attention backend options + + + `--speculative-dsa-topk-backend` + Choose the DSA indexer top-k backend for speculative draft workers. The `torch` backend currently requires `SGLANG_DSA_FUSE_TOPK=false`. + `sgl-kernel` + sgl-kernel, torch, flashinfer `--speculative-moe-runner-backend` diff --git a/docs/docs/advanced_features/speculative_decoding.mdx b/docs/docs/advanced_features/speculative_decoding.mdx index 52c89027b..084a2260e 100644 --- a/docs/docs/advanced_features/speculative_decoding.mdx +++ b/docs/docs/advanced_features/speculative_decoding.mdx @@ -202,6 +202,11 @@ To enable EAGLE speculative decoding the following parameters are relevant: Override attention backend for the draft model. None (same as target) + + --speculative-dsa-topk-backend + Select the DSA indexer top-k backend for speculative draft workers independently of --dsa-topk-backend. Options are sgl-kernel, torch, and flashinfer; torch requires SGLANG_DSA_FUSE_TOPK=false. + sgl-kernel + --speculative-draft-model-quantization Quantization method for the draft model. Use "unquant" to force no quantization even when the target model is quantized. @@ -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 None Override attention backend for the draft model + + --speculative-dsa-topk-backend + str + sgl-kernel + DSA indexer top-k backend for speculative draft workers, independent of --dsa-topk-backend (sgl-kernel, torch, or flashinfer) + --speculative-moe-runner-backend str diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index eaf5447c8..7f69d703c 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -523,12 +523,12 @@ SGLang supports various environment variables that can be used to configure its SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC - Use deterministic FlashInfer topk kernels when --dsa-topk-backend=flashinfer. + Use deterministic FlashInfer topk kernels when either --dsa-topk-backend=flashinfer or --speculative-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. + Tie-break mode for FlashInfer DSA topk when either --dsa-topk-backend=flashinfer or --speculative-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 diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index cb7871d7b..ac2345c7f 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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" ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py index 2d2122235..5be5e9f8f 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 2a0a993f3..3ac2ae68e 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1ee1ca36d..fb5c69cdf 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 72a37f5e6..3d05a7423 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -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 diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py index 00fc9864d..25f014347 100644 --- a/test/registered/kernels/ops/attention/test_dsa_indexer.py +++ b/test/registered/kernels/ops/attention/test_dsa_indexer.py @@ -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