[Perf] Add opt-in SM120 BF16 BA projection for small batches

[LingZ315]

Summary

Add an opt-in, shape-allowlisted BF16 Tensor Core projection for
in_proj_ba on SM120. This is an experimental performance improvement, not a
correctness fix or a claim of universally bitwise-identical inference.

Base inspected: 03d06a764e4a83268eefd1bafc676418f7269c89 (2026-09-09).

Scope and dispatch

• SGLANG_ENABLE_SM120_BA_GEMM=1, with the default --bf16-gemm-backend auto.
• Disabled by default; explicit BF16 backends are not overridden.
• Exact layer-name component in_proj_ba; SM120 only, contiguous BF16 input and
  weight, no bias or gradients, N=48, K=5120, 2 <= M <= 8.
• Preserve M=1 and unsupported shapes/layouts/dtypes/devices through F.linear.
• Keep deterministic, batch-invariant, autocast and compilation paths unchanged.
• Tensor-device-aware launch and current-stream execution. Cold CUDA Graph shapes
  fall back rather than compiling during capture; warm eager calls enable capture.
• No weight conversion, model configuration, scheduler, quantized projection,
  Mamba state update, or installed-environment changes.

Why Tensor Core dot instead of a row-wise reduction?

Earlier experiments with a row-wise reduction changed rounding and failed task
quality regression. This candidate preserves M=1 and uses BF16 tl.dot with
FP32 accumulation and K tiles of 128 for M=2..8. This numerical choice is integral
to the proposal; it is not interchangeable with the faster rejected reduction.

Validation of the source integration

Run from the repository root on SM120:

PYTHONPATH=python python -m unittest discover -s test/registered/kernel/gemm -p test_sm120_ba_gemm.py -v
PYTHONPATH=python python -m unittest discover -s test/registered/unit/layers/quantization -p test_sm120_ba_dispatch.py -v
PYTHONPATH=python python benchmark/kernels/benchmark_sm120_ba_gemm.py --device 0

The integration was imported directly from the isolated main-based source tree
using the existing test-machine dependency environment (not a newly installed
main wheel). Thirteen tests passed, including actual UnquantizedLinearMethod
dispatch, default/enable/reset gates, shape and layout fallbacks, gradient
fallbacks, autocast/deterministic/compiler fallbacks, cold/warm CUDA Graph replay,
and a noncurrent second GPU. torch.compile coverage uses backend="eager",
not an Inductor end-to-end serving test.

• Compute Sanitizer memcheck: 0 errors across the test suite.
• Compute Sanitizer racecheck: 0 hazards (0 errors, 0 warnings) across the suite.
• 148 saved real-activation snapshots, expanded to 1,184 batches / 255,744 output
  elements: packaged implementation exactly matched the frozen candidate.
  Expanded batches use repeated captured rows with seeded perturbations, not
  newly captured concurrent activations.
• Numerical unit test: 20 seeds for every M=2..8, exact same-environment F.linear
  parity; this is an empirical gate, not a cuBLAS-version-independent guarantee.
• Ruff 0.15.1 selected checks/format, isort 7.0.0 and diff whitespace checks pass.
  These are historical local checks. Upstream CI subsequently exposed test-taxonomy and registry errors, corrected in 9445e20; new CI results are pending.

Graph microbenchmark medians (100 calls/replay, 20 shuffled-order repetitions):

GPU	M	PyTorch (us)	Frozen candidate (us)	Packaged candidate (us)
RTX 5090 GPU0	1	2.817	2.811	2.814
RTX 5090 GPU0	4	19.784	7.251	7.251
RTX 5090 GPU0	8	19.847	7.109	7.110
RTX 5090 GPU1	1	2.851	2.851	2.850
RTX 5090 GPU1	4	20.225	7.408	7.411
RTX 5090 GPU1	8	20.245	7.302	7.302

Latest source-integrated validation (September 10)

The measurements in this section used cd2a21c7bc4e3846f68b6f3ab171f2d070c3e590, loaded from the isolated source tree with the existing dependency environment. They supersede the earlier two-boot performance estimate below. The subsequent 9445e20 commit changes only test registration/locations and documentation, not kernel arithmetic, dispatch or default enablement; it is not a new GPU measurement.

