From: opencoti Subject: [PATCH 0080] Gemma-4 assistant-MTP D512/gqa<=2 flash-attn — real f16 VEC kernel (#491+#493, bug-567) The Gemma-4 E4B assistant-MTP drafter's single GLOBAL (non-SWA) layer is head_count 4 / head_count_kv 2 -> gqa_ratio 2 at head_dim 512. On the 0.10.3 base NEITHER the MMA (switch_ncols2, fattn.cu) NOR the TILE (switch_ncols2, fattn-tile.cuh) D512 path instantiates ncols2 in {1,2} (gated DV<=256 for binary size), so a D512/gqa<=2 flash-attn op aborts at runtime (bug-567). Every other Gemma-4 config takes ncols2 in {4,8} and is unaffected (A4B/12B/31B and the E4B base are gqa>=4; A4B globals are head16/kv2 -> gqa 8). #491 (host stop-gap): a guard in build_attn_mha (llama-graph.cpp) routed the D512/gqa<=2 layer to the EXISTING non-FA explicit-softmax branch. Correct but not a real flash-attn — "fa off" on the hot path. #493 (this patch, the real fix): add a REAL D512 f16 flash-attn via the VEC kernel, and narrow the #491 guard so f16/f32 K/V take it; the guard now fires ONLY for the residual with no D512 kernel (quantized / bf16 K/V), which the assistant drafter never hits. - template-instances/fattn-vec-instance-f16-f16.cu + generate_cu_files.py: DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_F16) instance (the autogen line + the file). - ggml-cuda/fattn-vec.cuh: extern decl for the D512 f16 instance. - ggml-cuda/fattn.cu: (a) selector route in ggml_cuda_get_best_fattn_kernel — D512 && gqa_ratio<=2 && gqa_opt_applies && K/V f16 -> BEST_FATTN_KERNEL_VEC; (b) the matching VEC dispatch case FATTN_VEC_CASE(512, F16, F16) in ggml_cuda_flash_attn_ext_vec, BEFORE the GGML_ABORT. NOTE: fattn.cu has TWO independent head-size switches (selector + VEC dispatch); both need the D512 case or the selector picks VEC but the dispatch falls through to a decode-time abort. - llama.cpp/src/llama-graph.cpp: narrow the #491 guard — fa_no_gpu_kernel fires only when K/V are NOT f16/f32 (fa_d512_f16_vec), so f16 D512/gqa<=2 now flash-attns on-GPU. The VEC kernel tiles Q->ne[1] (cols_per_block 1 decode / 2 prefill+verify) so the one f16 D512 instance is a real flash-attn at ANY n_q (covers decode, prefill, and the spec-verify batch). logit_softcap is forced off for D!=128,256 in the VEC kernel; safe because Gemma-4 attn_soft_cap is false (softcap 0 -> use_logit_softcap false). Gate (.opencoti/m493-fa-d512-gate.sh): E4B -fa on assistant-MTP decodes 147.86 tps with NO abort; the verified-greedy output is BYTE-IDENTICAL to the #491 non-FA soft_max path and draft acceptance is unchanged (0.67910 both), so the new VEC FA is numerically correct (exact spec-decode validates every token, and a wrong FA would collapse acceptance). A4B (gqa 8, untouched MMA path) regresses nowhere. Both DSO paths rebuilt byte-identical (sha256 a8ddcd82). Marker convention: every changed file carries an `opencoti #491`/`opencoti #493` marker; see docs/protocols/UPSTREAM_SYNC.md. All hunks are nested (llama.cpp/); no outer-vendor delta. diff --git a/llama.cpp/src/llama-graph.cpp b/llama.cpp/src/llama-graph.cpp --- a/llama.cpp/src/llama-graph.cpp +++ b/llama.cpp/src/llama-graph.cpp @@ -2122,7 +2122,29 @@ ggml_tensor * cur; - const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr; + // opencoti F5 M6-S4 mtp (#441/#491): head_dim 512 has NO GPU flash-attn kernel for gqa_ratio ≤ 2. + // Both the MMA (ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2, fattn.cu) and TILE + // (launch_fattn_tile_switch_ncols2, fattn-tile.cuh) D512 paths only instantiate ncols2 ∈ {4,8} + // (gqa_ratio > 2); the ncols2 ∈ {1,2} variants are gated out for binary size (DV<=256 only), so a + // D512/gqa≤2 flash-attn op aborts at runtime (fattn.cu:132 / fattn-tile.cuh:1322). That is exactly + // the Gemma-4 E4B assistant-MTP drafter's single GLOBAL layer: head_count 4 / head_count_kv 2 → + // gqa_ratio 2. Route it through the EXISTING non-FA explicit-softmax branch below (any head_dim/gqa + // — see the "non-flash-attn branch" fallback note at the top of this fn). q/k here are post-permute + // (flash-attn layout), so q->ne[2]/k->ne[2] == the CUDA kernel's gqa_ratio and k->ne[0] == head_dim. + // NARROW: every working D512 model (A4B/E2B/12B/31B + the E4B base, gqa_ratio ≥ 4) keeps flash-attn, + // byte-unchanged; only the gqa≤2 case is special. + // opencoti #493 (F5 M6-S4): D512/gqa≤2 with f16 (or f32→f16-cast) K AND V now has a REAL GPU + // flash-attn — the f16 D512 VEC instance, which tiles Q->ne[1] (cols_per_block 1 decode / 2 + // prefill+verify) so it computes at ANY n_q (see ggml_cuda_get_best_fattn_kernel, fattn.cu). So + // -fa on COMPUTES on-GPU for the E4B drafter (its KV is f16) instead of diverting to soft_max; no + // global fa-off, no non-FA fallback for the real path. The guard now fires ONLY for the residual + // with no D512/gqa≤2 kernel: QUANTIZED or bf16 K/V (the VEC D512 instance is f16-only; MMA/TILE + // abort, bug-567) — an exotic config the assistant drafter never hits. f16/f32 K/V → flash-attn. + const bool fa_d512_f16_vec = (k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_F32) && + (v->type == GGML_TYPE_F16 || v->type == GGML_TYPE_F32); + const bool fa_no_gpu_kernel = k->ne[0] == 512 && k->ne[2] > 0 && (q->ne[2] / k->ne[2]) <= 2 && + !fa_d512_f16_vec; + const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr && !fa_no_gpu_kernel; if (use_flash_attn) { GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet"); diff --git a/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/llama.cpp/ggml/src/ggml-cuda/fattn.cu --- a/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -360,6 +360,13 @@ FATTN_VEC_CASE(512, GGML_TYPE_TURBO2_0, GGML_TYPE_TURBO2_0) FATTN_VEC_CASE(512, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0) + // opencoti #493 (F5 M6-S4): head_dim 512 f16 VEC dispatch — the runtime counterpart to the + // selector's D512/gqa<=2 f16 VEC route (ggml_cuda_get_best_fattn_kernel) + the f16-f16 D512 + // instance. Without this case the selector picks VEC but the dispatch falls through to the abort + // below (bug-567 reappeared as a *decode-time* fatal at fattn.cu:363). The macro accepts F16 and + // F32→F16 for both K and V, matching the graph guard. Gemma-4 E4B assistant-MTP drafter global layer. + FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_F16) + GGML_ABORT("fatal error"); } @@ -513,6 +520,20 @@ return (can_use_vector_kernel || turbo512_vec) ? BEST_FATTN_KERNEL_VEC : BEST_FATTN_KERNEL_NONE; } + // opencoti #493 (F5 M6-S4): D512 (Gemma-4 GLOBAL layer) at gqa_ratio <= 2 — the E4B assistant-MTP + // drafter (head_count 4 / head_count_kv 2 → gqa 2) — has NO MMA/TILE kernel: switch_ncols2 aborts + // for DKQ>256 && gqa<=2 (the ncols2 ∈ {1,2} D512 instances are if-constexpr'd out, bug-567). The + // f16 D512 VEC instance (#493) tiles Q->ne[1] (cols_per_block 1 for decode, 2 for prefill/verify — + // ggml_cuda_flash_attn_ext_vec_case) so it is a REAL flash-attn at ANY n_q. Route f16/f16 D512 here + // when gqa<=2 so -fa on COMPUTES instead of diverting to the #491 non-FA host guard. D512/gqa>2 + // (A4B / 12B / 31B / E4B base) keeps the faster MMA ncols2 ∈ {4,8} path — byte-unchanged. Past the + // case-512 check above, gqa_opt_applies already holds (so gqa_ratio>=2); gqa_ratio<=2 ⇒ exactly 2. + if (Q->ne[0] == 512 && gqa_ratio <= 2 && gqa_opt_applies && + K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && + K->ne[1] % FATTN_KQ_STRIDE == 0) { + return BEST_FATTN_KERNEL_VEC; + } + // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel) { diff --git a/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh b/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh --- a/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh @@ -631,3 +631,8 @@ // layers (key_length=512). Definitions in the turbo2/turbo3 instance .cu files. turbo4 stays ≤256. extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_TURBO2_0, GGML_TYPE_TURBO2_0); extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); + +// opencoti #493 (F5 M6-S4): head_dim 512 f16 VEC — Gemma-4 GLOBAL layer at gqa_ratio<=2 (the E4B +// assistant-MTP drafter), where MMA/TILE abort (DKQ>256 && gqa<=2, bug-567). The VEC kernel tiles +// Q->ne[1] so it is a real flash-attn at ANY n_q. Definition in fattn-vec-instance-f16-f16.cu. +extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_F16); diff --git a/llama.cpp/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu b/llama.cpp/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu --- a/llama.cpp/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu +++ b/llama.cpp/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu @@ -5,3 +5,4 @@ DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16); +DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_F16); // opencoti #493: Gemma-4 GLOBAL head_dim 512 f16, gqa<=2 diff --git a/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py --- a/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -74,6 +74,17 @@ with open(f"fattn-vec-instance-{get_short_name(type_k)}-{get_short_name(type_v)}.cu", "w") as f: f.write(SOURCE_FATTN_VEC.format(type_k=type_k, type_v=type_v)) +# opencoti #493 (F5 M6-S4): head_dim 512 f16 VEC for Gemma-4 GLOBAL layers at gqa_ratio<=2 (the E4B +# assistant-MTP drafter's single global layer is head_count 4 / head_count_kv 2 → gqa 2). The MMA and +# TILE D512 paths only instantiate ncols2 ∈ {4,8} (gqa>2); ncols2 ∈ {1,2} are if-constexpr'd out, so a +# D512/gqa<=2 op aborts at runtime (fattn.cu switch_ncols2, bug-567). The VEC kernel tiles Q->ne[1] +# (cols_per_block 1 for decode, 2 for prefill/verify — ggml_cuda_flash_attn_ext_vec_case) so it is a +# REAL flash-attn at ANY n_q for D512. Mirrors #411's turbo512 but for plain f16 K/V → -fa on COMPUTES +# instead of falling back to the #491 non-FA host guard. f16/f16 only (drafter KV is f16); quantized +# D512/gqa<=2 keeps the guard (no VEC D512 quant instance). +with open("fattn-vec-instance-f16-f16.cu", "a") as f: + f.write("DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_F16); // opencoti #493: Gemma-4 GLOBAL head_dim 512 f16, gqa<=2\n") + # opencoti-hook: turboquant perf-lift (Level B) — F5 M6-S2 #392. Fused turbo K/V vec instances # (K==V same tier only; the in-register readers live in fattn-turbo.cuh). for type_t in ["GGML_TYPE_TURBO2_0", "GGML_TYPE_TURBO3_0", "GGML_TYPE_TURBO4_0"]: