[Rui Ren (rui-ren)]

Description

Reduces CUDA packed decode tile/loader overhead above #34002. Narrows eligible
eight-row IQ2 FC1 tiles, adds aligned record-window reconstruction for eligible
IQ2 FC2 rows, and uses full-tile NVFP4 paths only when their shape/alignment guards
hold. Unsupported geometries retain the existing fallback.

Dependency: #34002 in
native stack #33967.
No schema, packed layout, cache/mask precision workaround, or GenAI source change.
The R8-prefill/R1-ordinary hybrid session/package work is in the companion
microsoft/onnxruntime-genai#2916, not implemented by this native-kernel PR.

Summary of Changes

Packed Decode Specializations

File	Change
onnxruntime/contrib_ops/cuda/moe/moe_packed.cu	Select narrower FC1 tiles and eligible aligned FC2 loader variants.
onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh, onnxruntime/contrib_ops/cuda/moe/moe_packed_iq2.cuh	Reuse the shared IQ2 decoder for aligned word windows, including two-byte-offset records.
onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu	Specialize eligible full NVFP4 output tiles while preserving guarded fallbacks.
onnxruntime/test/contrib_ops/moe_packed_routing_test.cc	Exercise eight-row width/routing boundaries and larger intermediate widths.

Testing

Review-fix snapshot: ed6e6e1991fd944f65f1d97472d3f29046d97361.

Bound aligned-window FC2 to four resident CTAs, kept the R8-specific narrow FC1 gate decode-only (!use_prefill_tiles), simplified the FC2 guard to its real row/base alignment invariant, and reused the lower shared decoder. Nonuniform regression protection is introduced in #33963/#34002. Mamba replay changes were moved into their owning PR #33969; this layer no longer changes Mamba numerics or tests.

Owning-layer validation: Integrated stack regression run and exact-production-kernel H100 A/B below.

Integrated tip regression: 51 C++ provider/schema tests, 5 actual internal CUDA tests and 2 Python IOBinding tests passed, zero failures/skips. All owning-layer provider cases are retained. The final overlay recompiles the changed CUDA/host/schema/test sources over ABI-matched preserved ORT support; not a clean full-source build, full-CI approval, independent HF parity or new end-to-end qualification. Targeted configured clang-format/Ruff and whitespace checks passed.

Review-requested kernel A/B (H100, CUDA 12.8.93)

Exact production kernel/helper text from the integrated head; random nonuniform IQ2, 512 experts, 16 rotating route sets, R=8. All variant outputs are bit-identical across all route sets. CUDA-event median: 20 warmups, 9 alternating batches x32 launches; no full-model/cold-cache guarantee. k=9 is decode, k=17 is four-subtile prefill; the latter remains excluded from the narrow dispatch.

Hidden	Inter	k	FC1 wide→narrow μs	FC2 plain / window / hinted μs
1024	1024	9	15.035 → 10.888	12.874 / 11.966 / 12.062
1024	1024	17	16.264 → 14.610	18.149 / 15.968 / 16.188
1024	2048	9	18.041 → 15.920	21.786 / 20.282 / 20.458
1024	2048	17	26.279 → 23.879	32.182 / 28.613 / 28.801
2048	1024	9	27.187 → 18.500	20.663 / 19.135 / 19.234
2048	1024	17	28.614 → 25.400	29.040 / 25.962 / 26.176
2048	2048	9	31.466 → 27.515	34.251 / 32.346 / 32.446
2048	2048	17	48.462 → 42.877	51.977 / 45.898 / 46.270
4096	1024	9	49.654 → 32.181	30.966 / 29.015 / 29.088
4096	1024	17	53.387 → 46.004	50.201 / 45.097 / 45.371
4096	2048	9	57.550 → 50.024	53.390 / 50.104 / 50.318
4096	2048	17	91.847 → 79.396	90.444 / 80.277 / 80.838

The local 12.8 compiler did not reproduce the reported 79-register/3-CTA window allocation: unhinted and hinted window kernels both use 64 registers/4 CTAs here. The selective hint is retained as a portability cap for the reviewer's newer-toolchain result, with its local timing tradeoff shown rather than hidden. These kernel-only measurements do not attribute the older hybrid end-to-end gains below to the later narrow-tile commit.

Historical measurements and qualification evidence (not current-head acceptance)

These previously reported results are retained for provenance. They have not been rerun on the rebased head above; earlier CI/build status and benchmark results do not establish current-head CI, accuracy, or merge readiness.

The following results include the companion GenAI runtime and hybrid package, not just this PR's native kernel changes. They use the original pinned binaries and package hashes below, not the current rebased ORT head.

Unprofiled frozen real30, six H100s, two reversed same-GPU passes, 240 complete measured
generations; original warmups, mode order, EOS/512 cap.

Actual input tokens	Before ordinary tok/s	Hybrid ordinary tok/s	Gain	Original R8/k7 tok/s
178-254	94.40	119.95	+27.1%	361.46 → 361.57
427-932	94.01	119.59	+27.2%	298.68 → 298.94
1028-1663	93.89	118.90	+26.6%	324.25 → 324.31