TP2 RTX 5090, tested NVFP4 checkpoint with BF16 BA, FP8 E4M3 KV, BF16 SSM, page size 256, chunked prefill 4096, memory fraction 0.85, max-running 8, Mamba cache 32, decode graph max batch 8, single-batch overlap, no speculative decoding. Exactly 50,000 WikiText input tokens / 1,000 output tokens, 16 requests per group; 45,056-token prefix warmup excluded from timing; measured KV hit 90.112%.

Six independent uninstrumented boots: A1/B1/B2/A2/A3/B3. Only the enable flag differed. No discarded or replaced runs.

Client concurrency	Baseline TPS	Candidate TPS	Change	Baseline TTFT P50/P90 (s)	Candidate TTFT P50/P90 (s)
1	116.37	116.20	-0.15%	0.733 / 0.754	0.733 / 0.755
4	329.77	349.59	+6.01%	2.091 / 2.609	2.146 / 2.619
8	396.53	414.44	+4.52%	3.394 / 10.494	3.377 / 10.198

TPS is total output tokens/group wall time, averaged over three boots per side. TTFT values are means of per-run quantiles, not pooled percentiles. All 288 requests succeeded; KV capacity remained 1,028,608 logical tokens. All three c4/c8 pairs improved, but c4 TTFT did not. The sampled running peak at eight-client concurrency was 7, not sustained eight-way decode.

Quality: distinguish controlled comparisons from asynchronous traffic

• Both-rank actual-dispatch audit confirmed candidate graph nodes at M=4/8, none in baseline; M=1 falls back.
• Same-input checks across all 48 BA layers and both ranks compared 110,592,000 active elements with zero differences, including candidate-driven inputs. Counts include repeated inputs/padding, not independent samples.
• Three controlled native-batch A/B pairs matched 12,096 teacher-forced full-vocabulary positions and 200 free-generation entries per pair (including repeated questions). An additional A/B/A audit matched all 8,000 logits positions and observed scheduler metadata on both ranks. Instrumented timing is not performance evidence.
• Independent HTTP generation had no correctness flips in three pairs, but token equality was 196/200, 200/200 and 200/200; baseline repeats also varied.
• Independent HTTP strict logits checks still failed, for candidate/baseline and baseline repeats. Mean/P99 KL: 0.0380393/0.359714 versus repeat 0.0227127/0.307314. The 0.001/0.01 gates were not relaxed. Controlled admission does not retroactively make that protocol pass.
• A fixed-input, two-request/two-rank replay isolated BF16 state rounding at differing prefill boundaries as one concrete mechanism. FP32 replay state removed those local differences; this does not attribute every model-level difference or justify universal losslessness. No FP32-state change is included.

CI integration correction and reproducibility

Commit 9445e20 fixes the registered-test taxonomy and incomplete CPU registration. Follow-up 8a4489b places GPU coverage in base-b-kernel-unit-test-1-gpu-small and adds a dedicated RTX 5090 job in the existing JIT-kernel workflow, preserving its caller authorization gate. The H100 job would skip SM120. The optional second-GPU test still requires the local two-GPU setup.

Local checks after this correction: repository registered-test lint passed; AST collection succeeded for all 2,244 registrations; whitespace checks passed; production kernel/dispatcher/environment files are unchanged. GPU/unit execution was not repeated on the current Windows host. After rebasing, upstream Base/Extra check-changes (including partition computation) passed. Full CPU/GPU jobs are currently blocked by the required run-ci label, which needs maintainer authorization. The latest kernel-suite correction passes local direct taxonomy checks and workflow YAML/suite-wiring validation; latest full lint is pending. Upstream CI completion remains required.

Validation/reproduction notes and all six boot results include exact operator commands, fixed dataset hash, serving protocol and retained numerical failures. The operator tests/microbenchmark are executable; the archived corpus and private full-model diagnostic harness are not redistributed, so this is not a claim of a fully standalone serving-quality reproducer.

Detailed validation update.

Historical end-to-end evidence for the frozen candidate

