Skip to content

[cleanup] Migrate kernels to fx.copy / .llvm_ptr; expand kernel-code-cleanup skill - #974

Merged
coderfeli merged 3 commits into
mainfrom
cleanup/fx-copy-and-llvm-ptr
Aug 6, 2026
Merged

[cleanup] Migrate kernels to fx.copy / .llvm_ptr; expand kernel-code-cleanup skill#974
coderfeli merged 3 commits into
mainfrom
cleanup/fx-copy-and-llvm-ptr

Conversation

@coderfeli

@coderfeli coderfeli commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Summary

Behavior-preserving cleanup migrating kernels to the current fx.* surface, plus documentation of the migrations in the kernel-code-cleanup skill. All changes are one-for-one API swaps; validated cold + benchmarked on gfx950 (CDNA4).

1. create_llvm_ptr(idx, address_space=3) → .llvm_ptr (flash_attn_utils.py)

The two init_lds methods that hold a real fly.ptr from SharedAllocator now use the to_llvm_ptr op (#921): lds.<field>.ptr.llvm_ptr. The backend resolves Shared → LLVM <3>, matching the old hardcoded value. Other create_llvm_ptr/get_llvm_ptr/IntToPtrOp sites hold raw int/index addresses or memrefs (no fly.ptr) and are correctly left on the legacy helpers.

2. per-tile fx.copy_atom_call → fx.copy (11 files)

Every single-tile fx.copy_atom_call(atom, src, dst) swapped to fx.copy across norm (rmsnorm/softmax/layernorm), attention (rope), topk gating, mxfp_moe (gemm1/gemm2/common), and preshuffle GEMM. fx.copy iterates the same atom over the single-tile partition = one atom call = behavior-preserving. The one genuinely fused copy-in-gemm-loop (preshuffle_gemm.py mma_kloop) was hoisted to a single whole-fragment fx.copy; the compiler folds it back into the MFMA loop. copy_atom_call_ssa and swizzle/ABI-coupled single-atom loads left as-is (no meaningful TV layout).

3. skill docs (kernel-code-cleanup/SKILL.md)

Added the three migrations to the skill: §3b fly.ptr→!llvm.ptr, §3c manual s_waitcnt bitfields → arch-dispatched fx.rocdl.s_waitcnt, §7a building tiled copies from a TV layout (make_layout_tv/make_tiled_copy/_A/_B/_C), plus quick-ref rows and grep patterns.

Validation (gfx950, cold FLYDSL_RUNTIME_ENABLE_CACHE=0)

  • Accuracy: flash_attn fwd 33; rmsnorm 30, softmax 1, layernorm 14, topk_gating 1, fused_rope_cache 78, preshuffle_gemm 50, mxfp_moe variants+stage2 24 — all pass. (qk_norm_rope_quant.py has no test → compile-validated only.)
  • Perf (median-of-7, ±14% gfx950 noise):
    • llvm_ptr flash-attn: bf16 0.994×, fp8 1.002× — parity.
    • preshuffle_gemm mma_kloop hoist (M=N=4096 K=8192 fp8): 145.0 → 144.5 µs (1895 → 1902 TFLOPS, +0.4%) — parity.

Notes

  • Conflict resolution on mxfp_moe/gemm2.py: kept origin/main's newer A-LDS swizzle addressing (byte_off // 16) and applied the fx.copy swap on top.
  • The buffer_ops → fx.copy port was investigated separately and found to be a no-op on CDNA: the clean dense-tile ports already landed in Port GEMM/MoE/conv kernels to the layout API #913; remaining buffer_ops sites are scalar/gather/ABI-coupled (non-portable), and the only dense-A/B candidates are RDNA kernels not runnable/benchmarkable on gfx950.

🤖 Generated with Claude Code

Felix Li and others added 3 commits August 6, 2026 07:46
Felix Li @claude
[cleanup] flash_attn_utils: use fx .llvm_ptr for LDS base pointers
248de62 
Replace buffer_ops.create_llvm_ptr(idx, address_space=3) with the
backend-resolved fx primitive lds.<field>.ptr.llvm_ptr in the two
init_lds paths (DualwaveKernelContext, DualwaveFp8KernelContext),
where a real fly.ptr (Shared address space) is in hand.

All other create_llvm_ptr / get_llvm_ptr / IntToPtrOp sites in kernels/
operate on raw int/index addresses or memref tensors with folded offset
arithmetic (no fly.ptr) and are left on the legacy helpers.

Behavior-preserving: Shared address space resolves to LLVM <3> on the
rocm backend, matching the previous hardcoded value. Verified cold on
gfx950: tests/kernels/test_flash_attn_fwd.py 33 passed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Felix Li @claude
[cleanup] Migrate per-tile fx.copy_atom_call to fx.copy in kernels
11977f6 
Replace all 35 fx.copy_atom_call single-atom sites across norm, attention,
topk-gating, mxfp_moe, and preshuffle GEMM kernels with fx.copy. fx.copy
issues the same atom over the (single-tile) partition, so the swap is
behavior-preserving; validated cold (FLYDSL_RUNTIME_ENABLE_CACHE=0) on
rmsnorm, softmax, layernorm, topk_gating, fused_rope_cache, preshuffle_gemm
(all dtypes) and mxfp_moe variants/stage2.

Also hoist the per-ki A LDS->register copy in preshuffle_gemm's mma_kloop
out of the fx.gemm loop into a single whole-fragment fx.copy. Benchmarked
median-of-7 (gfx950, M=N=4096 K=8192 fp8): 145.0us/1895 TFLOPS baseline vs
144.5us/1902 TFLOPS ported (+0.4%, within noise) -- no regression.

qk_norm_rope_quant.py has no test coverage; its swap is the identical
single-tile view->rmem pattern and is compile/import-validated only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Felix Li @claude
[docs] kernel-code-cleanup skill: prefer fx.copy/fx.gemm, add TV layo…
72df699 
…ut, to_llvm_ptr, s_waitcnt

Document the current fx.* surface for the migrations the cleanup agents used:
- §7b prefer fx.copy / fx.gemm over copy_atom_call / mma_atom_call, INCLUDING
  single-atom sites (textual one-for-one, perf-neutral; keep only *_atom_call_ssa)
- §3b fly.ptr -> !llvm.ptr via ptr.llvm_ptr / fx.to_llvm_ptr (backend-resolved AS)
- §3c manual s_waitcnt bitfields -> arch-dispatched fx.rocdl.s_waitcnt(vmcnt=/lgkmcnt=/expcnt=)
- §7a build tiled copies from a TV layout (make_layout_tv / make_tiled_copy / _A/_B/_C)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@coderfeli
coderfeli force-pushed the cleanup/fx-copy-and-llvm-ptr branch from 0e86e52 to 72df699 Compare August 6, 2026 08:08
@coderfeli
coderfeli merged commit 772f070 into main Aug 6, 2026
11 checks passed
@coderfeli
coderfeli deleted the cleanup/fx-copy-and-llvm-ptr branch August 6, 2026 08:17
This was referenced Aug 13, 2026
Merged
Merged
JohnQin-llm added a commit that referenced this pull request Aug 19, 2026
@JohnQin-llm @claude yanyuan.qin@amd.com
[Fix] Restore the gfx950 XCD mapping for bf16 dense attention
43081ba 
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, TFLOP/s at
S = 61,380 / 180,180 / 239,580, aiter measured in the same process as control:

    cf1db2f  7/31  #919     1273.2  1287.0  1287.4
    772f070  8/6   #974     1253.7  1074.5  1053.4
    494b2cb  8/9   #945     1251.9  1073.3  1052.6
    f6d0a41  8/14  #991     1247.2  1076.2  1061.2
    47009a3  8/14  #1004    1248.2  1076.2  1061.3
    daee90e  8/17  #930     1249.0  1076.4  1060.6
    96af91b  8/17  #1014    1247.2  1076.6  1064.5
    4521c8c  8/18  #960     1247.8  1074.8  1061.2
    aiter                   ~1250   ~1253   ~1252

One step, between 7/31 and 8/6, and #950 (8/3) is the only commit in that
window that touches the mapping. Everything after it is flat to within noise.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is still ~1.4% below the 7/31 measurement
(1268.8 against 1287.4 at S=239,580). The candidates are the
`MIN_Q_BLOCKS_XCD_SWIZZLE` threshold #950 introduced for fp8, which #919 did
not have, and the LLVM bump in #945.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@JohnQin-llm JohnQin-llm mentioned this pull request Aug 19, 2026
Merged
4 tasks done
JohnQin-llm added a commit that referenced this pull request Aug 19, 2026
@JohnQin-llm @claude yanyuan.qin@amd.com
[Fix] Restore the gfx950 XCD mapping for bf16 dense attention
6f825c4 
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, cache cold, TFLOP/s
at S = 61,380 / 180,180 / 239,580, aiter in the same process as control:

    cf1db2f  7/31  #919     1284.6  1288.2  1286.8
    772f070  8/6   #974     1256.7  1188.1  1184.3
    494b2cb  8/9   #945     1255.6  1187.9  1187.2
    f6d0a41  8/14  #991     1256.0  1182.8  1183.6
    47009a3  8/14  #1004    1260.5  1188.6  1183.4
    daee90e  8/17  #930     1254.1  1077.0  1060.9
    96af91b  8/17  #1014    1251.8  1074.9  1060.6
    4521c8c  8/18  #960     1249.3  1073.9  1059.0
    aiter                   ~1252   ~1255   ~1255

Two independent steps. The first, between 7/31 and 8/6 (-8.0%), is #950 and is
what this commit fixes. The second, between 47009a3 and daee90e (-10.4%), is
#930 dropping the fastmath flag when it moved these kernels onto fx.maxnumf,
which reaches the raw MLIR builder without passing through resolve_fastmath;
that is a separate fix, filed separately.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is 1268.3 against cf1db2f's 1286.8, a 1.4%
gap. That is #930's dropped fastmath, which costs 10.4% with the remap off but
only 1.4% with it on -- the two interact. Fixing #930 separately closes it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
JohnQin-llm added a commit that referenced this pull request Aug 20, 2026
@JohnQin-llm @claude yanyuan.qin@amd.com
[Fix] Restore the gfx950 XCD mapping for bf16 dense attention
0e41bfb 
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, cache cold, TFLOP/s
at S = 61,380 / 180,180 / 239,580, aiter in the same process as control:

    cf1db2f  7/31  #919     1284.6  1288.2  1286.8
    772f070  8/6   #974     1256.7  1188.1  1184.3
    494b2cb  8/9   #945     1255.6  1187.9  1187.2
    f6d0a41  8/14  #991     1256.0  1182.8  1183.6
    47009a3  8/14  #1004    1260.5  1188.6  1183.4
    daee90e  8/17  #930     1254.1  1077.0  1060.9
    96af91b  8/17  #1014    1251.8  1074.9  1060.6
    4521c8c  8/18  #960     1249.3  1073.9  1059.0
    aiter                   ~1252   ~1255   ~1255

Two independent steps. The first, between 7/31 and 8/6 (-8.0%), is #950 and is
what this commit fixes. The second, between 47009a3 and daee90e (-10.4%), is
#930 dropping the fastmath flag when it moved these kernels onto fx.maxnumf,
which reaches the raw MLIR builder without passing through resolve_fastmath;
that is a separate fix, filed separately.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is 1268.3 against cf1db2f's 1286.8, a 1.4%
gap. That is #930's dropped fastmath, which costs 10.4% with the remap off but
only 1.4% with it on -- the two interact. Fixing #930 separately closes it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
coderfeli added a commit that referenced this pull request Aug 20, 2026
@JohnQin-llm yanyuan.qin@amd.com
@claude @coderfeli
[Bugfix] Fix two A-tile loading bugs in the preshuffle GEMM (#1007)
4f527a2 
Both bugs are in how the A tile reaches and leaves LDS. Neither produces an
error today: the first returns wrong results, the second returns correct
results slowly.

1. Reject 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, so tile_m of 112/144/176/208
   load only 48/64/80/96 of the 56/72/88/104 bytes per thread they own. This is
   reachable by accident because _TILE_PRELOAD_TABLE advertises (48, 64, 128),
   (80, 128, 256) and (112, 64, 256) as tuned entries.

   Validate the exact condition on the tile size rather than on either
   truncated intermediate.

2. Restore per-k-step A-fragment LDS reads in mma_kloop.

   #974 replaced the per-ki copies 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. Measured on MI355X
   (gfx950), bf16 M=54560 K=8192 N=8192, sync path, single-hunk A/B on ac227c3:

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

   The regression is confined to specific large-tile_k shapes; everything else
   is within the ~1% run-to-run noise (128x256x64 is the control here, and the
   affected cell is stable to +-0.3% over five repeats).

Co-authored-by: yanyuan.qin@amd.com <yanyuan.qin@amd.com@mia1-p02-g56.mia.tensorwave.lan>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Felix Li <felix.li@amd.com>
coderfeli pushed a commit that referenced this pull request Aug 20, 2026
@JohnQin-llm @claude
[Fix] Restore the gfx950 XCD mapping for bf16 dense attention (#1009)
9d7accf 
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, cache cold, TFLOP/s
at S = 61,380 / 180,180 / 239,580, aiter in the same process as control:

    cf1db2f  7/31  #919     1284.6  1288.2  1286.8
    772f070  8/6   #974     1256.7  1188.1  1184.3
    494b2cb  8/9   #945     1255.6  1187.9  1187.2
    f6d0a41  8/14  #991     1256.0  1182.8  1183.6
    47009a3  8/14  #1004    1260.5  1188.6  1183.4
    daee90e  8/17  #930     1254.1  1077.0  1060.9
    96af91b  8/17  #1014    1251.8  1074.9  1060.6
    4521c8c  8/18  #960     1249.3  1073.9  1059.0
    aiter                   ~1252   ~1255   ~1255

Two independent steps. The first, between 7/31 and 8/6 (-8.0%), is #950 and is
what this commit fixes. The second, between 47009a3 and daee90e (-10.4%), is
#930 dropping the fastmath flag when it moved these kernels onto fx.maxnumf,
which reaches the raw MLIR builder without passing through resolve_fastmath;
that is a separate fix, filed separately.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is 1268.3 against cf1db2f's 1286.8, a 1.4%
gap. That is #930's dropped fastmath, which costs 10.4% with the remap off but
only 1.4% with it on -- the two interact. Fixing #930 separately closes it.

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@coderfeli