mikhail/llama.cpp-mtp-turboquant
1
0
Fork 0
Code Issues Pull Requests Actions 30 Packages Projects Releases Wiki Activity

merge: TurboQuant KV cache types onto upstream master (MTP built-in)

Browse Source
This commit is contained in:
root
2026-05-18 01:52:24 +03:00
parent 726704a160
commit b34ac655f3
72 changed files with 18946 additions and 257 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-graph.cpp
+119
View File
@@ -1958,6 +1958,10 @@ ggml_tensor * llm_graph_context::build_attn_mha(
k = ggml_permute(ctx0, k, 0, 2, 1, 3);
v = ggml_permute(ctx0, v, 0, 2, 1, 3);
// TurboQuant note: graph-side Q rotation (pre-rotate-queries) is implemented below
// in the flash-attn path. The VEC kernel bug (wrong Q/K stride in
// vec_dot_fattn_vec_KQ_turbo3_0) was fixed in fattn-common.cuh to match f16 pattern.
ggml_tensor * cur;
const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr;
@@ -1984,6 +1988,20 @@ ggml_tensor * llm_graph_context::build_attn_mha(
ggml_flash_attn_ext_add_sinks(cur, sinks);
ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);
// TurboQuant: inverse WHT on FA output when V values are WHT-rotated.
// For MLA, V is a view of K with different ne[0] (e.g. V=512, K=576).
// Group size must come from K (which determines the WHT rotation), not V.
if (v->type == GGML_TYPE_TURBO3_0 || v->type == GGML_TYPE_TURBO4_0 || v->type == GGML_TYPE_TURBO2_0) {
const bool k_is_turbo = (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0);
const ggml_tensor * group_src = k_is_turbo ? k : v;
const int turbo_group = (group_src->ne[0] % 128 == 0) ? 128 : 64;
if (cur->ne[0] % turbo_group == 0) {
if (!ggml_is_contiguous(cur)) { cur = ggml_cont(ctx0, cur); }
ggml_tensor * innerq_scale = mctx ? mctx->get_turbo_innerq_scale_inv() : nullptr;
cur = ggml_turbo_wht(ctx0, cur, 1, turbo_group, innerq_scale); // 1 = inverse
}
}
if (v_mla) {
#if 0
// v_mla can be applied as a matrix-vector multiplication with broadcasting across dimension 3 == n_tokens.
@@ -2050,6 +2068,18 @@ ggml_tensor * llm_graph_context::build_attn_mha(
ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq);
cb(kqv, "kqv", il);
// TurboQuant: inverse WHT on attention output (non-FA path)
if (v->type == GGML_TYPE_TURBO3_0 || v->type == GGML_TYPE_TURBO4_0 || v->type == GGML_TYPE_TURBO2_0) {
const bool k_is_turbo = (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0);
const ggml_tensor * group_src = k_is_turbo ? k : v;
const int turbo_group = (group_src->ne[0] % 128 == 0) ? 128 : 64;
if (kqv->ne[0] % turbo_group == 0) {
if (!ggml_is_contiguous(kqv)) { kqv = ggml_cont(ctx0, kqv); }
ggml_tensor * innerq_scale = mctx ? mctx->get_turbo_innerq_scale_inv() : nullptr;
kqv = ggml_turbo_wht(ctx0, kqv, 1, turbo_group, innerq_scale);
}
}
// for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA
if (v_mla) {
kqv = ggml_mul_mat(ctx0, v_mla, kqv);
@@ -2067,6 +2097,8 @@ ggml_tensor * llm_graph_context::build_attn_mha(
}
}
// TurboQuant: graph-side inverse WHT on attention output (undoes V rotation)
ggml_build_forward_expand(gf, cur);
return cur;
@@ -2227,12 +2259,44 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
// TurboQuant pre-rotate-queries: O(d log d) WHT rotation via custom op
// Q shape: (n_embd_head, n_head, n_tokens)
// For zero-padded models (head_dim not 128-aligned), pad Q to match padded K dim first.
if (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0) {
// Pad Q per-head to next multiple of 128 if needed
if (q->ne[0] % 128 != 0) {
const int64_t pad = ((q->ne[0] + 127) / 128) * 128 - q->ne[0];
q = ggml_pad(ctx0, q, pad, 0, 0, 0);
}
if (!ggml_is_contiguous(q)) { q = ggml_cont(ctx0, q); }
ggml_tensor * innerq_scale = mctx_cur->get_turbo_innerq_scale_inv();
q = ggml_turbo_wht(ctx0, q, 0, 0, innerq_scale); // 0 = forward, 0 = auto group size from q->ne[0]
}
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
cb(cur, "kqv_out", il);
if (inp->self_v_rot) {
cur = ggml_mul_mat_aux(ctx0, cur, inp->self_v_rot);
}
// TurboQuant: if V was padded, the output has padded dimensions.
// Extract original V head_dim after inverse WHT (applied inside build_attn_mha).
if (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0) {
const int64_t orig_v_head = hparams.n_embd_head_v(il);
// cur is 2D: (n_embd_head * n_head, n_tokens) after build_attn_mha
const int64_t padded_v_head = v->ne[0];
if (padded_v_head != orig_v_head) {
// Reshape to 4D, extract original head_dim, reshape back to 2D
const int64_t n_head_v = hparams.n_head_kv(il);
const int64_t n_tokens_cur = cur->ne[1];
cur = ggml_reshape_3d(ctx0, cur, padded_v_head, n_head_v, n_tokens_cur);
// ggml_view_3d to extract first orig_v_head elements per head
cur = ggml_view_3d(ctx0, cur, orig_v_head, n_head_v, n_tokens_cur,
cur->nb[1], cur->nb[2], 0);
cur = ggml_cont(ctx0, cur);
cur = ggml_reshape_2d(ctx0, cur, orig_v_head * n_head_v, n_tokens_cur);
}
}
if (wo) {
if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) {
@@ -2318,9 +2382,39 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
// TurboQuant: pre-rotate Q for K-only (MLA) attention
// For zero-padded models, pad Q to match padded K dim first.
if (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0) {
// Pad Q per-head to next multiple of 128 if needed
if (q->ne[0] % 128 != 0) {
const int64_t pad = ((q->ne[0] + 127) / 128) * 128 - q->ne[0];
q = ggml_pad(ctx0, q, pad, 0, 0, 0);
}
if (!ggml_is_contiguous(q)) { q = ggml_cont(ctx0, q); }
ggml_tensor * innerq_scale = mctx_cur->get_turbo_innerq_scale_inv();
q = ggml_turbo_wht(ctx0, q, 0, 0, innerq_scale); // 0 = forward, 0 = auto group size
}
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
cb(cur, "kqv_out", il);
// TurboQuant: if V was padded (MLA: V is view of K, may have padded dim),
// extract original V head_dim after inverse WHT.
if (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0) {
const int64_t orig_v_head = v_cur->ne[0]; // original V head_dim from model
const int64_t padded_v_head = v->ne[0]; // padded V head_dim in cache
if (padded_v_head != orig_v_head) {
// cur is 2D: (padded_v_head * n_head, n_tokens) after build_attn_mha
const int64_t n_head_v = hparams.n_head_kv(il);
const int64_t n_tokens_cur = cur->ne[1];
cur = ggml_reshape_3d(ctx0, cur, padded_v_head, n_head_v, n_tokens_cur);
cur = ggml_view_3d(ctx0, cur, orig_v_head, n_head_v, n_tokens_cur,
cur->nb[1], cur->nb[2], 0);
cur = ggml_cont(ctx0, cur);
cur = ggml_reshape_2d(ctx0, cur, orig_v_head * n_head_v, n_tokens_cur);
}
}
if (wo) {
if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) {
// GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators
@@ -2406,12 +2500,37 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
// TurboQuant: pre-rotate Q for ISWA attention (pad to 128-aligned if needed)
if (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0) {
if (q->ne[0] % 128 != 0) {
const int64_t pad = ((q->ne[0] + 127) / 128) * 128 - q->ne[0];
q = ggml_pad(ctx0, q, pad, 0, 0, 0);
}
if (!ggml_is_contiguous(q)) { q = ggml_cont(ctx0, q); }
ggml_tensor * innerq_scale = mctx_cur->get_turbo_innerq_scale_inv();
q = ggml_turbo_wht(ctx0, q, 0, 0, innerq_scale);
}
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
cb(cur, "kqv_out", il);
if (v_rot) {
cur = ggml_mul_mat_aux(ctx0, cur, v_rot);
}
// TurboQuant: if V was padded, extract original V head_dim after inverse WHT
if (k->type == GGML_TYPE_TURBO3_0 || k->type == GGML_TYPE_TURBO4_0 || k->type == GGML_TYPE_TURBO2_0) {
const int64_t orig_v_head = hparams.n_embd_head_v(il);
const int64_t padded_v_head = v->ne[0];
if (padded_v_head != orig_v_head) {
const int64_t n_head_v = hparams.n_head_kv(il);
const int64_t n_tokens_cur = cur->ne[1];
cur = ggml_reshape_3d(ctx0, cur, padded_v_head, n_head_v, n_tokens_cur);
cur = ggml_view_3d(ctx0, cur, orig_v_head, n_head_v, n_tokens_cur,
cur->nb[1], cur->nb[2], 0);
cur = ggml_cont(ctx0, cur);
cur = ggml_reshape_2d(ctx0, cur, orig_v_head * n_head_v, n_tokens_cur);
}
}
if (wo) {
cur = build_lora_mm(wo, cur, wo_s);