[Bugfix] Fix two A-tile loading bugs in the preshuffle GEMM

[JohnQinAMD]

Summary

Two bugs in how the A tile reaches and leaves LDS in preshuffle_gemm. The first returns wrong results, the second returns correct results 3.3× slower on some shapes.

Bug 1 — tiles whose A tile is only partially loaded

num_a_loads is a truncating division, so tile shapes whose per-thread A copy does not divide into whole 16B loads never fetch the tail of the A tile. The kernel then computes on stale LDS and returns wrong results at full speed with no diagnostic:

bytes_per_thread_a = (tile_m * tile_k * elem_bytes) // total_threads
num_a_loads        = bytes_per_thread_a // a_load_bytes   # remainder dropped

The requirement is tile_m * tile_k * elem_bytes % 4096 == 0. For bf16 at tile_k=64 that makes tile_m a multiple of 32:

tile_m	bytes/thread	loaded
112	56	48
144	72	64
176	88	80
208	104	96

Easy to hit by accident: _TILE_PRELOAD_TABLE advertises (48, 64, 128), (80, 128, 256) and (112, 64, 256) as tuned entries, so autotuning tile_m over that table on a 2-byte dtype silently produces garbage.

Reproduction — against main at ac227c3, tile_m=112, tile_n=256, tile_k=64, in_dtype="bf16" compiles, runs, and gives a relative error of ~1.0 against torch.mm. After this PR the same call raises ValueError naming the required multiple.

Bug 2 — whole-tile A-fragment read costs up to 3.3×

#974 replaced the per-ki A-fragment copies in mma_kloop with a single
whole-tile fx.copy. The two are semantically identical, but the single copy leaves the scheduler no room to interleave the reads with the MFMAs that consume them.

Single-hunk A/B against main at ac227c3, bf16, M=54560 K=8192 N=8192, sync path:

tile_m × tile_n × tile_k	whole-tile	per-k-step
128 × 256 × 256	303	1011	3.3×
128 × 128 × 256	687	774	1.13×
64 × 256 × 256	863	859	neutral
128 × 256 × 128	1106	1103	neutral
128 × 256 × 64	1388	1381	neutral (control)

The regression is confined to specific large-tile_k shapes. 128×256×64 is the control; the affected cell is stable to ±0.3% over five repeats, so the 3.3× is not measurement noise.

Testing

Environment: MI355X (gfx950, 256 CU), ROCm 7.2.3 / PyTorch 2.11.0 / FlyDSL 0.2.4, base commit ac227c3 (main as of #1006). This is older than CI's rocm/pytorch:rocm7.14_ubuntu24.04_py3.11_pytorch_release_2.12.0; the linux-flydsl-mi355-8 runner is the same GPU, so CI should confirm on the supported toolchain.

pytest tests/kernels/test_preshuffle_gemm.py -c tests/pytest.ini -q
#   126 passed, 34 skipped

ruff check kernels/gemm/preshuffle_gemm.py tests/kernels/test_preshuffle_gemm.py
#   All checks passed!

Adds test_preshuffle_rejects_partial_a_tile (48/80/112/144 must raise) and test_preshuffle_accepts_whole_a_tile (64/96/128/160 must still compile). The 34 skips are the pre-existing MXFP4 fp8-A cases plus the non-8-bit async_copy parametrization, both unchanged by this PR.

Not tested on gfx942 — bug 1's check is pure Python validation with no architecture dependence; bug 2 is a scheduling change and its measurements are gfx950 only.

Breaking Changes

Configurations that previously returned silently wrong results now raise at compile time. Callers autotuning tile_m should catch ValueError and skip, as they already must for the existing tile_k and lds_stage validation.

[Copilot]

Pull request overview

Fixes two correctness/performance issues in the preshuffle GEMM kernel’s A-tile movement through LDS: (1) prevents silently wrong results when the A tile cannot be evenly distributed into 16B per-thread loads, and (2) restores per-k-step A-fragment LDS reads to enable better MFMA/read interleaving for large tile_k shapes.

Changes:

• Add a compile-time validation in compile_preshuffle_gemm that rejects (tile_m, tile_k, dtype) combinations leaving a partial per-thread 16B A load (previously computed on stale LDS).
• Change mma_kloop to issue A-fragment LDS reads per k-iteration instead of one whole-tile copy, improving scheduling latitude and addressing the reported bf16 tile_k=256 slowdown.
• Add unit tests asserting rejection/acceptance of representative tile_m values for bf16 tile_k=64.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
kernels/gemm/preshuffle_gemm.py	Adds strict A-tile load-granularity validation and switches A-fragment reads back to per-k-step copies for better MFMA interleaving.
tests/kernels/test_preshuffle_gemm.py	Adds compile-time tests ensuring partial A-tile load shapes raise ValueError and whole-load shapes still compile.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.