[yibol9768-alt]

Motivation

On SM12x, ModelOpt static-FP8 linears can change reduction schedules between cold prefill and a shorter prefix-cache-hit batch despite deterministic inference being enabled. The same input sequence then receives different scores.

Related to #41351's ModelOpt path. Credit to furedericca-lab for the diagnosis and initial fallback, and twu3202 for identifying the small-M gap. This does not claim to resolve every hybrid/GDN issue in that thread.

Modifications

• In SM12x deterministic mode only, reuse Triton scaled MM with fixed N=64/K=128 reduction tiles and M-only tiles 16/32/128. Normal-mode and forced-Marlin dispatch remain unchanged.
• Zero masked A/B loads for partial K tiles.
• Extend the existing registered SM120 test file with 12 cached-suffix comparisons and two CUDA Graph replay cases covering bias, strided 3D inputs, BF16/FP16 and a K=80 tail.
• The 2026-10-22 refinement changes only the deterministic GEMM's pipeline from two stages to three, retaining the same tiles, warps 4, quantization, arithmetic and scaling. It overlaps loads while preserving the reduction schedule.

Accuracy Tests

Validated 2026-10-22 on one RTX 5090, 400W, TP=1, using the complete official SGLang 0.5.23.dev20261020+g55b2aef7ad nightly: PyTorch 2.14.1, Triton 3.8.0, FlashInfer 0.7.2, sglang-kernel0.4.10 and its official Rust radix cache. Complete source files are installed before each process starts; normal package imports are used, with no function/AST injection. This is official-nightly validation, not a latest-main rebuild or upstream CI.

Model: nvidia/Qwen3-8B-FP8@2cebc4c89e25abc17668c81b01dceaf3d8b914d5, BF16 activations, automatic FP8 KV, deterministic mode and Triton attention.

Check	Result
Complete existing SM120 test file	137 passed
Additional two-stage/three-stage tensor comparisons, BF16/FP16, two seeds, eight N/K shapes, M boundaries, bias, strided input and tails	288/288 bit-exact, also exact against full-batch suffixes
Real model: prefixes 128/1024/4096, batches 1/4/2, two repetitions; two Graph boots and one eager boot per PR arm	18 paired cache cases exact; cold/cached scores and warm repeats have zero delta; selected logprobs also match across pipeline versions
Complete fixed-work outputs, including warmups and measured repetitions	180/180 texts and token sequences exact between pipeline versions
Complete GSM8K test rows 32–159, cap 512, Graph mode	128/128 texts and token sequences exact between pipeline versions;120/128 correct in each, no missing final answers

The 128 questions reuse previous data and test preservation of the existing PR's outputs. They are not an independent quality evaluation, an accuracy gain, or evidence of statistical equivalence to upstream. Earlier upstream/candidate checks remain disclosed:15/16→16/16 on a16-question nightly subset, and119/128→117/128 on the older mixed runtime. See the prior nightly results and historical128-question results.

The unmodified nightly FP8 path still shows nonzero cold/cache selected-logprob deltas in the fresh Graph control. An independent RTX PRO 6000 Max-Q / Qwen3-14B-FP8 check also corroborates the fix at prior head 5f90880cc2; it predates this stage-only refinement. The original hybrid TP=2 model, SM121 and torch.compile have not been tested locally.

Raw paired results, full generated outputs, tests, source snapshots, scripts and hashes.

Speed Tests and Profiling

The isolated probe identifies GEMM as the main component cost in these cases; increasing pipeline depth is faster for several projection shapes while all 64 explored cases remain exact. These cache-evicted microbenchmarks are diagnostic; the end-to-end results below determine the performance claim.

32 forced output tokens per request, one warmup and five measured samples per cell per boot, cache flushed before requests, passive cooling to<=55C outside timing. PR-arm Graph boot order is A(two stages), B(three stages), B, A. Table medians pool ten samples per PR arm; per-boot medians, ranges and every sample are retained. The upstream Graph reference and each eager arm have one boot/five samples. Graph capture covers decode 1/2/4 and prefill 4/128/256/1024/2048/4096/8192; unsupported partial prefills may fall back to eager.

Prefix / batch	Graph two stages → three	Wall-time reduction	Upstream Graph reference	Eager two stages → three
128 / 1	0.3119s → 0.2796s	10.3%	0.2684s	1.4868s → 1.4839s
128 / 4	0.3539s → 0.3251s	8.1%	0.3322s	1.5320s → 1.5275s
2048 / 1	0.5314s → 0.5143s	3.2%	0.4766s	1.5629s → 1.5642s
2048 / 4	0.8900s → 0.8688s	2.4%	0.7953s	2.0843s → 2.0747s

The stage-only update reduces this Graph matrix's pooled median wall times by 2.4–10.3%, most consistently on the short prompts. The 2048/1 cell varies between boots: current medians0.5285/0.5895s, candidate0.4982/0.5434s. Eager differences are within about 0.5%, so no meaningful eager speedup is claimed.

Against the single fresh upstream Graph reference, the three-stage path ranges from 2.1% lower to9.2% higher wall time. This does not eliminate the determinism/cost tradeoff or establish a general speedup. Earlier 5.7–25.8% nightly and 5–20% mixed-runtime Graph costs describe the prior two-stage measurements and remain in their historical evidence. All ten successful jobs stayed below 72 C, peaking at 66C with the 400W limit unchanged.

Full samples, boot medians and measurement limits.

