From f947db93fe910f91c34867ae6386423117f573c4 Mon Sep 17 00:00:00 2001 From: juuso-oskari Date: Fri, 12 Jun 2026 08:18:21 +0000 Subject: [PATCH] CK-UA: fix wide-MMA FP8 P relayout (was cvt-only, missing QK-C->PV-A transpose) The wide 32x32x64 FP8 path shipped a "cvt-only, layouts coincide" P relayout (strategy C) that did no cross-lane movement, on the claim that the QK-C output and PV-A input per-thread layouts already match at K=64. They don't: QK-C holds one kv across many query rows while PV-A needs one query across many kv (a transpose), so the relayout MUST do the cross-lane permlane32_swap. Both the 32x32x16 and 32x32x64 MMAs share an identical 32x32 C-output distribution; only kABKPerLane changes (8->32), i.e. the per-lane chunk COUNT, not the per-chunk swap pattern. So strategy A's existing fused cvt+permlane32 relayout is correct for K=64 too -- the 8-fp8 loop just runs more iterations. The bug was masked by near-uniform softmax (transposing a flat P barely moves the row-sum), surfacing only as a few large-delta output lanes -> ~0.6-11% of elements over the loose fp8 tol on prefill_fp8, while bf16 and fp8 decode passed. Fix: gate strategy A for K=64 in addition to K=16; delete the cvt-only branch. prefill_fp8 + the full correctness matrix now PASS (standalone host-ref mismatch 0.96% -> 0.0000%); standalone perf holds ~1.66k TFLOPs (the permlane32 ops are cheap and overlap under softmax). Co-authored-by: Cursor --- .../pipeline/unified_attention_pipeline.hpp | 85 +++++++------------ 1 file changed, 31 insertions(+), 54 deletions(-) diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index 1b484e6009..5ae2c9a5a1 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp @@ -1788,10 +1788,26 @@ struct UnifiedAttentionPipeline // (A) Cross-lane in-register swap via // `__builtin_amdgcn_ds_bpermute` between paired // lanes (lane ^ 32). Cheap (one ds_bpermute_b32 - // per PV K-iter, no LDS traffic, no barrier), - // but ONLY works for the 32x32x16 MFMA shape: - // it assumes kABMLane=32 / kABKLane=2 with the - // paired-lane bit at position 5. + // per PV K-iter, no LDS traffic, no barrier). + // Works for the 32x32 MFMA shapes (both K=16 and + // K=64): both have kAMLane=32 / kABKLane=2 and an + // identical 32x32 C-output distribution, so the + // QK-C -> PV-A relayout is the SAME paired-lane + // (bit-5) swap-of-half regardless of K. (The wider + // K=64 A-operand only changes kABKPerLane 8->32, + // i.e. the per-lane chunk COUNT, not the per-chunk + // swap pattern -- the 8-fp8 loop below just runs + // more iterations. Verified byte-identical to the + // narrow path on hw.) + // + // NOTE: an earlier "cvt-only, layouts coincide" + // fast path for K=64 was WRONG -- QK-C holds one kv + // across many query rows while PV-A needs one query + // across many kv (a transpose), so skipping the + // swap silently corrupts P. The error was masked by + // near-uniform softmax (transposing a flat P barely + // moves the row-sum) and only surfaced as a few + // large-delta output lanes. // // (B) Layout-agnostic LDS roundtrip via // store_tile(QK-C dist) + s_barrier + @@ -1801,11 +1817,11 @@ struct UnifiedAttentionPipeline // On 4-warp decode_m128 this measured ~2-3x // worse end-to-end than (A). // - // We pick (A) for the 32x32x16 tiers (all of prefill, - // decode_m{32,64,128}) and (B) for the 16x16x32 m16 - // tiny-decode tier where (A) doesn't apply. This - // keeps the previously-tuned 32x32x16 perf intact - // while enabling FP8 on the m16 tier. + // We pick (A) for the 32x32 tiers -- 32x32x16 (decode + // m{32,64,128}) and 32x32x64 (wide-MMA prefill) -- and + // (B) for the 16x16x32 m16 tiny-decode tier where (A) + // doesn't apply. This keeps the previously-tuned 32x32x16 + // perf intact while enabling FP8 on the m16 tier. // // For strategy (A) the cvt and the cross-lane swap are // fused into a single 8-fp8-per-iter loop so that the @@ -1815,9 +1831,10 @@ struct UnifiedAttentionPipeline using PVWarpTile = typename UnifiedAttentionShape::Gemm1WarpTile; if constexpr(PVWarpTile::at(number<0>{}) == 32 && PVWarpTile::at(number<1>{}) == 32 && - PVWarpTile::at(number<2>{}) == 16) + (PVWarpTile::at(number<2>{}) == 16 || + PVWarpTile::at(number<2>{}) == 64)) { - // ---- (A) Fused cvt + cross-lane swap (32x32x16). ---- + // ---- (A) Fused cvt + cross-lane swap (32x32x16 / 32x32x64). ---- // // Per 8-fp8 K-chunk: // 1. cvt 8 fp32 -> 2 packed uint32 (lo_pack = slot[0..3], @@ -1834,9 +1851,10 @@ struct UnifiedAttentionPipeline // sub=1 | slot[0..3] | N=4..7 | K=8..11 BAD // sub=1 | slot[4..7] | N=12..15 | K=12..15 OK static_assert(sp(sp_reg_idx).p.thread_buf_.size() % 8 == 0, - "FP8 32x32x16 + Single cross-lane permute " + "FP8 32x32 (K=16/K=64) cross-lane permute " "expects PV per-thread buffer in chunks of 8 " - "fp8 (one warp-gemm K iteration)."); + "fp8 (one 32x32x16 warp-gemm K iteration worth " + "of the swap-of-half pattern)."); // On gfx950 the paired-lane (l^32) swap is a single VALU // op (v_permlane32_swap_b32), so the lane-id / ds_bpermute @@ -1918,47 +1936,6 @@ struct UnifiedAttentionPipeline p.thread_buf_[k_base + 7] = bit_cast(static_cast((out_hi >> 24) & 0xFFu)); }); -#pragma clang diagnostic pop - } - else if constexpr(PVWarpTile::at(number<0>{}) == 32 && - PVWarpTile::at(number<1>{}) == 32 && - PVWarpTile::at(number<2>{}) == 64) - { - // ---- (C) cvt-only, no cross-lane swap (32x32x64). ---- - // - // Under the wide v_mfma_f32_32x32x64 MMA the QK-C output - // and PV-A input per-thread layouts coincide (the K=64 - // CTransposed C-fragment is already in the A-operand order), - // so the relayout is just the fp32->fp8 pack — no permute, - // no LDS roundtrip, no barrier (matches the ASM kernel's - // _softmax_pack_P_fp8). Chained-`old` cvt pattern to match - // cast_tile_pk_fp8_fp32 byte-for-byte (see strategy A note). -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wuninitialized" - int dummy_old; - static_for<0, sp(sp_reg_idx).p.thread_buf_.size(), 4>{}([&](auto idx) { - const float a = - p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 0]); - const float b = - p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 1]); - const float c = - p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 2]); - const float d = - p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 3]); - - const uint32_t lo = - __builtin_amdgcn_cvt_pk_fp8_f32(a, b, dummy_old, /*hi=*/false); - const uint32_t packed = - __builtin_amdgcn_cvt_pk_fp8_f32(c, d, lo, /*hi=*/true); - sp(sp_reg_idx).p.thread_buf_[idx + 0] = - bit_cast(static_cast((packed >> 0) & 0xFFu)); - sp(sp_reg_idx).p.thread_buf_[idx + 1] = - bit_cast(static_cast((packed >> 8) & 0xFFu)); - sp(sp_reg_idx).p.thread_buf_[idx + 2] = - bit_cast(static_cast((packed >> 16) & 0xFFu)); - sp(sp_reg_idx).p.thread_buf_[idx + 3] = - bit_cast(static_cast((packed >> 24) & 0xFFu)); - }); #pragma clang diagnostic pop } else