[rapsealk]

Fixes #8831

Problem

For elems_per_row <= 128 the inference v2 norm launchers pack several rows into one 256-thread block, but the kernels have no row bound. When rows is larger than the rows per block and not a multiple of it, the last, partially filled block loads and stores rows past the end of every input and output. Reachable through CUDARMSNorm with a small channels, e.g. the per-head Q/K RMSNorm of EXAONE4 / EXAONE 4.5 (head_size 64 or 128, rows = tokens * heads).

Fix

• Pass rows to the rms_norm, pre_rms_norm, fused_ln and fused_residual_ln kernels.
• Derive the row index from the expression already used for block_offset, and mask the global loads and stores of rows past the end.
• Threads past the end still take part in the block reductions, so the grid and block shape stay the same. In-bounds outputs are bitwise unchanged.

Same change as #8830 makes for the v1 kernels in csrc/transformer/inference.

Tests

• Every kernel output in test_rms_norm.py, test_post_ln.py, test_pre_ln.py and the new test_ln.py (CUDAFPLN had no test in-tree) is allocated in front of 256 canary rows. A write past the last row fails under plain pytest, without compute-sanitizer.
• One sub-block shape per threads-per-row dispatch (1, 2, 4, 8, 16), each with one row more than fits in a block, in fp32, fp16 and bf16.
• assert_rows_independent: each row of a batched launch must match the same row run alone, within the dtype tolerance. The generic bf16 tolerance (rtol 0.48) is loose enough to accept a row normalized with its neighbour's statistics, so the norm tests use (3e-2, 1e-2) for bf16, measured with a 5x margin over all five kernels and the existing large shapes. No global tolerance change.

Test results

Hardware: 1x NVIDIA GB10 (sm_121, aarch64), driver 580.126.09, CUDA 13.0, torch 2.14.1+cu130, Python 3.13.2.

pytest tests/unit/inference/v2/kernels/core_ops/test_rms_norm.py \
       tests/unit/inference/v2/kernels/core_ops/test_post_ln.py \
       tests/unit/inference/v2/kernels/core_ops/test_pre_ln.py \
       tests/unit/inference/v2/kernels/core_ops/test_ln.py -m inference_v2_ops

kernels	passed	failed	skipped
master (67ec196) + this branch's tests	92	108	3
this branch	200	0	3

