From 8735fd3e71debe9327fc9415960d20d7993de5cd Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 19:12:52 -0400 Subject: [PATCH] Replace the per-MFMA V transpose with CDNA4's transposing LDS read The in-tree prefill kernel reaches 5.0% of matmul peak on gfx942 and 4.2% on gfx950. Profiling attributes that to one line rather than to the MFMA tile shape: compute_qk reads K with a plain load_fragment, while compute_sfm_v reads V through load_matrix_m16n16_trans, which appends transpose_mma_tile -- six cross-lane shuffles, once per PV MFMA. Both phases issue the same 16,777,216 MFMA instructions and PV costs 9x more on gfx942, 15x on gfx950. Ablating just that transpose: 4.134 -> 2.665 ms (gfx942), 2.812 -> 1.662 ms (gfx950). All six shuffles lower to ds_bpermute_b32 -- 2264 in the kernel ISA, with zero ds_swizzle_b32 and zero v_mov_b32_dpp -- so none of them are cheap. CDNA4 does the whole thing in one ds_read_b64_tr_b16. gfx942 rejects the builtin (needs target feature gfx950-insts) and keeps transpose_mma_tile; its device ISA is byte-identical to base, 43515 lines either way, so CDNA3 is proven untouched rather than merely measured-equal. The builtin is not a drop-in at the same address. It redistributes as out(16g+4a+b)[j] = in(16g+4j+a)[b], so at the load_fragment seed lane 0 gets {0,64,128,192} where the software path gives {0,16,32,48}. Solving for the required seed gives row = 4*(lane/16) + (lane%16)/4, column = lane%4 -- a different index expression into the same swizzled tile, so no LDS layout changes. Verified bit-identical to load_fragment + transpose_mma_tile across all 64 lanes before any kernel code was touched. Because the arches want different per-lane addresses, seeding the obvious way stays correct on CDNA3 and is silently wrong on CDNA4. The seeds therefore live on smem_t beside the load, and one constant (kHasTransposingLdsRead) gates the seed and the load together so they cannot drift apart. The read is 16-bit only, which KernelTraits::DTypeKVSmem already guarantees for every KV dtype: since #368 an fp8 cache is dequantized on the way into LDS and static_assert(sizeof(DTypeKVSmem) == 2) holds it there. The S-fragment transpose in compute_sfm_v is unchanged and still issues ds_bpermute_b32; S lives in registers, never LDS, so the hardware read cannot reach it. This removes the V-side transposes only. gfx950, three interleaved reps against base at this commit's parent: 2.813 -> 1.648 ms, 97.7 -> 166.8 TFLOP/s, 1.71x. That edges past the ablation ceiling because the hardware folds read and transpose into one op where the ablation still paid for a load. Correctness on gfx950, the only arch that compiles the new path: 1632 single-prefill fa2 cases against the independent naive_attention oracle, 162 in-tree fp8, 59 bf16 custom-mask, 24 new batch cases, 330 logits-cap, 622 POD/cascade/shared-prefix, 3492 batch-prefill, 6 forced-CTA_TILE_Q. gfx942: 1632 + 144 + 24. The new batch test adds head_dim 64/256, non-causal and fp16 to the independent batch oracle, which previously existed only for bf16 + causal + head_dim 128 in test_batch_prefill_bf16_custom_mask; every other batch test compares against single_prefill_with_kv_cache("fa2"), which reads V through the same path, so a uniform layout error cancels. A/B: with the CDNA4 seed deliberately wrong the error is 1.41-2.29 against 0.002 correct, so the test detects a layout defect by ~150x and the 1e-2 tolerance sits ~5x above the fp16 noise floor. Co-Authored-By: Claude --- .../rocm/attention/permuted_smem.cuh | 67 +++++----- include/flashinfer/rocm/attention/prefill.cuh | 16 +-- include/flashinfer/rocm/mma.h | 28 +++++ tests/rocm/test_batch_prefill_kernels.py | 116 ++++++++++++++++++ 4 files changed, 188 insertions(+), 39 deletions(-) diff --git a/include/flashinfer/rocm/attention/permuted_smem.cuh b/include/flashinfer/rocm/attention/permuted_smem.cuh index 4166a0fb7aa..b44802bd95f 100644 --- a/include/flashinfer/rocm/attention/permuted_smem.cuh +++ b/include/flashinfer/rocm/attention/permuted_smem.cuh @@ -204,42 +204,45 @@ struct smem_t { } /*! - * \brief Loads a fragment from shared memory and performs an in-register transpose across a quad. - * \details This function is designed to prepare the B-matrix operand for a CDNA3 MFMA - * instruction. - * It performs two actions in sequence for a quad of 4 threads: - * 1. Each thread loads a row-oriented fragment (e.g., 4 `half` values) from shared - * memory. - * 2. It then calls `transpose_intra_quad_fragments` to perform an in-register transpose - * of this data among the 4 threads. - * - * The result is that each thread's registers are populated with a column-oriented - * fragment, which is the required layout for the B-operand in a - * row-major(A) x col-major(B) MFMA. - * - * Visual Representation: - * If `[a,b,c,d]` are the 4 `half` values loaded by Thread 0: - * - * Data in Shared Memory (conceptually): - * Row 0: [a, b, c, d] - * Row 1: [e, f, g, h] - * Row 2: [i, j, k, l] - * Row 3: [m, n, o, p] - * - * After this function, registers hold: - * Thread 0: [a, e, i, m] (Column 0) - * Thread 1: [b, f, j, n] (Column 1) - * Thread 2: [c, g, k, o] (Column 2) - * Thread 3: [d, h, l, p] (Column 3) + * \brief Row within the 16x16 tile that `lane` must address for + * `load_matrix_m16n16_trans`. NOT the `load_fragment` layout. + */ + static __device__ __forceinline__ uint32_t trans_frag_row(uint32_t lane) { + if constexpr (mma::kHasTransposingLdsRead) { + return 4 * (lane / 16) + (lane % 16) / 4; + } else { + return lane % 16; + } + } + + /*! \brief Column, in `BasePtrTy` units, matching `trans_frag_row`. */ + static __device__ __forceinline__ uint32_t trans_frag_col(uint32_t lane) { + if constexpr (mma::kHasTransposingLdsRead) { + return lane % 4; + } else { + return lane / 16; + } + } + + /*! + * \brief Loads a 16x16 tile transposed, giving the column-oriented B operand an + * MFMA needs from row-oriented shared memory. * - * \tparam T The type of the register fragment (e.g., uint32_t). - * \param offset The starting offset in shared memory for the quad to begin loading. - * \param frag A pointer to the thread's local registers to store the resulting column fragment. + * \note Seed `offset` with `trans_frag_row`/`trans_frag_col`, not the + * `load_fragment` layout: CDNA4 does this in one `ds_read_b64_tr_b16` + * and CDNA3 in six `ds_bpermute_b32`, and the two want different + * per-lane addresses. Seeding the obvious way is silently wrong on CDNA4. */ template __device__ __forceinline__ void load_matrix_m16n16_trans(uint32_t offset, T* frag) { - load_fragment(offset, frag); - mma::transpose_mma_tile(frag); + static_assert(sizeof(T) == 4, "Only 32-bit fragment loading supported"); + if constexpr (mma::kHasTransposingLdsRead) { + static_assert(sizeof(BasePtrTy) == 8, "The transposing read consumes one 64-bit slot"); + mma::load_transposed_fragment(reinterpret_cast(frag), base + offset); + } else { + load_fragment(offset, frag); + mma::transpose_mma_tile(frag); + } } template diff --git a/include/flashinfer/rocm/attention/prefill.cuh b/include/flashinfer/rocm/attention/prefill.cuh index 54d96dceaac..3d25b369afc 100644 --- a/include/flashinfer/rocm/attention/prefill.cuh +++ b/include/flashinfer/rocm/attention/prefill.cuh @@ -1016,10 +1016,9 @@ __device__ __forceinline__ void compute_sfm_v( #pragma unroll for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { - // v_col_idx: current column j of *v_smem_offset_r before each advance_offset_by_column. - // Reset per KV row: each row's V fragment starts at column tid.x / WARP_THREAD_COLS. - // Needed by k128B_16Row; ignored by k128B. - uint32_t v_col_idx = tid.x / KTraits::WARP_THREAD_COLS; + // Current column j of *v_smem_offset_r before each advance_offset_by_column, + // reset per KV row. Needed by k128B_16Row; ignored by k128B. + uint32_t v_col_idx = v_smem->trans_frag_col(tid.x); #pragma unroll for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { uint32_t b_frag[INT32_ELEMS_PER_THREAD]; @@ -1520,7 +1519,8 @@ __device__ __forceinline__ void SinglePrefillWithKVCacheDevice( uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16)); uint32_t v_smem_offset_r = v_smem.template get_permuted_offset( - get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16); + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + v_smem.trans_frag_row(lane_idx), + v_smem.trans_frag_col(lane_idx)); uint32_t k_smem_offset_w = k_smem.template get_permuted_offset( warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, lane_idx % KV_THR_LAYOUT_COL), @@ -1932,7 +1932,8 @@ __global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchPrefillWithRaggedKV get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16)); uint32_t v_smem_offset_r = v_smem.template get_permuted_offset( - get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16); + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + v_smem.trans_frag_row(lane_idx), + v_smem.trans_frag_col(lane_idx)); uint32_t k_smem_offset_w = k_smem.template get_permuted_offset( warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, @@ -2200,7 +2201,8 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16)); uint32_t v_smem_offset_r = v_smem.template get_permuted_offset( - get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16); + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + v_smem.trans_frag_row(lane_idx), + v_smem.trans_frag_col(lane_idx)); uint32_t k_smem_offset_w = k_smem.template get_permuted_offset( warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, diff --git a/include/flashinfer/rocm/mma.h b/include/flashinfer/rocm/mma.h index 1d40673ee9b..cf8b6ed3ad9 100644 --- a/include/flashinfer/rocm/mma.h +++ b/include/flashinfer/rocm/mma.h @@ -141,6 +141,34 @@ __device__ __forceinline__ void load_fragment(uint32_t* R, const T* smem_ptr) { R[1] = reinterpret_cast(smem_ptr)[1]; } +// The one place the CDNA4 transposing LDS read is switched on. Everything that +// has to agree with it -- the per-lane seed and the load itself -- keys off this +// constant rather than repeating the guard. +#if defined(__HIP_DEVICE_COMPILE__) && defined(__gfx950__) +inline constexpr bool kHasTransposingLdsRead = true; +#else +inline constexpr bool kHasTransposingLdsRead = false; +#endif + +/// @brief CDNA4 transposing LDS read: one `ds_read_b64_tr_b16` replacing the six +/// `ds_bpermute_b32` of `transpose_mma_tile`. +/// +/// Redistributes as out(16g+4a+b)[j] = in(16g+4j+a)[b], so it needs its own +/// per-lane address -- see `smem_t::trans_frag_row`/`trans_frag_col`. +/// `smem_ptr` must be LDS-resident and 16-bit; `KernelTraits::DTypeKVSmem` +/// guarantees the latter for every KV dtype (prefill.cuh static_asserts it). +__device__ __forceinline__ void load_transposed_fragment(uint32_t* R, const void* smem_ptr) { +#if defined(__HIP_DEVICE_COMPILE__) && defined(__gfx950__) + // __fp16, not this header's _Float16 f16x4: the builtin rejects that spelling. + using h4 = __fp16 __attribute__((ext_vector_type(4))); + const h4 v = + __builtin_amdgcn_ds_read_tr16_b64_v4f16((h4 __attribute__((address_space(3)))*)(smem_ptr)); + __builtin_memcpy(R, &v, sizeof(v)); +#else + __builtin_trap(); // unreachable: callers gate on kHasTransposingLdsRead +#endif +} + // MMA operation for FP16 inputs with FP32 accumulator template __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, diff --git a/tests/rocm/test_batch_prefill_kernels.py b/tests/rocm/test_batch_prefill_kernels.py index 7157da0afb2..f0f3b8bca34 100644 --- a/tests/rocm/test_batch_prefill_kernels.py +++ b/tests/rocm/test_batch_prefill_kernels.py @@ -4,6 +4,7 @@ import pytest import torch +from attention_reference import naive_attention from jit_utils import gen_prefill_attention_modules import flashinfer @@ -52,6 +53,12 @@ def warmup_jit(): ), verbose=False, ) + flashinfer.jit.build_jit_specs( + gen_prefill_attention_modules( + [torch.float16], [torch.float16], [64], [0], [False], [False], [False] + ), + verbose=False, + ) yield @@ -1962,3 +1969,112 @@ def plan(page_size): if non_native is not None: with pytest.raises(ValueError, match="logits_soft_cap"): plan(non_native) + + +@pytest.mark.parametrize("head_dim", [64, 128, 256]) +@pytest.mark.parametrize("num_qo_heads", [4, 32]) +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("mode", ["ragged", "paged"]) +def test_batch_prefill_matches_independent_reference( + mode, head_dim, num_qo_heads, causal +): + """Batch fa2 prefill against a torch reference, at kv_len % CTA_TILE_KV != 0. + + The independent batch oracle that already exists covers only bf16 + causal + + head_dim 128 (test_batch_prefill_bf16_custom_mask); every other batch test + compares against single_prefill_with_kv_cache("fa2"), which reads V through the + same path, so a uniform V-layout error cancels. This adds head_dim 64/256, + non-causal and fp16. + """ + torch.manual_seed(0) + device = "cuda:0" + batch_size, qo_len, kv_len, num_kv_heads, page_size = 3, 37, 97, 4, 16 + + q = torch.randn( + batch_size * qo_len, num_qo_heads, head_dim, device=device, dtype=torch.float16 + ) + k = torch.randn( + batch_size * kv_len, num_kv_heads, head_dim, device=device, dtype=torch.float16 + ) + v = torch.randn( + batch_size * kv_len, num_kv_heads, head_dim, device=device, dtype=torch.float16 + ) + q_indptr = ( + torch.arange(0, batch_size + 1, device=device, dtype=torch.int32) * qo_len + ) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + if mode == "ragged": + kv_indptr = ( + torch.arange(0, batch_size + 1, device=device, dtype=torch.int32) * kv_len + ) + wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, "NHD", backend="fa2" + ) + wrapper.plan( + q_indptr, kv_indptr, num_qo_heads, num_kv_heads, head_dim, causal=causal + ) + o = wrapper.run(q, k, v) + else: + pages_per_seq = (kv_len + page_size - 1) // page_size + padded = pages_per_seq * page_size + kv_data = torch.zeros( + batch_size * pages_per_seq, + 2, + page_size, + num_kv_heads, + head_dim, + device=device, + dtype=torch.float16, + ) + for src, slot in ((k, 0), (v, 1)): + buf = torch.zeros( + batch_size, + padded, + num_kv_heads, + head_dim, + device=device, + dtype=torch.float16, + ) + buf[:, :kv_len] = src.view(batch_size, kv_len, num_kv_heads, head_dim) + kv_data[:, slot] = buf.view( + batch_size * pages_per_seq, page_size, num_kv_heads, head_dim + ) + kv_indptr = ( + torch.arange(0, batch_size + 1, device=device, dtype=torch.int32) + * pages_per_seq + ) + kv_indices = torch.arange( + 0, batch_size * pages_per_seq, device=device, dtype=torch.int32 + ) + last_len = torch.full( + (batch_size,), + kv_len - (pages_per_seq - 1) * page_size, + device=device, + dtype=torch.int32, + ) + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, "NHD", backend="fa2" + ) + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + last_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + causal=causal, + ) + o = wrapper.run(q, kv_data) + + for i in range(batch_size): + qs = slice(i * qo_len, (i + 1) * qo_len) + ks = slice(i * kv_len, (i + 1) * kv_len) + # fp32 reference: naive_attention does not upcast, and an fp16 one would + # eat the tolerance budget. + o_ref, _ = naive_attention( + q[qs].float(), k[ks].float(), v[ks].float(), causal=causal + ) + torch.testing.assert_close(o[qs].float(), o_ref.float(), rtol=1e-2, atol=1e-2)