Skip to content

Fix frozen kv MTP crash when bonus_tokens is None - #25204

Merged
hnyls2002 merged 1 commit into
sgl-project:mainfrom
pyc96:frozen-kv-mtp-crash
May 14, 2026
Merged

Fix frozen kv MTP crash when bonus_tokens is None#25204
hnyls2002 merged 1 commit into
sgl-project:mainfrom
pyc96:frozen-kv-mtp-crash

Conversation

@pyc96

@pyc96 pyc96 commented May 14, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Fix frozen kv MTP crash when bonus_tokens is None. Seems to be introduced recently in spec refactoring PRs.

Fix #24912

Modifications

Frozen KV MTP doesn't set bonus_tokens but we can get the last verified tokens from input_ids.

Tests

Server doesn't crash and accept rate looks normal.

python -m sglang.launch_server --model-path google/Gemma-4-26B-A4B-IT  --tp-size 1 --enable-torch-compile  --speculative-algorithm NEXTN --speculative-draft-model-path google/gemma-4-26B-A4B-it-assistant --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4

Checklist

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.

@gemini-code-assist

gemini-code-assist Bot commented May 14, 2026

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@pyc96
pyc96 requested a review from kpham-sgl May 14, 2026 01:20
@kpham-sgl

kpham-sgl commented May 14, 2026

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label May 14, 2026
@hnyls2002
hnyls2002 merged commit 78408d2 into sgl-project:main May 14, 2026
168 of 209 checks passed
@pyc96
pyc96 deleted the frozen-kv-mtp-crash branch May 14, 2026 23:38
@amd-bot amd-bot mentioned this pull request May 16, 2026
Open
SpencerGarnets added a commit to ai-blaise/optimization-playground that referenced this pull request May 16, 2026
@SpencerGarnets @claude
Merge upstream SGLang main into ai-blaise/optimization-playground main
0ea8076 
Merging upstream sgl-project/sglang main (d523ae1, 52 commits) into
ai-blaise/optimization-playground main, preserving the full custom stack
(IndexCache + TurboQuant 2.5-bit + SMC-SD, NVFP4 IndexCache+HISA 4:1,
HIGGS dense 2-bit MLA KV, HF config dispatcher with indexcache-hisa,
LayerSplit, NCCLX collective, Warp Decode, TokenSpeed/SMC/indexer
server_args defaults, W4A4 MoE WarpDecode).

Conflict resolutions (3 files):

1. python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
   COMPOSE - kept custom override_kv_cache_dim=compressed_mla_host_dim(kv)
   semantics; adopted upstream function rename (build_shared_anchor_stack
   -> build_anchor_sidecar_stack from PR sgl-project#25277 UnifiedRadixCache device
   match semantics). The renamed function still accepts override_kv_cache_dim
   (verified at hybrid_pool_assembler.py:583), so this is backwards
   compatible. The full call site at L946 already feeds the override into
   the kwarg.

2. sgl-kernel/cmake/flashmla.cmake
   TAKE UPSTREAM - upstream PR sgl-project#24130 (fix sm90 compile flashmla failed)
   renamed the FlashMLA-internal SM100 gate from SGL_KERNEL_ENABLE_SM100A
   to FLASHMLA_ENABLE_SM100 and re-introduced the SM100 source list under
   the new macro. HEADs older SGL_KERNEL_ENABLE_SM100A source-list block
   was a duplicate of upstreams new FLASHMLA_ENABLE_SM100 block; dropped
   the duplicate to avoid CMake emitting the same source files twice. The
   outer SGL_KERNEL_ENABLE_SM100A cmake option in CMakeLists.txt is
   preserved as it gates additional sgl-kernel targets beyond FlashMLA.

