[HIP] Fix rmsnorm_quant cross-row reads and writes for rows not a multiple of 4 bytes

[rk9595]

Summary

Fixes #5044.

add_rmsnorm_quant_kernel bounds each of its gmem descriptors at n rounded up to a whole dword:

const int oob_i = (n + ooba_i - 1) / ooba_i * ooba_i;   // 769 -> 770 for fp16
auto buffer_i   = opus::make_gmem<DTYPE_I>(input_ptr, oob_i * sizeof(DTYPE_I));

Rows are contiguous, so for any row whose byte length is not a multiple of 4 that window reaches into the next row. With one block per row and n = 769 fp16:

• element 769 is the next row's column 0, and it is inside the load window, so it is squared into that row's sum of squares;
• it is also inside the store window, so the block writes normalize(next_row[0]) * weight[769] over it — racing with the block that actually owns that row. Whichever store retires last wins, which is why only some rows show corruption and which ones varies between runs.

The same round-up reads past the end of the input tensor on the last row, and past the end of weight on every row.

This is not only the reported fp16 path. ooba_o = 4 / sizeof(DTYPE_O_STORE) is 4 for fp8/int8 output, so rmsnorm_quant / add_rmsnorm_quant clobber up to 3 bytes of the next row whenever n % 4 != 0, and residual_out in add_rmsnorm_quant has the 2-byte version of the same problem. Same class as #4467, which fixed it for the packed fp4 output only.

The existing sweeps run n in [4096, 8192, 16384, 32768, 65536] and [1024, 2048, 3584, 4096, 8192] — all dword-aligned — which is why this survived.

Technical details

• Every descriptor is now bounded at the row's exact byte length (row_bytes_i / row_bytes_o). fp4 keeps (n + 1) / 2, which was already exact. Reads past the row return 0, contributing nothing to the sum of squares or the abs-max; writes past it are dropped.
• A vectorized access moves whole dwords, so the one straddling the end of an unaligned row may now be dropped entirely. The row's trailing sub-dword elements are therefore reloaded and rewritten one at a time via opus::load<1> / opus::store<1>. Those are b16/b8 accesses that fit inside the exact bound, so they are performed whether the hardware range-checks per byte or per dword — the fix does not depend on which, and needs no per-arch probe.
• column_of(slot) maps a register slot back to its row column, mirroring the chunking in load_vector_nbytes / store_vector (num_load_inst reaches the store as num_repeat, so one mapping serves both). The quant tail reuses scaled_cast on a 2-lane vector and keeps lane 0, so its rounding is identical to the vector path rather than a reimplementation.
• Aligned rows are unchanged by construction: row_bytes equals the old rounded bound, the descriptors are bit-identical, and both tail loops are skipped on a uniform branch.

Test plan

• op_tests/test_rmsnorm2d.py: new test_rmsnorm2d_unaligned and test_rmsnorm2d_fuseAdd_unaligned with a guard-row helper, over n in [769, 1023, 2047, 4095, 6143, 8191] (one per dispatch bin; 769 is the issue's own shape) x m in [1, 7, 33] x {fp16, bf16}. Values go through torch.testing.assert_close, since the cross-row write lands inside the tensor and a guard-row check alone would pass it; the guard rows catch the separate write past the last row. These are separate cases rather than additions to l_n, to keep the file's CI time.
• op_tests/test_rmsnorm2dFusedAddQuant.py: generalized fix(module_rmsnorm_quant): bound packed FP4 output stores #4467's fp4-only guard to every kernel-written output (plain, quant, and residual_out), added n = 1027 and 2050 to the sweep, and made group quant skip shapes it cannot express so those n do not break modes 7/8.

Test result

gfx942 (MI325X), inside rocm/pytorch:rocm7.2.4_ubuntu22.04_py3.10_pytorch_release_2.10.0.

The issue's script verbatim:

before   hidden_size=768: mismatches=[],                        max_abs_diff=0.000977
before   hidden_size=769: mismatches=[[1,0],[3,0],[4,0],[6,0]], max_abs_diff=0.511719  -> AssertionError
after    hidden_size=768: mismatches=[],                        max_abs_diff=0.000977
after    hidden_size=769: mismatches=[],                        max_abs_diff=0.000000

(The corrupted row indices differ from the ones in the issue and move run to run, as expected for a race; column 0 and the 0.51 magnitude reproduce exactly.)

Both test files pass, including all 36 new unaligned cases. On test_rmsnorm2dFusedAddQuant.py (modes 1/2/5/6, fp8) the checkAllclose warning rate is identical before and after the patch at 2.0 per shape with no hard failures, so this introduces none — those are pre-existing bf16 tolerance noise on aligned shapes.

Perf, median of 7 reps x 200 iters, bf16, microseconds:

shape	fn	before	after	delta
4096x4096	rms_norm	15.891	15.805	-0.5%
4096x4096	fwd_with_add	28.840	29.012	+0.6%
8192x8192	rms_norm	54.757	54.912	+0.3%
8192x8192	fwd_with_add	106.377	106.410	+0.03%

Under half a percent, both directions. (m=256 shapes were also measured but are launch-overhead bound — 256x4096 and 256x8192 cost the same despite double the work — so their deltas are process noise.)

I only had gfx942; CI covers gfx950.

If you would rather keep this minimal, the exact-bound change alone fixes the reported corruption and I am happy to drop the tail path — I kept it because it makes the kernel correct regardless of the hardware's range-check granularity, but it is a clean separation if you prefer.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5290 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[valarLip]

hidden_size=769: this size will in real model?

[rk9595]

No — I don't know of a production model with an odd hidden size. Hidden dims are multiples of 64 in practice, and the reporter picked 769 as the smallest odd size next to an aligned one, so this is not on any hot path.

What made me write the fix is the blast radius rather than the shape. rms_norm sends every 2-D fp16/bf16 row with n <= 8192 to this kernel with no alignment check, and on an unaligned row the block writes its normalized values over the next row's column 0, racing with the block that owns that row. The damage lands in a different token's data, varies run to run, and raises nothing — which is what makes it expensive to attribute if anyone reaches it through the public op with a non-model tensor.

If you'd rather not carry the kernel change for a shape real models don't use, two smaller versions:

• Exact bounds only — about 10 lines, dropping the ~90-line tail path. That alone fixes the reported corruption and is what I verified on gfx942. I kept the tail because a bounded vector access straddling the row end may be dropped whole rather than partially performed, which would leave the last elements unwritten on an arch that range-checks per dword; I had no gfx950 to check that on.
• Reject unaligned n at the entry points. Worth noting this can't be done as a dispatch guard alone: the grouped/shuffle/e8m0 branch of rmsnorm2d_fwd_with_dynamicquant calls rmsnorm_quant with no opus fallback, so those paths would have to assert rather than fall back.

My preference is the full fix, then exact-bounds-only, but I'm happy to cut it down to whichever you want — say which and I'll push.

One other thing: the workflows on this PR are still sitting at action_required, so nothing but the title/label bots has run. Could you approve the run? gfx950 is the part I couldn't cover locally.