Replace the per-MFMA V transpose with CDNA4's transposing LDS read

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 <typename T = uint32_t>
   __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<uint32_t*>(frag), base + offset);
+    } else {
+      load_fragment(offset, frag);
+      mma::transpose_mma_tile(frag);
+    }
   }
 
   template <typename T = uint32_t>

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<UPCAST_STRIDE_K>(
       get_warp_idx_kv<KTraits>(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16));
   uint32_t v_smem_offset_r = v_smem.template get_permuted_offset<UPCAST_STRIDE_V>(
-      get_warp_idx_kv<KTraits>(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16);
+      get_warp_idx_kv<KTraits>(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<UPCAST_STRIDE_K>(
                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<KTraits>(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16));
 
   uint32_t v_smem_offset_r = v_smem.template get_permuted_offset<UPCAST_STRIDE_V>(
-      get_warp_idx_kv<KTraits>(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16);
+      get_warp_idx_kv<KTraits>(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<UPCAST_STRIDE_K>(
                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<UPCAST_STRIDE_K>(
       get_warp_idx_kv<KTraits>(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16));
   uint32_t v_smem_offset_r = v_smem.template get_permuted_offset<UPCAST_STRIDE_V>(
-      get_warp_idx_kv<KTraits>(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16);
+      get_warp_idx_kv<KTraits>(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<UPCAST_STRIDE_K>(
                warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL,

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<const uint32_t*>(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 <typename T, mma::MMAMode mma_mode = mma::MMAMode::kInplaceUpdate>
 __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A,

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)