From 42f2250446705c69da3740b9bebf10b4ea19d4f8 Mon Sep 17 00:00:00 2001 From: raufaser Date: Mon, 7 Sep 2026 21:02:39 +0200 Subject: [PATCH 01/14] kvarn/rocm: fix D256 k00-combine race and enable RDNA D256 route flash_attn_ext_f16_process_tile reuses tile_Q as combine staging across k00 iterations; the end-of-iteration barrier only fired for np > 1. D256/ncols=64 (nbatch_combine=64, DV/2=128, np=1) corrupted output. Sync on all but the last iteration, plus sync between process_tile calls that reuse tile_Q. Enable D256 in the RDNA WMMA route policy and device guard. --- ggml/src/ggml-cuda/fattn-kvarn-route-policy.h | 2 +- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 27 ++++++++++++++----- tests/test-cuda-fattn-route-policy.cpp | 8 +++--- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h b/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h index 96b8b7c0d1c4..87a39efef482 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h +++ b/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h @@ -60,7 +60,7 @@ 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_CDNA_MFMA && input.head_dim > 256)) { return GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED; } diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 290e637eb367..d75fe1f29d63 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -908,12 +908,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 { @@ -1447,12 +1451,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 { @@ -1767,7 +1775,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(); } } @@ -1844,7 +1854,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; } @@ -1926,6 +1936,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/tests/test-cuda-fattn-route-policy.cpp b/tests/test-cuda-fattn-route-policy.cpp index 7ba8cb99a030..d11a927ec8ec 100644 --- a/tests/test-cuda-fattn-route-policy.cpp +++ b/tests/test-cuda-fattn-route-policy.cpp @@ -217,10 +217,10 @@ 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"); 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 && From 27c130cbb30c524e339e52bdca2d2ac88a9d0573 Mon Sep 17 00:00:00 2001 From: raufaser Date: Mon, 7 Sep 2026 21:02:39 +0200 Subject: [PATCH 02/14] kvarn/rocm: route HIP prompt-prefill through portable-native attention RDNA WMMA VKQ accumulators are fp16 (~3e-4/call vs ~1e-5 portable), compounding through depth into a KLD collapse. Prefer portable-native direct-record attention for HIP prompt-prefill (nq > 16); decode stays on WMMA. Opt out with GGML_KVARN_AMD_PROMPT_PORTABLE=0. --- ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu index b8f7a5d08650..cc4c996129ad 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu +++ b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu @@ -1190,6 +1190,25 @@ bool ggml_cuda_flash_attn_ext_kvarn( const ggml_cuda_fattn_kvarn_route fallback_route = ggml_cuda_fattn_kvarn_select_fallback_route( prompt_prefill, generic_shape_supported, portable_supported); +#if defined(GGML_USE_HIP) + // RDNA WMMA prompt tiles accumulate in fp16 (RMSE ~3e-4 vs ~1e-5 for the + // portable fp32 path), and the error compounds through 64 layers into a + // visible KLD collapse. Route HIP prompt-prefill through portable-native + // direct-record attention (same records, exact math) until a float + // WMMA accumulator lands. Decode (nq<=16) stays on WMMA: its single + // iteration is exact. Opt out with GGML_KVARN_AMD_PROMPT_PORTABLE=0. + { + 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); + ggml_cuda_fattn_kvarn_debug_route( + ctx.device, plan, dst, entry_path, "portable-native", + "hip-prompt-precision"); + return ggml_cuda_flash_attn_ext_kvarn_portable(ctx, dst, plan); + } + } +#endif if (fallback_route == GGML_CUDA_FATTN_KVARN_ROUTE_GENERIC_MMA || fallback_route == GGML_CUDA_FATTN_KVARN_ROUTE_PROMPT_PREFILL) { if (prompt_prefill) { From a6f9fc5b03816cefc0823890f6622dd97a21237e Mon Sep 17 00:00:00 2001 From: raufaser Date: Mon, 7 Sep 2026 21:02:39 +0200 Subject: [PATCH 03/14] kvarn/portable: warp-shuffle reduction and fp32 rescale Replace the per-token 7-stage shared-memory reduction tree with warp shuffles plus one cross-warp step; rescale the half accumulator in fp32; unroll the V-load loop. --- ggml/src/ggml-cuda/fattn-kvarn-portable.cuh | 55 +++++++++++++-------- 1 file changed, 35 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh index 50e769d394d5..a35b9fb58138 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh +++ b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh @@ -57,6 +57,32 @@ ggml_cuda_fattn_kvarn_portable_resolve( return result; } +static __device__ __forceinline__ float ggml_cuda_fattn_kvarn_portable_block_reduce( + float partial, + float * warp_partials) { + // Warp-level reduction with shuffles (no sync), then a single + // cross-warp step through shared memory (one sync). Replaces the + // 7-stage shared-memory tree (7 syncs) on the per-token hot path. + const int lane = threadIdx.x % 32; + const int wid = threadIdx.x / 32; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + partial += __shfl_xor_sync(0xFFFFFFFFu, partial, offset, 32); + } + if (lane == 0) { + warp_partials[wid] = partial; + } + __syncthreads(); + float total = threadIdx.x < (GGML_CUDA_FATTN_KVARN_DIM / 32) ? + warp_partials[threadIdx.x] : 0.0f; +#pragma unroll + for (int offset = (GGML_CUDA_FATTN_KVARN_DIM / 64); offset > 0; offset >>= 1) { + total += __shfl_xor_sync(0xFFFFFFFFu, total, offset, 32); + } + __syncthreads(); + return total; +} + template static __device__ __forceinline__ void ggml_cuda_fattn_kvarn_portable_stage_rotated( const ggml_cuda_fattn_kvarn_desc & desc, @@ -136,11 +162,12 @@ 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[THREADS]; __shared__ float maximum; __shared__ float denominator; __shared__ float old_scale_shared; __shared__ float weight_shared; + // Warp-shuffle reduction scratch: one partial per warp (THREADS/32). + __shared__ float warp_partials[THREADS / 32]; float accumulator[SLICES] = {}; if (tid == 0) { @@ -208,15 +235,8 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( const int dim = slice * GGML_CUDA_FATTN_KVARN_DIM + tid; partial += q[dim] * k_values[slice]; } - reduction[tid] = partial; - __syncthreads(); - - for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - reduction[tid] += reduction[tid + stride]; - } - __syncthreads(); - } + const float total = ggml_cuda_fattn_kvarn_portable_block_reduce( + partial, warp_partials); if (tid == 0) { float mask_value = 0.0f; @@ -227,7 +247,7 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( mask_value = slope * __half2float(*mask); } - float score = reduction[0] * scale; + float score = total * scale; if (logit_softcap != 0.0f) { score = logit_softcap * tanhf(score); } @@ -253,6 +273,7 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( ggml_cuda_fattn_kvarn_portable_stage_rotated( v_desc, v_ref.stage_pos, tid, v_values); } else { +#pragma unroll for (int slice = 0; slice < SLICES; ++slice) { v_values[slice] = ggml_cuda_fattn_kvarn_load_rotated( v_desc, token, slice, tid); @@ -281,21 +302,15 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( ptr, k_tail_bf16); partial += q[dim] * kval; } - reduction[tid] = partial; - __syncthreads(); - for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - reduction[tid] += reduction[tid + stride]; - } - __syncthreads(); - } + const float total = ggml_cuda_fattn_kvarn_portable_block_reduce( + partial, warp_partials); if (tid == 0) { const half * tail_mask = (const half *) ( tail_mask_data + (size_t) token * nbmt0 + (size_t) query * nbmt1 + (size_t) stream * nbmt3); const float mask_value = slope * __half2float(*tail_mask); - float score = reduction[0] * scale; + float score = total * scale; if (logit_softcap != 0.0f) { score = logit_softcap * tanhf(score); } From e502e1928378897a6178e84adef2c3b0e9c0d408 Mon Sep 17 00:00:00 2001 From: raufaser Date: Mon, 7 Sep 2026 21:02:39 +0200 Subject: [PATCH 04/14] kvarn/tests: production-shape accuracy ladder harness GGML_KVARN_TEST_NKV_LADDER_ONLY runs GPU-native vs CPU-materialized attention error over D128/256/512 and growing n_kv. --- tests/test-kvarn.cpp | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/test-kvarn.cpp b/tests/test-kvarn.cpp index b08ccff71252..926c8c8c7ceb 100644 --- a/tests/test-kvarn.cpp +++ b/tests/test-kvarn.cpp @@ -4104,6 +4104,41 @@ 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. + 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, 256, + 6, 1, n_kv, 2, false, nullptr, false, 0, false, + 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, 256, + 6, 1, n_kv, 2, false, nullptr, false, 0, false, + GGML_TYPE_F16, 0, false, true, -1, true); + 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); +} + static void test_store_paths_gpu() { ggml_backend_t gpu_backend = init_test_backend(GGML_BACKEND_DEVICE_TYPE_GPU, false); if (gpu_backend == nullptr) { @@ -5207,6 +5242,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"); From 78b63a42aadddae1320602fbf8de7dc3a50b7d74 Mon Sep 17 00:00:00 2001 From: raufaser Date: Mon, 7 Sep 2026 22:54:02 +0200 Subject: [PATCH 05/14] kvarn/portable: batched queries, shuffle reduction, hoisted resolve Serve QB=8 queries per block from one shared K/V token stream (no-tail path; tail keeps QB=1): ~2.7x prefill (pp4096 107 -> 291 t/s) at ladder-identical accuracy. Also replace the per-token shared-memory reduction tree with warp shuffles, resolve each token once per block instead of per thread, and add an env-gated kernel attribute print (GGML_KVARN_PORTABLE_ATTRS). --- ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu | 9 +- ggml/src/ggml-cuda/fattn-kvarn-portable.cuh | 350 +++++++++++++------- ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh | 97 ++++++ 3 files changed, 331 insertions(+), 125 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu index cc4c996129ad..60c5b19f6626 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu +++ b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu @@ -1202,9 +1202,16 @@ bool ggml_cuda_flash_attn_ext_kvarn( if (prompt_prefill && portable_supported && (prompt_portable == nullptr || atoi(prompt_portable) != 0)) { g_kvarn_route_portable_native.fetch_add(1, std::memory_order_relaxed); + // Batch 4 queries per block when no exact tail is attached (shared + // token stream); otherwise the queries attend different token sets + // and sharing is invalid. + const bool batched = dst->src[5] == nullptr && dst->src[0]->ne[1] > 1; ggml_cuda_fattn_kvarn_debug_route( ctx.device, plan, dst, entry_path, "portable-native", - "hip-prompt-precision"); + batched ? "hip-prompt-precision-qb4" : "hip-prompt-precision"); + if (batched) { + return ggml_cuda_flash_attn_ext_kvarn_portable_batched(ctx, dst, plan); + } return ggml_cuda_flash_attn_ext_kvarn_portable(ctx, dst, plan); } } diff --git a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh index a35b9fb58138..730c341653f1 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh +++ b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh @@ -57,30 +57,40 @@ ggml_cuda_fattn_kvarn_portable_resolve( return result; } -static __device__ __forceinline__ float ggml_cuda_fattn_kvarn_portable_block_reduce( - float partial, +template +static __device__ __forceinline__ void ggml_cuda_fattn_kvarn_portable_block_reduce_qb( + float (&partials)[QB], float * warp_partials) { - // Warp-level reduction with shuffles (no sync), then a single - // cross-warp step through shared memory (one sync). Replaces the - // 7-stage shared-memory tree (7 syncs) on the per-token hot path. + // Reduce QB values concurrently: one warp-shuffle phase over all QB + // lanes (no sync), then a single cross-warp step (one sync) instead of + // QB sequential reductions (2 syncs each). const int lane = threadIdx.x % 32; const int wid = threadIdx.x / 32; #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - partial += __shfl_xor_sync(0xFFFFFFFFu, partial, offset, 32); +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { + partials[qb] += __shfl_xor_sync(0xFFFFFFFFu, partials[qb], offset, 32); + } } if (lane == 0) { - warp_partials[wid] = partial; +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { + warp_partials[qb * (GGML_CUDA_FATTN_KVARN_DIM / 32) + wid] = partials[qb]; + } } __syncthreads(); - float total = threadIdx.x < (GGML_CUDA_FATTN_KVARN_DIM / 32) ? - warp_partials[threadIdx.x] : 0.0f; #pragma unroll - for (int offset = (GGML_CUDA_FATTN_KVARN_DIM / 64); offset > 0; offset >>= 1) { - total += __shfl_xor_sync(0xFFFFFFFFu, total, offset, 32); + for (int qb = 0; qb < QB; ++qb) { + float total = threadIdx.x < (GGML_CUDA_FATTN_KVARN_DIM / 32) ? + warp_partials[qb * (GGML_CUDA_FATTN_KVARN_DIM / 32) + threadIdx.x] : 0.0f; +#pragma unroll + for (int offset = (GGML_CUDA_FATTN_KVARN_DIM / 64); offset > 0; offset >>= 1) { + total += __shfl_xor_sync(0xFFFFFFFFu, total, offset, 32); + } + partials[qb] = total; } __syncthreads(); - return total; } template @@ -98,7 +108,7 @@ static __device__ __forceinline__ void ggml_cuda_fattn_kvarn_portable_stage_rota } } -template +template static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( const char * q_data, const ggml_cuda_fattn_kvarn_desc * k_descs, @@ -148,31 +158,64 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( "portable KVarN attention supports 128/256/512-wide heads"); constexpr int THREADS = GGML_CUDA_FATTN_KVARN_DIM; constexpr int SLICES = D / GGML_CUDA_FATTN_KVARN_DIM; + static_assert(QB >= 1 && QB <= 8, "portable query batch out of range"); - const int query = (int) blockIdx.x; + const int query_base = (int) blockIdx.x * QB; const int query_head = (int) blockIdx.y; const int stream = (int) blockIdx.z; const int tid = (int) threadIdx.x; - if (query >= n_query || query_head >= n_query_heads) { + if (query_base >= n_query || query_head >= n_query_heads) { return; } + // QB>1 fast path shares one K/V token stream across QB queries. That is + // only valid without packed per-query descriptors (no exact tail). + // The launcher guarantees k_tail_data == nullptr for QB > 1. + if constexpr (QB > 1) { + if (k_tail_data != nullptr) { + return; + } + } const int gqa = n_query_heads / n_kv_heads; const int kv_head = query_head / gqa; - const float * q = (const float *) ( - q_data + query * nbq1 + query_head * nbq2 + stream * nbq3); + // Per-batch query data. q_valid handles the ragged last group. + const float * q_ptr[QB]; + bool q_valid[QB]; +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { + const int query = query_base + qb; + q_valid[qb] = query < n_query; + q_ptr[qb] = (const float *) ( + q_data + query * nbq1 + query_head * nbq2 + stream * nbq3); + } - __shared__ float maximum; - __shared__ float denominator; - __shared__ float old_scale_shared; - __shared__ float weight_shared; - // Warp-shuffle reduction scratch: one partial per warp (THREADS/32). - __shared__ float warp_partials[THREADS / 32]; + __shared__ float maximum[QB]; + __shared__ float denominator[QB]; + __shared__ float old_scale_shared[QB]; + __shared__ float weight_shared[QB]; + // Token resolution (index math, pointer setup) is identical for all + // threads: resolve once per token, broadcast via shared. + __shared__ ggml_cuda_fattn_kvarn_resolved_token k_rt_shared; + __shared__ ggml_cuda_fattn_kvarn_resolved_token v_rt_shared; + // Warp-shuffle reduction scratch: QB sets of one partial per warp. + __shared__ float warp_partials[8 * (THREADS / 32)]; - float accumulator[SLICES] = {}; + float accumulator[QB][SLICES] = {}; + float q_values[QB][SLICES]; +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { +#pragma unroll + for (int slice = 0; slice < SLICES; ++slice) { + const int dim = slice * GGML_CUDA_FATTN_KVARN_DIM + tid; + q_values[qb][slice] = q_valid[qb] ? q_ptr[qb][dim] : 0.0f; + } + } if (tid == 0) { - maximum = -FLT_MAX; - denominator = 0.0f; +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { + maximum[qb] = -FLT_MAX; + denominator[qb] = 0.0f; + } } __syncthreads(); @@ -190,21 +233,24 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( const int32_t * desc = nullptr; bool body_packed = false; int n_body = n_kv; - if (k_tail_data != nullptr) { - const int query_id = stream * n_query + query; - int active = -1; - for (int packed = 0; packed < query_order_nelements; ++packed) { - if (query_order[packed] == query_id) { - active = packed / query_order_ne0; - break; + if constexpr (QB == 1) { + if (k_tail_data != nullptr) { + const int query = query_base; + const int query_id = stream * n_query + query; + int active = -1; + for (int packed = 0; packed < query_order_nelements; ++packed) { + if (query_order[packed] == query_id) { + active = packed / query_order_ne0; + break; + } } + if (active < 0) { + return; + } + desc = run_desc + (size_t) active * run_desc_ne0; + body_packed = run_desc_ne0 > 6 + tail_mask_ne0; + n_body = body_packed ? desc[5] : n_kv; } - if (active < 0) { - return; - } - desc = run_desc + (size_t) active * run_desc_ne0; - body_packed = run_desc_ne0 > 6 + tail_mask_ne0; - n_body = body_packed ? desc[5] : n_kv; } for (int packed = 0; packed < n_body; ++packed) { @@ -216,82 +262,101 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( k_descs[(size_t) body_stream * n_kv_heads + kv_head]; const ggml_cuda_fattn_kvarn_desc & v_desc = v_descs[(size_t) body_stream * n_kv_heads + kv_head]; - const auto k_ref = ggml_cuda_fattn_kvarn_portable_resolve(k_desc, token); - const auto v_ref = ggml_cuda_fattn_kvarn_portable_resolve(v_desc, token); + if (tid == 0) { + k_rt_shared = ggml_cuda_fattn_kvarn_resolve_token(k_desc, token); + v_rt_shared = ggml_cuda_fattn_kvarn_resolve_token(v_desc, token); + } + __syncthreads(); float k_values[SLICES] = {}; - if (k_ref.stage) { + if (k_rt_shared.from_stage) { ggml_cuda_fattn_kvarn_portable_stage_rotated( - k_desc, k_ref.stage_pos, tid, k_values); + k_desc, k_rt_shared.stage_pos, tid, k_values); } else { #pragma unroll for (int slice = 0; slice < SLICES; ++slice) { - k_values[slice] = ggml_cuda_fattn_kvarn_load_rotated( - k_desc, token, slice, tid); + k_values[slice] = ggml_cuda_fattn_kvarn_load_resolved( + k_desc, k_rt_shared, slice, tid); } } - float partial = 0.0f; + // One shared K row serves all QB queries: dots, then a single + // combined reduction for all QB partials. + float totals[QB]; #pragma unroll - for (int slice = 0; slice < SLICES; ++slice) { - const int dim = slice * GGML_CUDA_FATTN_KVARN_DIM + tid; - partial += q[dim] * k_values[slice]; + for (int qb = 0; qb < QB; ++qb) { + totals[qb] = 0.0f; +#pragma unroll + for (int slice = 0; slice < SLICES; ++slice) { + totals[qb] += q_values[qb][slice] * k_values[slice]; + } } - const float total = ggml_cuda_fattn_kvarn_portable_block_reduce( - partial, warp_partials); + ggml_cuda_fattn_kvarn_portable_block_reduce_qb(totals, warp_partials); if (tid == 0) { - float mask_value = 0.0f; - if (mask_data != nullptr) { - const half * mask = (const half *) ( - mask_data + token * nbm0 + query * nbm1 + - (query_head % nmask2) * nbm2 + (body_stream % nmask3) * nbm3); - mask_value = slope * __half2float(*mask); - } +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { + if (!q_valid[qb]) { + continue; + } + const int query = query_base + qb; + float mask_value = 0.0f; + if (mask_data != nullptr) { + const half * mask = (const half *) ( + mask_data + token * nbm0 + query * nbm1 + + (query_head % nmask2) * nbm2 + (body_stream % nmask3) * nbm3); + mask_value = slope * __half2float(*mask); + } - float score = total * scale; - if (logit_softcap != 0.0f) { - score = logit_softcap * tanhf(score); - } - score += mask_value; - if (mask_value == -INFINITY) { - old_scale_shared = 1.0f; - weight_shared = 0.0f; - } else { - const float next_maximum = fmaxf(maximum, score); - const float old_scale = maximum == -FLT_MAX ? - 0.0f : expf(maximum - next_maximum); - const float weight = expf(score - next_maximum); - maximum = next_maximum; - denominator = denominator * old_scale + weight; - old_scale_shared = old_scale; - weight_shared = weight; + float score = totals[qb] * scale; + if (logit_softcap != 0.0f) { + score = logit_softcap * tanhf(score); + } + score += mask_value; + if (mask_value == -INFINITY) { + old_scale_shared[qb] = 1.0f; + weight_shared[qb] = 0.0f; + } else { + const float next_maximum = fmaxf(maximum[qb], score); + const float old_scale = maximum[qb] == -FLT_MAX ? + 0.0f : expf(maximum[qb] - next_maximum); + const float weight = expf(score - next_maximum); + maximum[qb] = next_maximum; + denominator[qb] = denominator[qb] * old_scale + weight; + old_scale_shared[qb] = old_scale; + weight_shared[qb] = weight; + } } } __syncthreads(); float v_values[SLICES] = {}; - if (v_ref.stage) { + if (v_rt_shared.from_stage) { ggml_cuda_fattn_kvarn_portable_stage_rotated( - v_desc, v_ref.stage_pos, tid, v_values); + v_desc, v_rt_shared.stage_pos, tid, v_values); } else { #pragma unroll for (int slice = 0; slice < SLICES; ++slice) { - v_values[slice] = ggml_cuda_fattn_kvarn_load_rotated( - v_desc, token, slice, tid); + v_values[slice] = ggml_cuda_fattn_kvarn_load_resolved( + v_desc, v_rt_shared, slice, tid); } } #pragma unroll - for (int slice = 0; slice < SLICES; ++slice) { - accumulator[slice] = accumulator[slice] * old_scale_shared + - v_values[slice] * weight_shared; + for (int qb = 0; qb < QB; ++qb) { +#pragma unroll + for (int slice = 0; slice < SLICES; ++slice) { + accumulator[qb][slice] = accumulator[qb][slice] * old_scale_shared[qb] + + v_values[slice] * weight_shared[qb]; + } } __syncthreads(); } + if constexpr (QB == 1) { if (k_tail_data != nullptr) { + const int query = query_base; const int n_tail = desc[4]; for (int token = 0; token < n_tail; ++token) { const int slot = desc[6 + token]; - float partial = 0.0f; + float tail_totals[1] = { 0.0f }; #pragma unroll for (int slice = 0; slice < SLICES; ++slice) { const int dim = slice * GGML_CUDA_FATTN_KVARN_DIM + tid; @@ -300,10 +365,10 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( (size_t) dim * sizeof(uint16_t); const float kval = ggml_cuda_fattn_kvarn_load_tail( ptr, k_tail_bf16); - partial += q[dim] * kval; + tail_totals[0] += q_values[0][slice] * kval; } - const float total = ggml_cuda_fattn_kvarn_portable_block_reduce( - partial, warp_partials); + ggml_cuda_fattn_kvarn_portable_block_reduce_qb<1>(tail_totals, warp_partials); + const float total = tail_totals[0]; if (tid == 0) { const half * tail_mask = (const half *) ( @@ -316,17 +381,17 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( } score += mask_value; if (mask_value == -INFINITY) { - old_scale_shared = 1.0f; - weight_shared = 0.0f; + old_scale_shared[0] = 1.0f; + weight_shared[0] = 0.0f; } else { - const float next_maximum = fmaxf(maximum, score); - const float old_scale = maximum == -FLT_MAX ? - 0.0f : expf(maximum - next_maximum); + const float next_maximum = fmaxf(maximum[0], score); + const float old_scale = maximum[0] == -FLT_MAX ? + 0.0f : expf(maximum[0] - next_maximum); const float weight = expf(score - next_maximum); - maximum = next_maximum; - denominator = denominator * old_scale + weight; - old_scale_shared = old_scale; - weight_shared = weight; + maximum[0] = next_maximum; + denominator[0] = denominator[0] * old_scale + weight; + old_scale_shared[0] = old_scale; + weight_shared[0] = weight; } } __syncthreads(); @@ -339,40 +404,55 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel( (size_t) dim * sizeof(uint16_t); const float vval = ggml_cuda_fattn_kvarn_load_tail( ptr, v_tail_bf16); - accumulator[slice] = - accumulator[slice] * old_scale_shared + vval * weight_shared; + accumulator[0][slice] = + accumulator[0][slice] * old_scale_shared[0] + vval * weight_shared[0]; } __syncthreads(); } } + } if (tid == 0) { - if (sinks != nullptr) { - const float score = sinks[query_head]; - const float next_maximum = fmaxf(maximum, score); - const float old_scale = maximum == -FLT_MAX ? - 0.0f : expf(maximum - next_maximum); - const float weight = expf(score - next_maximum); - denominator = denominator * old_scale + weight; - maximum = next_maximum; - old_scale_shared = old_scale; - } else { - old_scale_shared = 1.0f; - } - if (body_meta != nullptr) { - const size_t row = ((size_t) stream * n_query + query) * n_query_heads + query_head; - body_meta[row] = make_float2(maximum, denominator); +#pragma unroll + for (int qb = 0; qb < QB; ++qb) { + if (!q_valid[qb]) { + continue; + } + const int query = query_base + qb; + if (sinks != nullptr) { + const float score = sinks[query_head]; + const float next_maximum = fmaxf(maximum[qb], score); + const float old_scale = maximum[qb] == -FLT_MAX ? + 0.0f : expf(maximum[qb] - next_maximum); + const float weight = expf(score - next_maximum); + denominator[qb] = denominator[qb] * old_scale + weight; + maximum[qb] = next_maximum; + old_scale_shared[qb] = old_scale; + } else { + old_scale_shared[qb] = 1.0f; + } + if (body_meta != nullptr) { + const size_t row = ((size_t) stream * n_query + query) * n_query_heads + query_head; + body_meta[row] = make_float2(maximum[qb], denominator[qb]); + } + weight_shared[qb] = denominator[qb] > 0.0f ? 1.0f / denominator[qb] : 0.0f; } - weight_shared = denominator > 0.0f ? 1.0f / denominator : 0.0f; } __syncthreads(); - float * output = (float *) ( - dst_data + query_head * nbd1 + query * nbd2 + stream * nbd3); #pragma unroll - for (int slice = 0; slice < SLICES; ++slice) { - const int dim = slice * GGML_CUDA_FATTN_KVARN_DIM + tid; - output[dim] = accumulator[slice] * old_scale_shared * weight_shared; + for (int qb = 0; qb < QB; ++qb) { + if (!q_valid[qb]) { + continue; + } + const int query = query_base + qb; + float * output = (float *) ( + dst_data + query_head * nbd1 + query * nbd2 + stream * nbd3); +#pragma unroll + for (int slice = 0; slice < SLICES; ++slice) { + const int dim = slice * GGML_CUDA_FATTN_KVARN_DIM + tid; + output[dim] = accumulator[qb][slice] * old_scale_shared[qb] * weight_shared[qb]; + } } } @@ -416,7 +496,7 @@ static inline bool ggml_cuda_fattn_kvarn_portable_supported( tail_ok && body_meta_ok; } -template +template static void ggml_cuda_fattn_kvarn_portable_launch( ggml_backend_cuda_context & ctx, ggml_tensor * dst, @@ -454,8 +534,15 @@ static void ggml_cuda_fattn_kvarn_portable_launch( k_desc.actual_size + v_desc.actual_size); const dim3 blocks( - (uint32_t) q->ne[1], (uint32_t) q->ne[2], (uint32_t) q->ne[3]); - ggml_cuda_fattn_kvarn_portable_kernel + (uint32_t) ((q->ne[1] + QB - 1) / QB), (uint32_t) q->ne[2], (uint32_t) q->ne[3]); + if (getenv("GGML_KVARN_PORTABLE_ATTRS") != nullptr) { + hipFuncAttributes attrs = {}; + CUDA_CHECK(hipFuncGetAttributes( + &attrs, (const void *) ggml_cuda_fattn_kvarn_portable_kernel)); + fprintf(stderr, "portable-attrs D=%d QB=%d numRegs=%d shared=%zu\n", + D, QB, attrs.numRegs, (size_t) attrs.sharedSizeBytes); + } + ggml_cuda_fattn_kvarn_portable_kernel <<>>( (const char *) q->data, k_desc.get(), v_desc.get(), @@ -503,9 +590,24 @@ static bool ggml_cuda_flash_attn_ext_kvarn_portable( return false; } switch (dst->src[0]->ne[0]) { - case 128: ggml_cuda_fattn_kvarn_portable_launch<128>(ctx, dst, plan); return true; - case 256: ggml_cuda_fattn_kvarn_portable_launch<256>(ctx, dst, plan); return true; - case 512: ggml_cuda_fattn_kvarn_portable_launch<512>(ctx, dst, plan); return true; + case 128: ggml_cuda_fattn_kvarn_portable_launch<128, 1>(ctx, dst, plan); return true; + case 256: ggml_cuda_fattn_kvarn_portable_launch<256, 1>(ctx, dst, plan); return true; + case 512: ggml_cuda_fattn_kvarn_portable_launch<512, 1>(ctx, dst, plan); return true; + default: return false; + } +} + +static bool ggml_cuda_flash_attn_ext_kvarn_portable_batched( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + const ggml_cuda_fattn_kvarn_plan & plan) { + if (!ggml_cuda_fattn_kvarn_portable_supported(plan, dst)) { + return false; + } + switch (dst->src[0]->ne[0]) { + case 128: ggml_cuda_fattn_kvarn_portable_launch<128, 8>(ctx, dst, plan); return true; + case 256: ggml_cuda_fattn_kvarn_portable_launch<256, 8>(ctx, dst, plan); return true; + case 512: ggml_cuda_fattn_kvarn_portable_launch<512, 8>(ctx, dst, plan); return true; default: return false; } } diff --git a/ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh b/ggml/src/ggml-cuda/fattn-mma-kvarn-load.cuh index 0f8e029503ca..e8ec9ebb5e33 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, From 27639c97044caccd684fdcfd0769712e63b06f15 Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:19:45 +0200 Subject: [PATCH 06/14] kvarn/rocm: fp32 WMMA accumulator for DV=128/256 prefill (opt-in) RDNA3 WMMA VKQ accumulators are fp16 (~3e-4/call), compounding over depth into a KLD collapse (2.14 at 32k). Mirror the proven DV=80/112 fp32-PV tiles for DV=128/256: 16-wide A, one wmma_f32 per K step, persistent f32 accumulator. Also fix the RDNA3 VKQ_C entry count for 16-row A tiles (was sized for 32-row tiles); the count is unchanged for all pre-existing configs. gfx1100, Qwen3.6-27B-Q5_K_S, kvarn6: ladder RMSE ~1e-5, 32k KLD 0.022 vs portable 0.023 (was 2.14), same-top 97.0 percent, prefill pp4096 286 -> 497 t/s. Gated behind GGML_KVARN_AMD_PROMPT_PORTABLE=0; default serving path unchanged. --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 27 ++++++++++++++- tests/test-kvarn.cpp | 49 +++++++++++++++++++++++++--- 2 files changed, 71 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index d75fe1f29d63..2a3556df20c6 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1102,6 +1102,25 @@ 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 }; +// 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 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 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 struct mma_tile_sizes { using T_A_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR>; // row-major @@ -1235,7 +1254,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 diff --git a/tests/test-kvarn.cpp b/tests/test-kvarn.cpp index 926c8c8c7ceb..54bf06a36d4f 100644 --- a/tests/test-kvarn.cpp +++ b/tests/test-kvarn.cpp @@ -2298,6 +2298,13 @@ static std::vector test_native_flash_attention_output( std::vector k_data((size_t) 128 * 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) { @@ -2311,6 +2318,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; + } } } } @@ -4113,16 +4124,46 @@ static void test_kvarn_nkv_ladder() { // 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, 256, - 6, 1, n_kv, 2, false, nullptr, false, 0, false, + 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, 256, - 6, 1, n_kv, 2, false, nullptr, false, 0, false, + 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) { From 5ca2905e58595169188e8d0bf9e21d0ec117caae Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:07:41 +0200 Subject: [PATCH 07/14] kvarn/rocm: publish whole-tile final meta for KVarN (fix ub>64 garbage) The three direct final-meta stores in flash_attn_ext_f16_process_tile were gated '!is_kvarn_kv', so KVarN whole-tile blocks never wrote body_meta (softmax max/rowsum). Split tiles get meta from the stream-k fixup, but flash_attn_stream_k_fixup_general skips tiles whose K range aligns exactly to tile boundaries; those rows kept zero meta (den=0) and the tail merge silently discarded their correct body values. Symptom: KVarN WMMA prompt path clean at ub<=64, garbage at ub>=96. Removing the gate makes whole-tile blocks publish their (max, rowsum) like dense FA. Validated: ub512 KLD on 4B/27B/35B-MoE all at portable parity; 32k KLD 2.118 -> 0.029. --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 2a3556df20c6..04ef065d9e4d 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1540,7 +1540,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) { @@ -1582,7 +1587,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) { @@ -1658,7 +1668,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) { From bad6f86c169fd8e465babc8246b0e3d4a0ff27f6 Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:53:33 +0200 Subject: [PATCH 08/14] kvarn/rocm: default HIP KVarN prompt-prefill to F32-WMMA (fast path as standard) Make F32-WMMA direct-record attention the default for HIP KVarN prompt-prefill (was: portable-native by default, WMMA opt-in). The ub-geometry fix (whole-tile body_meta) plus the DV=128/256 fp32 accumulators make WMMA both the fast and the exact route (~1e-5 ladder RMSE, 32k KLD at portable parity), validated at the serving default (-ub 512): 4B KLD 0.006822 (bit-exact vs fp16 base), 27B 0.004814, 35B-MoE 0.003526. Portable-native remains for unsupported shapes and as an explicit opt-in (GGML_KVARN_AMD_PROMPT_PORTABLE=1). Decode (nq<=16) unchanged. HIP-gated; CUDA/Vulkan and all non-KVarN paths untouched. --- ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu index 7943e8279112..040f7f617b08 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu +++ b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu @@ -1232,16 +1232,18 @@ bool ggml_cuda_flash_attn_ext_kvarn( ggml_cuda_fattn_kvarn_select_fallback_route( prompt_prefill, generic_shape_supported, portable_supported); #if defined(GGML_USE_HIP) - // RDNA WMMA prompt tiles accumulate in fp16 (RMSE ~3e-4 vs ~1e-5 for the - // portable fp32 path), and the error compounds through 64 layers into a - // visible KLD collapse. Route HIP prompt-prefill through portable-native - // direct-record attention (same records, exact math) until a float - // WMMA accumulator lands. Decode (nq<=16) stays on WMMA: its single - // iteration is exact. Opt out with GGML_KVARN_AMD_PROMPT_PORTABLE=0. + // RDNA WMMA prompt tiles now accumulate in fp32 for DV=128/256 (mirroring + // the proven DV=80/112 fp32-PV tiles), so 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. Decode + // (nq<=16) stays on WMMA as before. Portable-native direct-record + // attention remains as the fallback for unsupported shapes, or opt in + // explicitly with GGML_KVARN_AMD_PROMPT_PORTABLE=1. CUDA, Vulkan, and all + // non-KVarN paths are untouched by this block. { const char * prompt_portable = getenv("GGML_KVARN_AMD_PROMPT_PORTABLE"); if (prompt_prefill && portable_supported && - (prompt_portable == nullptr || atoi(prompt_portable) != 0)) { + (prompt_portable != nullptr && atoi(prompt_portable) == 1)) { g_kvarn_route_portable_native.fetch_add(1, std::memory_order_relaxed); // Batch 4 queries per block when no exact tail is attached (shared // token stream); otherwise the queries attend different token sets @@ -1249,7 +1251,7 @@ bool ggml_cuda_flash_attn_ext_kvarn( const bool batched = dst->src[5] == nullptr && dst->src[0]->ne[1] > 1; ggml_cuda_fattn_kvarn_debug_route( ctx.device, plan, dst, entry_path, "portable-native", - batched ? "hip-prompt-precision-qb4" : "hip-prompt-precision"); + batched ? "hip-prompt-precision-optin-qb4" : "hip-prompt-precision-optin"); if (batched) { return ggml_cuda_flash_attn_ext_kvarn_portable_batched(ctx, dst, plan); } From edc1d2f6431f984b81552eb2f8829bfdcb42a688 Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:17:05 +0200 Subject: [PATCH 09/14] kvarn/portable: fix CUDA build (HIP-only func-attributes in debug block) The GGML_KVARN_PORTABLE_ATTRS debug print used hipFuncAttributes / hipFuncGetAttributes unconditionally, which do not exist on CUDA and broke the CUDA compile (even though the block only runs when the env var is set). Use the HIP API under GGML_USE_HIP and the CUDA API (cudaFuncAttributes / cudaFuncGetAttributes, same fields) otherwise. No behavior change on HIP; CUDA now compiles. Found while validating the F32-WMMA flip on thermis/4090. --- ggml/src/ggml-cuda/fattn-kvarn-portable.cuh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh index 7ee4ef8fda97..c52e713f2038 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh +++ b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh @@ -572,9 +572,15 @@ static void ggml_cuda_fattn_kvarn_portable_launch( (uint32_t) ((q->ne[1] + QB - 1) / QB), (uint32_t) q->ne[2], (uint32_t) q->ne[3]); constexpr int RECORD_DIM = D == 64 ? 64 : GGML_CUDA_FATTN_KVARN_DIM; if (getenv("GGML_KVARN_PORTABLE_ATTRS") != nullptr) { +#if defined(GGML_USE_HIP) hipFuncAttributes attrs = {}; CUDA_CHECK(hipFuncGetAttributes( &attrs, (const void *) ggml_cuda_fattn_kvarn_portable_kernel)); +#else + cudaFuncAttributes attrs = {}; + CUDA_CHECK(cudaFuncGetAttributes( + &attrs, (const void *) ggml_cuda_fattn_kvarn_portable_kernel)); +#endif fprintf(stderr, "portable-attrs D=%d QB=%d numRegs=%d shared=%zu\n", D, QB, attrs.numRegs, (size_t) attrs.sharedSizeBytes); } From 37ce9e3377a1ad7053576edf444486ca1a5dc0e5 Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:58:56 +0200 Subject: [PATCH 10/14] kvarn/rocm: drop portable QB-batching call after v0.4.7 merge Follow-up to the v0.4.7 (complete optimized D64) merge, which superseded our interim portable QB-batching: the flip's fallback dispatch still called the now-removed ggml_cuda_flash_attn_ext_kvarn_portable_batched (undeclared on the merged tree). Route the fallback through the standard portable kernel (correct, unbatched). Portable is fallback-only post-flip. --- ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu index ef0373fe47f5..49b951f5f7bf 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu +++ b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu @@ -1248,16 +1248,11 @@ bool ggml_cuda_flash_attn_ext_kvarn( if (prompt_prefill && portable_supported && (prompt_portable != nullptr && atoi(prompt_portable) == 1)) { g_kvarn_route_portable_native.fetch_add(1, std::memory_order_relaxed); - // Batch 4 queries per block when no exact tail is attached (shared - // token stream); otherwise the queries attend different token sets - // and sharing is invalid. - const bool batched = dst->src[5] == nullptr && dst->src[0]->ne[1] > 1; + // 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", - batched ? "hip-prompt-precision-optin-qb4" : "hip-prompt-precision-optin"); - if (batched) { - return ggml_cuda_flash_attn_ext_kvarn_portable_batched(ctx, dst, plan); - } + "hip-prompt-precision-optin"); return ggml_cuda_flash_attn_ext_kvarn_portable(ctx, dst, plan); } } From 69c616d80cdbdae085ccf519abfd829c4ed4998b Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:12:02 +0200 Subject: [PATCH 11/14] kvarn/rocm: validated RDNA (256,256,64) tile + portable shared OOB fix Retune the RDNA (256,256,64) WMMA tile config to the validated values (measured KLD-clean on gfx1100 prompt-generic-mma) and raise the AMD WMMA device guard accordingly, so KVarN D256 shapes route to the fast WMMA path instead of the portable fallback. Only the (256,256,64) row is reachable (the KVarN switch has 128/256/512, standard FA caps RDNA at 128, eligibility caps RDNA3 at 256); the 320/512/576 keys in this commit are unreachable and reverted by the follow-up. Also size the portable kernel shared scratch by head dim (reduction[D] / transform[D]): the D64/D256/D512 V-domain transform indexes full-head dims into arrays previously sized RECORD_DIM, an out-of-bounds access for D256/D512. Latent on HIP (always rotated domain) but reachable on CUDA portable with the original V domain. Validation (gfx1100, kvarn6, ub512 KLD vs fp16): 27B 0.0046/95.1% via prompt-generic-mma with zero kernel traps (was 10.8 garbage); 35B-MoE 0.0037/97.1%; 4B 0.0072/96.9% unchanged. CUDA 4090: 27B 0.0033, KVarN/FA ctest 16/16, backend-ops MUL_MAT + TOP_K clean. --- ggml/src/ggml-cuda/fattn-kvarn-portable.cuh | 4 ++-- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 21 ++++++++++++--------- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh index b57d9b9abc66..ef19eee26488 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; diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 40710fb4afd8..6559887675d0 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/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); @@ -1964,8 +1964,11 @@ 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 > 256) { + // RDNA WMMA D256 tiles are the validated configs in the RDNA table above + // (256, 320, 512, 576 keys); the host caps RDNA3_0/RDNA4 at 576 and + // RDNA3_5 at 320 (fattn.cu), so this bound must stay >= the largest + // selectable head. GGML_CUDA_FA_WMMA_MAX_HEAD overrides the host selection. + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 576) { NO_DEVICE_CODE; return; } From 3c690bc18c608dd8d1976fd8e88079957b3b87c0 Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:59:53 +0200 Subject: [PATCH 12/14] kvarn/rocm: address review - RDNA4 fail-closed, opt-in reorder, knob docs 1. RDNA4 stays fail-closed at D128: the fp32-accumulator tiles that justify the raised D256 limit compile under RDNA3 (gfx11) only, so a new RDNA4_WMMA arch keeps the 128 head-dim cap (D64-style fail-closed) while RDNA3/3.5 admit 256. Dispatch comment corrected to RDNA3 scope. Route-policy test covers RDNA4-256 rejection. 2. Portable opt-in checked before the generic probe, so GGML_KVARN_AMD_PROMPT_PORTABLE=1 no longer pays for a discarded WMMA pass. Boolean aligned to nonzero-means-set (was strict ==1). 3. Knob documented in docs/beellama-args.md with exact accepted values. 4. D256 WMMA prefill with attached tail is covered on HIP via the ladder (GGML_KVARN_TEST_LADDER_TAIL=128: D256 n_kv 256-8192 RMSE ~1e-5 on the WMMA route); D512+tail hits a pre-existing fail-closed identical on the pre-change build. --- docs/beellama-args.md | 8 ++++ ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu | 44 +++++++++++-------- ggml/src/ggml-cuda/fattn-kvarn-route-policy.h | 5 ++- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 7 +-- tests/test-cuda-fattn-route-policy.cpp | 5 +++ 5 files changed, 47 insertions(+), 22 deletions(-) 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/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu b/ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu index 49b951f5f7bf..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,28 +1231,21 @@ 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 (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) { - g_kvarn_route_generic_shape_rejected.fetch_add(1, std::memory_order_relaxed); - } - } - const ggml_cuda_fattn_kvarn_route fallback_route = - ggml_cuda_fattn_kvarn_select_fallback_route( - prompt_prefill, generic_shape_supported, portable_supported); #if defined(GGML_USE_HIP) - // RDNA WMMA prompt tiles now accumulate in fp32 for DV=128/256 (mirroring - // the proven DV=80/112 fp32-PV tiles), so 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. Decode - // (nq<=16) stays on WMMA as before. Portable-native direct-record - // attention remains as the fallback for unsupported shapes, or opt in - // explicitly with GGML_KVARN_AMD_PROMPT_PORTABLE=1. CUDA, Vulkan, and all - // non-KVarN paths are untouched by this block. + // 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) == 1)) { + (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. @@ -1257,6 +1256,15 @@ bool ggml_cuda_flash_attn_ext_kvarn( } } #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) { + g_kvarn_route_generic_shape_rejected.fetch_add(1, std::memory_order_relaxed); + } + } + const ggml_cuda_fattn_kvarn_route fallback_route = + ggml_cuda_fattn_kvarn_select_fallback_route( + prompt_prefill, generic_shape_supported, portable_supported); if (fallback_route == GGML_CUDA_FATTN_KVARN_ROUTE_GENERIC_MMA || fallback_route == GGML_CUDA_FATTN_KVARN_ROUTE_PROMPT_PREFILL) { if (prompt_prefill) { diff --git a/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h b/ggml/src/ggml-cuda/fattn-kvarn-route-policy.h index a09bca6ffc6c..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, }; @@ -62,13 +63,15 @@ inline ggml_cuda_fattn_kvarn_mma_eligibility ggml_cuda_fattn_kvarn_amd_mma_eligi } if (input.head_dim <= 0 || (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 6559887675d0..c0cc9a634e97 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1965,9 +1965,10 @@ static __global__ void flash_attn_ext_f16( #if defined(AMD_WMMA_AVAILABLE) // Mirrored by ggml_cuda_fattn_kvarn_amd_mma_eligibility on the host. // RDNA WMMA D256 tiles are the validated configs in the RDNA table above - // (256, 320, 512, 576 keys); the host caps RDNA3_0/RDNA4 at 576 and - // RDNA3_5 at 320 (fattn.cu), so this bound must stay >= the largest - // selectable head. GGML_CUDA_FA_WMMA_MAX_HEAD overrides the host selection. + // (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; diff --git a/tests/test-cuda-fattn-route-policy.cpp b/tests/test-cuda-fattn-route-policy.cpp index 073505d2132e..11d53465c09f 100644 --- a/tests/test-cuda-fattn-route-policy.cpp +++ b/tests/test-cuda-fattn-route-policy.cpp @@ -221,6 +221,11 @@ int main(int argc, char ** argv) { 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 && From c5f2c022ecb29f7e28d4c6de4d3d5d40a1181345 Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Fri, 11 Sep 2026 09:11:55 +0200 Subject: [PATCH 13/14] kvarn/rocm: review round 2 - reachable rows only, fp32 scoping, committed tests - Drop the unreachable 320/512/576 retunes (no route selects those keys: KVarN switch has 128/256/512, standard FA caps RDNA at 128, eligibility caps RDNA3 at 256) and restore the device guard to DKQ > 256, mirrored by host eligibility again. Only the measured (256,256,64) row stays. - Scope the RDNA3 D128/D256 fp32 PV-accumulator tile to the KVarN path (mma_tile_sizes kvarn_accum gate): dense HIP attention keeps the qualified half2 tile, closing the dense-D128 evidence gap structurally. - Committed regression tests: D256 WMMA prompt prefill with attached exact tail vs CPU reference (ub-fix coverage), and portable with original_value_domain=true at D256/D512 vs CPU (shared-scratch fix coverage, force-pinned to the portable route). - AMD route-boundary expectation updated for admitted D256; feature matrix splits RDNA3/3.5 (WMMA to D256) from RDNA4 (D128 + portable). --- docs/beellama-features.md | 3 +- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 51 ++++++++++++------------ tests/test-kvarn.cpp | 59 +++++++++++++++++++++++++++- 3 files changed, 85 insertions(+), 28 deletions(-) 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-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index c0cc9a634e97..f5f829c15f8d 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -167,17 +167,17 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co 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, 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(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(512, 512, 8, 128, 3, 64, 96, 64, 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, 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, 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, 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, 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, 64, 128, 2, 32, 160, 128, 128, 1, true); return fattn_mma_config(32, 1, 0, 0, 0, 0, 0, false); @@ -1121,7 +1121,7 @@ template struct mma_tile_sizes { }; #elif defined(AMD_WMMA_AVAILABLE) #ifdef RDNA3 -template struct mma_tile_sizes { +template struct mma_tile_sizes { 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 @@ -1147,22 +1147,27 @@ template struct mma_tile_sizes<112, ncols> { }; // 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 struct mma_tile_sizes<128, ncols> { +// DV=80/112 fp32-PV tiles here for the KVarN path; the dense path keeps the +// qualified half2 tile until its fp32 variant is measured. +template struct mma_tile_sizes<128, ncols, kvarn_accum> { 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 + using T_C_VKQ = typename std::conditional, + tile<16, 16, half2, DATA_LAYOUT_I_MAJOR>>::type; // column-major }; -template struct mma_tile_sizes<256, ncols> { +template struct mma_tile_sizes<256, ncols, kvarn_accum> { 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 + using T_C_VKQ = typename std::conditional, + tile<16, 16, half2, DATA_LAYOUT_I_MAJOR>>::type; // column-major }; #else template struct mma_tile_sizes { @@ -1243,12 +1248,13 @@ 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 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 int cols_per_warp = T_B_KQ::I; constexpr int cols_per_thread = get_cols_per_thread(); @@ -1258,7 +1264,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"); @@ -1964,12 +1969,8 @@ static __global__ void flash_attn_ext_f16( #if defined(AMD_WMMA_AVAILABLE) // Mirrored by ggml_cuda_fattn_kvarn_amd_mma_eligibility on the host. - // 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) { + // Keep this final invariant for callers outside the KVarN dispatcher. + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) { NO_DEVICE_CODE; return; } diff --git a/tests/test-kvarn.cpp b/tests/test-kvarn.cpp index 52b8e4968849..fc6cf6899dc8 100644 --- a/tests/test-kvarn.cpp +++ b/tests/test-kvarn.cpp @@ -3777,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, @@ -4385,6 +4384,60 @@ static void test_kvarn_nkv_ladder() { 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. +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); + const std::vector expected = test_native_flash_attention_output( + cpu_backend, false, false, 256, 6, 6, 256, + 6, 1, 512, 2, false, nullptr, false, 128, false, + GGML_TYPE_F16, 0, false, true, -1, true); + const std::vector actual = test_native_flash_attention_output( + gpu_backend, true, true, 256, 6, 6, 256, + 6, 1, 512, 2, false, nullptr, false, 128, false, + GGML_TYPE_F16, 0, false, true, -1, true); + require_close_f32_rmse(actual, expected, 1e-2f, + "D256 WMMA prompt prefill with attached exact tail differs from CPU reference"); + 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); + require_close_f32_rmse(actual, expected, 1e-2f, + "portable original-V KVarN attention differs from CPU reference"); + } + 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) { @@ -5630,6 +5683,8 @@ int main() { test_odd_offset_record_decode_gpu(); test_d64_materialized_body_exact_tail_gpu(); test_native_flash_attention_prefill_route_parity(); + test_kvarn_d256_prompt_tail_regression(); + test_native_flash_attention_portable_original_v(); test_dflash_non_causal_attention_parity(); test_rotated_decode_transform_consistency(GGML_BACKEND_DEVICE_TYPE_CPU, true); test_rotated_decode_transform_consistency(GGML_BACKEND_DEVICE_TYPE_GPU, false); From f7f3350070db992a9ac3610c1068bd91e1efefdd Mon Sep 17 00:00:00 2001 From: raufaser <72879802+raufaser@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:04:15 +0200 Subject: [PATCH 14/14] kvarn/rocm: review round 2b - MIXED portable enable, fp32 scoping, committed tests - Enable portable for ROTATED_K_ORIGINAL_V at every head dim (was D64-only): the full-head V inverse transform inverts the stored forward transform, validated vs CPU at D256/D512 RMSE ~5e-4. Previously the portable+MIXED combination fell through to abort. - Scope the RDNA3 D128/D256 fp32 PV-accumulator tile to the KVarN path via mma_tile_sizes_kvarn (dense keeps its qualified half2 tiles). - Committed regression tests: D256 WMMA prompt prefill with attached exact tail (route-pinned, meta-denominator assertions, serving-geometry parity) and portable with original_value_domain=true at D256/D512 vs CPU (route- pinned via force-portable env). AMD route-boundary expectation updated for admitted D256. --- ggml/src/ggml-cuda/fattn-kvarn-portable.cuh | 3 +- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 51 ++++++++++------- tests/test-kvarn.cpp | 61 ++++++++++++++++++--- 3 files changed, 87 insertions(+), 28 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh index ef19eee26488..c7e5d4a483d9 100644 --- a/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh +++ b/ggml/src/ggml-cuda/fattn-kvarn-portable.cuh @@ -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-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index f5f829c15f8d..036b93dc7229 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1102,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 @@ -1121,7 +1135,7 @@ template struct mma_tile_sizes { }; #elif defined(AMD_WMMA_AVAILABLE) #ifdef RDNA3 -template struct mma_tile_sizes { +template struct mma_tile_sizes { 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 @@ -1145,29 +1159,26 @@ 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 }; -// 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 for the KVarN path; the dense path keeps the -// qualified half2 tile until its fp32 variant is measured. -template struct mma_tile_sizes<128, ncols, kvarn_accum> { +// 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 = typename std::conditional, - tile<16, 16, half2, DATA_LAYOUT_I_MAJOR>>::type; // column-major + using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major }; -template struct mma_tile_sizes<256, ncols, kvarn_accum> { +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 = typename std::conditional, - tile<16, 16, half2, DATA_LAYOUT_I_MAJOR>>::type; // column-major + using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major }; #else template struct mma_tile_sizes { @@ -1249,12 +1260,14 @@ 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; constexpr bool is_kvarn_kv = ggml_cuda_fattn_kvarn_template_type(type_K) || ggml_cuda_fattn_kvarn_template_type(type_V); - 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; + 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(); diff --git a/tests/test-kvarn.cpp b/tests/test-kvarn.cpp index fc6cf6899dc8..a29cf1b8711c 100644 --- a/tests/test-kvarn.cpp +++ b/tests/test-kvarn.cpp @@ -3800,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, @@ -4389,22 +4389,56 @@ static void test_kvarn_nkv_ladder() { // 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, - 6, 1, 512, 2, false, nullptr, false, 128, false, + 24, 4, 512, 2, false, nullptr, false, 128, true, GGML_TYPE_F16, 0, false, true, -1, true); - const std::vector actual = test_native_flash_attention_output( + const std::vector tailed = test_native_flash_attention_output( gpu_backend, true, true, 256, 6, 6, 256, - 6, 1, 512, 2, false, nullptr, false, 128, false, + 24, 4, 512, 2, false, nullptr, false, 128, true, GGML_TYPE_F16, 0, false, true, -1, true); - require_close_f32_rmse(actual, expected, 1e-2f, + 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); } @@ -4431,8 +4465,18 @@ static void test_native_flash_attention_portable_original_v() { 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); @@ -5676,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(); @@ -5683,8 +5732,6 @@ int main() { test_odd_offset_record_decode_gpu(); test_d64_materialized_body_exact_tail_gpu(); test_native_flash_attention_prefill_route_parity(); - test_kvarn_d256_prompt_tail_regression(); - test_native_flash_attention_portable_original_v(); test_dflash_non_causal_attention_parity(); test_rotated_decode_transform_consistency(GGML_BACKEND_DEVICE_TYPE_CPU, true); test_rotated_decode_transform_consistency(GGML_BACKEND_DEVICE_TYPE_GPU, false);