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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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;

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;