diff --git a/include/flashinfer/rocm/attention/permuted_smem.cuh b/include/flashinfer/rocm/attention/permuted_smem.cuh index 4166a0fb7a..b44802bd95 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 54d96dceaa..3d25b369af 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 1d40673ee9..cf8b6ed3ad 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 7157da0afb..f0f3b8bca3 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)