Make KVarN usable on ROCm RDNA3 (toward #122: D256 race, portable prefill, f32 WMMA prototype)

docs/beellama-args.md

@@ -51,6 +51,14 @@ peak transient scratch for concurrent long prompts but adds partial-softmax
 merges and changes floating-point reduction order. It does not alter context or
 persistent KV-cache capacity.
 
+On HIP/ROCm, KVarN prompt prefill defaults to the F32-accumulator WMMA route
+on arches whose tiles accumulate in fp32 (RDNA3/gfx11); RDNA4 stays on the
+portable route until its fp32 tiles qualify. `GGML_KVARN_AMD_PROMPT_PORTABLE`
+opts a prompt back into portable-native direct-record attention: any nonzero
+value (conventionally `1`) selects portable, while unset, `0`, or
+non-numeric values keep the WMMA default. The check runs before the generic
+probe, so opting in does not pay for a discarded WMMA pass.
+
 ## KV cache precision tail for quantized caches
 
 The KV cache precision tail (KVCPT) makes the newest attention-visible entries exact in F16 or BF16 for

ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu

@@ -880,6 +880,12 @@ static bool ggml_cuda_flash_attn_ext_kvarn_decode(
 }
 
 static ggml_cuda_fattn_kvarn_amd_mma_arch ggml_cuda_fattn_kvarn_amd_arch(int cc) {
+    if (GGML_CUDA_CC_IS_RDNA4(cc)) {
+        // RDNA4 compiles the half2 WMMA tiles only: the fp32-accumulator
+        // tiles that justify the raised D256 limit are RDNA3 (gfx11) builds.
+        // Keep RDNA4 fail-closed at D128 until its fp32 tiles are qualified.
+        return GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA;
+    }
     if (amd_wmma_available(cc)) {
         return GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA;
     }
@@ -1225,6 +1231,31 @@ bool ggml_cuda_flash_attn_ext_kvarn(
         ggml_cuda_fattn_kvarn_portable_supported(plan, dst);
     bool generic_shape_supported = false;
     bool wide_mma = false;
+#if defined(GGML_USE_HIP)
+    // RDNA3 (gfx11) WMMA prompt tiles accumulate in fp32 for DV=128/256
+    // (mirroring the proven DV=80/112 fp32-PV tiles), so on fp32-tile arches
+    // the WMMA path is both the fast and the exact route (~1e-5 ladder RMSE,
+    // 32k KLD at portable parity). It is therefore the default for HIP KVarN
+    // prompt-prefill. RDNA4 compiles the half2 tiles only and stays
+    // fail-closed on portable (see the RDNA4 eligibility gate). Decode
+    // (nq<=16) stays on WMMA as before.
+    // Checked BEFORE the generic probe below: the probe launches the WMMA
+    // kernel to test the shape, so diverting first avoids running prompt
+    // prefill twice and discarding the WMMA pass.
+    {
+        const char * prompt_portable = getenv("GGML_KVARN_AMD_PROMPT_PORTABLE");
+        if (prompt_prefill && portable_supported &&
+                (prompt_portable != nullptr && atoi(prompt_portable) != 0)) {
+            g_kvarn_route_portable_native.fetch_add(1, std::memory_order_relaxed);
+            // QB-batching was superseded by upstream's complete optimized D64
+            // rewrite (v0.4.7); the fallback uses the standard portable kernel.
+            ggml_cuda_fattn_kvarn_debug_route(
+                ctx.device, plan, dst, entry_path, "portable-native",
+                "hip-prompt-precision-optin");
+            return ggml_cuda_flash_attn_ext_kvarn_portable(ctx, dst, plan);
+        }
+    }
+#endif
     if (capabilities.generic_mma && Q->ne[0] != 64) {
         generic_shape_supported = ggml_cuda_flash_attn_ext_mma_kvarn(ctx, dst, wide_mma);
         if (!generic_shape_supported) {

ggml/src/ggml-cuda/fattn-kvarn-portable.cuh

@@ -182,8 +182,8 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel(
     const float * q = (const float *) (
         q_data + query * nbq1 + query_head * nbq2 + stream * nbq3);
 
-    __shared__ float reduction[RECORD_DIM];
-    __shared__ float transform[RECORD_DIM];
+    __shared__ float reduction[D];
+    __shared__ float transform[D];
     __shared__ float maximum;
     __shared__ float denominator;
     __shared__ float old_scale_shared;

ggml/src/ggml-cuda/fattn-kvarn-route-policy.h

@@ -30,6 +30,7 @@ enum ggml_cuda_fattn_kvarn_route {
 enum ggml_cuda_fattn_kvarn_amd_mma_arch {
     GGML_CUDA_FATTN_KVARN_AMD_NONE,
     GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA,
+    GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA,
     GGML_CUDA_FATTN_KVARN_AMD_CDNA_MFMA,
 };
 
@@ -61,14 +62,16 @@ inline ggml_cuda_fattn_kvarn_mma_eligibility ggml_cuda_fattn_kvarn_amd_mma_eligi
         return GGML_CUDA_FATTN_KVARN_MMA_INVALID_COLUMNS;
     }
     if (input.head_dim <= 0 ||
-            (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA && input.head_dim > 128) ||
+            (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA && input.head_dim > 256) ||
+            (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA && input.head_dim > 128) ||
             (input.arch == GGML_CUDA_FATTN_KVARN_AMD_CDNA_MFMA && input.head_dim > 256)) {
         return GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED;
     }
     if (input.ncols1 * input.ncols2 < 16) {
         return GGML_CUDA_FATTN_KVARN_MMA_TILE_TOO_SMALL;
     }
-    if (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA && input.ncols2 == 1) {
+    if ((input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA ||
+            input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA) && input.ncols2 == 1) {
         return GGML_CUDA_FATTN_KVARN_MMA_RDNA_SINGLE_GQA_COLUMN;
     }
     return GGML_CUDA_FATTN_KVARN_MMA_ELIGIBLE;

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -164,20 +164,20 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 2,  32, 128, 128, 128, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 2,  32, 128, 128, 128, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  64, 128, 128,  64, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  64, 128, 128,  64, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 2,  32, 128, 128,  32, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 128, 256, 1,  64, 128, 128,  64, 1, true);
 
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 160, 128, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 128, 2,  32, 160, 128, 128, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 256, 2,  64,  96,  16,  16, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 256, 2,  64,  96,  16,  16, 1, true);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8, 128, 3,  64,  96,  64, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 128, 3,  64,  96,  64, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2,  32, 128, 128, 128, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 128, 2,  64,  96,  16,  16, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 256, 2, 128,  96,  16,  16, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 128, 2,  32, 128, 128, 128, 1, true);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8, 128, 3,  64,  96,  64, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16, 128, 3,  64,  96,  64, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 128, 2,  32, 160, 128, 128, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16, 128, 2,  64,  96,  16,  16, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 256, 2, 128,  96,  64,  16, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 64, 128, 2,  32, 160, 128, 128, 1, true);
 
     return fattn_mma_config(32, 1, 0, 0, 0, 0, 0, false);
@@ -950,12 +950,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
         }
 #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
         if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
-            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
+            // Rescale in fp32 to avoid double-rounding the scale to half first.
+            const float scale_f32 = KQ_max_scale[0];
 #pragma unroll
             for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
 #pragma unroll
                 for (int l = 0; l < T_C_VKQ::ne; ++l) {
-                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
+                    float2 acc_f32 = __half22float2(VKQ_C[i].x[l]);
+                    acc_f32.x *= scale_f32;
+                    acc_f32.y *= scale_f32;
+                    VKQ_C[i].x[l] = make_half2(acc_f32.x, acc_f32.y);
                 }
             }
         } else {
@@ -1141,6 +1145,25 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
     using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
     using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
 };