Measured on 5e2f6ea. The current revision (3f70aeb) trims the shape matrix and relaxes the row check to tolerance; the re-run on the same machine is pending and this table will be updated. Every failure on master is a sub-block case failing on the canary assertion. The skips are fp32 with 65-128 channels (#8827, see below).

compute-sanitizer memcheck (reads as well as writes)

Each of the four kernels on exactly sized [rows, channels] tensors with PYTORCH_NO_CUDA_MEMORY_CACHING=1 under compute-sanitizer --tool memcheck, 11 shapes covering every threads-per-row dispatch and the in-place path:

	result
master kernels	Invalid __global__ read of size 16 bytes in rms_norm / fused_ln / fused_residual_ln, 262 errors over the 11 shapes, then a sticky launch failure
this branch	ERROR SUMMARY: 0 errors

Throughput (GB10, median of 200 CUDA-event timings, master -> branch, microseconds)

dtype	channels	rows	rms_norm	fused_ln	post_ln	pre_ln
fp16	4096	1	7.62 -> 7.58	7.65 -> 7.62	7.74 -> 7.71	7.84 -> 7.74
fp16	4096	4096	279.9 -> 281.1	280.3 -> 281.3	419.7 -> 421.6	566.7 -> 571.1
fp16	8192	4096	561.7 -> 558.9	552.0 -> 552.9	843.2 -> 842.6	1151 -> 1143
bf16	8192	4096	563.9 -> 558.8	553.4 -> 551.9	841.4 -> 840.6	1144 -> 1136
fp16	128	2560 (EXAONE4 Q, 64 tokens)	7.46 -> 7.30

All within +-1%, inside round-to-round noise. In-range outputs of all five wrappers were compared bitwise against master over 2430 (kernel, dtype, channels, rows) combinations: no differences.

Compatibility matrix

• Validated: inference v2 core ops on CUDA; fp16 and bf16 for 8-128 channels (sub-block schedule) and the existing 1024-14432 channel shapes; fp32 for 4-64 channels and the large shapes.
• Not validated: ROCm builds, pre-sm_80 GPUs, end-to-end inference v2 model runs (e.g. EXAONE4 through the ragged engine).
• Not applicable: ZeRO stage, offload, training parallelism. Inference-only kernels.

Not a duplicate

No open PR references #8831. #8830 (OOB rows), #8743 (rows == 0) and #8827 (fp32 65-128 channels launches no kernel) change only the v1 kernels under csrc/transformer/inference/csrc/. The v1 kernels, the v2 rms and layer norm kernels, and the other v2 kernels under kernels/{core_ops,ragged_ops} were checked; only the v1 and v2 norm kernels pack rows per block this way.

Related v2 defects not fixed here (same on master)

• fp32 with 65-128 channels: threads_per_group == 32 falls through the sub-block dispatch and no kernel is launched (v2 sibling of Fix inference layer_norm/rms_norm launching no kernel for fp32 hidden sizes 65-128 #8827). The LN tests skip that cell with an explicit reason.
• rows == 0 divides by zero in all four v2 launchers (v2 sibling of Return early from inference norm launchers when rows == 0 #8743).
• rms_norm.cpp has no dim() == 2 or contiguity check; layer_norm.cpp flattens 3-D inputs, so the two bindings disagree.
• block_offset is int in all four kernels and overflows past 2^31 elements. The new row index adds no overflow of its own.
• csrc/quantization/quantize.cu (cached_quantization) uses the same row packing with bounded stores but unbounded loads.
• get_dtypes() in the v2 test utils caches the first caller's filtered list, so under the directory-level CI invocation later callers lose fp32 (found by the Codex review here). Split out into its own PR so this one stays a kernel fix.

AI assistance

This change was written with AI assistance (Claude Code). The submitter reviewed every changed line and ran the tests above on the hardware listed.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-22T11:05:09.868386Z	05e6419	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 05e64191f3

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Prevent the dtype cache from skipping fp32 cases

Under the full inference-v2 CI invocation in .github/workflows/nv-a6000.yml:61-63, collection reaches test_bias_activation.py:66-68 before this module, and that decorator calls get_dtypes(include_float=False). Since inference_test_utils.py:25-37 caches the resulting half-only list globally, this check skips both explicit DtypeEnum.fp32 cases even on fp32-capable hardware; only the targeted three-file command happens to seed the cache with fp32 first. This leaves normal CI without fp32 coverage of the new row bounds, so test optional bf16 support directly (or make the cache respect include_float) instead of using the shared cached list here.

Useful? React with 👍 / 👎.

[rapsealk]

Confirmed, and it is wider than this test: under the directory-level CI invocation test_bias_activation.py seeds the cache first with include_float=False, so every later get_dtypes() caller (test_post_ln, test_pre_ln, test_rms_dtypes, test_gated_activation, test_post_ln_module, test_ragged_embed, test_logits_gather) lost float32 as well. Fixed in dd9f108 by caching the full list in get_dtypes() and filtering per call, so this skip check now only drops bf16 where it is unsupported. Locally kernels/core_ops goes from 110 passed + 8 skipped to 136 passed. Note for CI: the fp32 cases of test_ragged_embed and test_logits_gather will run for the first time; I could not run them here (no aarch64 deepspeed-kernels).

[rapsealk]

Pushed bbef33f after an adversarial review pass (two independent reviewers plus the Codex finding above). Kernel diff is unchanged; the evidence around it was the weak part:

• get_dtypes() cache fix so fp32 cases actually run under the directory-level CI invocation (Codex finding; it affected every get_dtypes() caller, not just this file).
• fused_ln / CUDAFPLN had no test anywhere in-tree: added test_ln.py.
• Sub-block shapes now cover every threads-per-row dispatch (1-16): 17 / 33 / 65 / 129 rows, 16-128 channels.
• assert_rows_independent: each row of a batched launch must be bitwise equal to the row run alone. The bf16 tolerance (rtol 0.48) alone cannot see one row's reduction leaking into a neighbour, so the bf16 cases were effectively canary-only before.
• test_exaone4.py exercises the real _apply_qk_norm with the real kernels as the inference-level path that motivated the fix.
• Description now lists the v2 siblings that are not fixed here (rows == 0 as in Return early from inference norm launchers when rows == 0 #8743, fp32 65-128 channels as in Fix inference layer_norm/rms_norm launching no kernel for fp32 hidden sizes 65-128 #8827, and the unbounded loads in quantize.cu), narrows the fp32 claim to <= 64 channels, and adds before/after throughput numbers (all within +-1%).

Branch: 157 passed. Same tests on master's kernels: 78 failed / 79 passed.

[rapsealk]

Pushed 5e2f6ea after a second independent review pass. Kernel diff still unchanged.

• Dropped the stub-based test_exaone4.py: whether it failed on master depended on where the caching allocator placed K relative to Q (2-3 of 6 layouts per run), and it coupled to private attribute names. The EXAONE4 shape ((tokens * heads, 128) fp16 / bf16) is covered deterministically by the kernel-level canary tests; the description says why no model-level test is included.
• The bf16 oracle was the real gap: at the generic bf16 tolerance a neighbour-row-variance leak passes allclose 50/50. Norm tests now use (rtol 3e-2, atol 1e-2), measured 5x margin across all five kernels and 13 shapes; no global tolerance change. assert_rows_independent stays, documented as pinning the current schedule property rather than the norm contract.
• The 'every threads-per-row dispatch' claim was false for 1 thread per row: added fp32/4 and fp16/8 at 257 rows, plus (33, 128) in the LN tests (fp32 skipped there with an explicit Fix inference layer_norm/rms_norm launching no kernel for fp32 hidden sizes 65-128 #8827 reason). Canary depth raised to a full 256-row block.
• 'Not fixed here' now also lists the rms_norm.cpp 2-D / contiguity gap and the int32 block_offset overflow.

Branch: 200 passed, 3 skipped. Same tests on master's kernels: 108 failed / 92 passed, all on the canary. Note that no CI workflow has run on this PR yet (action_required).