Ordinary TTFT medians: short 357.35→359.00ms, medium 826.97→831.29ms,
long 2153.15→2161.92ms. The earlier static-R1 version's prefill regression is not
the package being qualified here.

Separate same-GPU synthetic-2048 CUDA profiles, 64 captured target passes:
IQ2 expert work 3.259→1.496ms/pass, Mamba 0.824→0.423ms/pass.
These are device-work attribution, not an additional throughput claim.

Correctness: all240 exact same-mode sequences and complete counters match;
388 complete live committed-state boundaries (66.43GB compared), capture off/on,
two resets, prefix through130; original R8 ordinary/k7 retains1782 complete finite
tensor hashes. Three native kernel tests, seven Mamba operator tests, fifteen
GenAI tests execute without skips. The new replay test also passes CUDA memcheck
with zero errors (100 repetitions).

Scope: hybrid is ordinary-only. Both speculative arms use the original R8 package.
The benchmark explicitly keeps two per-mode models resident in each arm; it is not
a unified hybrid-speculative package or a k7 speedup. No precision/cache/mask workaround,
no weight repack, no independent HF rerun, no clean-build/full-CI claim.

Reproduce the exact tested graph/config with:

python examples/python/nemotron-bolt-hybrid.py --input <packed-r8> --output <fresh-same-filesystem-hybrid>

Requires onnx-ir/numpy/ml-dtypes and both updated runtimes. External-data validation
stays enabled; payloads are hardlinked inside the output package, not copied.

Pins:

• ORT: 0357a7dfe47d3c383bf36be62015162b6a612d5e
• Compiled/common GenAI: d73a6e627d6d20a5a1688b9fbd46237069a7be89
• Published GenAI: db6ffd29fc96e319f645ff461efa05c643235f29 (converter/docs-only child)
• CUDA provider: fac33c75d0336236736aa96456268057b628d93d0b9901d49598e275bf335275
• GenAI shared library: f977660a408bc0a01a31151c7280282ba443fa256785680ae3e841654eec9be7
• Tested hybrid graph: aaf5d6ae06b6740ab5b3af523359581930d2b6c42c38a18529a22737c2cfe530
• Config: 77f2d920d2b04e17884c3d2ee691d9053225b957c047454ebd0a4ea3649dac44
• Unchanged packed payload: 06d3d508e8d11235d0b770484b6b59764c2f297c3383b6d0de75648d26c11867

Full commands, per-case exact tokens/counters, raw kernel traces, provider maps,
UUID/admission receipts, and rejection evidence remain on the validation host under
/opt/bolt-runtime-fixes-20261108-v1/evidence/ordinary-r1-v1. Sealed manifest SHA256: 412b3898f09565557b01e82906358f017f9df3227db2bf330850245a52d32bbc.
Root-disk binaries were reloaded and verified with full ordinary/k7 generation;
the large model payload still resides on the existing tmpfs, not reboot-persistent.
Original package/deployment untouched; no merge.

Checklist

• Review feedback addressed with focused regressions
• Documentation and supported contracts updated
• Native stack order and immediate-parent dependencies preserved
• Clean full-source/full-CI qualification (not claimed by focused overlays)

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Rui Ren (rui-ren)]

Stack #33967 review — parent 99f4ab7c63a5 -> head bbac6c023794.

Reviewed all parent-relative changes: full NVFP4 tile guards, narrow eight-row FC1 tiles, aligned FC2 record-window selection, and multi-slot single-row Mamba readout ordering. The Mamba dispatch change preserves the four-part multi-row reduction for multi-slot replay while retaining the single-slot fast path; the added replay regression checks outputs and both snapshot types.

One additional regression-coverage gap, detailed inline: the newly eligible aligned FC2 record loader is exercised with repeated records/scales, so the tests cannot distinguish certain wrong-record selections. Add a nonuniform packed-row reference case for this new loader.

No wrong kernel result or out-of-row read demonstrated. Host checks covered 132,096 record windows across even-block row sizes and confirmed that the eligibility guard keeps the two-byte lookahead inside the row. This is not a fresh CUDA execution, capture/replay run, or throughput remeasurement.

[Tianlei Wu (tianleiwu)]

