[justinchuby]

Summary

Closes the §8.5 blocker from #900: adds a byte-exact bf16 fused skip-RMSNorm kernel and an optimizer fold that collapses the Add → SimplifiedLayerNormalization residual seam. Outcome: SHIP the kernel (proven 0-ulp), NO-SHIP the fold as default (measured −1.5% regression under graph replay) — the fold is retained opt-in behind an env flag, so the default binary is unchanged.

Refs #900. Do NOT self-merge — Chew gates numerics.

What

1. skip_rmsnorm_bf16 NVRTC kernel (kernels/normalization.rs): sum = __float2bfloat16_rn(f32(residual)+f32(x)) (bit-for-bit a standalone bf16 Add), stores the bf16-rounded sum as the next layer's residual, then the identical rmsnorm_bf16 block-tree reduction (fp32 accumulate over the rounded sum, NORM_BLOCK=256). Guarded native dispatch with run_bf16_via_f32 fallback for non-dense/bias/header-unavailable (Rule 11 portability).
2. CudaSkipRmsNormFusion optimizer pass: collapses the seam into one com.microsoft::SkipSimplifiedLayerNormalization, deleting the standalone Add+norm. bf16 (or f32-gamma-over-bf16) only; fp16 left to CudaSkipRmsNormMatMulFusion. Opt-in via ONNX_GENAI_CUDA_ENABLE_SKIP_RMSNORM_FUSION (default OFF).

Numeric fidelity (Chew gate): BYTE-EXACT — 0-ulp

• GPU unit tests bf16_native_skip_rmsnorm_is_byte_exact_with_{bf16,f32}_gamma pass (bit-identical vs standalone Add(bf16)→rmsnorm_bf16 at H=6656).
• Real-model greedy stream fold OFF vs ON: bit-identical (128/128 tokens). No reduction reordered.

Perf A/B (H200, ONNX_GENAI_CUDA_GRAPH=1, --steady --warmups 1 --runs 3 --tokens 128, interleaved)

metric	fold OFF (default)	fold ON
SkipSimplifiedLayerNormalization nodes	4	104 (2/layer × 52)
decode	20.93 ms/token	21.25 ms/token
throughput	47.77 tok/s	47.06 tok/s (−1.5%)
greedy tokens	reference	byte-identical