Duplicate check: inspected #34429 (small-M dispatch), #42865 (block/MXFP8), #38921 (SM89 prequantized inputs) and #44076 (GDN chunk alignment); these do not implement this static per-tensor deterministic path. Original pinned diff audit.

AI assistance was used for diagnosis, implementation, validation and this description. The human contributor reviewed the original fix and authorized this performance investigation and PR update.

Checklist

• Applicable pre-commit hooks pass.
• Regressions extend the existing registered file.
• Real pretrained-model correctness, complete outputs and performance checks are recorded with limitations.
• No new public interface; deterministic-mode behavior and tradeoffs are documented here.

Review and Merge Process

Maintainer review and upstream CI follow the normal project process. Local validation and independent reproductions do not constitute upstream CI or maintainer approval.

---

CI States

Latest PR Test (Base): ❌ Run #39746823806
Latest PR Test (Extra): ❌ Run #39746823720
Latest PR Test (AMD ROCm 10): ❌ Run #39746823779

[waynehacking8]

On a second SM12x part, an RTX PRO 6000 Max-Q (SM120) with nvidia/Qwen3-14B-FP8, head 5f90880 makes cold and prefix-cache-hit selected-token logprobs identical in all 12 comparisons, where main differs in all 12 (max abs delta 0.5 to 2.375). The CUDA graph cost matches the 5090 numbers the body gave for that head at batch 1 and is higher at batch 4. All of this is for 5f90880; I haven't run the current three-stage head 6a7695d.

CUDA graph, wall time for 32 output tokens, prefix / batch	main, s	PR, s	PR / main	PR / main on the body's 5090 (body for 5f90880)
128 / 1	0.422	0.504	1.193	1.189
128 / 4	0.493	0.577	1.171	1.057
2048 / 1	0.711	0.901	1.268	1.258
2048 / 4	1.167	1.477	1.265	1.169

The comparisons are the body's cases (prefixes 128/1024/4096 at batches 1/4/2, two repetitions) in one eager and one graph boot per arm. Times are medians of five, and the ranges don't overlap in any graph cell. In eager mode PR / main runs from 0.89 to 1.055, but those boots shared the card with another user's job; the graph boots had it alone. The PR's test_sm120_fp8_linear.py passes 137 of 137, and on main's code its cached-suffix cases fail 5 of 12 (4 of 12 on the body's 5090).

Setup: the body's engine settings (deterministic inference, Triton attention, TP1, BF16), fp8 KV, the PR diff applied to main 4fc86f7. Logs and scripts: https://gist.github.com/waynehacking8/e8474d2b9449f657396df5a040566e16

[adcostanza]

Validated on a hybrid GDN model, which the description leaves out of scope: it fixes this case.

Setup. One RTX PRO 6000 (SM120), TP=1. Image lmsysorg/sglang:nightly-dev-cu13-20261101-cda284e4 with the PR head 6a7695ddc4 merged on top (no conflicts), bind-mounted over the installed package. Model ukisai/Swift-1.5-Qwen3.8-27b-NVFP4 @ 25482027, a ModelOpt mixed checkpoint of a Qwen3.8-27B fine-tune: the GDN in_proj_qkv / in_proj_z / out_proj and the full-attention q/k/v/o_proj are static per-tensor FP8, the MLPs are NVFP4. Flags: --enable-deterministic-inference --attention-backend triton --kv-cache-dtype auto --mamba-ssm-dtype bfloat16 --mamba-radix-cache-strategy extra_buffer_lazy --chunked-prefill-size 2048 --page-size 64 --enable-hierarchical-cache --hicache-size 64 --hicache-write-policy write_through --schedule-policy lpm --enable-session-radix-cache --mem-fraction-static 0.41 --max-running-requests 4.

Result, python3 -m sglang.test.test_deterministic --test-mode radix_cache, cached vs uncached prefill:

	cached	uncached
nightly cda284e4	-5.0687	-4.8716	FAIL
+ this PR, prefill CUDA graph on	-3.7109196186	-3.7109196186	PASS
+ this PR, eager prefill	-3.7109196186	-3.7109196186	PASS

A prefill-only probe (cold, prefix hit with a 37-token suffix, cold again; 101 tokens) gives -6.1011090279 all three times with the PR.

Where the difference came from on this model (diagnosed independently before finding this PR, by dumping per-layer tensors for the hit tokens, cold vs hit):

• The first tensor that differs is layer 0's GDN out_proj output (1.2e-4, from bit-equal inputs). It accumulates in the residual stream and reaches layer 3's norm input as 1.3e-3. NVFP4 MLP outputs stay bit-equal.
• These FP8 projections go through ModelOptFp8LinearMethod.apply → apply_fp8_linear_bmm_flashinfer → flashinfer bmm_fp8 (use_flashinfer_bmm=True); try_sm120_fp8_linear returns None for these shapes.
• Standalone, on the real layer-0 out_proj weight (K=6144, N=5120), rows computed in a call of M rows vs the same rows inside a 101-row call: bmm_fp8 differs by 0.0039 at M = 8, 16, 37, 64; torch._scaled_mm and sgl_kernel.fp8_scaled_mm are exact at M = 1, 8, 16, 37, 64, 100.
• With the PR, all 128 FP8 linear calls per forward (M=101 cold and M=37 hit) take the new use_sm120_fp8 and deterministic branch into apply_fp8_linear_deterministic, which sits before the bmm_fp8 branch, so bmm_fp8 is never reached.

Not measured: throughput of the deterministic Triton path on this model.