+// Prototype (stew675 f32-VKQ guidance): DV=128/256 with fp16 PV accumulator show
+// ~3e-4/tile error compounding over 64 layers on gfx1100. Mirror the proven
+// DV=80/112 fp32-PV tiles here; generic path stays fp16 until qualified.
+template<int ncols> struct mma_tile_sizes<128, ncols> {
+    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+    using T_A_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+};
+template<int ncols> struct mma_tile_sizes<256, ncols> {
+    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+    using T_A_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+};
 #else
 template<int DV, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // row-major
@@ -1279,7 +1302,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
 #if defined(TURING_MMA_AVAILABLE)
     T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)];
 #elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
-    T_C_VKQ VKQ_C[DV % 32 != 0       ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J)];
+    // Entry count mirrors the rescale loops: half2 accumulators fold two
+    // stacked K-halves per entry via the opsel pair (DV/32 for DV%32==0),
+    // float accumulators keep one 16-row tile per entry (DV/16 always).
+    static constexpr int VKQ_C_COUNT = std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>
+        ? DV/T_C_VKQ::J
+        : (DV % 32 != 0 ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J));
+    T_C_VKQ VKQ_C[VKQ_C_COUNT];
 #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
     T_C_VKQ VKQ_C[                                     DV/(2*T_C_VKQ::J)];
 #else // Volta
@@ -1496,12 +1525,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
         }
 #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
         if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