Reviewed this layer (stack #33967, layer 6, the top) at bbac6c0237.

Kernel review. I extracted this head's IQ2 FC1/FC2 and NVFP4 kernels into a standalone harness and ran them against the #34002 versions on random, non-uniform data. Every specialization is bit-identical across 10 shapes x 5 weight-base offsets (decode and 4-subtile prefill tiles), and compute-sanitizer memcheck is clean with exact-size weight allocations (a deliberately 4-byte-short bank is flagged, so the check is byte-exact). A standalone replica of the new Mamba replay test passes on this head and fails with the previous rows > 1 dispatch. I found no correctness defect.

Requests (inline)

1. The FC2 window variant compiles to 79 registers on sm_90 (75 on sm_80), dropping occupancy from 4 to 3 CTAs/SM, and measures 3-8% slower than the plain loader at 2-9 rows on an H200; a min-CTAs __launch_bounds__ hint on that instantiation makes it 3-5% faster. I reproduced both.
2. The narrow FC1 tile has no measurement for this head and is mixed at prefill-tile route counts.
3. The Mamba2Verify dispatch change is unrelated to the PR title, changes numerics of an op owned by #33969, and leaves the doc and the test references stale.
4. The new Mamba regression test is not selected by the documented FusedStateSpaceKernelTest.Mamba* filter.
5. Nit: redundant guard terms for the FC2 window.

Evidence in the description. The performance numbers pin ORT commit 0357a7df, which is the pre-rebase of 42ebcad5e8 ("Preserve Mamba readout order ...") and predates bbac6c0237 (narrow FC1). The only R8 row in the table is flat (361.46 -> 361.57 tok/s), the +27% ordinary-decode result is attributed to companion runtime/package changes, and the detailed evidence lives on a host path maintainers cannot inspect. The description is upfront that these results are historical, but a PR titled "Reduce ... overhead" should carry per-kernel before/after numbers (nsys/ncu or a microbenchmark) for the narrow FC1 tile, the FC2 window and the NVFP4 full tile on the target GPU, or be retitled as specializations pending measurement. For reference, my standalone H200 numbers: narrow FC1 -2..-8% at 72 routes but +4..+11% at 136 routes for hidden <= 2048; FC2 window +3..+8% slower as compiled; NVFP4 full tile -0..-7% for five of six chunk-128 shapes (+3.2% at n=8192, k=2048) and -11% / -40% at chunk 64.

Minor

• NVFP4 full-tile coverage: Nvfp4Bf16ProductionProjectionShapes reaches the new specialization only for chunk 128 (its single k % 128 != 0 shape uses n=35), so chunk-64 + full tile has no test; a rows=8 case with n a multiple of 32 and k=192 or 576 would cover it. kFullNvfp4Tile also silently assumes m_rt == 8 and n % (MT*16) == 0 (enforced only by the launcher guard); a one-line comment at the template declaration would stop a future caller reading rows past m_rt.
• moe_packed_iq2.cuh: no kFc1 = true, kAlignedRecordWindow = true instantiation exists any more, yet code = (kCodesWord || kAlignedRecordWindow) ? ... : ... still carries that combination, and the template's precondition (4-byte-aligned row base plus an even block count, otherwise the 2-byte look-ahead reads past the last row) is only documented in the function body. Consider static_assert(kCodesWord || !kAlignedRecordWindow) (or deleting the dead branch) and stating the precondition at the template parameter.
• CI: Kernel Documentation Generation fails here; it is inherited from #33994 (see my comment there).

[Tianlei Wu (tianleiwu)]

Re-review of 623a725f5c.

Verdict: COMMENT — no defects found in this layer; it is blocked only by #33994's compile errors.

I checked each new guard against the invariant it relies on:

• narrow FC1 tile: num_rows == 8, 16-byte aligned input with hidden_size % 8 == 0, decode tiling only (!use_prefill_tiles); the MT = 1 instantiation derives every index from MT, so it changes no partition boundary or rounding point;
• aligned FC2 record window: down_row_bytes % 4 == 0 ⇒ even number of 66-byte blocks ⇒ the last block's records are word-aligned ⇒ the two-byte look-ahead of the misaligned ones always stays inside the row (correct, and confirmed by memcheck below);
• full NVFP4 tile: rows == 8, N % (MT * 16) == 0, 16-byte aligned A, K % 8 == 0.

Independent verification (stack tip = this head; Linux, CUDA 13.0, sm_80;sm_90a, Release; A100-SXM4-80GB)

• Build OK; the stack's 60 provider tests and 5 internal CUDA tests pass (also 3x shuffled). Iq2NonuniformSingleProductOracleAndEightNineRowEquivalence (rows 8 take the new narrow FC1 tile, rows 9 do not; 42 channels of nonuniform data with one surviving FC1 channel, so the oracle is exact) passes. The pre-existing *Packed* suites (655 tests) show 0 failures.
• The aligned IQ2 record-window decoder is bit-exact against the CPU dequantizer (K = 512/1536/2048/2560) with 0 compute-sanitizer memcheck errors on exact-size allocations; synccheck, racecheck and initcheck on the tests are clean.
• Not covered by anything I could run: the Hopper-only branches that the NVFP4 specializations extend (my GPUs are A100s).

CI is red

• Windows GPU CUDA CI Pipeline, Windows CUDA Plugin EP Build, Windows GPU TensorRT CI Pipeline: moe_group_router_impl.cu(157), C4244 → C2220 — from #33994.
• wasm_Debug / wasm_Release / wasm_Release_static_library: linear_attention_gates.cc:300,302, -Wshorten-64-to-32 — from #33994.
• build_x64_debug: onnxruntime_provider_test passed; the only failure is onnxruntime_shared_lib_cpuinfo_dlopen_test, which this stack does not touch.