3. sgl-kernel/csrc/flashmla_extension.cc
   TAKE UPSTREAM - upstream PR sgl-project#24130 switched the FlashMLA-internal
   #ifdef gate at L94 from SGL_KERNEL_ENABLE_SM100A to FLASHMLA_ENABLE_SM100.
   flashmla.cmake (post-resolution 2) emits target_compile_definitions
   PRIVATE FLASHMLA_ENABLE_SM100 at L173, so the new macro is the one
   actually defined at FlashMLA extension compile time. The unrelated
   #ifndef SGL_KERNEL_ENABLE_SM100A at L21 (different code path) is left
   untouched.

Verification gates:

  GATE 1 git diff --check                                PASS
  GATE 2 no conflict markers                             PASS
  GATE 3 python3.12 -m compileall srt/jit_kernel/ci      PASS
  GATE 4 cargo check sgl-model-gateway -j1               PASS
  GATE 5 test_quantization_config_dispatch (33 tests)    PASS 33/33
  GATE 6 test_nsa_layersplit (26 tests)                  PASS 26/26
  GATE 7 test_higgs_dense_2bit_kv unit                   BLOCKED (GPU
         saturation: 8x H200 fully occupied by sglang::scheduler_TP0-7
         from another job; CUDA error: devices busy or unavailable.
         File compiles cleanly, imports succeed, not merge-induced.)
  GATE 8 test_higgs_dense_2bit_kv_integration            BLOCKED (same)
  GATE 9 test_torchcomms_ncclx CPU subset                PASS 4/4
         (remaining 8 are CUDA-only, expected skip on busy GPU)
  GATE 10 focused custom-stack CPU subset                PASS
         GPU-dependent tests (HiCache transfer, pool host integration)
         BLOCKED on same GPU saturation - reproduced identical
         cudaErrorDevicesUnavailable across all GPU-touching tests.

Custom-stack file integrity check (all 10 sentinel files present):
  layersplit.py, higgs_dense_2bit_kv.cuh, higgs_dense_2bit.py,
  torchcomms_ncclx.py, nvfp4_indexer_quant.cuh, nvfp4_indexer.py,
  hisa_tilelang_kernels/hisa.py, indexer_policy.py,
  indexer_quantization.py, turboquant_dense_kv.cuh.

Custom server_args defaults preserved (verified by grep): tokenspeed_mla,
nsa_indexer_mode, nsa_indexcache_freq, enable_nsa_nvfp4_hisa,
enable_turboquant_dense_kv_cache, enable_higgs_dense_2bit_kv_cache,
enable_hisparse, torchcomms_ncclx_*, enable_return_indexer_topk,
enable_nsa_prefill_context_parallel.

Upstream additions accepted into shared infra (no custom regression):
  - Multi-detokenizer worker support (PR sgl-project#24944, server_args.py).
  - SGLANG_OPT_FP8_WO_A_GEMM auto-disable on non-Blackwell
    (PR sgl-project#25181, server_args.py).
  - SpeculativeAlgorithm.supports_target_verify_for_draft +
    create_future_map extension points (PR sgl-project#24999, spec_info.py).
  - Frozen-KV MTP bonus_tokens=None fix (PR sgl-project#25204,
    frozen_kv_mtp_utils.py - uses existing draft_input.input_ids
    fallback; spec-decoding naming rule still satisfied).
  - HiCache UnifiedRadixCache device match semantics
    (PR sgl-project#25277, hybrid_pool_assembler function rename).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026
@pyc96
Fix frozen kv MTP crash when bonus_tokens is None (sgl-project#25204)
d5ce6a5
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kpham-sgl kpham-sgl kpham-sgl approved these changes
@Ying1123 Ying1123 Awaiting requested review from Ying1123
@merrymercy merrymercy Awaiting requested review from merrymercy
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@Qiaolin-Yu Qiaolin-Yu Awaiting requested review from Qiaolin-Yu Qiaolin-Yu is a code owner

Assignees

No one assigned

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Bug] TypeError: 'NoneType' object is not subscriptable in frozen_kv_mtp_utils.py when using Gemma 4 Assistant draft model

3 participants

@pyc96 @kpham-sgl @hnyls2002