-            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
+            // Rescale in fp32 to avoid double-rounding the scale to half first.
+            const float scale_f32 = KQ_max_scale[0];
 #pragma unroll
             for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
 #pragma unroll
                 for (int l = 0; l < T_C_VKQ::ne; ++l) {
-                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
+                    float2 acc_f32 = __half22float2(VKQ_C[i].x[l]);
+                    acc_f32.x *= scale_f32;
+                    acc_f32.y *= scale_f32;
+                    VKQ_C[i].x[l] = make_half2(acc_f32.x, acc_f32.y);
                 }
             }
         } else {
@@ -1562,7 +1595,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
                 dstk_fixup_meta[jc_cwm] = KQ_cmr;
             }
-            if (!is_kvarn_kv && !needs_fixup && !is_fixup && dst_final_meta && threadIdx.x < T_B_KQ::I) {
+            // KVarN whole-tile blocks must publish final (max, rowsum) too: the
+            // tail merge reads body_meta for every row, and the stream-k fixup
+            // skips tiles whose K range aligns exactly to tile boundaries, so
+            // without this store those rows keep zero meta and their (correct)
+            // body values are silently discarded by the merge.
+            if (!needs_fixup && !is_fixup && dst_final_meta && threadIdx.x < T_B_KQ::I) {
                 const int j = jc_cwm / ncols2;
                 const int c = jc_cwm % ncols2;
                 if (jt*ncols1 + j < int(ne01.z) && zt_gqa*ncols2 + c < gqa_ratio) {
@@ -1608,7 +1646,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
                 dstk_fixup_meta[jc_cwm] = KQ_cmr;
             }
-            if (!is_kvarn_kv && !needs_fixup && !is_fixup && dst_final_meta && thread_should_write) {
+            // KVarN whole-tile blocks must publish final (max, rowsum) too: the
+            // tail merge reads body_meta for every row, and the stream-k fixup
+            // skips tiles whose K range aligns exactly to tile boundaries, so
+            // without this store those rows keep zero meta and their (correct)
+            // body values are silently discarded by the merge.
+            if (!needs_fixup && !is_fixup && dst_final_meta && thread_should_write) {
                 const int j = jc_cwm / ncols2;
                 const int c = jc_cwm % ncols2;
                 if (jt*ncols1 + j < int(ne01.z) && zt_gqa*ncols2 + c < gqa_ratio) {
@@ -1684,7 +1727,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
             float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
             dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs);
         }
-        if (!is_kvarn_kv && !needs_fixup && !is_fixup && dst_final_meta &&
+        // KVarN whole-tile blocks must publish final (max, rowsum) too: the
+        // tail merge reads body_meta for every row, and the stream-k fixup
+        // skips tiles whose K range aligns exactly to tile boundaries, so
+        // without this store those rows keep zero meta and their (correct)
+        // body values are silently discarded by the merge.
+        if (!needs_fixup && !is_fixup && dst_final_meta &&
                 (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) {
             const int jc = (threadIdx.y/np)*cols_per_warp + threadIdx.x;
             if (jc < ncols) {
@@ -1826,7 +1874,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 }
             }
         }
-        if (np > 1) {
+        // The tile_Q buffer is reused for the next k00 iteration, so all warps must sync here
+        // before its data is overwritten. With np > 1 only some warps read back, but they all write.
+        if (np > 1 || k00 + nbatch_combine < DV/2) {
             __syncthreads();
         }
     }
@@ -1914,8 +1964,12 @@ static __global__ void flash_attn_ext_f16(
 
 #if defined(AMD_WMMA_AVAILABLE)
     // Mirrored by ggml_cuda_fattn_kvarn_amd_mma_eligibility on the host.
-    // Keep this final invariant for callers outside the KVarN dispatcher.
-    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) {
+    // RDNA WMMA D256 tiles are the validated configs in the RDNA table above
+    // (256, 320, 512, 576 keys). The KVarN dispatcher admits D256 only where
+    // the fp32 tiles compile (RDNA3/gfx11; RDNA4 stays fail-closed at D128),
+    // and standard FA keeps upstream's D128 cap, so this bound is reachable
+    // only through qualified shapes.
+    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 576) {
         NO_DEVICE_CODE;
         return;
     }
@@ -1998,6 +2052,11 @@ static __global__ void flash_attn_ext_f16(
                  ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
         }
 
+        // The next process_tile call reuses the tile_Q buffer for its Q/K tiles, so all warps must
+        // have finished reading the combined results before any of them starts the next call.
+        // (With np == 1 the end-of-k00 barrier does not fire, so this is required for correctness.)
+        __syncthreads();
+
         kbc += iter_k;
         kbc -= kbc % iter_k;