cuda: fix Turbo K classic V FA prefill

Anbeeld · Anbeeld · commit cb83fe74f21e · 2026-06-13T13:04:28.000+02:00
Route Turbo K with classic non-q8 V through the Turbo prefill MMA path for Gemma-sized D256/D512 batch attention, while preserving native vec fallback for decode paths and prior classic-K/Turbo-V safeguards. Add a route-policy regression test covering the mixed-pair gaps.
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
@@ -1026,6 +1026,53 @@ static __global__ void k_turbo_fwht_forward(
     }
 }
 
+static inline bool ggml_cuda_fattn_is_classic_non_q8_type(const ggml_type type) {
+    return type == GGML_TYPE_BF16 ||
+           type == GGML_TYPE_Q4_0 ||
+           type == GGML_TYPE_Q4_1 ||
+           type == GGML_TYPE_Q5_0 ||
+           type == GGML_TYPE_Q5_1 ||
+           type == GGML_TYPE_Q6_0 ||
+           type == GGML_TYPE_Q6_1 ||
+           type == GGML_TYPE_Q3_0 ||
+           type == GGML_TYPE_Q3_1 ||
+           type == GGML_TYPE_Q2_0 ||
+           type == GGML_TYPE_Q2_1;
+}
+
+static void ggml_cuda_fattn_materialize_to_f16(
+        const ggml_tensor * src, half * dst, cudaStream_t stream, ggml_tensor & src_f16) {
+    const size_t bs = ggml_blck_size(src->type);
+    const size_t ts = ggml_type_size(src->type);
+
+    src_f16 = *src;
+    src_f16.type = GGML_TYPE_F16;
+    src_f16.data = dst;
+    src_f16.nb[0] = sizeof(half);
+
+    if (ggml_is_contiguously_allocated(src)) {
+        const to_fp16_cuda_t to_fp16 = ggml_get_to_fp16_cuda(src->type);
+        GGML_ASSERT(to_fp16 != nullptr);
+        to_fp16(src->data, dst, ggml_nelements(src), stream);
+
+        src_f16.nb[1] = src->nb[1] * bs * sizeof(half) / ts;
+        src_f16.nb[2] = src->nb[2] * bs * sizeof(half) / ts;
+        src_f16.nb[3] = src->nb[3] * bs * sizeof(half) / ts;
+    } else {
+        GGML_ASSERT(src->nb[0] == ts);
+        const to_fp16_nc_cuda_t to_fp16 = ggml_get_to_fp16_nc_cuda(src->type);
+        GGML_ASSERT(to_fp16 != nullptr);
+        const int64_t s01 = src->nb[1] / ts;
+        const int64_t s02 = src->nb[2] / ts;
+        const int64_t s03 = src->nb[3] / ts;
+        to_fp16(src->data, dst, src->ne[0], src->ne[1], src->ne[2], src->ne[3], s01, s02, s03, stream);
+
+        src_f16.nb[1] = src->ne[0] * sizeof(half);
+        src_f16.nb[2] = src->ne[1] * src_f16.nb[1];
+        src_f16.nb[3] = src->ne[2] * src_f16.nb[2];
+    }
+}
+
 static void ggml_cuda_turbo_prefill_attend(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     load_tcq_decode_alpha(ctx.device);
     cudaStream_t stream = ctx.stream();
@@ -1034,12 +1081,17 @@ static void ggml_cuda_turbo_prefill_attend(ggml_backend_cuda_context & ctx, ggml
 
     const bool turbo_k = K->type == GGML_TYPE_TURBO2_0 || K->type == GGML_TYPE_TURBO3_0 || K->type == GGML_TYPE_TURBO4_0 || K->type == GGML_TYPE_TURBO4_TCQ || K->type == GGML_TYPE_TURBO3_TCQ || K->type == GGML_TYPE_TURBO2_TCQ;
     const bool turbo_v = V->type == GGML_TYPE_TURBO2_0 || V->type == GGML_TYPE_TURBO3_0 || V->type == GGML_TYPE_TURBO4_0 || V->type == GGML_TYPE_TURBO4_TCQ || V->type == GGML_TYPE_TURBO3_TCQ || V->type == GGML_TYPE_TURBO2_TCQ;
+    const bool classic_non_q8_v = !turbo_v && ggml_cuda_fattn_is_classic_non_q8_type(V->type);
 
     int device;
     CUDA_CHECK(cudaGetDevice(&device));
 
     half * k_fp16 = nullptr;
     half * v_fp16 = nullptr;
+    bool v_f16_layout_set = false;
+
+    ggml_tensor K_f16 = *K;
+    ggml_tensor V_f16 = *V;
 
     // Allocate and dequant K to fp16 (turbo2, turbo3, or turbo4)
     if (turbo_k) {
@@ -1115,8 +1167,8 @@ static void ggml_cuda_turbo_prefill_attend(ggml_backend_cuda_context & ctx, ggml
         }
     }
 
-    // Allocate and dequant V to fp16 (turbo2, turbo3, or turbo4)
-    if (turbo_v) {
+    // Allocate and materialize V to fp16 (turbo2, turbo3, turbo4, or classic non-q8).
+    if (turbo_v || classic_non_q8_v) {
         // Size for full cache (kv_size from root) so we never realloc mid-session.
         const ggml_tensor * v_root = V;
         while (v_root->view_src) v_root = v_root->view_src;
@@ -1182,16 +1234,15 @@ static void ggml_cuda_turbo_prefill_attend(ggml_backend_cuda_context & ctx, ggml
             }
             k_turbo2_tcq_dequant_f16<<<grid_v, V->ne[0], 0, stream>>>(
                 (const char *)V->data, v_fp16, V->ne[0], V->ne[1], V->ne[2], V->nb[1], V->nb[2], V->nb[3], tcq_compute_alpha_v(V->type, V->ne[1]));
-        } else {
+        } else if (V->type == GGML_TYPE_TURBO4_0) {
             k_turbo4_dequant_f16<<<grid_v, V->ne[0], 0, stream>>>(
                 (const char *)V->data, v_fp16, V->ne[0], V->ne[1], V->ne[2], V->nb[1], V->nb[2], V->nb[3]);
+        } else {
+            ggml_cuda_fattn_materialize_to_f16(V, v_fp16, stream, V_f16);
+            v_f16_layout_set = true;
         }
     }
 
-    // Create fp16 tensor copies on stack
-    ggml_tensor K_f16 = *K;
-    ggml_tensor V_f16 = *V;
-
     if (k_fp16) {
         K_f16.type = GGML_TYPE_F16;
         K_f16.data = k_fp16;
@@ -1201,7 +1252,7 @@ static void ggml_cuda_turbo_prefill_attend(ggml_backend_cuda_context & ctx, ggml
         K_f16.nb[3] = K->ne[0] * K->ne[1] * K->ne[2] * sizeof(half);
     }
 
-    if (v_fp16) {
+    if (v_fp16 && !v_f16_layout_set) {
         V_f16.type = GGML_TYPE_F16;
         V_f16.data = v_fp16;
         V_f16.nb[0] = sizeof(half);
@@ -1380,18 +1431,21 @@ static inline bool ggml_cuda_fattn_pair_compiled(const ggml_type type_K_in, cons
 #endif
 }
 
-static inline bool ggml_cuda_fattn_is_classic_non_q8_type(const ggml_type type) {
-    return type == GGML_TYPE_BF16 ||
-           type == GGML_TYPE_Q4_0 ||
-           type == GGML_TYPE_Q4_1 ||
-           type == GGML_TYPE_Q5_0 ||
-           type == GGML_TYPE_Q5_1 ||
-           type == GGML_TYPE_Q6_0 ||
-           type == GGML_TYPE_Q6_1 ||
-           type == GGML_TYPE_Q3_0 ||
-           type == GGML_TYPE_Q3_1 ||
-           type == GGML_TYPE_Q2_0 ||
-           type == GGML_TYPE_Q2_1;
+static inline bool ggml_cuda_fattn_prefers_native_vec_for_turbo_k_classic_v(
+        const ggml_tensor * Q, const ggml_tensor * K, const ggml_tensor * V) {
+    return ggml_cuda_fattn_is_turbo_kv_type(K->type) &&
+           !ggml_cuda_fattn_is_turbo_kv_type(V->type) &&
+           ggml_cuda_fattn_is_classic_non_q8_type(V->type) &&
+           Q->ne[0] <= 512 &&
+           Q->ne[0] % 64 == 0 &&
+           ggml_cuda_fattn_pair_compiled(K->type, V->type);
+}
+
+static inline bool ggml_cuda_fattn_prefill_mma_can_materialize_turbo_k_classic_v(
+        const ggml_tensor * K, const ggml_tensor * V) {
+    return ggml_cuda_fattn_is_turbo_kv_type(K->type) &&
+           !ggml_cuda_fattn_is_turbo_kv_type(V->type) &&
+           ggml_cuda_fattn_is_classic_non_q8_type(V->type);
 }
 
 // Shape guard for the effective K/V pair after Turbo V decode-dequant.
@@ -2317,6 +2371,8 @@ static ggml_cuda_fattn_route_plan ggml_cuda_fattn_make_route_plan(const int devi
 #endif
 
     const bool turbo_decode_native = getenv("GGML_TURBO_DECODE_NATIVE") != nullptr;
+    const bool prefer_native_vec =
+        ggml_cuda_fattn_prefers_native_vec_for_turbo_k_classic_v(Q, K, V);
 
     const bool turbo_k_only = ggml_cuda_fattn_is_turbo_kv_type(K->type);
     const bool turbo_v_only = ggml_cuda_fattn_is_turbo_kv_type(V->type);
@@ -2327,6 +2383,8 @@ static ggml_cuda_fattn_route_plan ggml_cuda_fattn_make_route_plan(const int devi
         !turbo_k_only;
 
     // Decode-dequant policy:
+    // - Compiled Turbo K + classic non-q8 V fallback routes stay native vec
+    //   instead of taking generic f16 materialization after Turbo K decode.
     // - D <= 256: Turbo K/V sides may be decoded to f16.
     // - D = 512: f16/q8_0/Turbo pairs keep the existing f16 route.
     // - D = 512 classic non-q8 K + Turbo V decodes Turbo V to f16 and keeps
@@ -2340,6 +2398,7 @@ static ggml_cuda_fattn_route_plan ggml_cuda_fattn_make_route_plan(const int devi
     plan.decode_dequant =
         !hip_native_tcq_decode &&
         !turbo_decode_native &&
+        !prefer_native_vec &&
         turbo_kv &&
         (Q->ne[0] <= 256 ||
          (Q->ne[0] <= 512 && k_f16_q8_or_turbo && v_f16_q8_or_turbo) ||
@@ -2425,6 +2484,7 @@ static ggml_cuda_fattn_route_plan ggml_cuda_fattn_make_route_plan(const int devi
             "Kshape=[%lld,%lld,%lld,%lld] Vshape=[%lld,%lld,%lld,%lld] "
             "Keff=%s Veff=%s "
             "turbo_kv=%d "
+            "prefer_native_vec=%d "
             "decode=%d dk=%d dv=%d "
             "unsafe_vec_after_turbo_v_decode=%d allow_vec=%d "
             "kernel=%s "
@@ -2438,6 +2498,7 @@ static ggml_cuda_fattn_route_plan ggml_cuda_fattn_make_route_plan(const int devi
             (long long) V->ne[0], (long long) V->ne[1], (long long) V->ne[2], (long long) V->ne[3],
             ggml_type_name(plan.effective_type_K), ggml_type_name(plan.effective_type_V),
             (int) turbo_kv,
+            (int) prefer_native_vec,
             (int) plan.decode_dequant,
             (int) plan.decode_dequant_K,
             (int) plan.decode_dequant_V,
@@ -2588,12 +2649,19 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
         return;
     }
 
+    const bool turbo_k_classic_v_prefill =
+        ggml_cuda_fattn_prefill_mma_can_materialize_turbo_k_classic_v(K, V);
+    const bool turbo_prefill_can_make_f16_K = ggml_cuda_turbo_prefill_mma_can_make_f16(K->type);
+    const bool turbo_prefill_can_make_f16_V = ggml_cuda_turbo_prefill_mma_can_make_f16(V->type);
     const bool turbo_prefill_mma_safe =
         turbo_kv &&
-        ggml_cuda_turbo_prefill_mma_can_make_f16(K->type) &&
-        ggml_cuda_turbo_prefill_mma_can_make_f16(V->type);
+        ((turbo_prefill_can_make_f16_K &&
+          turbo_prefill_can_make_f16_V) ||
+         turbo_k_classic_v_prefill);
+    const bool turbo_prefill_turing_mma =
+        turing_mma_available(ggml_cuda_info().devices[ggml_cuda_get_device()].cc);
 
-    if (turbo_prefill_mma_safe && !turbo_prefill_vec && Q->ne[1] > 1 && Q->ne[0] <= 256 && turing_mma_available(ggml_cuda_info().devices[ggml_cuda_get_device()].cc)) {
+    if (turbo_prefill_mma_safe && !turbo_prefill_vec && Q->ne[1] > 1 && Q->ne[0] <= 512 && turbo_prefill_turing_mma) {
         if (turbo_fa_debug) {
             fprintf(stderr,
                 "GGML_TURBO_FA_DEBUG: path=prefill-dequant K=%s V=%s Q=[%lld,%lld,%lld,%lld] K=[%lld,%lld,%lld,%lld] V=[%lld,%lld,%lld,%lld]\n",
@@ -2608,11 +2676,19 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
     } else {
         if (turbo_fa_debug && turbo_kv) {
             fprintf(stderr,
-                "GGML_TURBO_FA_DEBUG: path=decode-dequant-or-vec K=%s V=%s Q=[%lld,%lld,%lld,%lld] K=[%lld,%lld,%lld,%lld] V=[%lld,%lld,%lld,%lld]\n",
+                "GGML_TURBO_FA_DEBUG: path=decode-dequant-or-vec K=%s V=%s Q=[%lld,%lld,%lld,%lld] K=[%lld,%lld,%lld,%lld] V=[%lld,%lld,%lld,%lld] prefill_safe=%d can_make_f16_K=%d can_make_f16_V=%d turbo_k_classic_v_prefill=%d vec_override=%d batch_ok=%d dim_ok=%d turing_mma=%d\n",
                 ggml_type_name(K->type), ggml_type_name(V->type),
                 (long long) Q->ne[0], (long long) Q->ne[1], (long long) Q->ne[2], (long long) Q->ne[3],
                 (long long) K->ne[0], (long long) K->ne[1], (long long) K->ne[2], (long long) K->ne[3],
-                (long long) V->ne[0], (long long) V->ne[1], (long long) V->ne[2], (long long) V->ne[3]);
+                (long long) V->ne[0], (long long) V->ne[1], (long long) V->ne[2], (long long) V->ne[3],
+                (int) turbo_prefill_mma_safe,
+                (int) turbo_prefill_can_make_f16_K,
+                (int) turbo_prefill_can_make_f16_V,
+                (int) turbo_k_classic_v_prefill,
+                (int) turbo_prefill_vec,
+                (int) (Q->ne[1] > 1),
+                (int) (Q->ne[0] <= 512),
+                (int) turbo_prefill_turing_mma);
         }
         load_tcq_decode_alpha(ctx.device);
 
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -286,6 +286,7 @@ llama_build_and_test(test-dflash-plumbing.cpp ARGS ${PROJECT_SOURCE_DIR})
 target_include_directories(test-dflash-plumbing PRIVATE ${PROJECT_SOURCE_DIR}/src)
 llama_build_and_test(test-mtmd-plumbing.cpp ARGS ${PROJECT_SOURCE_DIR})
 llama_build_and_test(test-perplexity-plumbing.cpp ARGS ${PROJECT_SOURCE_DIR})
+llama_build_and_test(test-cuda-fattn-route-policy.cpp ARGS ${PROJECT_SOURCE_DIR})
 
 if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC)
   # TODO: repair known memory leaks
diff --git a/tests/test-cuda-fattn-route-policy.cpp b/tests/test-cuda-fattn-route-policy.cpp
@@ -0,0 +1,97 @@
+#include <cstdio>
+#include <cstdlib>
+#include <fstream>
+#include <sstream>
+#include <string>
+
+static std::string read_file(const std::string & path) {
+    std::ifstream file(path);
+    if (!file.good()) {
+        std::fprintf(stderr, "failed to open %s\n", path.c_str());
+        std::exit(1);
+    }
+
+    std::ostringstream ss;
+    ss << file.rdbuf();
+    return ss.str();
+}
+
+static bool expect(bool ok, const char * message) {
+    if (!ok) {
+        std::fprintf(stderr, "%s\n", message);
+    }
+    return ok;
+}
+
+static std::string slice_between(const std::string & text, const std::string & begin, const std::string & end) {
+    const size_t b = text.find(begin);
+    if (b == std::string::npos) {
+        return {};
+    }
+    const size_t e = text.find(end, b);
+    if (e == std::string::npos) {
+        return text.substr(b);
+    }
+    return text.substr(b, e - b);
+}
+
+int main(int argc, char ** argv) {
+    bool ok = true;
+
+    ok &= expect(argc == 2, "expected repo root argument");
+    if (!ok) {
+        return 1;
+    }
+
+    const std::string root = argv[1];
+    const std::string fattn = read_file(root + "/ggml/src/ggml-cuda/fattn.cu");
+    const std::string helper = slice_between(fattn,
+            "static inline bool ggml_cuda_fattn_prefers_native_vec_for_turbo_k_classic_v",
+            "struct ggml_cuda_fattn_route_plan");
+    const std::string planner = slice_between(fattn,
+            "static ggml_cuda_fattn_route_plan ggml_cuda_fattn_make_route_plan",
+            "size_t ggml_cuda_flash_attn_ext_get_alloc_size");
+    const std::string prefill = slice_between(fattn,
+            "static void ggml_cuda_turbo_prefill_attend",
+            "#define FATTN_VEC_CASE");
+    const std::string prefill_policy = slice_between(fattn,
+            "static inline bool ggml_cuda_fattn_prefill_mma_can_materialize_turbo_k_classic_v",
+            "// Shape guard for the effective K/V pair after Turbo V decode-dequant.");
+    const std::string exec = slice_between(fattn,
+            "void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst)",
+            "bool ggml_cuda_flash_attn_ext_support");
+
+    ok &= expect(!helper.empty(),
+        "CUDA FA routing must have an explicit native vec policy for Turbo K + classic V pairs");
+    ok &= expect(helper.find("ggml_cuda_fattn_is_turbo_kv_type(K->type)") != std::string::npos &&
+                 helper.find("ggml_cuda_fattn_is_turbo_kv_type(V->type)") != std::string::npos,
+        "native vec policy must distinguish Turbo K from Turbo V");
+    ok &= expect(helper.find("ggml_cuda_fattn_is_classic_non_q8_type(V->type)") != std::string::npos,
+        "native vec policy must be limited to classic non-q8 V types");
+    ok &= expect(helper.find("ggml_cuda_fattn_pair_compiled(K->type, V->type)") != std::string::npos,
+        "native vec policy must require a compiled raw K/V vec pair");
+    ok &= expect(helper.find("Q->ne[0] <= 512") != std::string::npos &&
+                 helper.find("Q->ne[0] % 64 == 0") != std::string::npos,
+        "native vec policy must preserve the vec kernel head-dimension guard");
+
+    ok &= expect(planner.find("const bool prefer_native_vec =") != std::string::npos &&
+                 planner.find("ggml_cuda_fattn_prefers_native_vec_for_turbo_k_classic_v(Q, K, V)") != std::string::npos,
+        "route planner must compute the Turbo K + classic V native vec preference");
+    ok &= expect(planner.find("!prefer_native_vec &&") != std::string::npos,
+        "decode-dequant policy must not override preferred native vec Turbo K + classic V pairs");
+
+    ok &= expect(prefill.find("classic_non_q8_v") != std::string::npos &&
+                 prefill.find("ggml_cuda_fattn_materialize_to_f16(V, v_fp16, stream, V_f16)") != std::string::npos,
+        "Turbo K + classic non-q8 V prefill must materialize classic V to f16 for the MMA path");
+    ok &= expect(exec.find("turbo_k_classic_v_prefill") != std::string::npos &&
+                 exec.find("ggml_cuda_fattn_prefill_mma_can_materialize_turbo_k_classic_v(K, V)") != std::string::npos,
+        "Turbo K + classic non-q8 V batch prefill must be eligible for the prefill MMA path");
+    ok &= expect(exec.find("Q->ne[0] <= 512") != std::string::npos,
+        "Turbo K + classic non-q8 V batch prefill must cover Gemma D512 layers");
+    ok &= expect(prefill_policy.find("ggml_cuda_fattn_is_turbo_kv_type(K->type) &&") != std::string::npos &&
+                 prefill_policy.find("!ggml_cuda_fattn_is_turbo_kv_type(V->type)") != std::string::npos &&
+                 prefill_policy.find("ggml_cuda_fattn_is_classic_non_q8_type(V->type)") != std::string::npos,
+        "Turbo K + classic V prefill eligibility must not broaden classic-K/Turbo-V routing");
+
+    return ok ? 0 : 1;
+}
diff --git a/tests/test-dflash-plumbing.cpp b/tests/test-dflash-plumbing.cpp
@@ -410,7 +410,7 @@ int main(int argc, char ** argv) {
         "CUDA FlashAttention all-quant dispatch must include D=512 TCQ mixed q8/turbo3 pairs");
     ok &= expect(cuda_fattn.find("hip_native_tcq_decode") != std::string::npos &&
                  cuda_fattn.find("#if defined(GGML_USE_HIP)") != std::string::npos &&
-                 contains_ws(cuda_fattn, "!hip_native_tcq_decode && !turbo_decode_native && turbo_kv"),
+                 contains_ws(cuda_fattn, "!hip_native_tcq_decode && !turbo_decode_native && !prefer_native_vec && turbo_kv"),
         "HIP TCQ decode must stay on the native VEC path instead of dequantizing into generic tile/MMA FlashAttention");
     ok &= expect(cuda_fattn.find("turbo_mma_fused && turbo_mma_supported && Q->ne[1] <= 4") != std::string::npos &&
                  cuda_fattn.find("K->type == GGML_TYPE_TURBO4_0 ||") != std::string::npos &&
