CUDA: MLA flash-attention decode on Pascal (vec_f32 K=576/V=512), incl. Q8_0 KV (#2079)

On GPUs without FP16 tensor cores (Pascal / sm_60, e.g. Tesla P100) MLA
flash-attention decode falls back to the CPU. The !fp16_mma_available path
routes decode to the f16 vector kernel, whose is_supported check requires
K == V head sizes; MLA's absorbed head sizes are 576/512 (asymmetric), so it is
rejected and attention runs on the CPU. With --cpu-moe that recomputes the full
MLA attention on the CPU every decoded token, which dominates decode at long
context.

Route Pascal MLA decode (Q->ne[1] <= 8 && K == 576 && V == 512) to the f32
vector kernel and enable that kernel for the 576/512 case, including Q8_0 KV.

Scope: decode only (batch <= 8). Prefill (batch > 8) and -fa 0 are untouched;
tensor-core GPUs never reach this branch. Aligned head sizes are byte-identical
(the asymmetric/Q8_0 work folds to a no-op at compile time), so no other model
or configuration is affected.

- fattn.cu: route 576/512 decode to vec_f32 in the !fp16_mma dispatch and its
  is_supported mirror.
- fattn-vec-f32.cu/.cuh: accept + instantiate 576/512 (F16 and Q8_0); fix latent
  issues exposed by the first asymmetric/large-head use (KQ-row granularity uses
  FATTN_KQ_STRIDE not Dv; guard the dst store to tid < Dv; guard the softmax exp
  on the KV tail; size Q_i32 by ceil; only convert K/V to F16 when the type is
  F16). All are no-ops for the previously-exercised symmetric cases.
- fattn-vec-f32.cuh / fattn-vec-common.cuh: guard the Q8_0 ragged tail (Dk=576 is
  144 int32 lanes = 4.5 warps) with the ragged-dim idiom; compile-time-constant
  for aligned head dims, so it folds away.

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

ggml/src/ggml-cuda/fattn-vec-common.cuh

@@ -326,6 +326,12 @@ static __device__ __forceinline__ T vec_dot_fattn_vec_KQ_q8_0(
     for (int k_KQ_0 = 0; k_KQ_0 < Dk/sizeof(int); k_KQ_0 += WARP_SIZE) {
         const int k_KQ = k_KQ_0 + threadIdx.x;
 
+        // Ragged-tail guard (mirrors mainline llama.cpp fattn-vec.cuh:169). For a chunk that fully
+        // fits (k_KQ_0 + WARP_SIZE <= Dk/sizeof(int)) the predicate is a compile-time constant true
+        // and folds away -> byte-identical for aligned Dk (128/256). For the ragged tail of Dk=576
+        // (144 int32 = 4.5*WARP_SIZE) the out-of-bounds lanes (k_KQ >= 144) skip the K-block read
+        // (ib would be 18/19 > last valid block 17 -> OOB -> NaN) and contribute 0 to the warp sum.
+        if (k_KQ_0 + WARP_SIZE <= Dk/(int)sizeof(int) || k_KQ < Dk/(int)sizeof(int)) {
             const int ib  = k_KQ / QI8_0;
             const int iqs = k_KQ % QI8_0;
 
@@ -342,6 +348,7 @@ static __device__ __forceinline__ T vec_dot_fattn_vec_KQ_q8_0(
 
             sum += vec_dot_q8_0_q8_1_impl<T, 1>(&v, &Q_q8[k_KQ_0/WARP_SIZE], K_q8_0[ib].d, Q_d);
         }
+    }
 
     return sum;
 }
@@ -388,14 +395,21 @@ static __device__ __forceinline__ T vec_dot_fattn_vec_KQ_f16(
     return sum;
 }
 
-template <typename Tds>
+// ni = number of valid lanes for this chunk (compile-time). For a chunk that fully covers
+// WARP_SIZE int32 words ni == WARP_SIZE and every guard below is provably always-true, so the
+// compiler eliminates it (byte-identical to the unguarded version). For a ragged tail chunk
+// (e.g. Dk=576 -> the i0=128 chunk has only 144-128=16 valid int32 words) ni < WARP_SIZE and the
+// guards engage: the out-of-bounds lanes read 0 (no OOB Q load), and skip the yq32 / yds writes
+// (no clobber of the trailing q8_1 scale region / no OOB scale write). Mirrors mainline
+// llama.cpp fattn-common.cuh quantize_q8_1_to_shared<Tds, ni> (lane guards :338 / :366).
+template <typename Tds, int ni = WARP_SIZE>
 static __device__ __forceinline__ void quantize_q8_1_to_shared(
     const float * __restrict__ x, const float scale, int * __restrict__ yq32, void * __restrict__ yds) {
 
     float vals[sizeof(int)] = {0.0f};
 #pragma unroll
     for (int l = 0; l < sizeof(int); ++l) {
-        vals[l] = scale * x[4*threadIdx.x + l];
+        vals[l] = (ni == WARP_SIZE || threadIdx.x < ni) ? scale * x[4*threadIdx.x + l] : 0.0f;
     }
 
     float amax = fabsf(vals[0]);
@@ -422,8 +436,10 @@ static __device__ __forceinline__ void quantize_q8_1_to_shared(
         }
     }
 
+    if (ni == WARP_SIZE || threadIdx.x < ni) {
         yq32[threadIdx.x] = q32;
-    if (threadIdx.x % QI8_1 == 0) {
+    }
+    if (threadIdx.x % QI8_1 == 0 && (ni == WARP_SIZE || threadIdx.x < ni)) {
         if (std::is_same<Tds, half2>::value) {
             ((half2  *) yds)[threadIdx.x/QI8_1] =  make_half2(d, sum);
         } else {

ggml/src/ggml-cuda/fattn-vec-f32.cu

@@ -79,6 +79,9 @@ void ggml_cuda_flash_attn_ext_vec_f32(ggml_backend_cuda_context & ctx, ggml_tens
 
     FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16)
     FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
+
+    FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16)
+    FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
 #else
     FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)
 
@@ -97,6 +100,9 @@ void ggml_cuda_flash_attn_ext_vec_f32(ggml_backend_cuda_context & ctx, ggml_tens
     FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16)
     FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
 
+    FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16)
+    FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
+
 #endif // GGML_CUDA_FA_ALL_QUANTS
 
     on_no_fattn_vec_case(Q->ne[0], V->ne[0]);
@@ -106,8 +112,17 @@ bool ggml_cuda_fattn_vec_f32_is_supported([[maybe_unused]] ggml_backend_cuda_con
     auto K = dst->src[1];
     auto V = dst->src[2];
     if (K->ne[0] != V->ne[0]) {
-        if (K->ne[0] != 192 || V->ne[2] != 128) return false;
+        const bool is_192_128 = K->ne[0] == 192 && V->ne[0] == 128;
+        const bool is_576_512 = K->ne[0] == 576 && V->ne[0] == 512;
+        if (!is_192_128 && !is_576_512) return false;
         if (K->type != V->type) return false;
+        // Asymmetric (Dk != Dv) vec_f32 path. F16 KV is always fine (its dot loops stride Dk/2,
+        // a clean multiple of 32). Q8_0 KV is now supported: the per-warp Q->q8_1 quantization,
+        // the shared->register load, and the quantized-K dot all carry a compile-time-gated
+        // per-lane bounds guard (mirroring mainline llama.cpp's fattn-vec ragged-tail idiom) so
+        // the ragged final warp iteration for Dk=576 (144 int32 = 4.5*WARP_SIZE) and Dk=192
+        // (48 int32 = 1.5*WARP_SIZE) no longer reads Q/K out of bounds. The guard folds to a
+        // no-op for aligned Dk (128/256), keeping those instances byte-identical.
         return K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0;
     }
 #ifdef GGML_CUDA_FA_ALL_QUANTS

ggml/src/ggml-cuda/fattn-vec-f32.cuh

@@ -115,7 +115,7 @@ static __global__ void flash_attn_vec_ext_f32(
 
     // Convert Q to float2 (f16 K) or q8_1 (quantized K) and store in registers:
     float2  Q_f2[ncols][Dk/(2*WARP_SIZE)];
-    int    Q_i32[ncols][Dk/(sizeof(int)*QK8_1) == 0 ? 1 : Dk/(sizeof(int)*QK8_1)];
+    int    Q_i32[ncols][(Dk/(int)sizeof(int) + WARP_SIZE - 1)/WARP_SIZE];
     float2  Q_ds[ncols][Dk/QK8_1 == 0 ? 1 : Dk/QK8_1];
     if (Q_q8_1) {
 #pragma unroll
@@ -145,9 +145,21 @@ static __global__ void flash_attn_vec_ext_f32(
             }
 
             const float * Q_f = (const float *) (Q + j*nb01);
+            // Full WARP_SIZE-wide int32 chunks: ni == WARP_SIZE so quantize_q8_1_to_shared's lane
+            // guards fold away -> byte-identical to the original unguarded loop for aligned Dk.
+            constexpr int Dk_i32  = Dk/(int)sizeof(int);
+            constexpr int Dk_full = (Dk_i32/WARP_SIZE)*WARP_SIZE;
 #pragma unroll
-            for (int i0 = 0; i0 < Dk/sizeof(int); i0 += WARP_SIZE) {
+            for (int i0 = 0; i0 < Dk_full; i0 += WARP_SIZE) {
-                quantize_q8_1_to_shared<float2>(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1);
+                quantize_q8_1_to_shared<float2, WARP_SIZE>(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1);
+            }
+            // Ragged tail: only compiled for Dk whose int32 count is not a multiple of WARP_SIZE
+            // (Dk=576 -> 16 valid lanes, Dk=192 -> 16 valid lanes). `if constexpr` drops this
+            // entirely for aligned Dk (128/256), so the aligned path stays byte-identical.
+            if constexpr (Dk_i32 % WARP_SIZE != 0) {
+                constexpr int i0 = Dk_full;
+                constexpr int ni = Dk_i32 - Dk_full;
+                quantize_q8_1_to_shared<float2, ni>(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1);
             }
         }
 
@@ -162,8 +174,17 @@ static __global__ void flash_attn_vec_ext_f32(
             for (int i0 = 0; i0 < Dk/sizeof(int); i0 += WARP_SIZE) {
                 const int i = i0 + threadIdx.x;
 
+                // Ragged-tail guard (mirrors mainline fattn-vec.cuh:169). Aligned Dk: predicate
+                // folds to always-true -> byte-identical. Dk=576 ragged tail: OOB lanes (i >= 144)
+                // load 0 rather than reading the trailing shared region (tmp_q_i32[144..159] is now
+                // never written by the guarded quantize, so it is stale) and feed 0 to the dot.
+                if (i0 + WARP_SIZE <= Dk/(int)sizeof(int) || i < Dk/(int)sizeof(int)) {
                     Q_i32[j][i0/WARP_SIZE] = tmp_q_i32[i];
                     Q_ds[j][i0/WARP_SIZE]  = tmp_q_ds[i/QI8_1];
+                } else {
+                    Q_i32[j][i0/WARP_SIZE] = 0;
+                    Q_ds[j][i0/WARP_SIZE]  = make_float2(0.0f, 0.0f);
+                }
             }
         }
 
@@ -257,7 +278,7 @@ static __global__ void flash_attn_vec_ext_f32(
             const float KQ_max_scale = expf(kqmax[j] - kqmax_new_j);
             kqmax[j] = kqmax_new_j;
 
-            const float val = expf(KQ[j*Dk + tid] - kqmax[j]);
+            const float val = (FATTN_KQ_STRIDE % Dk != 0 && k_VKQ_0 + tid >= ne11) ? 0.0f : expf(KQ[j*Dk + tid] - kqmax[j]);
             kqsum[j] = kqsum[j]*KQ_max_scale + val;
             KQ[j*Dk + tid] = val;
 
@@ -344,8 +365,10 @@ static __global__ void flash_attn_vec_ext_f32(
         if (gridDim.y == 1) {
             dst_val /= kqsum[j_VKQ];
         }
+        if (tid < Dv) {
             dst[(((sequence*ne01 + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y)*Dv + tid] = dst_val;
         }
+    }
 
     if (gridDim.y != 1 && tid < ncols && (ncols <= 2 || ic0 + tid < ne01)) {
         dst_meta[((sequence*ne01 + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = make_float2(kqmax[tid], kqsum[tid]);
@@ -359,10 +382,13 @@ template <int Dk, int Dv, int cols_per_block, ggml_type type_K, ggml_type type_V
 void ggml_cuda_flash_attn_ext_vec_f32_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     constexpr int nwarps = Dk/WARP_SIZE;
     fattn_kernel_t fattn_kernel = flash_attn_vec_ext_f32<Dk, Dv, cols_per_block, type_K, type_V, use_logit_softcap>;
-    constexpr bool need_f16_K = Dk != 128 && Dk != 256;
+    constexpr bool need_f16_K = (Dk != 128 && Dk != 256) && type_K == GGML_TYPE_F16;
-    constexpr bool need_f16_V = Dv != 64 && Dv != 128 && Dv != 256;
+    constexpr bool need_f16_V = (Dv != 64 && Dv != 128 && Dv != 256) && type_V == GGML_TYPE_F16;
     constexpr size_t nbytes_shared = 0;
-    launch_fattn<Dv, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, Dv, need_f16_K, need_f16_V);
+    // MLA (Dv=512): the KV cache is padded to FATTN_KQ_STRIDE (256), not to Dv, so use 256 as the KQ-row
+    // granularity. Passing Dv=512 asserts n_kv % 512 == 0, which fails for odd multiples of 256.
+    constexpr int kq_row_granularity = (Dv == 512) ? FATTN_KQ_STRIDE : Dv;
+    launch_fattn<Dv, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, kq_row_granularity, need_f16_K, need_f16_V);
 }
 
 template <int Dk, int Dv, ggml_type type_K, ggml_type type_V>
@@ -488,3 +514,6 @@ extern DECL_FATTN_VEC_F32_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0);
 
 extern DECL_FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16);
 extern DECL_FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0);
+
+extern DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16);
+extern DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0);

ggml/src/ggml-cuda/fattn.cu

@@ -23,6 +23,14 @@ static inline bool mma_better_than_turing(const int cc) {
     return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) > CC_TURING;
 }
 
+// On GPUs without fp16 MMA (Pascal/sm_60), the MLA absorbed head sizes 576/512 are asymmetric, so the f16
+// vector kernel (which requires K==V) rejects them and decode falls back to the CPU. Route MLA decode
+// (batch <= 8) to the f32 vector kernel to keep it on the GPU. Single-sourced so the dispatch and the
+// is_supported check cannot drift.
+static inline bool is_pascal_mla_absorbed_decode(const ggml_tensor * Q, const ggml_tensor * K, const ggml_tensor * V) {
+    return Q->ne[1] <= 8 && K->ne[0] == 576 && V->ne[0] == 512;
+}
+
 void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * KQV  = dst;
     const ggml_tensor * Q    = dst->src[0];
@@ -76,6 +84,10 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
     }
 
     if (!fp16_mma_available(cc)) {
+        if (is_pascal_mla_absorbed_decode(Q, K, V)) {
+            ggml_cuda_flash_attn_ext_vec_f32(ctx, dst);
+            return;
+        }
         if (precision == GGML_PREC_DEFAULT) {
             if (Q->ne[1] <= 8 || Q->ne[0] == 256) {
                 ggml_cuda_flash_attn_ext_vec_f16(ctx, dst);
@@ -171,6 +183,9 @@ bool ggml_cuda_fattn_is_supported(ggml_backend_cuda_context & ctx, const ggml_te
     }
 
     if (!fp16_mma_available(cc)) {
+        if (is_pascal_mla_absorbed_decode(Q, K, V)) {
+            return ggml_cuda_fattn_vec_f32_is_supported(ctx, dst);
+        }
         if (precision == GGML_PREC_DEFAULT) {
             if (Q->ne[1] <= 8 || Q->ne[0] == 256) {
                 return ggml_cuda_fattn_vec_f16_is_supported(ctx, dst);

ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-f16-f16.cu

@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../fattn-vec-f32.cuh"
+
+DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16);

ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-q8_0-q8_0.cu

@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../fattn-vec-f32.cuh"
+
+DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0);