[ThomasNing]

ISSUE ID : ROCm/composable_kernel#3766

Ported from ROCm/composable_kernel#3765 because that repository is now a read-only mirror. This PR applies the same change to the canonical projects/composablekernel tree.

Problem

PermuteNEpilogue's ScaleM/ScaleN path gathers both scale tiles at the output thread-buffer index:

const auto sm = static_cast<float>(sm_tile.get_thread_buffer()[dst]);
const auto sn = static_cast<float>(sn_tile.get_thread_buffer()[dst]);

Both tiles are loaded through the store distribution, but scale_m is an (M, 1) view with stride 0 in N, so stepping the n2 Y-coordinate carries into its M coordinate: the value delivered at slot dst is scale_m[m + n2], not scale_m[m]. Each lane owns NRepeat consecutive columns, so NRepeat-1 of them get a neighbouring row's scale.

This is invisible in the configurations upstream exercises (a whole warp shares one B scale, BQuantGroupSize::kN = 128). At per-channel granularity (kN = 1) it corrupts (NRepeat-1)/NRepeat of all elements.

Second defect in the same path: the output and D windows advance by MWave*MPerXdl per MRepeat iteration, but scale_m's window never moves — so every iteration past the first re-reads the first M block's scales whenever MTile > MWave*MPerXdl. That includes the upstream flatmm example configs (M_Tile=128, MWave*MPerXdl=32).

Reproducer — stock CK, no downstream code

example/ck_tile/18_flatmm/flatmm_basic -prec=fp8 -warp_tile=2 -scale=1
  before: GPU verification fails, 13.77% of elements wrong
  after : verification result: correct

Fix

scale_m does not depend on n and scale_n does not depend on m, so gather sm at this lane's n2=0 slot and sn at its m2=0 slot; advance the scale_m window with the others. The store mapping is untouched.

Additional validation

Through AITER's a8w8_bpreshuffle tuner on gfx950 (MI350X): 23 kernel configurations × 5 GEMM shapes with per-token A scales and per-channel B scales, all at errRatio 0.0000; MRepeat=8 probes confirm the window advance. Identity-encoding probes were used to derive the actual lane → (m, n) mapping; that derivation is added as a comment block, including its correspondence with the paired host-side shuffle_b_permuteN layout.

Found while enabling TilePermuteN for a per-token/per-channel FP8 GEMM on MI350X.

🤖 Generated with Claude Code

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[ThomasNing]

Addressed the unit-test policy feedback in 16e133ac64:

• added regression coverage for the PermuteN per-token/per-channel scale buffer coordinates;
• covered the per-token scale-window advance across multiple M repeats;
• wired the test into the gfx950 epilogue test target.

The policy table now reports Unit Test: Pass. The remaining pre-commit gate is awaiting maintainer approval because this is a fork PR: https://github.com/ROCm/rocm-libraries/actions/runs/33921371027

[ThomasNing]

Strengthened the regression in 8bb48678b0:

• replaced the helper-only arithmetic test with an actual gfx950 GPU kernel that instantiates PermuteNEpilogue;
• uses non-periodic per-token and per-channel scales, plus MRepeat=8, so it exercises both corrected scale-buffer coordinates and the scale-M window advance;
• compares scaled device output against the canonical-coordinate result and removes the test-only helper API from the production header.

The new fork CI jobs are awaiting maintainer approval, including pre-commit and Component CI.

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-09-14 — Included in TheRock via ROCm/TheRock#8188.