[triton] fix transpose_scale in fused_rms_fp8_group_quant (was silently row-major)

[karverma-amd]

Summary

fused_rms_fp8_group_quant(transpose_scale=True) — and its sibling
fused_reduce_rms_fp8_group_quant — is supposed to return the per-group
activation scale in column-major ([num_groups, M]) byte-order so a CK
bpreshuffle w8a8 GEMM consumer can read the transposed layout directly, without
a .t().contiguous() relayout copy.

The kernel allocates the scale buffer as [num_bs_cols, M] and writes it
column-major (swapped strides), but the wrapper then returned it with:

out1_bs = out1_bs.view(M, num_bs_cols)   # BUG: reinterprets the buffer ROW-major

.view() on the [num_bs_cols, M]-contiguous buffer silently produces a
row-major (M, num_bs_cols) tensor — so the values are scrambled and
transpose_scale=True becomes a no-op (identical layout to
transpose_scale=False). Same bug in both functions.

Fix

out1_bs = out1_bs.transpose(0, 1)   # true column-major view of the [num_groups, M] buffer

A genuine transposed view: shape (M, num_bs_cols), strides (1, M), same
logical values as the default path. Pure-Python triton wrapper — no rebuild.

Tests

The two existing *_transpose_scale op-tests were themselves encoding the bug:
they asserted transpose_scale=True had the same row-major strides as the
default path and matched a .transpose().contiguous().view() (scrambled)
reference — i.e. they were green because the flag did nothing.

Both are rewritten to pin the correct contract, matching the already-correct
test_fused_flatten_fp8_group_quant_transpose_scale:

• same logical values as the default path (assert_close(..., atol=0)),
• column-major layout (stride() == (1, M), .T.is_contiguous()),
• bit-exact with the materialize (.t().contiguous().t()) copy path.

The degenerate num_bs_cols == 1 / M == 1 cases (where row- and column-major
coincide, so the bug is invisible) are guarded out.

Validation

Ran the affected op-tests on MI355X (gfx950), rocm/sgl-dev:v0.5.16-rocm720-mi35x-20260731
(torch 2.9.1+rocm7.2), across the full existing parametrization
(M, N1/N2[/N3], SPK, dtype ∈ {fp16, bf16}).

With the fix — all pass:

$ pytest test_fused_fp8_quant.py -k transpose_scale -q
93 passed, 210 deselected in 15.91s

Reverting only the two-line fix (transpose(0, 1) → view(M, num_bs_cols)),
same tests — the new assertions catch the bug:

$ pytest test_fused_fp8_quant.py -k transpose_scale
40 failed, 53 passed
  36 FAILED test_fused_reduce_rms_fp8_group_quant_transpose_scale
   4 FAILED test_fused_rms_fp8_group_quant_transpose_scale

E  AssertionError: Tensor-likes are not equal!
E  Mismatched elements: 1790 / 1792 (99.9%)          # transpose_scale=True vs default

The 40 failures are exactly the num_bs_cols >= 2, M > 1 cases (N1 ∈ {1536,
7168}); the 53 that still pass are the degenerate num_bs_cols == 1 cases and
the unaffected fused_flatten_fp8_group_quant transpose test — confirming the
guard boundaries and that the regression is genuinely exercised.

Impact (DeepSeek-V4 MI355X, gfx950)

Unlocks eliminating the largest remaining bpreshuffle scale copies on the MLA
wqkv_a (K=7168) / wq_b projections (×61 layers/step). GSM8K unchanged.

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 4506 --add-label <label>

[k50112113]

Hi @karverma-amd, I am a little bit confused as to why your transposed scale is now trans -> contg -> trans, according to AITER's CK gemm test https://github.com/karverma-amd/aiter/blob/50250c1eb350f192669baac688779e9c7d2930d7/op_tests/test_gemm_a8w8_blockscale.py#L112
it's doing trans -> contg -> view explicitly. As far as I understand it, the stride information after the view is not going to be meaningful anymore, you should still treat the x_scale as column major.

[k50112113]

update: I think #4406 adds the strided support for different triton version, so would you be able to help support both? we can default it to the trans -> contg -> trans and support trans -> contg -> view if compiler version <3.7

[k50112113]

LGTM