Repository navigation
[Feat] nemotron: Port Nemotron 3.5 Super serving changes - #44518
Merged
Merged
Conversation
Splitting the global --max-running-requests across attention DP workers used bare floor division, so `--max-running-requests 1 --dp 4` gave every worker 0 slots. On hybrid models this surfaced as a misleading "mamba state cache is too small" startup failure even though the mamba pool had plenty of slots; the same floor division also produced 0-sized mamba pools on the disable-radix-cache path. The MLX resolver, which mirrors the same split, returned 0 unchecked outside its hybrid branch. Floor every per-worker split at 1, matching the guard the DeepEP v2 dispatch budget check already had. Requests are load balanced across workers, so a worker with 0 slots can never serve anything. The division itself stays floored, so the request pool, the mamba pools, the PD in-transfer reserve, the GDN ReplaySSM ring reserve, the LoRA MoE CUDA graph buffers and the MLX resolver keep sizing from the same per-worker count. (cherry picked from commit f284c66)
Under MAX_LEN padding, prepare_mlp_sync_batch relabels a decoding DP rank
as EXTEND and pads each row to a 1-token extend, so every rank presents
one shape to the collectives.
The relabel only decouples tokens from rows for the DP collectives; the
rows stay decode rows. Consumers that need the decode semantics behind
the label now read back through the new ForwardBatch.logical_forward_mode
property:
- mamba metadata, which would otherwise treat the rows as prefills,
skipping the decode state tracking (num_decodes == 0) and
dereferencing mamba_track_seqlens, which a decode batch never builds;
- the KDA gate preprocessing in kimi_linear and kimi_k3, which would
otherwise take the prefill branch and hand the decode kernel a
[1, T, H, K] forget gate and a pre-sigmoided beta;
- the GDN input layout in qwen3_5, which would otherwise hand the
decode kernels the strided views built for prefill.
[Port note: resolved a comment-only conflict in kimi_linear.py. The base no
longer carries the "Prefill passes raw gates" comment, so only the new
logical-mode comment is kept, matching the upstream PR head.]
(cherry picked from commit 5e1c25d)
Cover KVCacheConfigurator.resolve_max_num_reqs and the SWA chunk cap with --max-running-requests below attn_dp_size. Both sized each worker for zero requests before the split was floored at one. (cherry picked from commit 98c7a89)
Run a relabeled decode batch on a mamba track boundary through MambaAttnBackendBase._forward_metadata, where the scheduler crash fired; the existing cases only exercised Mamba2Metadata.prepare_mixed. Move those cases into test_mamba_prefill_track_metadata.py alongside it. (cherry picked from commit 28220ac)
Jiminator
marked this pull request as draft
October 15, 2026 16:35
Keep the measured kernels and launch parameters while removing unused Mamba megakernel and L2 prefetch experiments. Consolidate related wrappers, move phase-layout ownership into the model, reuse derived-cache update guards, and retain numerical/graph tests. Fix the prefill overlap TP lookup against the current runtime context.
BBuf
changed the title
Use the upstream TRTLLM workspace semaphore isolation instead of a separate fmha_v2 allocation, retaining the prefill/decode regression test.
BBuf
marked this pull request as ready for review
October 17, 2026 14:07
BBuf
added
run-ci
The per-worker max_running_requests floor, ForwardBatch.logical_forward_mode and its Mamba/KDA/GDN consumers, and their unit tests belong to #34535, which is reviewed separately. Restore these files to main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
main's NemotronHMoEDecoderLayer.forward, added with Nemotron-H EPLB support, already passes forward_batch to NemotronHMoE, so the prefill-overlap gate keeps its input. The MLP-like helper's MoE branch is unreachable; restore the base forward. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
BBuf
approved these changes
Oct 18, 2026
4 tasks done
Keep the upstream removal of unused JIT MoE gate wrappers while preserving the SM120 small-M router and its opt-in dispatch.
Collaborator
Author
brandonfzhang
added a commit
to brandonfzhang/sglang
that referenced
this pull request
Oct 19, 2026
Conflict: main's sgl-project#44518 added an opt-in host-indptr draft plan that reuses the token-block launch with 512-token blocks (<=32). Resolved to one launch sized by spec_kv_index_token_blocks, never below the host-plan path's tiling, and a single NUM_STEPS kwarg covering both triggers. main absorbed the kernel-side token-block support, so this branch now carries only the sizing helper, call-site wiring, and tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Merged
3 of 5 tasks
detain
pushed a commit
to detain/sglang
that referenced
this pull request
Oct 29, 2026
…project#44901) cherry-pick -x 3a059de. PORT ORDER: landed FIRST in this family because its merged mainline chronology makes it the strict precondition of sgl-project#44904 (main merges: sgl-project#44901 16:54 -> sgl-project#44904 16:56 -> sgl-project#44912 16:57 -> sgl-project#44913 16:59; sgl-project#44904's diff preimage literally consumes fused_shared_expert_weight / is_routed_scaling_factor_fused_in_topk introduced here, and sgl-project#44913's preimage deletes the helpers sgl-project#44912 introduces -- the only cherry-pickable order is the mainline one; later commits in this family build on earlier ones). Conflict resolutions (2): * models/deepseek_v2.py: the pick's hunk anchors on `class MergedMoEOutput(MoEOutput)` as CONTEXT, but that class is upstream sgl-project#44657 content this fork never took (review-record skip family), so the region conflicted empty-vs-context. Kept ONLY sgl-project#44901's own addition: _scale_routed_and_add_shared (verified byte-identical to the upstream post-image function text) plus all nine routed_scaling_factor_fused_in_topk sites (two __init__ sets, the aiter shared-scaling compensation block, two non-CUDA post-MoE-multiply guards, and the forward_deepep/state-based forward replacements). Nothing in this tree references MergedMoEOutput; did not smuggle it in. * test/registered/unit/layers/moe/test_mega_moe_decode_capacity.py: modify/delete conflict (the file was created by sgl-project#44657, absent here; sgl-project#44901 merely touched 1 line inside it). Kept it deleted. Fidelity: 7/12 python files are BYTE-IDENTICAL to the upstream post-image (topk.py, hash_topk.py, waterfill.py, moe_ep_setup.py, glm4_moe.py, glm4_moe_lite.py, vl_routing.py). mega_moe/mega_moe_flydsl/bailing_moe_v3/ dots3 landed inside files carrying nemotron-era (sgl-project#44518) upstream drift we do not have; their fold changes were diff-reviewed line by line (old `should_fuse_routed_scaling_factor_in_topk or _use_aiter` predicates replaced by the new config-level predicate; residual old-form reads match upstream's own post-image -- legitimately kept). WaterfillBalancer's routed_scaling_factor kwarg became shared_weight (one production constructor, moe_ep_setup, updated by the pick; no other call sites in the tree). Battery (CVD=9) vs base 0df18e9: test_fused_shared_expert_scaling.py 6P -> 12P/25sub (upstream's +343-line fold matrix, green); test_hash_topk.py 3P -> 4P; test_jit_grouped_topk 10S -> 12S (CUDA-gated adds); every pin held (standard_dispatcher_runner_backend 12P, smallm 28P/66sub, b12x 67P/46sub, kernels_namespace 76P, model_overrides 111P/1S/179sub, modelopt suites exact). Pre-existing classes unchanged (trtllm-gen 8F, fused_gate 46F, round-robin 37F, triton-fused-moe 96F, mem_pool_ep 1F, mega_moe_deepgemm 1F). ruff --select=F401,F821,F823,UP037 + ruff format clean on all 18 files. Fold-predicate e2e on real NVFP4 weights stays GPU-window debt (the standalone CPU contract probe follows this family as its own commit). (cherry picked from commit 3a059de) Co-authored-by: port-moe-contracts subagent (cherry picked from commit c4511d515077ea018200eeccfff256cc25ff18e5)
detain
pushed a commit
to detain/sglang
that referenced
this pull request
Oct 29, 2026
…ts only (sgl-project#44904) cherry-pick -x 48ea486. Landed second of the renorm pair: mainline makes sgl-project#44904 build ON TOP of sgl-project#44901 (its diff preimage contains fused_shared_ expert_weight/is_routed_scaling_factor_fused_in_topk), so the family follows the mainline order (see sgl-project#44901's note here). Wrong-output class: before this, the fused per-rank shared-slot paths asked the gate for top_k while num_fused_shared_experts markers also went through, renormalizing over one routed expert too many (routed sum 0.85-0.97 on CUDA). Conflict resolution (1): test/registered/kernels/ops/moe/test_moe_fused_gate.py carries pre-sgl-project#44690 docstring drift here (this fork never took sgl-project#44518's nemotron rework of that file's header); took the upstream post-image's new one-line docstring, kept this fork's file otherwise. The _reference_gate renorm-order fix, the renorm x apply_scale parametrization, and the production-impl comparison changes all landed. Fidelity: topk.py, hash_topk.py, state_capturer/routed_experts.py, test_fused_append_remap_per_rank_shared_slots.py and test_fused_shared_expert_scaling.py are BYTE-IDENTICAL to the upstream post-image; _router_triton_kernel (kernels/ops/moe/moe_fused_gate.py) verified byte-identical function body against upstream (the file's OTHER regions carry this fork's jit-loader drift and were not touched); fused_moe_triton_kernels.py hunks are docstring-only. Battery (CVD=9) vs sgl-project#44901 commit: test_fused_shared_expert_scaling 12P -> 13P/25sub; waterfill_eplb 4P -> 5P; scattered_a2a 3P -> 5P/2sub (the new routed-width capture case); hash_topk 4P held; append 6S -> 9S, padded 10S -> 11S (CUDA-gated adds); moe_fused_gate 46F -> 52F = the file's pre-existing no-CUDA class +6 new parametrized members (collects clean; needs hardware); kernels_namespace 76P (inventory guard, nothing added/removed); every other pin exact. ruff + format clean on all 11 files. Remap-on-CUDA numerics (torch.cat width change + padded-row remask) remain GPU-window debt, pinned at contract level by the fold-predicate probe commit that follows. (cherry picked from commit 48ea486) Co-authored-by: port-moe-contracts subagent (cherry picked from commit 0432a8fce0b4919b4a85fc6d3e724c6de30189be)
3 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Bring Nemotron 3.5 Super MTP serving optimizations for SM120 into SGLang: small-M NVFP4 MoE execution, SM120 GEMV, FlashInfer host planning, and hybrid prefill CUDA graphs.
Scope
SGLANG_SM120_NVFP4_PHASE_LAYOUT). Nemotron owns the layout, and the runner reaches it through an optionalprepare_for_forwardhook. Weight updates reuse the existing derived-cache guard, and weight export is rejected while the layout is enabled.--bf16-gemm-backend gemvis accepted on SM12x.SGLANG_SM120_MTP_ROUTER_TOP22), with unchanged device arithmetic.SGLANG_FLASHINFER_DRAFT_HOST_INDPTR,SGLANG_FLASHINFER_VERIFY_HOST_PLAN).forward_batchplumbing instead of duplicating it.Validation
For this revision (
4631e828b8), only CPU and static checks were run, on macOS:git diff --check, the repository's pinned ruff 0.15.1 (hook rule selection and format check), isort 7.0.0, codespell 2.4.1, clang-format 20.1.7,check-ast,debug-statements, and the registered-test lint scripts all pass on the changed files.forward_batchto the MoE core (the input the SM120 prefill-overlap gate reads). Results match the previous headd3f59dd9a6.No GPU tests, E2E serving runs or benchmarks were run for this revision.
Historical E2E evidence (not rerun)
The measurements below come from earlier heads, before main was merged (
d3f59dd9a6) and before the carried upstream changes were removed. Both measured heads still included the carried changes. They compare the earlier public headaf2628444ewith the cleanup headce427d4d9b. Both ran on the same RTX PRO 6000 (SM120), EA1008 NVFP4 checkpoint, pinned container, launch settings and request token IDs. Settings: TP=1, concurrency=1, real EAGLE acceptance (3 steps / 4 draft tokens). Each run has two warmups and 80 measured requests of 8,192 input / 1,024 output tokens, with the prefix cache flushed before every measured request.af2628444ece427d4d9bThroughput changed by -0.42% and TTFT by +13.07 ms (+1.55%). All 80 complete output-token sequences and all 100 GSM8K answer texts match exactly. Both runs had zero cached prompt tokens, zero request retractions, and no competing GPU process during measurement. The workload uses repeated text and has high acceptance. These results do not establish general conversational throughput or performance on other GPU architectures.
Earlier hardware recipes and model-support validation are recorded in cookbook PR #44510.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): 🚫 Run #39201116873
Latest PR Test (Extra): 🚫 Run #39201116823
Latest PR Test (AMD ROCm 10): ⏳ Run #39201116860