From a8cf53fd69bada5450bd653eb0e32d1113fbd7fe Mon Sep 17 00:00:00 2001 From: mb8565 <244351746+mb8565@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:59:29 -0500 Subject: [PATCH] CUDA: fix flash attention for gpt-oss (SWA + attention sinks) on the tile kernels (no-tensor-core GPUs) (#2087) * CUDA: use the mask tensor stride (nb31) in the tile FA kernels, not ne11 The tile flash-attention kernels (no-tensor-core GPUs, e.g. Pascal/sm_60) indexed the KQ mask using ne11 (= K->ne[1]) as the row stride. The SWA windowing in ggml_cuda_flash_attn_ext (the n_swa branch) re-points K/V/mask to the last nton tokens, setting ne11 = nton while the mask keeps its original row stride nb31. Indexing by ne11 then reads across mask rows and yields NaN once the window slice engages (context past nton). Use the mask's own stride nb31/sizeof(half); in the non-sliced case this equals ne11, so it is a no-op there. Matches how the vec kernel already computes its mask offset, and how mainline llama.cpp fixed the same latent bug in its unified tile kernel. Co-Authored-By: Claude Opus 4.8 (1M context) * CUDA: apply attention sinks in the tile FA kernels (gpt-oss) The tile flash-attention kernels took a sinks argument but never applied it, so gpt-oss (which has a per-head sink logit) got a wrong softmax denominator with -fa 1 while the -fa 0 path (ggml_soft_max_add_sinks) matched CPU. Apply the sink after the KV loop, mirroring the vec kernel: the sink joins the running max and adds exp(sink - max) to the denominator once, rescaling kqsum and VKQ. Only ip==0 adds it so it is counted once across the parallel_blocks KV split (the epilogue writes the sink-inclusive max/denominator into dst_meta before the combine). The per-head index is blockIdx.y (the same index the slope uses). Guarded by non-null sinks, so non-sink models are unaffected. This is the tile-kernel equivalent of mainline llama.cpp PR #15178; the wmma kernel is left untouched (no Turing/Volta hardware to validate here, but the same defect likely applies). Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- ggml/src/ggml-cuda/fattn-tile-f16.cu | 36 +++++++++++++++++++++++++-- ggml/src/ggml-cuda/fattn-tile-f32.cu | 37 ++++++++++++++++++++++++++-- 2 files changed, 69 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-tile-f16.cu b/ggml/src/ggml-cuda/fattn-tile-f16.cu index a2fe6bd3d..9d72b44f7 100644 --- a/ggml/src/ggml-cuda/fattn-tile-f16.cu +++ b/ggml/src/ggml-cuda/fattn-tile-f16.cu @@ -68,7 +68,12 @@ static __global__ void flash_attn_tile_ext_f16( const float2 * Q_f2 = (const float2 *) (Q + nb02* blockIdx.y + nb01*ic0); const half2 * K_h2 = (const half2 *) (K + nb12*(blockIdx.y / gqa_ratio)); const half2 * V_h2 = (const half2 *) (V + nb12*(blockIdx.y / gqa_ratio)); // K and V have same shape - const half * maskh = (const half *) mask + ne11*ic0; + // Mask row stride must come from the mask tensor (nb31), not ne11 (= K->ne[1]). For SWA models the + // fattn.cu n_swa windowing re-points K/V/mask to the last nton tokens (ne11 = nton) while the mask keeps + // its original row stride, so indexing the mask by ne11 reads garbage and yields NaN on the tile kernels. + const int stride_mask = nb31 / sizeof(half); + const half * maskh = (const half *) mask + stride_mask*ic0; + const float * sinksf = (const float *) sinks; const int stride_KV2 = nb11 / sizeof(half2); @@ -176,7 +181,7 @@ static __global__ void flash_attn_tile_ext_f16( } else { sum = __low2half(sum2[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]) + __high2half(sum2[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]); } - sum += mask ? slopeh*maskh[j_KQ*ne11 + k_VKQ_0 + i_KQ] : __float2half(0.0f); + sum += mask ? slopeh*maskh[j_KQ*stride_mask + k_VKQ_0 + i_KQ] : __float2half(0.0f); kqmax_new[j_KQ_0/nwarps] = ggml_cuda_hmax(kqmax_new[j_KQ_0/nwarps], sum); @@ -258,6 +263,33 @@ static __global__ void flash_attn_tile_ext_f16( __syncthreads(); } + // Apply attention sinks (e.g. gpt-oss): the sink is a per-head extra softmax logit whose value + // contribution is zero, so it only joins the running max and rescales the denominator/value + // accumulator. Only ip==0 adds the sink term so it is counted exactly once across the + // parallel_blocks KV split. Ported from fattn-vec-f16.cuh. kqmax is warp-uniform here (already + // warp_reduce_max'd in the KV loop), so no shared-mem reduction is needed. + if (sinksf && ip == 0) { + const half sink = __float2half(sinksf[blockIdx.y]); + +#pragma unroll + for (int j0 = 0; j0 < ncols; j0 += nwarps) { + const half kqmax_new_j = ggml_cuda_hmax(kqmax[j0/nwarps], sink); + const half2 KQ_max_scale = __half2half2(hexp(kqmax[j0/nwarps] - kqmax_new_j)); + kqmax[j0/nwarps] = kqmax_new_j; + + kqsum[j0/nwarps] = kqsum[j0/nwarps]*KQ_max_scale; + if (threadIdx.x == 0) { + // Add the sink term to only the low half; the epilogue sums low+high once per lane. + kqsum[j0/nwarps].x += hexp(sink - kqmax[j0/nwarps]); + } + +#pragma unroll + for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) { + VKQ[j0/nwarps][i0/WARP_SIZE] *= KQ_max_scale; + } + } + } + #pragma unroll for (int j_VKQ_0 = 0; j_VKQ_0 < ncols; j_VKQ_0 += nwarps) { const int j_VKQ = j_VKQ_0 + threadIdx.y; diff --git a/ggml/src/ggml-cuda/fattn-tile-f32.cu b/ggml/src/ggml-cuda/fattn-tile-f32.cu index 675f22e05..7a1d36e5a 100644 --- a/ggml/src/ggml-cuda/fattn-tile-f32.cu +++ b/ggml/src/ggml-cuda/fattn-tile-f32.cu @@ -68,7 +68,12 @@ static __global__ void flash_attn_tile_ext_f32( const float2 * Q_f2 = (const float2 *) (Q + nb02* blockIdx.y + nb01*ic0); const half2 * K_h2 = (const half2 *) (K + nb12*(blockIdx.y / gqa_ratio)); const half2 * V_h2 = (const half2 *) (V + nb12*(blockIdx.y / gqa_ratio)); // K and V have same shape - const half * maskh = (const half *) mask + ne11*ic0; + // Mask row stride must come from the mask tensor (nb31), not ne11 (= K->ne[1]). For SWA models the + // fattn.cu n_swa windowing re-points K/V/mask to the last nton tokens (ne11 = nton) while the mask keeps + // its original row stride, so indexing the mask by ne11 reads garbage and yields NaN on the tile kernels. + const int stride_mask = nb31 / sizeof(half); + const half * maskh = (const half *) mask + stride_mask*ic0; + const float * sinksf = (const float *) sinks; const int stride_KV2 = nb11 / sizeof(half2); @@ -171,7 +176,7 @@ static __global__ void flash_attn_tile_ext_f32( sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] = softcap * tanhf(sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]); } - sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += mask ? slope*__half2float(maskh[j_KQ*ne11 + k_VKQ_0 + i_KQ]) : 0.0f; + sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += mask ? slope*__half2float(maskh[j_KQ*stride_mask + k_VKQ_0 + i_KQ]) : 0.0f; kqmax_new[j_KQ_0/nwarps] = fmaxf(kqmax_new[j_KQ_0/nwarps], sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]); @@ -256,6 +261,34 @@ static __global__ void flash_attn_tile_ext_f32( __syncthreads(); } + // Apply attention sinks (e.g. gpt-oss): the sink is a per-head extra softmax logit whose value + // contribution is zero, so it only joins the running max and rescales the denominator/value + // accumulator. Only ip==0 adds the sink term so it is counted exactly once across the + // parallel_blocks KV split. Ported from fattn-vec-f16.cuh. kqmax is warp-uniform here (already + // warp_reduce_max'd in the KV loop); kqsum holds per-lane partials reduced in the epilogue, so + // the sink term is added on a single lane. + if (sinksf && ip == 0) { + const float sink = sinksf[blockIdx.y]; + +#pragma unroll + for (int j0 = 0; j0 < ncols; j0 += nwarps) { + const float kqmax_new_j = fmaxf(kqmax[j0/nwarps], sink); + const float KQ_max_scale = expf(kqmax[j0/nwarps] - kqmax_new_j); + kqmax[j0/nwarps] = kqmax_new_j; + + kqsum[j0/nwarps] = kqsum[j0/nwarps]*KQ_max_scale; + if (threadIdx.x == 0) { + kqsum[j0/nwarps] += expf(sink - kqmax[j0/nwarps]); + } + +#pragma unroll + for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) { + VKQ[j0/nwarps][i0/WARP_SIZE].x *= KQ_max_scale; + VKQ[j0/nwarps][i0/WARP_SIZE].y *= KQ_max_scale; + } + } + } + #pragma unroll for (int j_VKQ_0 = 0; j_VKQ_0 < ncols; j_VKQ_0 += nwarps) { const int j_VKQ = j_VKQ_0 + threadIdx.y;