Mechanism (honest): eager per-op timer shows the fused glue is faster (Add+norm+skip 355.6 → 349.0 ms — real launch saving), but graph replay already amortizes launches (~0.9 µs/node floor). At M=1 the RMS reduction is single-CTA; folding the residual add into it serializes work the standalone Add had spread across all 132 SMs. Same structural reason the multi-CTA GEMV megakernel was NO-GO (§7/#898) and glue collapse paid only +0.9% (§8.5/#900): decode is at its launch-amortized latency floor.

Verdict

• Kernel: SHIP — 0-ulp, portability-gated, no production change. Prerequisite for the only winning bf16 path: folding the norm into the neighbouring multi-CTA int4 GEMV prologue/epilogue (bf16 analogue of CudaSkipRmsNormMatMulFusion) — a larger GEMV job, not this fold.
• Fold: NO-SHIP as default — opt-in only; default binary unchanged (47.77 tok/s == baseline).

Validation

• cargo fmt --all -- --check clean; clippy clean on ep-cuda (pre-existing optimizer.rs:3347 + standard_attention.rs warnings are not mine).
• 61 optimizer + 20 normalization (incl. 2 new byte-exact) unit tests pass.
• CUDA EP cdylib ldd-confirmed not linking libonnxruntime.
• Doc docs/research/dense-decode-megakernel-feasibility.md §8.6 + decision drop .squad/decisions/inbox/sebastian-bf16-skip-rmsnorm.md.

🤖 Generated with Copilot. Co-authored-by: Copilot.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.79%. Comparing base (e9c9fc4) to head (05aa127).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #903   +/-   ##
=======================================
  Coverage   77.79%   77.79%           
=======================================
  Files         354      354           
  Lines      148003   148003           
  Branches   148003   148003           
=======================================
+ Hits       115133   115134    +1     
+ Misses      28378    28377    -1     
  Partials     4492     4492           

Flag	Coverage Δ
cli-ort-linux	83.42% <ø> (ø)
cli-ort-windows	83.01% <ø> (ø)
mlas	80.96% <ø> (+0.04%)	⬆️
offline	77.53% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.
see 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	52.64 µs	119.92 µs	+127.8%
🔴	gather/medium_f32_threads=1-internal/32768	5.18 µs	7.65 µs	+47.5%
🔴	matmul/large_generic_f16_threads=8/32x1024x1024	83.80 µs	110.24 µs	+31.6%
⚠️	matmul/large_generic_bf16_threads=8/32x1024x1024	2.18 ms	2.72 ms	+25.2%
⚠️	gather/large_f16_threads=1-internal/131072	17.55 µs	21.82 µs	+24.3%
⚠️	sampling_latency/top_k_per_token	58.32 µs	71.48 µs	+22.6%
⚠️	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.10 ms	7.36 ms	+20.7%
⚠️	qwen3_sampling_processors/top_k_top_p_fast	677.39 µs	807.51 µs	+19.2%
✅	sampling_latency/top_p_per_token	434.77 µs	485.36 µs	+11.6%
✅	matmul/medium_generic_f16_threads=8/32x512x512	41.76 µs	45.53 µs	+9.0%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	5.69 ms	6.16 ms	+8.3%
✅	qwen3_sampling_processors/top_k_partial_selection	153.89 µs	165.51 µs	+7.5%
✅	sampling_latency/greedy_per_token	3.55 µs	3.79 µs	+7.0%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	76.54 µs	81.84 µs	+6.9%
✅	kv_cache/alloc_dealloc_pages	40.61 µs	43.33 µs	+6.7%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.22 ms	2.35 ms	+5.9%
✅	logit_processing/seven_processor_chain_per_step	324.00 µs	341.67 µs	+5.5%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	464.48 µs	478.59 µs	+3.0%
✅	add/large_f16_threads=1-internal/4194304	41.85 ms	42.55 ms	+1.7%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.54 ms	2.55 ms	+0.2%
✅	gather/small_f16_threads=1-internal/4096	514.2 ns	514.7 ns	+0.1%
✅	reduce_mean/small_f32_threads=1-internal/4096	16.64 µs	16.31 µs	-2.0%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.75 ms	10.46 ms	-2.7%
✅	gather/medium_f16_threads=1-internal/32768	2.66 µs	2.58 µs	-3.0%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.94 ms	3.78 ms	-4.2%
✅	reduce_mean/medium_f32_threads=1-internal/65536	265.44 µs	254.05 µs	-4.3%
✅	grammar_masking/llguidance_compute_mask/32	78.22 µs	74.63 µs	-4.6%
✅	gather/large_f32_threads=1-internal/131072	42.41 µs	40.42 µs	-4.7%
✅	matmul/small_generic_f16_threads=1/1x256x256	37.74 µs	35.85 µs	-5.0%
✅	add/large_f32_threads=1-internal/4194304	43.63 ms	41.41 ms	-5.1%
✅	matmul/medium_generic_f16_threads=1/32x512x512	39.82 µs	37.76 µs	-5.2%
✅	add/medium_f16_threads=1-internal/262144	2.81 ms	2.66 ms	-5.2%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	579.81 µs	547.21 µs	-5.6%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	575.08 µs	540.03 µs	-6.1%
✅	tokenization/encode_tokens_per_second	468.64 µs	428.87 µs	-8.5%
✅	add/medium_f32_threads=1-internal/262144	2.90 ms	2.65 ms	-8.5%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	91.06 µs	82.99 µs	-8.9%
✅	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	87.45 µs	79.54 µs	-9.0%
✅	add/large_bf16_threads=1-internal/4194304	47.91 ms	43.20 ms	-9.8%
✅	add/small_bf16_threads=1-internal/1024	14.88 µs	13.35 µs	-10.3%
✅	add/small_f16_threads=1-internal/1024	15.24 µs	13.44 µs	-11.8%
✅	matmul/medium_generic_f32_threads=8/32x512x512	1.58 ms	1.39 ms	-11.9%
✅	sampling_latency/min_p_per_token	259.01 µs	224.65 µs	-13.3%
✅	gather/small_bf16_threads=1-internal/4096	604.2 ns	523.0 ns	-13.4%
✅	add/medium_bf16_threads=1-internal/262144	3.14 ms	2.67 ms	-14.9%
🟢	gather/large_bf16_threads=1-internal/131072	18.43 µs	15.60 µs	-15.3%
🟢	add/small_f32_threads=1-internal/1024	255.1 ns	213.5 ns	-16.3%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	37.27 µs	31.15 µs	-16.4%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	3.02 ms	2.43 ms	-19.8%
🟢	gather/small_f32_threads=1-internal/4096	940.5 ns	747.7 ns	-20.5%
🟢	tokenization/decode_tokens_per_second	8.94 ms	7.06 ms	-21.0%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	44.81 µs	34.78 µs	-22.4%
🟢	gather/medium_bf16_threads=1-internal/32768	3.46 µs	2.66 µs	-23.1%
🟢	matmul/small_generic_f32_threads=8/1x256x256	59.99 µs	45.96 µs	-23.4%
🟢	matmul/small_generic_f32_threads=1/1x256x256	55.40 µs	41.76 µs	-24.6%
🟢	matmul/small_generic_f16_threads=8/1x256x256	42.80 µs	31.50 µs	-26.4%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	673.46 µs	458.59 µs	-31.9%
🟢	reduce_mean/large_f32_threads=1-internal/262144	1.51 ms	1.01 ms	-33.3%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.69 ms	652.72 µs	-61.3%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 4.22 4.01 7.87 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)