From 9373fab5537f5d53935647ef4ef291a4cf66d654 Mon Sep 17 00:00:00 2001 From: juuso-oskari Date: Mon, 1 Jun 2026 11:41:52 +0000 Subject: [PATCH] CK-UA: replace FP8 repack ds_bpermute with v_permlane32_swap_b32 (gfx950) The FP8 32x32x16 fmha_alu1 repack trades each lane's "bad" 4-byte pack with the lane 32 away (QK-C vs PV-A layouts differ by an l^32 swap). Since 3431615ff this used an LDS-crossbar ds_bpermute plus an is_sub_0 v_cndmask mux. gfx950 exposes v_permlane32_swap_b32 (__builtin_amdgcn_permlane32_swap, feature permlane32-swap) which does the l^32 exchange in a single VALU op with no LDS round-trip. Verified on-device that permlane32_swap( lo_pack, hi_pack) returns {out_lo, out_hi} for every lane, folding both the cross-lane swap and the per-lane sub-block muxing into one instruction. Guarded #if defined(__gfx950__); ds_bpermute kept as the #else fallback (gfx942 lacks the feature). ISA (prefill_d128 fp8 instance): 12 ds_bpermute -> 0, replaced by 12 v_permlane32_swap_b32; v_cndmask muxing removed. FP8 prefill + decode PASS vs torch reference. Clean A/B (median of 3, b=4 FP8 prefill): sq=sk 2048/5000/10000 -> 1.6% / 1.9% / 2.1% faster, scaling with the per-iter repack count. Co-authored-by: Cursor --- .../pipeline/unified_attention_pipeline.hpp | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) 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 87628763c5..a5032684bd 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 @@ -1429,9 +1429,15 @@ struct UnifiedAttentionPipeline "expects PV per-thread buffer in chunks of 8 " "fp8 (one warp-gemm K iteration)."); + // On gfx950 the paired-lane (l^32) swap is a single VALU + // op (v_permlane32_swap_b32), so the lane-id / ds_bpermute + // address machinery below is only needed for the + // ds_bpermute fallback on older arches (e.g. gfx942). +#if !defined(__gfx950__) const int lane_id = ck_tile::get_lane_id(); const int paired_addr = (lane_id ^ 32) << 2; // bytes const bool is_sub_0 = (lane_id & 32) == 0; +#endif #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wuninitialized" @@ -1458,6 +1464,23 @@ struct UnifiedAttentionPipeline const uint32_t hi_pack = __builtin_amdgcn_cvt_pk_fp8_f32(g, h, hi_tmp, /*hi=*/true); +#if defined(__gfx950__) + // gfx950: do the paired-lane (l^32) swap in a single + // VALU instruction instead of an LDS round-trip. + // v_permlane32_swap_b32 exchanges operand0's high half + // (lanes 32..63) with operand1's low half (lanes 0..31) + // and keeps the diagonal halves, so feeding + // (lo_pack, hi_pack) returns {out_lo, out_hi} directly + // for every lane -- the swap and the per-lane sub-block + // muxing the ds_bpermute path needs are both folded into + // the instruction. (Semantics verified on hw.) + const auto swapped = + __builtin_amdgcn_permlane32_swap(lo_pack, hi_pack, + /*fi=*/false, + /*bound_ctrl=*/false); + const uint32_t out_lo = swapped[0]; + const uint32_t out_hi = swapped[1]; +#else // Issue ds_bpermute as early as possible so its LDS-DMA // latency overlaps with the byte writes below (and with // the next K-chunk's cvts after this iter unrolls). @@ -1467,6 +1490,7 @@ struct UnifiedAttentionPipeline const uint32_t out_lo = is_sub_0 ? lo_pack : recv; const uint32_t out_hi = is_sub_0 ? recv : hi_pack; +#endif p.thread_buf_[k_base + 0] = bit_cast(static_cast((out_lo >> 0) & 0xFFu));