diff --git a/docs/beellama-args.md b/docs/beellama-args.md index 811ba6c57bfb..ddc1aef2b970 100644 --- a/docs/beellama-args.md +++ b/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 diff --git a/docs/beellama-features.md b/docs/beellama-features.md index 2f2fe7d076b0..496a62e95fa7 100644 --- a/docs/beellama-features.md +++ b/docs/beellama-features.md @@ -190,7 +190,8 @@ correctness, memory behavior, or performance on that GPU. | HIP architecture | Physical wave | Native KVarN route | |---|---:|---| -| RDNA3, RDNA3.5, RDNA4 | 32 | WMMA generic/prefill and occupancy-selected split decode | +| RDNA3, RDNA3.5 | 32 | WMMA generic/prefill (D256 on fp32-accumulator tiles, qualified on gfx1100) and occupancy-selected split decode | +| RDNA4 | 32 | WMMA generic/prefill up to D128; D256+ stays on portable direct-record attention until its fp32 tiles qualify | | CDNA1-CDNA4 | 64 | MFMA generic/prefill and physical-wave split decode | | Older GCN, RDNA1, RDNA2 | device default | Portable direct-record attention | diff --git a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu index 2fe12cabee1f..53916d09f0da 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu +++ b/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) { diff --git a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh index b57d9b9abc66..c7e5d4a483d9 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh +++ b/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; @@ -828,8 +828,7 @@ static inline bool ggml_cuda_fattn_kvarn_portable_supported( body_meta->ne[1] == q->ne[2] && body_meta->ne[2] == q->ne[1] && body_meta->ne[3] == q->ne[3] && ggml_is_contiguous(body_meta)); const bool domain_ok = ggml_cuda_fattn_kvarn_rotated_decode_domain(dst) || - (q->ne[0] == 64 && - ggml_cuda_fattn_kvarn_domain(dst) == GGML_FLASH_ATTN_EXT_KVARN_DOMAIN_ROTATED_K_ORIGINAL_V); + ggml_cuda_fattn_kvarn_domain(dst) == GGML_FLASH_ATTN_EXT_KVARN_DOMAIN_ROTATED_K_ORIGINAL_V; return domain_ok && (q->ne[0] == 64 || q->ne[0] == 128 || q->ne[0] == 256 || q->ne[0] == 512) && q->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h b/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h index df5bd63e0644..58cbcb6eedf5 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h +++ b/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; diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index e19e81e3ec7b..036b93dc7229 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -164,7 +164,7 @@ 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); @@ -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) { - 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 { @@ -1098,6 +1102,20 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } +// KVarN-only tile selector. Forwards to mma_tile_sizes except RDNA3 D128/D256, +// which use the fp32-accumulator specializations below. Lets the dense path +// keep its qualified tiles while KVarN uses the qualified fp32 ones. The +// member lookups are dependent and resolve after the arch regions below. +template struct mma_tile_sizes; +template struct mma_tile_sizes_kvarn { + using T_A_KQ = typename mma_tile_sizes::T_A_KQ; + using T_B_KQ = typename mma_tile_sizes::T_B_KQ; + using T_C_KQ = typename mma_tile_sizes::T_C_KQ; + using T_A_VKQ = typename mma_tile_sizes::T_A_VKQ; + using T_B_VKQ = typename mma_tile_sizes::T_B_VKQ; + using T_C_VKQ = typename mma_tile_sizes::T_C_VKQ; +}; + #if defined(TURING_MMA_AVAILABLE) template struct mma_tile_sizes { using T_A_KQ = tile<16, 8, half2>; // row-major @@ -1141,6 +1159,27 @@ template 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 }; +// KVarN-only fp32-accumulator tiles (stew675 f32-VKQ guidance): DV=128/256 +// with fp16 PV accumulator show ~3e-4/tile error compounding over 64 layers +// on gfx1100. Selected explicitly for the KVarN path via mma_tile_sizes_kvarn +// (forwarding primary declared above the region chain); dense keeps the +// primary half2 tile. +template struct mma_tile_sizes_kvarn<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 struct mma_tile_sizes_kvarn<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 struct mma_tile_sizes { using T_A_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR>; // row-major @@ -1220,12 +1259,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int ncols = ncols1 * ncols2; - using T_A_KQ = typename mma_tile_sizes::T_A_KQ; - using T_B_KQ = typename mma_tile_sizes::T_B_KQ; - using T_C_KQ = typename mma_tile_sizes::T_C_KQ; - using T_A_VKQ = typename mma_tile_sizes::T_A_VKQ; - using T_B_VKQ = typename mma_tile_sizes::T_B_VKQ; - using T_C_VKQ = typename mma_tile_sizes::T_C_VKQ; + constexpr bool is_kvarn_kv = ggml_cuda_fattn_kvarn_template_type(type_K) || ggml_cuda_fattn_kvarn_template_type(type_V); + using tile_sizes_sel = typename std::conditional, mma_tile_sizes>::type; + using T_A_KQ = typename tile_sizes_sel::T_A_KQ; + using T_B_KQ = typename tile_sizes_sel::T_B_KQ; + using T_C_KQ = typename tile_sizes_sel::T_C_KQ; + using T_A_VKQ = typename tile_sizes_sel::T_A_VKQ; + using T_B_VKQ = typename tile_sizes_sel::T_B_VKQ; + using T_C_VKQ = typename tile_sizes_sel::T_C_VKQ; constexpr int cols_per_warp = T_B_KQ::I; constexpr int cols_per_thread = get_cols_per_thread(); @@ -1235,7 +1277,6 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr bool is_kvarn_kv = ggml_cuda_fattn_kvarn_template_type(type_K) || ggml_cuda_fattn_kvarn_template_type(type_V); constexpr int nstages = is_kvarn_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, use_sparse); static_assert(!is_kvarn_kv || !use_sparse, "sparse KVarN record loads are not qualified"); @@ -1279,7 +1320,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 + ? 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 +1543,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) { - 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 +1613,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 +1664,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 +1745,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 +1892,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(); } } @@ -1915,7 +1983,7 @@ 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) { + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) { NO_DEVICE_CODE; return; } @@ -1998,6 +2066,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; diff --git a/ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh b/ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh index 54ed9e915393..680a4f36615e 100644 --- a/ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh @@ -59,6 +59,103 @@ static __device__ __forceinline__ float ggml_cuda_fattn_kvarn_load_stage_rotated return __half2float(desc.stage[base + dim]); } +// Block-shared resolution of a token to its storage location. All fields +// depend only on (desc, token), so one thread can resolve per token and +// broadcast to the block instead of all 128 threads repeating the index +// math (64-bit div/mod, branches). +struct ggml_cuda_fattn_kvarn_resolved_token { + bool from_stage; + bool from_record; + int pos; + int stage_pos; + const uint8_t * record; + const half * scale_axis; + const half * zp_axis; + const half * other_axis; +}; + +static __device__ __forceinline__ ggml_cuda_fattn_kvarn_resolved_token +ggml_cuda_fattn_kvarn_resolve_token( + const ggml_cuda_fattn_kvarn_desc & desc, + const int token) { + ggml_cuda_fattn_kvarn_resolved_token out = {}; + int group; + int pos; + bool from_stage; + bool from_record; + int stage_pos; + int record_group; + + if (desc.swa || desc.read_indirect) { + const int64_t encoded = desc.indices[token]; + if (encoded == -1) { + return out; + } + bool explicitly_staged; + int assigned_slot = -1; + const int64_t abs_pos = ggml_cuda_fattn_kvarn_read_cell( + desc, encoded, explicitly_staged, &assigned_slot); + group = (int) (abs_pos / GGML_CUDA_FATTN_KVARN_DIM); + pos = (int) (abs_pos - (int64_t) group * GGML_CUDA_FATTN_KVARN_DIM); + from_stage = explicitly_staged || + (!(desc.read_indirect && !desc.swa) && ggml_cuda_fattn_kvarn_group_from_stage(desc, group)); + from_record = !explicitly_staged && (desc.read_indirect && !desc.swa ? true : + ggml_cuda_fattn_kvarn_group_from_record(desc, group)); + stage_pos = ggml_cuda_fattn_kvarn_stage_pos( + desc, group, pos, assigned_slot); + record_group = desc.swa ? group % desc.groups_per_stream : + desc.stream * desc.groups_per_stream + group; + } else { + group = token / GGML_CUDA_FATTN_KVARN_DIM; + pos = token - group * GGML_CUDA_FATTN_KVARN_DIM; + from_stage = ggml_cuda_fattn_kvarn_group_from_stage(desc, group); + from_record = ggml_cuda_fattn_kvarn_group_from_record(desc, group); + const int stage_base = desc.stream * GGML_CUDA_FATTN_KVARN_DIM * desc.stage_groups; + stage_pos = stage_base + (group == 0 ? pos : + GGML_CUDA_FATTN_KVARN_DIM + ((group - 1) % desc.tail_groups) * GGML_CUDA_FATTN_KVARN_DIM + pos); + record_group = desc.stream * desc.groups_per_stream + group; + } + + out.pos = pos; + out.from_stage = from_stage; + out.from_record = from_record; + out.stage_pos = stage_pos; + if (from_record) { + // NOTE: record_head (slice) is applied by the caller. + out.record = desc.records + (int64_t) record_group * desc.n_record_heads * desc.record_bytes; + const int payload_bytes = GGML_CUDA_FATTN_KVARN_DIM * GGML_CUDA_FATTN_KVARN_DIM * desc.bits / 8; + out.scale_axis = (const half *) (out.record + payload_bytes); + out.zp_axis = out.scale_axis + GGML_CUDA_FATTN_KVARN_DIM; + out.other_axis = out.zp_axis + GGML_CUDA_FATTN_KVARN_DIM; + } + return out; +} + +static __device__ __forceinline__ float ggml_cuda_fattn_kvarn_load_resolved( + const ggml_cuda_fattn_kvarn_desc & desc, + const ggml_cuda_fattn_kvarn_resolved_token & rt, + const int slice, + const int dim) { + const int record_head = desc.head_base + slice; + if (rt.from_stage) { + return ggml_cuda_fattn_kvarn_load_stage_rotated(desc, rt.stage_pos, record_head, dim); + } + if (!rt.from_record) { + return 0.0f; + } + const uint8_t * record = rt.record + (int64_t) record_head * desc.record_bytes; + const int payload_bytes = GGML_CUDA_FATTN_KVARN_DIM * GGML_CUDA_FATTN_KVARN_DIM * desc.bits / 8; + const half * scale_axis = (const half *) (record + payload_bytes); + const half * zp_axis = scale_axis + GGML_CUDA_FATTN_KVARN_DIM; + const half * other_axis = zp_axis + GGML_CUDA_FATTN_KVARN_DIM; + const int row = desc.value ? rt.pos : dim; + const int col = desc.value ? dim : rt.pos; + const uint8_t q = ggml_cuda_fattn_kvarn_unpack_record( + record, row * GGML_CUDA_FATTN_KVARN_DIM + col, desc.bits); + return (float(q) * __half2float(scale_axis[row]) + __half2float(zp_axis[row])) * + __half2float(other_axis[col]); +} + static __device__ __forceinline__ float ggml_cuda_fattn_kvarn_load_rotated( const ggml_cuda_fattn_kvarn_desc & desc, const int token, diff --git a/tests/test-cuda-fattn-route-policy.cpp b/tests/test-cuda-fattn-route-policy.cpp index c2109939733b..11d53465c09f 100644 --- a/tests/test-cuda-fattn-route-policy.cpp +++ b/tests/test-cuda-fattn-route-policy.cpp @@ -217,10 +217,15 @@ int main(int argc, char ** argv) { GGML_CUDA_FATTN_KVARN_MMA_ELIGIBLE, "RDNA WMMA must reject ncols2=1 and admit the same-width ncols2=2 tile"); ok &= expect(mma_eligibility(GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA, 256, 8, 2) == - GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED && - mma_eligibility(GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA, 512, 8, 2) == - GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED, - "RDNA WMMA must reject D256 and D512 before template launch"); + GGML_CUDA_FATTN_KVARN_MMA_ELIGIBLE && + mma_eligibility(GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA, 512, 8, 2) == + GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED, + "RDNA WMMA must admit D256 and reject D512 before template launch"); + ok &= expect(mma_eligibility(GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA, 128, 8, 2) == + GGML_CUDA_FATTN_KVARN_MMA_ELIGIBLE && + mma_eligibility(GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA, 256, 8, 2) == + GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED, + "RDNA4 WMMA must stay fail-closed at D128 until its fp32 tiles qualify"); for (int head_dim : {128, 256}) { ok &= expect(mma_eligibility(GGML_CUDA_FATTN_KVARN_AMD_CDNA_MFMA, head_dim, 5, 3) == GGML_CUDA_FATTN_KVARN_MMA_TILE_TOO_SMALL && diff --git a/tests/test-kvarn.cpp b/tests/test-kvarn.cpp index 1b9c71e69743..a29cf1b8711c 100644 --- a/tests/test-kvarn.cpp +++ b/tests/test-kvarn.cpp @@ -2378,6 +2378,13 @@ static std::vector test_native_flash_attention_output( std::vector k_data((size_t) record_dim * record_heads * n_kv * n_stream); std::vector v_data(k_data.size()); + // GGML_KVARN_TEST_LADDER_OUTLIER=S replicates LLM outlier channels + // (Qwen K/V have a few channels at 10-100x the typical magnitude) to + // test half-narrowing sensitivity of WMMA tile loaders. + float ladder_outlier = 0.0f; + if (const char * outlier_env = std::getenv("GGML_KVARN_TEST_LADDER_OUTLIER")) { + ladder_outlier = strtof(outlier_env, nullptr); + } for (int t = 0; t < n_kv; ++t) { for (int h = 0; h < n_kv_heads; ++h) { for (int slice = 0; slice < slices; ++slice) { @@ -2391,6 +2398,10 @@ static std::vector test_native_flash_attention_output( v_data[off] = 0.75f * std::cos(float(full_d) * 0.013f - float(t) * 0.019f) + 0.08f * std::sin(float(t) * 0.015f + float(h) * 0.23f); + if (ladder_outlier != 0.0f && (full_d % 64) == 0) { + k_data[off] *= ladder_outlier; + v_data[off] *= ladder_outlier; + } } } } @@ -3766,8 +3777,7 @@ static void test_native_flash_attention_gpu() { require(stats.decode_split == 0 && stats.amd_decode_split == 0 && stats.decode_vector == 0 && stats.amd_decode_vector == 0, "AMD route-boundary case entered a CUDA-only specialized decode route"); - const bool known_invalid_generic = hip_physical_wave_size == 32 ? - head_dim > 128 : head_dim > 256; + const bool known_invalid_generic = head_dim > 256; if (known_invalid_generic) { require(stats.generic_shape_rejected > 0 && stats.portable_native > 0 && stats.generic_mma == 0 && stats.prompt_prefill == 0, @@ -3790,7 +3800,7 @@ static void test_native_flash_attention_gpu() { require_amd_case(128, 17, gqa, 0, GGML_TYPE_F16, "AMD D128 GQA route-boundary output differs from the materialized oracle"); } - for (int head_dim : { 256, 512 }) { + for (int head_dim : { 256, 512 }) { for (int n_q : { 17, 256 }) { for (ggml_type exact_type : { GGML_TYPE_F16, GGML_TYPE_BF16 }) { require_amd_case(head_dim, n_q, 6, 128, exact_type, @@ -4309,6 +4319,169 @@ static void test_native_flash_attention_prefill_route_parity() { ggml_backend_free(gpu_backend); } +static void test_kvarn_nkv_ladder() { + ggml_backend_t gpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_GPU, false); + if (gpu_backend == nullptr) { + return; + } + ggml_backend_t cpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_CPU, true); + // Production proxy: D256, k6/v6, GQA 6 (24q/4kv Qwen), nq=256 prompt tile, + // production query layout + eager records (op_params[9]=1 in serving). + // Head-dim ladder decides which dims need the portable prompt route. + // GGML_KVARN_TEST_LADDER_NQ overrides the prompt tile (32/128/256) to + // isolate ncols-dependent prefill paths. + int ladder_nq = 256; + if (const char * nq_env = std::getenv("GGML_KVARN_TEST_LADDER_NQ")) { + ladder_nq = std::atoi(nq_env); + } + int ladder_q_heads = 6, ladder_kv_heads = 1; + if (const char * heads_env = std::getenv("GGML_KVARN_TEST_LADDER_HEADS")) { + if (std::sscanf(heads_env, "%d,%d", &ladder_q_heads, &ladder_kv_heads) != 2) { + ladder_q_heads = 6; + ladder_kv_heads = 1; + } + } + int ladder_tail = 0; + if (const char * tail_env = std::getenv("GGML_KVARN_TEST_LADDER_TAIL")) { + ladder_tail = std::atoi(tail_env); + } + auto [route_reset, route_get] = get_kvarn_route_stats_fns(gpu_backend); + for (int head_dim : { 128, 256, 512 }) { + for (int n_kv : { 256, 512, 1024, 2048, 4096, 8192 }) { + const std::vector expected = test_native_flash_attention_output( + cpu_backend, false, false, head_dim, 6, 6, ladder_nq, + ladder_q_heads, ladder_kv_heads, n_kv, 2, false, nullptr, false, ladder_tail, ladder_tail > 0, + GGML_TYPE_F16, 0, false, true, -1, true); + if (route_reset != nullptr) { + route_reset(); + } + const std::vector actual = test_native_flash_attention_output( + gpu_backend, true, true, head_dim, 6, 6, ladder_nq, + ladder_q_heads, ladder_kv_heads, n_kv, 2, false, nullptr, false, ladder_tail, ladder_tail > 0, + GGML_TYPE_F16, 0, false, true, -1, true); + if (route_get != nullptr && head_dim == 256 && n_kv == 512) { + test_kvarn_route_stats stats = make_test_kvarn_route_stats(); + route_get(&stats); + std::printf("kvarn-ladder-routes: generic_mma=%llu prompt_prefill=%llu portable_native=%llu amd_generic_mma=%llu materialize=%llu vec=%llu\n", + (unsigned long long) stats.generic_mma, (unsigned long long) stats.prompt_prefill, + (unsigned long long) stats.portable_native, (unsigned long long) stats.amd_generic_mma, + (unsigned long long) stats.materialize_fallback, (unsigned long long) stats.decode_vector); + std::fflush(stdout); + } + double sum = 0.0; + double mx = 0.0; + for (size_t i = 0; i < actual.size(); ++i) { + const double d = double(actual[i]) - double(expected[i]); + sum += d * d; + mx = std::max(mx, std::fabs(d)); + } + std::printf("kvarn-ladder: D=%d n_kv=%d rmse=%g maxabs=%g n=%zu\n", + head_dim, n_kv, std::sqrt(sum / actual.size()), mx, actual.size()); + std::fflush(stdout); + } + } + ggml_backend_free(cpu_backend); + ggml_backend_free(gpu_backend); +} + +// Committed regression coverage for the ub>64 whole-tile body_meta fix: +// D256 k6/v6 prompt prefill (nq=256, whole-tile K blocks) with an attached +// 128-candidate exact tail, GPU native vs CPU materialized reference. With +// the `!is_kvarn_kv` gate restored on the whole-tile dst_final_meta stores, +// tail-merge rows keep zero meta and this diverges catastrophically. +// Committed regression coverage for the ub>64 whole-tile body_meta fix: +// D256 k6/v6 prompt prefill (nq=256, whole-tile K blocks) pins the WMMA +// prompt route and asserts (a) every published body denominator is positive +// and finite, and (b) the attached-exact-tail pass matches the CPU reference. +// With the `!is_kvarn_kv` gate restored on the whole-tile dst_final_meta +// stores, tail-merge rows keep zero meta: (a) fails deterministically. +static void test_kvarn_d256_prompt_tail_regression() { + ggml_backend_t gpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_GPU, false); + if (gpu_backend == nullptr) { + return; + } + ggml_backend_t cpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_CPU, true); + auto [route_reset, route_get] = get_kvarn_route_stats_fns(gpu_backend); + if (route_reset != nullptr) { + route_reset(); + } + std::vector body_meta; + const std::vector actual = test_native_flash_attention_output( + gpu_backend, true, true, 256, 6, 6, 256, + 24, 4, 512, 2, false, &body_meta, false, 0, true, + GGML_TYPE_F16, 0, false, true, -1, true); + if (route_get != nullptr) { + test_kvarn_route_stats stats = make_test_kvarn_route_stats(); + route_get(&stats); + std::printf("kvarn-tail-regression-routes: generic_mma=%llu prompt_prefill=%llu portable_native=%llu amd_generic_mma=%llu materialize=%llu\n", + (unsigned long long) stats.generic_mma, (unsigned long long) stats.prompt_prefill, + (unsigned long long) stats.portable_native, (unsigned long long) stats.amd_generic_mma, + (unsigned long long) stats.materialize_fallback); + std::fflush(stdout); + require(stats.prompt_prefill > 0, + "D256 tail regression did not execute the WMMA prompt-prefill route it guards"); + } + require(body_meta.size() % 2 == 0 && !body_meta.empty(), + "D256 tail regression did not publish body softmax metadata"); + for (size_t i = 0; i < body_meta.size(); i += 2) { + require(std::isfinite(body_meta[i + 1]) && body_meta[i + 1] > 0.0f, + "D256 WMMA whole-tile body row kept zero denominator"); + } + const std::vector expected = test_native_flash_attention_output( + cpu_backend, false, false, 256, 6, 6, 256, + 24, 4, 512, 2, false, nullptr, false, 128, true, + GGML_TYPE_F16, 0, false, true, -1, true); + const std::vector tailed = test_native_flash_attention_output( + gpu_backend, true, true, 256, 6, 6, 256, + 24, 4, 512, 2, false, nullptr, false, 128, true, + GGML_TYPE_F16, 0, false, true, -1, true); + require_close_f32_rmse(tailed, expected, 1e-2f, + "D256 WMMA prompt prefill with attached exact tail differs from CPU reference"); + std::printf("test-kvarn: D256 prompt-tail regression OK\n"); + std::fflush(stdout); + ggml_backend_free(cpu_backend); + ggml_backend_free(gpu_backend); +} + +// Committed coverage for the portable shared-scratch sizing fix: portable +// attention with original_value_domain=true at D256/D512 exercises the +// full-head V-domain transform (reduction/transform indexed to D-1), which +// silently ran out of bounds when the arrays were sized RECORD_DIM. HIP +// stays rotated by policy, so the force-portable env pins the route here; +// the CPU materialized reference is route-independent. +static void test_native_flash_attention_portable_original_v() { + ggml_backend_t gpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_GPU, false); + if (gpu_backend == nullptr) { + return; + } + ggml_backend_t cpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_CPU, true); + scoped_test_env force_portable("GGML_KVARN_TEST_FORCE_PORTABLE_FATTN", "1"); + for (int head_dim : { 256, 512 }) { + const std::vector expected = test_native_flash_attention_output( + cpu_backend, false, false, head_dim, 6, 6, 32, + 6, 1, 512, 2, false, nullptr, false, 0, true, + GGML_TYPE_F16, 0, false, true, -1, true); + const std::vector actual = test_native_flash_attention_output( + gpu_backend, true, true, head_dim, 6, 6, 32, + 6, 1, 512, 2, false, nullptr, false, 0, true, + GGML_TYPE_F16, 0, false, true, -1, true); + double sum = 0.0; + for (size_t i = 0; i < actual.size(); ++i) { + const double d = double(actual[i]) - double(expected[i]); + sum += d * d; + } + std::printf("test-kvarn: portable original-V D%d rmse=%g n=%zu\n", + head_dim, std::sqrt(sum / actual.size()), actual.size()); + std::fflush(stdout); + require_close_f32_rmse(actual, expected, 1e-2f, + "portable original-V KVarN attention differs from CPU reference"); + std::printf("test-kvarn: portable original-V D%d parity OK\n", head_dim); + std::fflush(stdout); + } + ggml_backend_free(cpu_backend); + ggml_backend_free(gpu_backend); +} + static void test_store_paths_gpu() { ggml_backend_t gpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_GPU, false); if (gpu_backend == nullptr) { @@ -5412,6 +5585,12 @@ int main() { return 0; } + if (std::getenv("GGML_KVARN_TEST_NKV_LADDER_ONLY") != nullptr) { + test_kvarn_nkv_ladder(); + std::printf("test-kvarn: nkv ladder OK\n"); + return 0; + } + if (std::getenv("GGML_KVARN_TEST_DFLASH_NONCAUSAL_ONLY") != nullptr) { test_dflash_non_causal_attention_parity(); std::printf("test-kvarn: DFlash non-causal attention parity OK\n"); @@ -5541,6 +5720,11 @@ int main() { test_cache_ops_swa(GGML_BACKEND_DEVICE_TYPE_GPU, false, 1); // CUDA SWA ring parity test_cache_ops_swa(GGML_BACKEND_DEVICE_TYPE_CPU, true, 2); test_cache_ops_swa(GGML_BACKEND_DEVICE_TYPE_GPU, false, 2); // multi-slot SWA ring parity + // Placed before the store-route gauntlet below: the head-wide store + // assertion aborts on some HIP devices (pre-existing), and these two + // regression cases must execute on every backend regardless. + test_kvarn_d256_prompt_tail_regression(); + test_native_flash_attention_portable_original_v(); test_store_paths_gpu(); test_native_flash_attention_support_gates(); test_native_flash_attention_cpu();