Commit 358a420ec for llama.cpp
commit 358a420ec4ed59b6b0be08b1d1f71dd478b99a99
Author: Pascal <admin@serveurperso.com>
Date: Mon Oct 12 23:00:27 2026 +0200
graph: pad the KV dimension of the no-cache attention to 256 (#30416)
The no-cache path builds a [n_tokens, n_tokens] mask and passes K and V
unpadded, so FLASH_ATTN_EXT sees K->ne[1] == n_tokens. The CUDA kernels
for head size 512 require a multiple of FATTN_KQ_STRIDE, so the op falls
back to the CPU for most lengths, and the -fa auto probe (1 token)
disables flash attention for such models.
Pad the mask like llama_kv_cache::get_n_kv and pad K and V with zeros up
to the KV dimension of the mask. clef builds its own causal mask and
pads it the same way.
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index fe3d03c02..f06b94d15 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -462,7 +462,7 @@ static void print_mask(const T * data, int64_t n_tokens, int64_t n_kv, int64_t n
}
void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {
- const int64_t n_kv = ubatch->n_tokens;
+ const int64_t n_kv = self_kq_mask->ne[0];
const int64_t n_tokens = ubatch->n_tokens;
const auto fill_mask = [&](auto * data, int64_t ne, int n_swa, llama_swa_type swa_type) {
@@ -3000,15 +3000,17 @@ llm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() con
// flash attention requires an f16 mask
const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
- // note: there is no KV cache, so the number of KV values is equal to the number of tokens in the batch
- inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
+ // note: there is no KV cache, so the KV values are the tokens of the batch, padded like the KV cache (see llama_kv_cache::get_n_kv)
+ const int64_t n_kv = GGML_PAD(n_tokens, 256);
+
+ inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1);
ggml_set_input(inp->self_kq_mask);
cb(inp->self_kq_mask, "self_kq_mask", -1);
inp->self_kq_mask_cnv = inp->self_kq_mask;
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
- inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
+ inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1);
ggml_set_input(inp->self_kq_mask_swa);
inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
@@ -3054,6 +3056,13 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = k_cur;
ggml_tensor * v = v_cur;
+ // the KV values are padded with zeros up to the KV dimension of the mask
+ const int64_t n_kv_pad = kq_mask->ne[0] - k->ne[2];
+ if (n_kv_pad > 0) {
+ k = ggml_pad(ctx0, k, 0, 0, n_kv_pad, 0);
+ v = ggml_pad(ctx0, v, 0, 0, n_kv_pad, 0);
+ }
+
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
diff --git a/src/llama-graph.h b/src/llama-graph.h
index e15d78450..fbbe67f25 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -322,11 +322,11 @@ public:
ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
- // n_tokens == n_batch
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_tokens, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_tokens, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]
+ // n_tokens == n_batch, n_kv == n_tokens padded to 256
+ ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
const llama_hparams hparams;
const llama_cparams cparams;
diff --git a/src/models/clef.cpp b/src/models/clef.cpp
index 81e4ad65e..cd2889b22 100644
--- a/src/models/clef.cpp
+++ b/src/models/clef.cpp
@@ -320,13 +320,14 @@ public:
using llm_graph_input_attn_no_cache::llm_graph_input_attn_no_cache;
void set_input(const llama_ubatch * ubatch) override {
+ const int64_t n_kv = self_kq_mask->ne[0];
const int64_t n_tokens = ubatch->n_tokens;
const auto fill_mask = [&](auto * data, auto zero, auto ninf) {
for (int64_t i1 = 0; i1 < n_tokens; ++i1) {
- for (int64_t i0 = 0; i0 < n_tokens; ++i0) {
+ for (int64_t i0 = 0; i0 < n_kv; ++i0) {
const bool visible = i0 <= i1 && ubatch->seq_id[i0][0] == ubatch->seq_id[i1][0];
- data[i1 * n_tokens + i0] = visible ? zero : ninf;
+ data[i1 * n_kv + i0] = visible ? zero : ninf;
}
}
};
@@ -346,7 +347,7 @@ llm_graph_input_attn_no_cache * llama_model_clef::graph::build_attn_inp_causal()
const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
- inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
+ inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, GGML_PAD(n_tokens, 256), n_tokens, 1, 1);
ggml_set_input(inp->self_kq_mask);
cb(inp->self_kq_mask, "self_kq_mask", -1);