[Spec][DSA] Add --speculative-dsa-topk-backend

[zianglih]

Motivation

@HumansAnd

• Scope: Split DSA indexer top-k selection between the target model and
  speculative draft workers.
• --dsa-topk-backend currently reaches both roles because target and draft
  runners share one ServerArgs. A target override such as flashinfer
  therefore also changes repeated MTP draft top-k.
• Add --speculative-dsa-topk-backend, defaulting to sgl-kernel, so the
  draft stays on the usual fast path unless it is explicitly overridden.
• Default launches are unchanged. To use the same non-default backend for both
  roles, pass both selectors.

Modifications

• Keep --dsa-topk-backend as the target-model selector.
• Add --speculative-dsa-topk-backend for speculative draft workers with the
  same sgl-kernel, torch, and flashinfer choices and a
  sgl-kernel default.
• Resolve the backend per runner from the post-publish config namespaces:
  • target: exec.kernel.dsa_topk_backend;
  • draft: spec.speculative_dsa_topk_backend.
• Use the common resolver in both DeepseekSparseAttnBackend (DSA/GLM) and
  DeepseekV4AttnBackend without mutating the shared ServerArgs.
• Make the existing DSA test runners explicit target workers so their mocks
  satisfy the same role contract as production ModelRunner.
• Document the independent selector in the server-argument and speculative
  decoding references, including the DSA MTP path and FlashInfer environment
  switches.
• Compatibility: Existing launches that set only
  --dsa-topk-backend=<non-default> now apply that override only to the
  target. Pass the new flag as well to restore the previous target/draft
  pairing.
• Non-goals: No kernel implementation changes. SGLANG_DSA_FUSE_TOPK,
  SGLANG_OPT_USE_TOPK_V2, and the FlashInfer deterministic/tie-break
  settings remain process-wide. Selecting torch for either role still
  requires SGLANG_DSA_FUSE_TOPK=false, which disables fused top-k for both
  roles.

Accuracy Tests

Tested commits and environment:

• PR: e06cd057c3cb319970a283aafb3e23facbf7f559
• Base: 99c02d71b170673f97676f097fc928de38d847e0
• C2 bare devbox with 8 x NVIDIA B300 SXM6 AC allocated; tests were constrained
  to GPU 0 with CUDA_VISIBLE_DEVICES=0
• Image:
  lmsysorg/sglang:nightly-dev-cu13-20260825-1ec20fd2
  (linux/amd64 digest
  sha256:65376f9f5c317da614be2d6e2c683d8c4d7964cf8fa63f663a1cd058929598ca)
• NVIDIA driver 590.48.01; CUDA 13.0.3; Python 3.12.3;
  PyTorch 2.13.0+cu130; Triton 3.7.1
• Exact PR and base sources were downloaded from GitHub codeload tarballs for
  the SHAs above and selected with PYTHONPATH.

Existing DSA and DeepSeek-V4 attention unit suites on the PR:

cd /sgl-workspace/sglang-pr36313
export PYTHONPATH=/sgl-workspace/sglang-pr36313/python
export CUDA_VISIBLE_DEVICES=0
python3 -m pytest -q -ra \
  test/registered/attention/unittests/dsa/test_dsa.py \
  test/registered/attention/unittests/dsv4/test_deepseek_v4.py

32 passed, 23 skipped, 21 warnings, 73 subtests passed in 15.24s

The skips are capability-gated on SM 10.3: FA3 is limited to SM 9.x, the image
does not contain a matching TileLang B300 template, TensorRT-LLM covers SM
10.0 rather than SM 10.3 here, AITER is AMD-only, and the remaining cases are
unsupported FP8/backend combinations.

Existing DSA metadata/transform and DeepSeek-V4 indexer unit suites on the PR:

cd /sgl-workspace/sglang-pr36313
export PYTHONPATH=/sgl-workspace/sglang-pr36313/python
export CUDA_VISIBLE_DEVICES=0
python3 -m pytest -q -ra \
  test/registered/kernels/ops/attention/test_dsa_metadata.py \
  test/registered/kernels/ops/attention/test_dsa_transform_index.py \
  test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py \
  test/registered/unit/layers/test_dsv4_nonpaged_indexer.py

41 passed, 15 warnings, 18 subtests passed in 37.91s

Direct selector smoke against the resolved runtime namespaces:

from types import SimpleNamespace

from sglang.srt.layers.attention.dsa.dsa_topk_backend import (
    DSATopKBackend as Backend,
)
from sglang.srt.runtime_context import get_context

target = SimpleNamespace(is_draft_worker=False)
draft = SimpleNamespace(is_draft_worker=True)

override = get_context().override_server_args(dsa_topk_backend="flashinfer")
server_args = override.install()
assert Backend.resolve(target) is Backend.FLASHINFER
assert Backend.resolve(draft) is Backend.SGL_KERNEL
assert server_args.speculative_dsa_topk_backend == "sgl-kernel"
override.restore()

override = get_context().override_server_args(
    dsa_topk_backend="sgl-kernel",
    speculative_dsa_topk_backend="torch",
)
override.install()
assert Backend.resolve(target) is Backend.SGL_KERNEL
assert Backend.resolve(draft) is Backend.TORCH
print("target/draft DSA top-k selection: PASS")