These numbers are not a full-service rerun of the rebased integration.
They were measured on SGLang 0.5.19.dev20260901+g783af667fb, TP2 RTX 5090,
with the tested NVFP4 checkpoint, BF16 BA projections and FP8 KV cache.

Fixed WikiText-103 token inputs, 50,000 input / exactly 1,000 output, 16 requests
per group, 45,056-token prefix warmup outside timing, measured KV hit 90.112%.
Separate uninstrumented performance boots used A1/B1/B2/A2 order. TPS is total
output tokens divided by elapsed group wall time, then averaged over two boots.

Client concurrency	Baseline TPS	Candidate TPS	Change
1	115.85	116.22	+0.31% (essentially unchanged)
4	327.45	351.33	+7.30%
8	392.77	421.59	+7.34%

All 192 performance requests succeeded. In the 8-client performance groups,
the recorded server decode peak was 7, for both implementations. Do not
interpret this as sustained eight-way decode throughput. Two runs per side do
not establish statistical significance or a universal fixed speedup.

Quality A1/B1/A2/B2: 128-question math score stayed 121/128 in all four rounds;
structured output stayed 16/16 and long retrieval 8/8. No new stable correct-to-
incorrect case appeared. Short teacher-forced logits matched exactly. Long
teacher-forced logits did not: mean KL at the long c8 group was 0.010914 for
baseline repeat, 0.027757 / 0.026138 for candidates, within the predeclared 3x
repeat-noise gate but approximately 2.54x / 2.39x its measured baseline mean.
This uncertainty is explicitly retained; passing the gate is not losslessness.

Related work and remaining review gates

FlashInfer #4223 tracks
small-N BF16 GEMV, and #4250
proposed a CuTe-DSL SM120/121 implementation. FlashInfer #4250 was closed (not shown merged)
at inspection. Its published shapes include N=96/K=5120 and N=64/K=2048, while
this proposal targets TP-sharded N=48/K=5120 and a different reduction strategy.
No head-to-head benchmark of that alternative is claimed here, and this PR does
not claim the optimization direction is novel.

Remaining review gates (Ready for review does not imply ready to merge):

• Complete actual eight-way decode performance coverage; the completed independent-HTTP rerun below reached a sampled peak of seven.
• Retain the independent-HTTP strict logits failures and extend numerical coverage across scheduling patterns and dependency versions.
• Resolve whether to reuse an existing small-N implementation, after same-shape
  numerical and performance comparison, and agree on supported SM120 CI coverage.
• Extend workload and stability coverage; current results do not validate
  97k/512k, multimodal input, other checkpoints or other GPU architectures.

All reported measurements above are actual recorded tests; unrun
acceptance gates are listed explicitly.

---

CI States

Latest PR Test (Base): ❌ Run #34552524249
Latest PR Test (Extra): ❌ Run #34552524072
Latest PR Test (AMD ROCm 10): ❌ Run #34552524231

[LingZ315]

Validation update for cd2a21c7bc4e3846f68b6f3ab171f2d070c3e590

Additional serving validation is complete. This update separates operator correctness, controlled-trajectory model comparisons, and uninstrumented performance. It does not claim bitwise-identical inference under arbitrary asynchronous traffic.

The tested integration is the current PR head at inspection. The kernel and unquant.py dispatch files match the PR files after normalizing line endings. Tests used the isolated source tree with the existing dependency environment, not a newly installed wheel or a newer upstream main.

1. Actual candidate execution and same-input correctness

• An independent dispatch audit confirmed candidate CUDA Graph node replay on both TP ranks for M=4 and M=8. Baseline replay contained no candidate nodes. M=1 follows the existing fallback.
• Same-input checks covered all 48 BA layers and both TP ranks, including inputs encountered while the candidate output was driving the model, not only baseline-driven inputs.
• Across baseline/candidate/baseline diagnostic runs, 110,592,000 active output elements were compared: zero baseline/candidate differences, zero baseline repeat differences, and zero differences between live output and the corresponding checked path. Counts include graph padding/repeated workload coverage; they are not independent samples.
• Clean and deliberately corrupted shadow-output detector tests passed on both live paths. The injected error stayed in the diagnostic shadow and was not propagated through the model.