target/draft DSA top-k selection: PASS

Controlled PR/base comparison of the attention and B300 indexer suites, using
the same node, image, GPU, and command:

cd /sgl-workspace/sglang-pr36313
PYTHONPATH=/sgl-workspace/sglang-pr36313/python CUDA_VISIBLE_DEVICES=0 \
  python3 -m pytest -q -ra \
    test/registered/attention/unittests/dsa/test_dsa.py \
    test/registered/attention/unittests/dsv4/test_deepseek_v4.py \
    test/registered/kernels/ops/attention/test_dsa_indexer.py

2 failed, 42 passed, 23 skipped, 21 warnings, 106 subtests passed in 224.45s

cd /sgl-workspace/sglang-pr36313-base
PYTHONPATH=/sgl-workspace/sglang-pr36313-base/python CUDA_VISIBLE_DEVICES=0 \
  python3 -m pytest -q -ra \
    test/registered/attention/unittests/dsa/test_dsa.py \
    test/registered/attention/unittests/dsv4/test_deepseek_v4.py \
    test/registered/kernels/ops/attention/test_dsa_indexer.py

2 failed, 42 passed, 23 skipped, 21 warnings, 106 subtests passed in 17.77s

Both commits fail the same two B300-only cases:

FAILED test/registered/kernels/ops/attention/test_dsa_indexer.py::TestDSAIndexer::test_forward_decode_mode
FAILED test/registered/kernels/ops/attention/test_dsa_indexer.py::TestDSAIndexer::test_forward_extend_mode
AttributeError: 'MockModelRunner' object has no attribute 'max_running_requests'

The failing workspace allocation reads
model_runner.max_running_requests; the unchanged upstream mock does not
define that field. The identical base and PR outcomes bound this as a
pre-existing B300 test-mock limitation rather than a regression from this
change. The wall times are not a performance comparison because the second run
used warm JIT/caches.

No new standalone test cases are included; validation uses the existing DSA
suites plus the direct selector smoke above.

Static validation:

pre-commit run --files \
  docs/docs/advanced_features/server_arguments.mdx \
  docs/docs/advanced_features/speculative_decoding.mdx \
  docs/docs/references/environment_variables.mdx \
  python/sglang/srt/layers/attention/deepseek_v4_backend.py \
  python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py \
  python/sglang/srt/layers/attention/dsa_backend.py \
  python/sglang/srt/server_args.py \
  python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py \
  test/registered/kernels/ops/attention/test_dsa_indexer.py

check for broken symlinks..............................(no files to check)Skipped
detect destroyed symlinks..................................................Passed
trim trailing whitespace...................................................Passed
fix end of files...........................................................Passed
check yaml.............................................(no files to check)Skipped
check toml.............................................(no files to check)Skipped
check python ast...........................................................Passed
check for added large files................................................Passed
check for merge conflicts..................................................Passed
check that scripts with shebangs are executable............................Passed
detect private key.........................................................Passed
debug statements (python)..................................................Passed
don't commit to branch.....................................................Passed
isort......................................................................Passed
ruff (legacy alias)........................................................Passed
black-jupyter..............................................................Passed
codespell..................................................................Passed
clang-format...........................................(no files to check)Skipped
nbstripout.............................................(no files to check)Skipped
forbid docs_new/ (renamed to docs/)....................(no files to check)Skipped
check chinese characters in multimodal_gen.............(no files to check)Skipped
sort CI_PERMISSIONS.json...............................(no files to check)Skipped
check for duplicate workflow job names.................(no files to check)Skipped
check rust-ext cache key prefix defaults match.........(no files to check)Skipped
reject bare pytest.main calls in __main__ blocks...........................Passed
unit tests for lint checkers...........................(no files to check)Skipped
validate registered test CI registries.....................................Passed
reject CI-registered tests inside the sglang package.......................Passed
rustfmt sgl-model-gateway (nightly)....................(no files to check)Skipped
rustfmt experimental/sgl-router........................(no files to check)Skipped
clippy rust/ workspace (auto-fix)......................(no files to check)Skipped
rustfmt rust/ workspace................................(no files to check)Skipped

python3 -m compileall -q \
  python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py \
  python/sglang/srt/layers/attention/dsa_backend.py \
  python/sglang/srt/layers/attention/deepseek_v4_backend.py \
  python/sglang/srt/server_args.py \
  python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py \
  test/registered/kernels/ops/attention/test_dsa_indexer.py
git diff --check upstream/main..HEAD

# Both commands exited 0 with no output.

Speed Tests and Profiling

• Not run. This PR changes backend selection configuration, not a kernel
  implementation, and makes no quantitative speedup claim.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #32886603742
Latest PR Test (Extra): ❌ Run #32886603233
Latest PR Test (AMD ROCm 7.2): ⏳ Run #32886603588

[zianglih]

/tag-and-rerun-ci

[zianglih]

/tag-and-rerun-ci

[zianglih]

/tag-and-rerun-ci