2. Controlled-trajectory full-model comparison

Using native pretokenized batch admission in fixed-order waves, the previously frozen full regression passed three A/B pairs: 12,096 teacher-forced positions and 200 free-generation test entries per pair, with identical logits and output token sequences. These entries include repeated questions across concurrency groups, not 200 independent benchmark questions.

An additional A/B/A audit now verifies the actual scheduler metadata for the 50,000-input / 1,000-output, eight-request case:

• Each arm/rank recorded 2,011 batches: 11 extend and 2,000 decode batches.
• Ordered request IDs, sequence lengths, prefix/extend lengths, logical batch sizes, and CUDA Graph-use flags matched exactly across all three arms and both ranks.
• B/A and repeated A/A each matched at all 8,000 complete-vocabulary logits positions, including the first output position: maximum absolute difference, mean/P99 KL, and mean target-NLL delta were all zero.
• Original BF16 SSM, FP8 KV, attention, CUDA Graph and overlap settings were retained. No shape-padding, FP32-state, or attention-control diagnostic changes were used.

This audits matching observed trajectories; it is not a scheduler replay feature, and atomic admission is not the independent-HTTP performance workload below. Instrumented timings are not used as performance evidence.

3. Uninstrumented serving performance

TP2 RTX 5090; tested NVFP4 checkpoint with BF16 BA projections; FP8 E4M3 KV; BF16 SSM; page size 256; chunked prefill 4096; max-running 8; Mamba cache 32; single-batch overlap; no speculative decoding.

Fixed WikiText token inputs: exactly 50,000 input / 1,000 output, 16 requests/group, 45,056-token prefix warmup outside timing, measured cache hit 90.112%. Six independent boots in A1/B1/B2/A2/A3/B3 order, with no diagnostic hooks and no discarded/replaced runs. The only A/B treatment was the candidate enable flag.

Client concurrency	Baseline output TPS	Candidate output TPS	Change	Baseline TTFT P50/P90 (s)	Candidate TTFT P50/P90 (s)
1	116.37	116.20	-0.15%	0.733 / 0.754	0.733 / 0.755
4	329.77	349.59	+6.01%	2.091 / 2.609	2.146 / 2.619
8	396.53	414.44	+4.52%	3.394 / 10.494	3.377 / 10.198

TPS is aggregate output tokens/group wall time, averaged over three runs per side. TTFT values are means of per-run quantiles, not pooled percentiles. All three c4 and c8 pairs improved, but this is limited hardware/workload coverage, not a universal speedup. c4 TTFT did not improve.

All 288 timed requests succeeded. KV capacity was unchanged at 1,028,608 logical tokens. The eight-client performance groups had a sampled server running peak of 7, not sustained eight-way decode. Independent-HTTP free-generation checks showed no correctness flips in any of three 200-entry pairs; token sequences matched in 196/200, 200/200 and 200/200 entries. Baseline repeats also varied in output tokens without correctness flips.

4. Numerical limitations retained

The independent-HTTP strict logits gate still failed for both candidate-vs-baseline and baseline-repeat comparisons. Those failures are retained; controlled admission does not retroactively make that protocol pass.

A separate fixed-input replay isolated one concrete source in two requests on both ranks: different prefill boundaries with identical complete layer-0 inputs and starting state reproduced all 12 recorded GDN outputs exactly with BF16 state. Changing only the replay state storage to FP32 removed the between-plan differences in those tested aligned plans. This is evidence for that local mechanism, not an attribution of every historical/final-model difference. Candidate execution timing may still affect scheduling indirectly.

No numerical thresholds were relaxed. No claim is made that FP32 state alone fixes full-model reproducibility, that arbitrary chunk boundaries are invariant, or that the limited tests establish universal arithmetic equivalence.

Remaining scope

This supports review of the narrow, disabled-by-default optimization. It does not close actual eight-way decode performance coverage, broader workload/hardware coverage, alternative-kernel comparisons, upstream CI, or default enablement. No new claims for multimodal inputs, 97k/512k contexts, or other checkpoints are made.

Raw traces, per-position metrics, individual benchmark results and scripts have been retained in the experiment archive; they are not attached to this comment.