From 925792c693d7254d93a02d35786653206976982c Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 14:17:45 -0600 Subject: [PATCH 01/16] qwen4exp: catch up to upstream through 36b101543 Ports every upstream qwen4exp (Qwen3.8-Flash-Next) commit from the past 10 days that this fork's manual PR port had not received: - reduce graph splits by hoisting the PLE embedding gather out of the per-layer loop (#27880) - sum indexer heads via strided adds instead of transpose+sum_rows (#28023) - support recurrent state rollback for MTP speculative decoding (#28123) - rewrite QSA sparse-attention block/bias selection: fixes NaN-producing bias rows for short sequences, fixes cross-sequence block pooling in a unified KV cache, adds mrope duplicate-position ranking, and fixes a CUDA rms_norm gridDim.y overflow (#27941) - indexer cache seq_cp staleness fix, ext.x/ext.y state-restore fix, PLE-must-be-linear-attention validation, correct -sm tensor disablement (#27941) - Hadamard k_rot context-shift crash fix, shared with other archs (#27967) Also replaces raw GGML_ASSERT aborts in hparams loading with proper error messages, and adds test coverage: a PLE fixture in test-llama-archs (which required porting the per_layer_token_embd row-count-from-metadata fix to make it loadable) and a state round-trip test in test-save-load-state. Verified against the real Qwen3.8-Flash-Next model: correct generation at short and long (~66k token) context, and test-llama-archs passes qwen4exp on both CUDA and CPU. --- src/llama-arch.cpp | 2 + src/llama-kv-cache.cpp | 18 +- src/llama-kv-cells.h | 11 +- src/llama-memory-hybrid-idx.cpp | 293 ++++++++++++++++++++++++++++---- src/llama-memory-hybrid-idx.h | 3 +- src/models/models.h | 12 +- src/models/qwen4exp.cpp | 261 ++++++++++++++++++++++------ tests/test-llama-archs.cpp | 24 +++ tests/test-save-load-state.cpp | 63 +++++++ 9 files changed, 591 insertions(+), 96 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index ba98fde5daf1..6df93bcd90e1 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1089,6 +1089,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { switch (arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: return true; default: @@ -1124,6 +1125,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_QWEN3TTS: + case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs return false; default: return true; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index b397abde40c6..d90c38013ad3 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2245,7 +2245,7 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) { kv_self->set_input_k_shift(k_shift); } - if (k_rot) { + if (k_rot && k_rot->buffer) { kv_self->set_input_k_rot(k_rot); } } @@ -2592,6 +2592,12 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 ubatch.seq_id_unq[0] = dest_seq_id; + // the ext as it was saved, to put back after apply_ubatch() + std::vector exts; + if (hparams.n_pos_per_embd() > 1) { + exts.resize(cell_count); + } + for (uint32_t i = 0; i < cell_count; ++i) { llama_pos pos; uint32_t n_seq_id; @@ -2610,6 +2616,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 ubatch.pos[i + ubatch.n_tokens] = ext.y; ubatch.pos[i + ubatch.n_tokens*2] = ext.x; + + exts[i] = ext; } // read the sequence id, but directly discard it - we will use dest_seq_id instead @@ -2662,6 +2670,14 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 // see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350 apply_ubatch(sinfo, ubatch); + // apply_ubatch() takes the 2D position from the ubatch, and that ubatch is built with this + // cache's own n_pos_per_embd. a cache that does not use M-RoPE itself but mirrors one that + // does (the qwen4exp QSA indexer) would drop x and y. put the saved ext back instead, which + // is what the whole-context path below already does. + for (uint32_t i = 0; i < (uint32_t) exts.size(); ++i) { + cells.ext_set(sinfo.idxs[0][i], exts[i]); + } + LLAMA_LOG_DEBUG("%s: cell_count = %d, dest_seq_id = %d\n", __func__, cell_count, dest_seq_id); // DEBUG CHECK: verify that all cells were allocated and have correct seq_id and pos values diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index fddd31a0b219..903a1aeb5c8c 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -31,6 +31,8 @@ struct llama_kv_cell_ext { // TODO: add unit tests class llama_kv_cells { public: + using seq_set_t = std::bitset; + void reset() { for (uint32_t i = 0; i < pos.size(); ++i) { pos[i] = -1; @@ -297,6 +299,13 @@ class llama_kv_cells { return seq[i].count(); } + // the full set of sequences this cell is visible to + const seq_set_t & seq_get_all(uint32_t i) const { + assert(i < pos.size()); + + return seq[i]; + } + // check if the cell contains seq_id bool seq_has(uint32_t i, llama_seq_id seq_id) const { assert(i < pos.size()); @@ -483,8 +492,6 @@ class llama_kv_cells { // std::vector shift; - using seq_set_t = std::bitset; - // the bitset seq[i] tells us which sequences are currently occupying the i-th cell std::vector seq; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 3dfaa517fdc1..31e391da9a8b 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -563,7 +563,10 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( llama_context * lctx, bool optimize) : llama_memory_hybrid_context(mem, lctx, optimize), - mem(mem) {} + mem(mem), + // update() applies a pending cross-stream seq_cp, else the copy keeps stale indexer keys + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + mem->get_mem_idx()->init_update(lctx, optimize)) {} llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( llama_memory_hybrid_idx * mem, @@ -619,7 +622,8 @@ void llama_memory_hybrid_idx_context::set_input_qsa( ggml_tensor * blk_pos, ggml_tensor * bias, const llama_ubatch * ubatch, - uint32_t ratio) const { + uint32_t ratio, + bool blk_bias) const { GGML_ASSERT(ratio > 0); GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); @@ -639,19 +643,29 @@ void llama_memory_hybrid_idx_context::set_input_qsa( int32_t * dst_blk_pos = (int32_t *) blk_pos->data; float * dst_bias = (float *) bias->data; - // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio - // all mrope sections carry it: exact for text, approximate for images - for (int64_t sec = 0; sec < 4; ++sec) { - for (int64_t s = 0; s < n_ns; ++s) { - for (int64_t b = 0; b < n_blocks; ++b) { - dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r); - } - } - } - - // one pass per stream: cell j is a different token in each, so no mapping is shared - std::vector blk_of(n_kv); - std::vector filled(n_blocks); + // a block is keyed on (sequence set, index bucket): a unified cache counts every sequence + // from zero, so the bucket alone would pool two sequences into one block + GGML_ASSERT(r <= 64); + const uint64_t slots_full = r == 64 ? ~uint64_t(0) : ((uint64_t(1) << r) - 1); + + // TODO: this runs per ubatch and is O(n_kv) per stream, about 865 us at 33k context. the cost + // is the per-cell scan rather than these allocations, so hoisting them buys nothing + std::vector blk_of(n_kv); + std::vector cell_grp(n_kv); + std::vector grp_head(n_blocks); + std::vector grp_next; + std::vector grp_first; + std::vector grp_slot0; + std::vector grp_slots; + std::vector grp_bid; + std::vector bid_idx; + std::vector bid_cell; + std::vector bid_slot0; + + std::vector order; + std::vector rank; + + std::fill(dst_blk_pos, dst_blk_pos + 4*n_blocks*n_ns, 0); for (int64_t s = 0; s < n_ns; ++s) { // ubatch index s*n_tps belongs to this stream; ask which cells array it uses @@ -661,52 +675,257 @@ void llama_memory_hybrid_idx_context::set_input_qsa( int32_t * cur_cell_blk = dst_cell_blk + s*n_kv; int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks); - // an incomplete block cannot be pooled; the bias below forces those tail cells in - // -1 means no usable block, and block 0 only keeps the gather in range - std::fill(blk_of.begin(), blk_of.end(), -1); - std::fill(filled.begin(), filled.end(), 0); std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0); - for (int64_t j = 0; j < n_kv; ++j) { - if (cells.is_empty(j)) { - continue; + bid_idx .clear(); + bid_cell .clear(); + bid_slot0.clear(); + + int n_seq_present = 0; + + for (int sq = 0; sq < LLAMA_MAX_SEQ && n_seq_present < 2; ++sq) { + if (cells.seq_pos_min(sq) >= 0) { + n_seq_present++; } + } - const llama_pos p = cells.pos_get(j); - const int64_t b = p/r; + const bool one_seq = n_seq_present <= 1; - if (b >= n_blocks) { - continue; + // a cell no block covers needs its own -inf, which a per-block bias cannot carry + // every cache path keeps the position below the cell window, so this stays false + bool oor = false; + + bool dup = false; + + bool ranked = false; + + auto group_cells = [&]() { + // -1 means no usable block: an incomplete or short group cannot be pooled + std::fill(blk_of.begin(), blk_of.end(), -1); + std::fill(cell_grp.begin(), cell_grp.end(), -1); + std::fill(grp_head.begin(), grp_head.end(), -1); + + grp_next .clear(); + grp_first.clear(); + grp_slot0.clear(); + grp_slots.clear(); + grp_bid .clear(); + + oor = false; + dup = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + + const int64_t idx = ranked ? rank[j] : cells.pos_get(j); + const int64_t pb = idx/r; + + if (pb >= n_blocks) { + oor = true; + continue; + } + + int32_t g = -1; + + for (int32_t c = grp_head[pb]; c >= 0; c = grp_next[c]) { + if (one_seq || cells.seq_get_all((uint32_t) grp_first[c]) == cells.seq_get_all((uint32_t) j)) { + g = c; + break; + } + } + + if (g < 0) { + g = (int32_t) grp_first.size(); + + grp_next .push_back(grp_head[pb]); + grp_first.push_back((int32_t) j); + grp_slot0.push_back(-1); + grp_slots.push_back(0); + grp_bid .push_back(-1); + + grp_head[pb] = g; + } + + const uint64_t bit = uint64_t(1) << (idx%r); + + dup |= (grp_slots[g] & bit) != 0; + + cell_grp[j] = g; + grp_slots[g] |= bit; + + if (idx%r == 0) { + grp_slot0[g] = (int32_t) j; + } + } + }; + + group_cells(); + + // mrope repeats one position across an image, so rank cells instead of using the position + if (dup && ubatch->is_pos_2d() && one_seq) { + order.clear(); + order.reserve(n_kv); + + for (int64_t j = 0; j < n_kv; ++j) { + if (!cells.is_empty(j)) { + order.push_back((int32_t) j); + } + } + + // same total order the mrope causal mask uses: pos, then ext.y, then ext.x + std::sort(order.begin(), order.end(), [&cells](int32_t a, int32_t b) { + const llama_pos pa = cells.pos_get(a); + const llama_pos pb = cells.pos_get(b); + + if (pa != pb) { + return pa < pb; + } + + const auto & ea = cells.ext_get(a); + + return cells.ext_get(b).is_2d_gt(ea.x, ea.y); + }); + + rank.assign(n_kv, -1); + + for (int64_t k = 0; k < (int64_t) order.size(); ++k) { + rank[order[k]] = (int32_t) k; } - blk_of[j] = (int32_t) b; - cur_blk_cells[b*r + (p%r)] = (int32_t) j; - filled[b]++; + ranked = true; + + group_cells(); } + GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); + + int32_t n_bid = 0; + + for (int64_t pb = 0; pb < n_blocks; ++pb) { + for (int32_t g = grp_head[pb]; g >= 0; g = grp_next[g]) { + if (grp_slots[g] != slots_full) { + continue; + } + + grp_bid[g] = n_bid++; + + bid_idx .push_back((int32_t) (pb*r)); + bid_cell .push_back(grp_first[g]); + bid_slot0.push_back(grp_slot0[g]); + } + } + + GGML_ASSERT(n_bid <= n_blocks); + + for (int32_t b = 0; b < n_bid; ++b) { + int32_t sec_pos[4] = { bid_idx[b], bid_idx[b], bid_idx[b], bid_idx[b] }; + + if (ranked) { + const int32_t c = bid_slot0[b]; + const llama_pos p = cells.pos_get(c); + const auto & e = cells.ext_get(c); + + sec_pos[0] = p; + sec_pos[1] = e.y; + sec_pos[2] = e.x; + sec_pos[3] = p; + } + + for (int64_t sec = 0; sec < 4; ++sec) { + dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = sec_pos[sec]; + } + } + + // unpooled cells all point at one spare block. a spare block exists only when some + // cell is unpooled: n_bid == n_blocks means every cell sits in a full block. + const bool have_dead = n_bid < n_blocks; + const int32_t dead_bid = have_dead ? n_bid : n_blocks - 1; + for (int64_t j = 0; j < n_kv; ++j) { - if (blk_of[j] >= 0 && filled[blk_of[j]] < r) { - blk_of[j] = -1; + const int32_t g = cell_grp[j]; + + blk_of[j] = g < 0 ? -1 : grp_bid[g]; + + if (blk_of[j] >= 0) { + const int64_t idx = ranked ? rank[j] : cells.pos_get(j); + + cur_blk_cells[blk_of[j]*r + (idx%r)] = (int32_t) j; } - cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; + + cur_cell_blk[j] = blk_of[j] < 0 ? dead_bid : blk_of[j]; } for (int64_t ii = 0; ii < n_tps; ++ii) { const int64_t i = s*n_tps + ii; const llama_seq_id seq_id = ubatch->seq_id[i][0]; - const llama_pos q = ubatch->pos[i]; + + int64_t q = ubatch->pos[i]; + + if (ranked) { + const llama_pos qt = ubatch->pos[i]; + const llama_pos qy = ubatch->pos[i + n_tokens]; + const llama_pos qx = ubatch->pos[i + n_tokens*2]; + + int64_t lo = 0; + int64_t hi = (int64_t) order.size(); + + while (lo < hi) { + const int64_t mid = (lo + hi)/2; + const int32_t c = order[mid]; + const llama_pos pc = cells.pos_get(c); + + if (pc < qt || (pc == qt && !cells.ext_get(c).is_2d_gt(qx, qy))) { + lo = mid + 1; + } else { + hi = mid; + } + } + + q = lo - 1; + } // the tail is an incomplete block and is always visible, as in the reference - const llama_pos tail_start = (q + 1)/r*r; + const int64_t tail_start = (q + 1)/r*r; + + if (blk_bias) { + // a block sits wholly inside or outside the tail, so one value covers it + // the caller adds the attention mask, which drops empty, foreign and future cells + float * cur_blk_bias = dst_bias + i*n_blocks; + + for (int64_t b = 0; b < n_blocks; ++b) { + if (b >= n_bid || !cells.seq_has((uint32_t) bid_cell[b], seq_id)) { + cur_blk_bias[b] = -INFINITY; + continue; + } + + // finite, so it can never meet a -inf and produce a nan + cur_blk_bias[b] = bid_idx[b] >= tail_start ? 1e9f : 0.0f; + } + + // the spare block holds the unpooled cells, which are the incomplete tail, so + // it gets the tail value. it must stay finite: a sequence with fewer than + // `ratio` cells owns no full block, and a row of -inf only gives a nan. + if (have_dead) { + cur_blk_bias[dead_bid] = 1e9f; + } + + continue; + } float * cur_bias = dst_bias + i*n_kv; for (int64_t j = 0; j < n_kv; ++j) { float v = -INFINITY; - if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { - // finite, so it can never meet a -inf and produce a nan - v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + if (!cells.is_empty(j) && cells.seq_has(j, seq_id)) { + const int64_t idx = ranked ? rank[j] : cells.pos_get(j); + + if (idx <= q) { + // finite, so it can never meet a -inf and produce a nan + v = idx >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + } } cur_bias[j] = v; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 9bc0748a0bbe..ddccb1ee7fce 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -174,7 +174,8 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, - ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio) const; + ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, + bool blk_bias) const; private: const llama_memory_hybrid_idx * mem = nullptr; diff --git a/src/models/models.h b/src/models/models.h index 995ff88aad04..ddbba9fa2e3a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2125,6 +2125,8 @@ struct llama_model_qwen35 : public llama_model_base { }; +class llm_graph_input_qsa; + struct llama_model_qwen4exp : public llama_model_base { llama_model_qwen4exp(const struct llama_model_params & params) : llama_model_base(params) {} @@ -2168,11 +2170,16 @@ struct llama_model_qwen4exp : public llama_model_base { float kq_scale, int il); + // the QSA cache layout inputs do not depend on the layer, only on its compress ratio, + // so the layers sharing a ratio share one input set + std::map qsa_inps; + // QSA: token indices this layer's queries may attend to, or nullptr for dense ggml_tensor * build_qsa_top_k( const llama_memory_hybrid_idx_context * mctx_hyb, ggml_tensor * cur, ggml_tensor * inp_pos, + ggml_tensor * kq_mask, int * sections, int il); @@ -2201,9 +2208,12 @@ struct llama_model_qwen4exp : public llama_model_base { int64_t channels, int il); + ggml_tensor * build_inp_ple( + const llama_memory_hybrid_idx_context * mctx_hyb); + ggml_tensor * build_ple( llm_graph_input_rs * inp, - const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * emb, ggml_tensor * hidden, int il); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 13da092416f1..335d139d9351 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -4,6 +4,24 @@ #include "llama-memory-recurrent.h" #include +#include + +// bad metadata must be catchable: GGML_ASSERT aborts the whole process +static void qwen4exp_require_nonzero(const llama_model_loader & ml, llm_kv kid, uint32_t value) { + if (value == 0) { + throw std::runtime_error(format("%s must be greater than zero, got %u", ml.llm_kv(kid).c_str(), value)); + } +} + +// get_arr() copies a short array as-is, leaving a zero tail the n-gram hash silently drops +static void qwen4exp_require_arr_len(llama_model_loader & ml, llm_kv kid, uint32_t n_min) { + uint32_t n_arr = 0; + ml.get_arr_n(kid, n_arr, true); + if (n_arr < n_min) { + throw std::runtime_error(format("%s has %u entries, but at least %u are required", + ml.llm_kv(kid).c_str(), n_arr, n_min)); + } +} void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); @@ -17,17 +35,29 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + qwen4exp_require_nonzero(ml, LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + qwen4exp_require_nonzero(ml, LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); + qwen4exp_require_nonzero(ml, LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); + qwen4exp_require_nonzero(ml, LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); + qwen4exp_require_nonzero(ml, LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank) ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); - GGML_ASSERT(hparams.dsv4_hc_mult > 0 && "qwen4exp needs a hyper-connection count"); - GGML_ASSERT(hparams.hc_low_rank > 0 && "qwen4exp needs a hyper-connection low rank"); + // a count of 1 has nothing to mix + if (hparams.dsv4_hc_mult <= 1) { + throw std::runtime_error(format("%s must be greater than one, got %u", + ml.llm_kv(LLM_KV_HYPER_CONNECTION_COUNT).c_str(), hparams.dsv4_hc_mult)); + } + qwen4exp_require_nonzero(ml, LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false); std::fill(hparams.is_ple_impl.begin(), hparams.is_ple_impl.end(), 0); @@ -38,8 +68,15 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { if (n_ple > 0) { std::vector ple_layers; ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers); + if (n_ple != 1) { + // hparams holds one set of hash constants, so several PLE modules cannot be represented + throw std::runtime_error(format("%s lists %u layers, but only one PLE layer is supported", + ml.llm_kv(LLM_KV_PLE_LAYERS).c_str(), n_ple)); + } for (uint32_t il : ple_layers) { - GGML_ASSERT(il < hparams.n_layer_all); + if (il >= hparams.n_layer_all) { + throw std::runtime_error(format("PLE layer %u is out of range", il)); + } hparams.is_ple_impl[il] = 1; } @@ -50,12 +87,19 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { // optional: files written before this key fall back to the EOS token ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false); ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); + qwen4exp_require_nonzero(ml, LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel); + qwen4exp_require_nonzero(ml, LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; hparams.ple_head_dim = hparams.n_embd_per_layer; GGML_ASSERT(hparams.ple_ngram_size >= 2 && hparams.ple_ngram_size <= LLAMA_MAX_PLE_NGRAM); GGML_ASSERT(hparams.ple_n_heads > 0 && hparams.ple_n_heads <= LLAMA_MAX_PLE_HEADS); + // get_arr() copies a short array as-is, leaving a zero tail the n-gram hash silently drops + qwen4exp_require_arr_len(ml, LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_ngram_size); + qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_OFFSETS, hparams.ple_n_heads); + qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_VOCAB_SIZES, hparams.ple_n_heads); + ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, hparams.ple_head_offsets); ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, hparams.ple_head_vocab_sizes); @@ -65,11 +109,19 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { uint32_t full_attn_interval = 4; ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + qwen4exp_require_nonzero(ml, LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval); for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0); } } + // the PLE conv history is a row of the recurrent cache, which linear layers alone have + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + if (hparams.is_ple(i) && !hparams.is_recr(i)) { + throw std::runtime_error(format("PLE layer %u is not a linear attention layer", i)); + } + } + switch (hparams.n_layer()) { case 48: type = LLM_TYPE_A3B; break; default: type = LLM_TYPE_UNKNOWN; @@ -95,12 +147,24 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); } - // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back + // flat [ple_head_dim, n_rows] gather target if (hparams.ple_n_heads > 0) { + // the head ranges are what the gather indexes, so they set the minimum row count + int64_t ple_rows = 0; + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + ple_rows = std::max(ple_rows, (int64_t) (hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h])); + } + + // the converter pads the table; a model synthesised from metadata has no tensor to ask const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); - const auto * ple_w = ml.get_weight(ple_name.c_str()); - GGML_ASSERT(ple_w != nullptr && "qwen4exp is missing the PLE n-gram table"); - const int64_t ple_rows = ple_w->tensor->ne[1]; + if (const auto * ple_w = ml.get_weight(ple_name.c_str())) { + if (ple_w->tensor->ne[1] < ple_rows) { + throw std::runtime_error(format("%s has %" PRId64 " rows, too few for the PLE head ranges (%" PRId64 ")", + ple_name.c_str(), ple_w->tensor->ne[1], ple_rows)); + } + ple_rows = ple_w->tensor->ne[1]; + } + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), { hparams.ple_head_dim, ple_rows }, 0); } @@ -261,6 +325,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inpL = build_inp_embd(model.tok_embd); cb(inpL, "model.input_embed", -1); + ggml_build_forward_expand(gf, inpL); auto * inp = build_inp_mem_hybrid(); @@ -277,6 +342,13 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + ggml_tensor * ple_emb = nullptr; + if (hparams.ple_n_heads > 0) { + ple_emb = build_inp_ple(mctx_hyb); + // make sure ple_emb and build_inp_embd are in the same graph split + ggml_build_forward_expand(gf, ple_emb); + } + // the wide residual starts as hc identical copies of the embedding ggml_tensor * res_hc = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), @@ -287,7 +359,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa res->t_layer_inp[il] = res_hc; if (hparams.is_ple(il)) { - res_hc = build_ple(inp->get_recr(), mctx_hyb, res_hc, il); + res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il); } ggml_tensor * inject = nullptr; @@ -375,13 +447,40 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( // mean-pooled indexer key, plus the incomplete tail. set_input resolves the cache layout. class llm_graph_input_qsa : public llm_graph_input_i { public: - llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio) : - mctx(mctx), ratio(ratio) {} + llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) : + mctx(mctx), ratio(ratio), blk_bias(blk_bias) {} virtual ~llm_graph_input_qsa() = default; void set_input(const llama_ubatch * ubatch) override { mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio); + mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); + } + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx); + + const auto * idx = mctx->get_idx(); + if (idx == nullptr) { + return false; + } + + const int64_t n_kv = idx->get_n_kv(); + const int64_t n_stream = mctx->get_n_stream(); + const int64_t n_blocks = (n_kv + ratio - 1)/ratio; + + bool res = true; + + res &= params.ubatch.n_tokens % n_stream == 0; + + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= cell_blk->ne[0] == n_kv; + res &= cell_blk->ne[1] == n_stream; + res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks; + res &= blk_pos->ne[0] == 4*n_blocks*n_stream; + res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv); + res &= bias->ne[1] == params.ubatch.n_tokens/n_stream; + + return res; } // per stream: a cell index names a different token in each stream @@ -389,16 +488,20 @@ class llm_graph_input_qsa : public llm_graph_input_i { ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream] ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] - ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens/n_stream, n_stream] + ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream] const llama_memory_hybrid_idx_context * mctx; const uint32_t ratio; + + // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded + const bool blk_bias; }; ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( const llama_memory_hybrid_idx_context * mctx_hyb, ggml_tensor * cur, ggml_tensor * inp_pos, + ggml_tensor * kq_mask, int * sections, int il) { const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); @@ -417,21 +520,38 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( GGML_ASSERT(n_tokens % n_stream == 0); const int64_t n_tps = n_tokens/n_stream; - auto qsa = std::make_unique(mctx_hyb, (uint32_t) r); + // only the "which block is visible" half of the bias varies per block + // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells + // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out + // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there + const bool blk_bias = kq_mask != nullptr && + kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && + cparams.causal_attn && !hparams.use_alibi; - qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); - qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); - qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); - qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); - qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream); + // nothing above depends on the layer, so the layers sharing a ratio share one input set + llm_graph_input_qsa * inp = nullptr; - ggml_set_input(qsa->cell_blk); - ggml_set_input(qsa->blk_cells); - ggml_set_input(qsa->blk_pos); - ggml_set_input(qsa->bias); - - llm_graph_input_qsa * inp = qsa.get(); - res->add_input(std::move(qsa)); + const auto it = qsa_inps.find((uint32_t) r); + if (it != qsa_inps.end()) { + inp = it->second; + } else { + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + + qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); + qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); + qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); + qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream); + + ggml_set_input(qsa->cell_blk); + ggml_set_input(qsa->blk_cells); + ggml_set_input(qsa->blk_pos); + ggml_set_input(qsa->bias); + + inp = qsa.get(); + res->add_input(std::move(qsa)); + qsa_inps.emplace((uint32_t) r, inp); + } // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); @@ -459,9 +579,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); cb(pooled, "indexer_k_pooled", il); + // count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536 + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks*n_stream, 1); + pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after. pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); - pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); @@ -479,20 +602,38 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // rectify each head dot product before the sum, as in the DeepSeek lightning indexer // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s ggml_tensor * score = ggml_mul_mat(ctx0, pooled, - ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream)); + ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h*n_tps, n_stream)); score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream); score = ggml_relu(ctx0, score); - score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); - score = ggml_sum_rows(ctx0, score); - score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); + + // the heads sit side by side on ne[1] and there are only a few of them + ggml_tensor * summed = nullptr; + for (int64_t h = 0; h < n_idx_h; ++h) { + ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream, + score->nb[2], score->nb[3], h*score->nb[1]); + summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice); + } + + score = summed; cb(score, "indexer_score", il); - // give every token of a block the block score; the budget is a whole number of - // blocks, so the top-k cut still lands on a block boundary + // one value per block, so it is cheaper to bias here than after the cells are expanded + if (blk_bias) { + score = ggml_add(ctx0, score, inp->bias); + } + + // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary ggml_tensor * expanded = ggml_get_rows(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); - expanded = ggml_add(ctx0, expanded, inp->bias); + + if (blk_bias) { + // flash attention keeps the mask in f16; the scores are f32 + ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); + expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream)); + } else { + expanded = ggml_add(ctx0, expanded, inp->bias); + } cb(expanded, "indexer_score_tokens", il); // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail @@ -602,7 +743,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( // indexer reads the same block input as q/k/v; no cache or no ratio means dense const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; - ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, sections, il) : nullptr; + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; // Qwen3Next uses a single Q projection that outputs query + gate ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] @@ -1007,33 +1148,34 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at( ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0); - // keep the last state_cols columns for the next ubatch + // [TAG_RECURRENT_ROLLBACK_SPLITS] keep the last state_cols columns once per rollback slot, + // slot s ending s tokens earlier so a rollback of s tokens reads a history that never saw them const size_t row_size = ggml_row_size(conv_states_all->type, row_total); + const uint32_t mem_size = mctx_cur->get_size(); - ggml_tensor * tail = ggml_view_3d(ctx0, conv_input, - state_cols, channels, n_seqs, - conv_input->nb[1], conv_input->nb[2], - ggml_row_size(conv_input->type, conv_input->ne[0] - state_cols)); + const int64_t n_slots = (int64_t) cparams.n_rs_seq + 1; - ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all, - state_cols * channels, n_seqs, - conv_states_all->nb[1], - kv_head * row_size); + for (int64_t slot = 0; slot < n_slots; ++slot) { + const int64_t s_idx = std::max(0, conv_input->ne[0] - state_cols - slot); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst)); + ggml_tensor * tail = ggml_view_3d(ctx0, conv_input, + state_cols, channels, n_seqs, + conv_input->nb[1], conv_input->nb[2], + ggml_row_size(conv_input->type, s_idx)); + + ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all, + state_cols * channels, n_seqs, + conv_states_all->nb[1], + (slot * mem_size + kv_head) * row_size); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst)); + } return conv_input; } -ggml_tensor * llama_model_qwen4exp::graph::build_ple( - llm_graph_input_rs * inp, - const llama_memory_hybrid_idx_context * mctx_hyb, - ggml_tensor * hidden, - int il) { - GGML_UNUSED(inp); - - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t hc_dim = hc * n_embd; +ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple( + const llama_memory_hybrid_idx_context * mctx_hyb) { const int64_t n_heads = hparams.ple_n_heads; auto ple_inp = std::make_unique( @@ -1047,7 +1189,18 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple( // gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows); emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens); - cb(emb, "ple_embd", il); + cb(emb, "ple_embd", -1); + + return emb; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_ple( + llm_graph_input_rs * inp, + ggml_tensor * emb, + ggml_tensor * hidden, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; ggml_tensor * key = build_lora_mm(model.layers[il].ple_key, emb); ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 30489016c703..24231015512a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -210,6 +210,30 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8)); // without this the QSA layers fall back to dense and go uncovered ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); + + // has_cell_ext() needs ple_n_heads here: the indexer cache serializes no ext without it + const uint32_t ple_ngram_size = 3; + const uint32_t ple_heads_per_ngram = 2; + const uint32_t ple_n_heads = (ple_ngram_size - 1)*ple_heads_per_ngram; + GGML_ASSERT(n_embd % ple_n_heads == 0); + const uint32_t ple_head_dim = n_embd/ple_n_heads; + + std::vector ple_head_offsets(ple_n_heads); + std::vector ple_head_vocab_sizes(ple_n_heads, n_vocab); + for (uint32_t h = 0; h < ple_n_heads; h++) { + ple_head_offsets[h] = uint64_t(h)*n_vocab; + } + + // the PLE history lives in the recurrent cache, so it must sit on a linear attention layer + ms.add_kv(LLM_KV_PLE_LAYERS, std::vector({ 0 })); + ms.add_kv(LLM_KV_PLE_NGRAM_SIZE, ple_ngram_size); + ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, ple_heads_per_ngram); + ms.add_kv(LLM_KV_PLE_CONV_KERNEL, uint32_t(4)); + ms.add_kv(LLM_KV_PLE_EOS_TOKEN_ID, uint32_t(0)); + ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, ple_head_dim); + ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector({ 1, 3, 5 })); + ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS, ple_head_offsets); + ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, ple_head_vocab_sizes); } ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1)); diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6e93ce6fb8da..5aea842d97eb 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -347,6 +347,64 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p } +// Test 6: state blob round-trip +// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens +static bool test_state_roundtrip(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { + auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + + LOG("\n=== Test 6: state blob round-trip ===\n"); + + if (llama_decode(ctx.get(), llama_batch_get_one(const_cast(tokens.data()), (int32_t) tokens.size()))) { + LOG_ERR("\n%s: failed to decode prompt\n", __func__); + return false; + } + + std::vector blob_a(llama_state_seq_get_size(ctx.get(), 0)); + const size_t n_a = llama_state_seq_get_data(ctx.get(), blob_a.data(), blob_a.size(), 0); + if (n_a != blob_a.size()) { + LOG_ERR("\n%s: saved %zu bytes, expected %zu\n", __func__, n_a, blob_a.size()); + return false; + } + + if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) { + LOG_ERR("\n%s: failed to erase seq 0\n", __func__); + return false; + } + + if (llama_state_seq_set_data(ctx.get(), blob_a.data(), blob_a.size(), 0) != blob_a.size()) { + LOG_ERR("\n%s: failed to restore seq 0\n", __func__); + return false; + } + + std::vector blob_b(llama_state_seq_get_size(ctx.get(), 0)); + const size_t n_b = llama_state_seq_get_data(ctx.get(), blob_b.data(), blob_b.size(), 0); + if (n_b != n_a) { + LOG_ERR("\n%s: re-saved %zu bytes, expected %zu\n", __func__, n_b, n_a); + return false; + } + + size_t n_diff = 0; + size_t i_diff = 0; + for (size_t i = 0; i < n_a; i++) { + if (blob_a[i] != blob_b[i]) { + if (n_diff == 0) { + i_diff = i; + } + n_diff++; + } + } + + if (n_diff > 0) { + LOG_ERR("\n%s: state changed across a restore: %zu of %zu bytes differ, first at offset %zu\n", + __func__, n_diff, n_a, i_diff); + return false; + } + + LOG("\nPASS\n"); + return true; +} + + int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -434,6 +492,11 @@ int main(int argc, char ** argv) { return 1; } + // Test 6: state blob round-trip + if (!test_state_roundtrip(model, params, tokens)) { + return 1; + } + LOG("\nAll tests passed.\n"); return 0; From 7fd6a7cc019e48d63410d27146db2887ec8de2d4 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 14:39:46 -0600 Subject: [PATCH 02/16] CUDA: XOR swizzle flash attn K,V smem fp16 tiles Prerequisite for the upcoming sparse-fa flash attention path: adds the smem swizzle layout (fattn-swizzle.cuh) that the sparse gather load tiles need. Reconciled against TurboQuant's turbo2/3/4 SRAM tile loaders in fattn-mma-f16.cuh, which are untouched by this change. Cherry-picked from upstream e4b9af007. --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 72 ++++++++++----- ggml/src/ggml-cuda/fattn-swizzle.cuh | 126 +++++++++++++++++++++++++++ tests/test-backend-ops.cpp | 28 +++++- 3 files changed, 200 insertions(+), 26 deletions(-) create mode 100644 ggml/src/ggml-cuda/fattn-swizzle.cuh diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 2cd9053c9bf0..32ad9f4892fd 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -4,6 +4,7 @@ #include "cp-async.cuh" #include "mma.cuh" #include "fattn-common.cuh" +#include "fattn-swizzle.cuh" using namespace ggml_cuda_mma; @@ -68,7 +69,7 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2, 32, 96, 64, 64, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2, 32, 96, 64, 64, 2, true); - GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 8, 64, 4, 64, 128, 128, 128, 2, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 8, 128, 2, 64, 128, 128, 128, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 64, 4, 32, 128, 128, 128, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2, 32, 128, 128, 128, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2, 32, 128, 128, 128, 2, true); @@ -391,7 +392,7 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, c // ------------------------------------------------------------------------------------------------------------------ -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); @@ -428,7 +429,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk); + if constexpr (swz) { + const int smem_offs_b = ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk); + cp_async_cg_16(tile_KV_32 + smem_offs_b, KV + i*stride_KV + k*h2_per_chunk); + } else { + cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk); + } } } }; @@ -463,8 +469,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, - !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + if constexpr (swz) { + ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk), + !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + } else { + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, + !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + } } } }; @@ -802,9 +813,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool is_turbo_kv = (type_K != GGML_TYPE_F16 || type_V != GGML_TYPE_F16); constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); - constexpr int stride_tile_K = nbatch_K2 + 4; - - constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4; + // swizzle the tile stride for K and V based on the batch size. + constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2); + constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2); + constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2); + constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2); const int k_VKQ_0 = kb0 * nbatch_fa; #if defined(TURING_MMA_AVAILABLE) @@ -822,7 +835,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool use_cp_async = true; cp_async_wait_all(); __syncthreads(); - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); } else { constexpr bool use_cp_async = nstages == 1; @@ -861,7 +874,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } else if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); if (use_cp_async) { cp_async_wait_all(); @@ -877,7 +890,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( #pragma unroll for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) { T_A_KQ K_A; - load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K); + ggml_cuda_fattn_smem_swizzle::load_ldmatrix(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start); if constexpr (cols_per_warp == 8) { mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]); } else { @@ -903,7 +916,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I; T_A_KQ K_A; - load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K); + ggml_cuda_fattn_smem_swizzle::load_ldmatrix(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start); if constexpr (cols_per_warp == 8) { mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[0]); @@ -1197,7 +1210,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( flash_attn_ext_f16_load_mask (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); } - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); } } @@ -1234,7 +1247,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); if (use_cp_async) { cp_async_wait_all(); @@ -1253,7 +1266,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::J; T_A_VKQ A; // Transposed in SRAM but not in registers, gets transposed on load. - load_ldmatrix_trans(A, tile_V_i + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V); + ggml_cuda_fattn_smem_swizzle::load_ldmatrix_trans(A, tile_V, (int)(tile_V_i - tile_V) + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2); if constexpr (T_B_KQ::I == 8) { mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]); } else { @@ -1279,7 +1292,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::I; T_A_VKQ A; // Transposed in both SRAM and registers, load normally. - load_ldmatrix(A, tile_V_i + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V); + ggml_cuda_fattn_smem_swizzle::load_ldmatrix(A, tile_V, (int)(tile_V_i - tile_V) + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2); mma(VKQ_C[i_VKQ_0/i0_stride], B[k00/(np*T_A_VKQ::I)], A); } } @@ -1446,10 +1459,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( static_assert(nwarps * (cols_per_warp/ncols2) % ncols1 == 0, "bad nwarps"); constexpr int stride_tile_Q = DKQ/2 + 4; - constexpr int stride_tile_K = nbatch_K2 + 4; - - constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4; + // swizzle the tile stride for K and V based on the batch size. + constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2); + constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2); constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V; + constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2); + constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2); extern __shared__ half2 tile_Q[]; half2 * tile_K = Q_in_reg ? tile_Q : tile_Q + ncols * stride_tile_Q; @@ -1543,7 +1558,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_load_mask (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); } - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); } @@ -1708,11 +1723,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int tile_stride = nbatch_combine + 4; static_assert((DV/2) % nbatch_combine == 0, "bad nbatch_combine"); + constexpr bool combine_needs_sync = swz_K || swz_V; + if constexpr (cols_per_warp == 8) { const int jc_cwmo = (threadIdx.x % (2*T_C_VKQ::J)) / T_C_VKQ::J; // jc combine write meta offset const int jc_cwm = threadIdx.y*(2*T_C_VKQ::J) + 2*T_C_VKQ::get_j(-1) + jc_cwmo; // jc combine write meta const float2 KQ_cmr = make_float2(KQ_max[jc_cwmo], KQ_rowsum[jc_cwmo]); // KQ combine max rowsum + if constexpr (combine_needs_sync) { + __syncthreads(); + } + if (((!needs_fixup && !is_fixup) || np > 1) && threadIdx.x < 2*T_C_VKQ::J) { // Use the 16 bytes of padding in each row to store the meta data: KQ max, KQ rowsum, KQ max scale. ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr; @@ -1749,6 +1770,10 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const bool thread_should_write = T_C_KQ::J == 8 || T_C_KQ::get_j(threadIdx.x & 2) < 8; #endif // defined(TURING_MMA_AVAILABLE) + if constexpr (combine_needs_sync) { + __syncthreads(); + } + if (((!needs_fixup && !is_fixup) || np > 1) && thread_should_write) { ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr; } @@ -2198,8 +2223,11 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu - const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(nbatch_K2 + 4, nbatch_V2 + 4) * sizeof(half2); - const size_t nbytes_shared_KV_2stage = nbatch_fa * (nbatch_K2 + 4 + nbatch_V2 + 4) * sizeof(half2); + // KV tile strides must match flash_attn_ext_f16_iter / _process_tile. + const int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2, cc); + const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2, cc); + const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(stride_tile_K, stride_tile_V) * sizeof(half2); + const size_t nbytes_shared_KV_2stage = nbatch_fa * (stride_tile_K + stride_tile_V) * sizeof(half2); const size_t nbytes_shared_Q = ncols * (DKQ/2 + 4) * sizeof(half2); const size_t nbytes_shared_mask = ncols1 * (nbatch_fa/2 + 4) * sizeof(half2); const size_t nbytes_shared_combine = nwarps*cols_per_warp * (nbatch_combine + 4) * sizeof(half2); diff --git a/ggml/src/ggml-cuda/fattn-swizzle.cuh b/ggml/src/ggml-cuda/fattn-swizzle.cuh new file mode 100644 index 000000000000..44338c8db08d --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-swizzle.cuh @@ -0,0 +1,126 @@ +#pragma once + +#include "common.cuh" +#include "mma.cuh" + +// XOR swizzle for K/V SMEM tiles to avoid bank conflicts without row padding (Turing+ only). +// Stride must be a multiple of 32 half2 columns, otherwise we keep +4 row padding. + +namespace ggml_cuda_fattn_smem_swizzle { + +static __host__ __device__ constexpr bool bank_aligned(const int nbatch_2) { + return nbatch_2 >= 32 && nbatch_2 % 32 == 0; +} + +static __device__ constexpr bool enabled(const int nbatch_2) { +#if defined(TURING_MMA_AVAILABLE) + return bank_aligned(nbatch_2); +#else + GGML_UNUSED(nbatch_2); + return false; +#endif // defined(TURING_MMA_AVAILABLE) +} + +static __host__ bool enabled(const int nbatch_2, const int cc) { +#ifdef GGML_USE_HIP + GGML_UNUSED(nbatch_2); + GGML_UNUSED(cc); + return false; +#else + return turing_mma_available(cc) && bank_aligned(nbatch_2); +#endif // GGML_USE_HIP +} + +static __device__ constexpr int tile_stride(const int nbatch_2) { + return enabled(nbatch_2) ? nbatch_2 : nbatch_2 + 4; +} + +static __host__ int tile_stride(const int nbatch_2, const int cc) { + return enabled(nbatch_2, cc) ? nbatch_2 : nbatch_2 + 4; +} + +// Swizzled byte offset for tile element (row, col_h2), same map used for writes and reads. +template +static __device__ __forceinline__ int bytes_rc(const int row, const int col_h2) { + static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32"); + return ((row * stride_h2 + col_h2) * (int) sizeof(half2)) ^ ((row & 7) << 4); +} + +// ldmatrix.x4 via 64-bit generic pointer. +static __device__ __forceinline__ void ldmatrix_x4(int * xi, const half2 * addr) { +#if defined(TURING_MMA_AVAILABLE) + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) + : "l"(addr)); +#else + GGML_UNUSED_VARS(xi, addr); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) +} + +static __device__ __forceinline__ void ldmatrix_x4_trans(int * xi, const half2 * addr) { +#if defined(TURING_MMA_AVAILABLE) + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) + : "l"(addr)); +#else + GGML_UNUSED_VARS(xi, addr); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) +} + +// Per-lane swizzled address for one tile<16, 8, half2> ldmatrix: 16 rows, 4 half2 columns per lane. +template +static __device__ __forceinline__ const half2 * lane_addr( + const half2 * tile_base, const int base_row, const int base_col_h2, const int I, const int J) { + static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32"); + const int lane_row = threadIdx.x % I; + const int lane_col = (threadIdx.x / I) * (J / 2); + uint32_t byte_off = (uint32_t) ((base_row + lane_row)*stride_h2 + base_col_h2 + lane_col) * (uint32_t) sizeof(half2); + byte_off ^= (uint32_t) (((base_row + lane_row) & 7) << 4); + return (const half2 *) ((const char *) tile_base + byte_off); +} + +template +static __device__ __forceinline__ void load_ldmatrix( + TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) { + if constexpr (swz) { + static_assert(std::is_same_v>, + "the swizzled layout is only supported for tile<16, 8, half2>"); + ldmatrix_x4((int *) t.x, lane_addr(tile_base, base_row, base_col_h2, TileT::I, TileT::J)); + } else { + ggml_cuda_mma::load_ldmatrix(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2); + } +} + +template +static __device__ __forceinline__ void load_ldmatrix(TileT & t, const half2 * tile_base, const int off_h2) { + if constexpr (swz) { + load_ldmatrix(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2); + } else { + ggml_cuda_mma::load_ldmatrix(t, tile_base + off_h2, stride_h2); + } +} + +template +static __device__ __forceinline__ void load_ldmatrix_trans( + TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) { + if constexpr (swz) { + static_assert(std::is_same_v>, + "the swizzled layout is only supported for tile<16, 8, half2>"); + ldmatrix_x4_trans((int *) t.x, lane_addr(tile_base, base_row, base_col_h2, TileT::I, TileT::J)); + } else { + ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2); + } +} + +template +static __device__ __forceinline__ void load_ldmatrix_trans(TileT & t, const half2 * tile_base, const int off_h2) { + if constexpr (swz) { + load_ldmatrix_trans(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2); + } else { + ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + off_h2, stride_h2); + } +} + +} // namespace ggml_cuda_fattn_smem_swizzle diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index bc97cc656aba..92eebedebe99 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10269,6 +10269,22 @@ static std::vector> make_test_cases_eval() { GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); } + // dense-allocated (non-view) quant K/V at batch >= 64, in cache and native layouts + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false)); + + // FLASH_ATTN_EXT MMA: non-pow2 head size and MLA K/V view. + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 8, {8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + + // FLASH_ATTN_EXT MMA, swizzled K/V tiles, power-of-two stride: nbatch_K2 = 32, 64, 128, 256. + test_cases.emplace_back(new test_flash_attn_ext( 64, 64, 8, {8, 1}, 4096, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 4096, 8, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); @@ -10643,10 +10659,14 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext_turbo4_vec(128)); test_cases.emplace_back(new test_flash_attn_ext_turbo4_vec(256)); - for (int kv : { 4096, 8192, 16384, }) { - for (int hs : { 64, 128, }) { - for (int nr : { 1, 4, }) { - test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {nr, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + for (int kv : { 4096, 8192, 16384,32768, 65536, }) { + for (int hs : { 64, 128, 256, 576, }) { + const int hsv = hs == 576 ? 512 : hs; + const bool v_view = hs == 576; + for (int nr : { 1, 4, 8, }) { + for (int nb : { 1, 4096, }) { + test_cases.emplace_back(new test_flash_attn_ext(hs, hsv, 8, {nr, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, v_view)); + } } } } From f3b50395070b793a3107b01b3b3d9feb06a1e025 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 14:44:51 -0600 Subject: [PATCH 03/16] test: make the FA V-is-view-of-K case a test case parameter Prerequisite for the upcoming sparse-fa flash attention path: the new sparse-fa test cases need kv_view/v_is_view_of_k as named constructor parameters instead of hardcoded booleans. Also wires kv_view through to the K/V tensor creation (previously hardcoded true), matching an earlier upstream commit's intent that our tree never received. Cherry-picked from upstream 5fff12845, plus the kv_view wiring fix. --- tests/test-backend-ops.cpp | 46 +++++++++++++++++++++++++++++--------- 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 92eebedebe99..a040e82c12bc 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7224,9 +7224,11 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_K; const ggml_type type_V; std::array permute; + const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) + const bool v_is_view_of_k; std::string vars() override { - return VARS_TO_STR14(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute); + return VARS_TO_STR16(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k); } double max_nmse_err() override { @@ -7242,9 +7244,10 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, - ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}) + ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, + bool kv_view = true, bool v_is_view_of_k = false) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7272,21 +7275,21 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false); ggml_set_name(q, "q"); - ggml_tensor * k = create_permuted(type_K, hsk_padded, kv, nh, nr23[1], true); // the K tensor is usually a view of the K cache + ggml_tensor * k = create_permuted(type_K, hsk_padded, kv, nh, nr23[1], kv_view); // the K tensor is usually a view of the K cache ggml_set_name(k, "k"); ggml_tensor * v = nullptr; - if (type_K == type_V && hsk_padded == 576 && hsv_padded == 512) { - // TODO: this branch should become a separate test case parameter instead of hardcoding this for these head shapes - - // in this branch, the V cache is sub-view of the K cache. this is used by some MLA-based models + if (v_is_view_of_k) { + // the V cache is a sub-view of the K cache. this is used by some MLA-based models // for more info: // - https://github.com/ggml-org/llama.cpp/pull/13435 // - https://github.com/ggml-org/llama.cpp/pull/18953#issuecomment-3774948392 // - https://github.com/ggml-org/llama.cpp/pull/18986 + GGML_ASSERT(type_K == type_V && hsv_padded <= hsk_padded); + v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0); } else { - v = create_permuted(type_V, hsv_padded, kv, nh, nr23[1], true); // the V tensor is usually a view of the V cache + v = create_permuted(type_V, hsv_padded, kv, nh, nr23[1], kv_view); // the V tensor is usually a view of the V cache } ggml_set_name(v, "v"); @@ -10222,12 +10225,14 @@ static std::vector> make_test_cases_eval() { for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO4_0}) { if ((type_KV == GGML_TYPE_TURBO3_0 || type_KV == GGML_TYPE_TURBO4_0) && hsk < 128) continue; if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72 && hsk != 128) continue; + // DeepSeek MLA: the V cache is a sub-view of the K cache + const bool v_is_view_of_k = hsk == 576; test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV)); + hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 1, 2, 3}, true, v_is_view_of_k)); // run fewer test cases permuted if (mask == true && max_bias == 0.0f && logit_softcap == 0 && kv == 512) { test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3})); + hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3}, true, v_is_view_of_k)); } } } @@ -10260,6 +10265,25 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_turbo4_vec(128)); test_cases.emplace_back(new test_flash_attn_ext_turbo4_vec(256)); + // q8_0 KV cases: decode and prompt batches, KV pad, permuted KV, feature flags, and long context + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 2}, 1025, 1, true, true, 8, 30, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + + // MLA shape: the V cache is a sub-view of the K cache, with quantized KV + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + + // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes + test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 512, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). for (int64_t kv : { 4096, 16384 }) { From 3ddea31f7d10c0b4f9020398fe256c1f157915d0 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 15:04:24 -0600 Subject: [PATCH 04/16] CUDA + ggml: add sparse-fa for DSV4/GLM Cherry-picked from upstream 8e93a9773 on top of the swizzle and V-is-view-of-K prerequisites. Adds ggml_flash_attn_ext_set_n_kv_max(): the CUDA flash-attention kernel can treat the mask's finite entries as a sparse K/V set and skip the rest, instead of computing dense attention and masking it out. DeepSeek-V4's own top-k sparse-attention path is wired to it. Reconciled against TurboQuant's turbo2/3/4 tile loaders and dedicated fattn-mma-turbo.cuh kernel, which the upstream diff has no knowledge of. Two real bugs were caught and fixed in that reconciliation, not just merge conflicts: - fattn-mma-turbo.cuh instantiated flash_attn_ext_f16<...> with the old positional template argument list; type_K was landing in the new use_sparse slot (a nonzero ggml_type implicitly converts to true), and type_K/type_V were shifted off the end entirely. - Its launch_fattn call had the same problem one level up: the new use_sparse parameter was inserted before warp_size, so the trailing warp_size_host argument would have silently become use_sparse=true and warp_size would have silently fallen back to its default. - qwen35.cpp's MTP draft-head graph calls build_attn_mha directly and wasn't part of upstream's diff at all, so it still had the pre-change 9-argument signature. qwen4exp's own build_attn_mha call is updated for the new signature but left at n_kv_max=0 (disabled): the sparse kernel path is compile-time gated to DeepSeek-V4/GLM's specific MLA head shapes (512/512 or 576/512, GQA 8 or 16). qwen4exp's actual shape is 256/256 with GQA 12, which doesn't match, so enabling it would take no effect today. Extending the shape gate to cover qwen4exp is separate, higher-risk kernel work not attempted here. Verified: full build, test-llama-archs passes on CUDA and CPU with no regressions (qwen4exp, qwen35, qwen35moe, and all other architectures). --- ggml/include/ggml.h | 6 + ggml/src/ggml-cuda/fattn-common.cuh | 23 ++- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 223 +++++++++++++++++-------- ggml/src/ggml-cuda/fattn-mma-turbo.cuh | 6 +- ggml/src/ggml-cuda/fattn-tile.cuh | 12 +- ggml/src/ggml-cuda/fattn-vec.cuh | 2 +- ggml/src/ggml-cuda/fattn.cu | 133 +++++++++++++++ ggml/src/ggml.c | 9 + src/llama-graph.cpp | 17 +- src/llama-graph.h | 1 + src/models/deepseek4.cpp | 7 +- src/models/qwen35.cpp | 2 +- src/models/qwen4exp.cpp | 2 +- tests/test-backend-ops.cpp | 71 +++++++- 14 files changed, 410 insertions(+), 104 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 65eaebfe0873..5761f8226074 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2445,6 +2445,12 @@ extern "C" { GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a); + // Use finite mask entries as a sparse K/V set. Set 0 to disable. + // n_kv_max must bound the number of finite entries in every mask row. + GGML_API void ggml_flash_attn_ext_set_n_kv_max( + struct ggml_tensor * a, + int32_t n_kv_max); + GGML_API void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks); diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 58efb7633f73..b0dabf86c1ca 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1117,6 +1117,9 @@ static __global__ void flash_attn_mask_to_KV_max( KV_max[sequence*ne31 + jt] = KV_max_sj; } +void ggml_cuda_flash_attn_ext_compact_mask( + const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream); + template // D == head size __launch_bounds__(D, 1) static __global__ void flash_attn_stream_k_fixup_uniform( @@ -1359,7 +1362,8 @@ static __global__ void flash_attn_combine_results( template void launch_fattn( ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared, - const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE + const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse, + const int warp_size = WARP_SIZE ) { constexpr int ncols = ncols1 * ncols2; @@ -1520,10 +1524,20 @@ void launch_fattn( const int ntiles_z_gqa = ((gqa_ratio + ncols2 - 1) / ncols2); const int ntiles_dst = ntiles_x * ntiles_z_gqa * K->ne[2] * Q->ne[3]; + const int32_t n_kv_max = use_sparse ? ggml_get_op_params_i32(KQV, 4) : 0; + if (use_sparse) { + GGML_ASSERT(mask != nullptr); + GGML_ASSERT(n_kv_max > 0); + const size_t mask_rows = size_t(mask->ne[1]) * mask->ne[3]; + + KV_max.alloc(size_t(n_kv_max) * mask_rows); + ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, n_kv_max, main_stream); + } + // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); @@ -1545,7 +1559,8 @@ void launch_fattn( GGML_ASSERT(max_blocks_per_sm > 0); int parallel_blocks = max_blocks_per_sm; - const int ntiles_KV = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. + const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1]; + const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. dim3 blocks_num; if (stream_k) { @@ -1647,7 +1662,7 @@ void launch_fattn( !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], - K->ne[0], K->ne[1], K->ne[2], K->ne[3], nb11, nb12, nb13, + K->ne[0], n_kv, K->ne[2], K->ne[3], nb11, nb12, nb13, nb21, nb22, nb23, mask ? mask->ne[1] : 0, mask ? mask->ne[2] : 0, mask ? mask->ne[3] : 0, mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0 diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 32ad9f4892fd..dcf5d381df82 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -381,20 +381,24 @@ static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, return cp_async_available(cc) && ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2, cc) : 0; } -static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2) { +static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages( + const int DKQ, const int DV, const int ncols1, const int ncols2, const bool use_sparse) { #ifdef CP_ASYNC_AVAILABLE - return ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; + const int nstages_target = ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; + // sparse gather is not implemented for multi-stage loading + return use_sparse && nstages_target > 1 ? 1 : nstages_target; #else - GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2); + GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2, use_sparse); return 0; #endif // CP_ASYNC_AVAILABLE } // ------------------------------------------------------------------------------------------------------------------ -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( - const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) { + const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, + const int k_VKQ_0, const int i_sup, const int32_t * const __restrict__ indices) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); // K/V data is loaded with decreasing granularity for D for better memory bandwidth. // The minimum granularity is 16 bytes. @@ -402,7 +406,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( const int chunks_per_row = D2 / h2_per_chunk; if constexpr (use_cp_async) { static_assert(warp_size == 32, "bad warp_size"); - static_assert(!oob_check, "OOB check not compatible with cp_async"); + static_assert(!oob_check || use_sparse, "OOB check not compatible with cp_async"); constexpr int preload = 64; const unsigned int tile_KV_32 = ggml_cuda_cvta_generic_to_shared(tile_KV); @@ -425,15 +429,24 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( break; } + int64_t i_KV; + if constexpr (use_sparse) { + // padded slots gather row 0, the -inf mask removes their contribution + const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0; + i_KV = index >= 0 ? index : 0; + } else { + i_KV = k_VKQ_0 + i; + } + #pragma unroll for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); if constexpr (swz) { const int smem_offs_b = ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk); - cp_async_cg_16(tile_KV_32 + smem_offs_b, KV + i*stride_KV + k*h2_per_chunk); + cp_async_cg_16(tile_KV_32 + smem_offs_b, KV + i_KV*stride_KV + k*h2_per_chunk); } else { - cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk); + cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i_KV*stride_KV + k*h2_per_chunk); } } } @@ -469,12 +482,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); + const half2 * src; + if constexpr (use_sparse) { + const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1; + src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero; + } else { + src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero; + } if constexpr (swz) { - ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk), - !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk), src); } else { - ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, - !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, src); } } } @@ -687,14 +705,16 @@ static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile( } } -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const half * const __restrict__ mask_h, half * const __restrict__ tile_mask, - const int stride_mask, const int i_sup, const int j0, const uint3 ne01) { + const int stride_mask, const int k_VKQ_0, const int i_sup, const int j0, const uint3 ne01, + const int32_t * const __restrict__ indices) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); if constexpr (use_cp_async) { static_assert(nbatch_fa <= 8*warp_size && nbatch_fa % 8 == 0, "bad nbatch_fa"); static_assert(!oob_check, "OOB check incompatible with cp_async"); + static_assert(!use_sparse, "sparse gather incompatible with cp_async"); constexpr int preload = nbatch_fa >= 32 ? nbatch_fa * sizeof(half) : 64; constexpr int cols_per_warp = 8*warp_size/nbatch_fa; constexpr int stride_j = nwarps * cols_per_warp; @@ -712,9 +732,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const int i = 8 * (threadIdx.x % (nbatch_fa/8)); - cp_async_cg_16(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + i); + cp_async_cg_16(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + i); } - } else if constexpr (oob_check) { + } else if constexpr (oob_check || use_sparse) { #pragma unroll for (int j1 = 0; j1 < ncols1; j1 += nwarps) { const int j_sram = j1 + threadIdx.y; @@ -728,7 +748,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( for (int i0 = 0; i0 < nbatch_fa; i0 += warp_size) { const int i = i0 + threadIdx.x; - tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + i] : half(0.0f); + if constexpr (use_sparse) { + const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1; + tile_mask[j_sram*(nbatch_fa + 8) + i] = index >= 0 ? mask_h[int64_t(j_vram)*stride_mask + index] : half(-INFINITY); + } else { + tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + k_VKQ_0 + i] : half(0.0f); + } } } } else if constexpr (nbatch_fa < 2*warp_size) { @@ -745,7 +770,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const int i = threadIdx.x % (warp_size/cols_per_warp); - ggml_cuda_memcpy_1(tile_mask + j_sram*(nbatch_fa + 8) + 2*i, mask_h + int64_t(j_vram)*stride_mask + 2*i); + ggml_cuda_memcpy_1(tile_mask + j_sram*(nbatch_fa + 8) + 2*i, mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + 2*i); } } else { #pragma unroll @@ -761,14 +786,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( for (int i0 = 0; i0 < nbatch_fa; i0 += 2*warp_size) { const int i = i0 + 2*threadIdx.x; - ggml_cuda_memcpy_1(tile_mask + j_sram*(nbatch_fa + 8) + i, mask_h + int64_t(j_vram)*stride_mask + i); + ggml_cuda_memcpy_1(tile_mask + j_sram*(nbatch_fa + 8) + i, mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + i); } } } } template static __device__ __forceinline__ void flash_attn_ext_f16_iter( @@ -776,6 +801,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ indices, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, @@ -811,7 +837,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( // multi-stage pipeline (nstages>1) would copy raw turbo bytes as half2 => garbage. // Force single-stage synchronous loading for the turbo path. constexpr bool is_turbo_kv = (type_K != GGML_TYPE_F16 || type_V != GGML_TYPE_F16); - constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); + constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, use_sparse); // swizzle the tile stride for K and V based on the batch size. constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2); @@ -835,13 +861,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool use_cp_async = true; cp_async_wait_all(); __syncthreads(); - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr); } else { - constexpr bool use_cp_async = nstages == 1; + // the sparse mask values are gathered per element, always load them synchronously + constexpr bool use_cp_async = nstages == 1 && !use_sparse; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, k_VKQ_0, k_VKQ_sup, jt*ncols1, ne01, indices); } } @@ -874,8 +901,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } else if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices); if (use_cp_async) { cp_async_wait_all(); } @@ -1200,6 +1227,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } if constexpr (nstages > 1) { + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading"); // Preload K tile for next iteration: constexpr bool use_cp_async = true; @@ -1207,11 +1235,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( __syncthreads(); if (!last_iter) { if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr); } } @@ -1247,8 +1275,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_0, k_VKQ_sup, indices); if (use_cp_async) { cp_async_wait_all(); } @@ -1303,7 +1331,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, @@ -1401,13 +1429,14 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) -template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ indices, const float * const __restrict__ sinks_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, @@ -1449,7 +1478,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); // Force single-stage synchronous loading for the turbo path (see iter for rationale). constexpr bool is_turbo_kv = (type_K != GGML_TYPE_F16 || type_V != GGML_TYPE_F16); - constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); + constexpr int nstages = is_turbo_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, use_sparse); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1550,37 +1579,38 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); constexpr bool use_cp_async = true; constexpr bool oob_check = false; constexpr int k_VKQ_sup = nbatch_fa; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr); } // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. - if constexpr (ncols2 == 1) { + if constexpr (ncols2 == 1 || use_sparse) { constexpr bool oob_check = true; for (; kb0 < kb0_stop-1; ++kb0) { constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } constexpr bool last_iter = true; const int k_VKQ_sup = ne11 - kb0*nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } else { @@ -1589,18 +1619,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } constexpr bool last_iter = true; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } @@ -1995,7 +2025,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, kb0_start, kb0_stop); @@ -2003,7 +2033,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } -template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( @@ -2030,14 +2066,15 @@ static __global__ void flash_attn_ext_f16( const int32_t nb31, const int32_t nb32, const int64_t nb33) { ggml_cuda_pdl_sync(); // TODO optimize placement #if defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) - const char * GGML_CUDA_RESTRICT Q = Q_ptr; - const char * GGML_CUDA_RESTRICT K = K_ptr; - const char * GGML_CUDA_RESTRICT V = V_ptr; - const char * GGML_CUDA_RESTRICT mask = mask_ptr; - const char * GGML_CUDA_RESTRICT sinks = sinks_ptr; - const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; - float * GGML_CUDA_RESTRICT dst = dst_ptr; - float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; + const char * GGML_CUDA_RESTRICT Q = Q_ptr; + const char * GGML_CUDA_RESTRICT K = K_ptr; + const char * GGML_CUDA_RESTRICT V = V_ptr; + const char * GGML_CUDA_RESTRICT mask = mask_ptr; + const char * GGML_CUDA_RESTRICT sinks = sinks_ptr; + const int * GGML_CUDA_RESTRICT KV_max = use_sparse ? nullptr : KV_max_ptr; + const int * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr; + float * GGML_CUDA_RESTRICT dst = dst_ptr; + float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; // Skip unused kernel variants for faster compilation: if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) { @@ -2048,6 +2085,11 @@ static __global__ void flash_attn_ext_f16( NO_DEVICE_CODE; return; } + + if (!ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2) && use_sparse) { + NO_DEVICE_CODE; + return; + } #ifdef VOLTA_MMA_AVAILABLE if (ncols1*ncols2 < 32) { NO_DEVICE_CODE; @@ -2129,6 +2171,7 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr; const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; @@ -2138,13 +2181,13 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } @@ -2175,6 +2218,7 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr; const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; @@ -2184,8 +2228,8 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, @@ -2201,6 +2245,8 @@ static __global__ void flash_attn_ext_f16( #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) } +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; @@ -2247,20 +2293,49 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml using fattn_kernel_ptr_t = fattn_kernel_t; #endif // defined(GGML_USE_HIP) fattn_kernel_t fattn_kernel; + bool use_sparse = false; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - fattn_kernel = flash_attn_ext_f16; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) { + constexpr bool use_sparse_kernel = true; + fattn_kernel = flash_attn_ext_f16; + use_sparse = true; + + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } + } else { + constexpr bool use_sparse_kernel = false; + fattn_kernel = flash_attn_ext_f16; + + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } + } + } else +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + { + constexpr bool use_sparse_kernel = false; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) - static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; - if (!shared_memory_limit_raised[id]) { - CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); - shared_memory_limit_raised[id] = true; - } + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } #endif // !defined(GGML_USE_MUSA) + } } else { constexpr bool use_logit_softcap = true; - fattn_kernel = flash_attn_ext_f16; + constexpr bool use_sparse_kernel = false; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -2272,7 +2347,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml } launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); } diff --git a/ggml/src/ggml-cuda/fattn-mma-turbo.cuh b/ggml/src/ggml-cuda/fattn-mma-turbo.cuh index f727e22ff825..2dfbab3c39d6 100644 --- a/ggml/src/ggml-cuda/fattn-mma-turbo.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-turbo.cuh @@ -65,7 +65,7 @@ void ggml_cuda_flash_attn_ext_mma_turbo_case(ggml_backend_cuda_context & ctx, gg fattn_kernel_t fattn_kernel; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -76,7 +76,7 @@ void ggml_cuda_flash_attn_ext_mma_turbo_case(ggml_backend_cuda_context & ctx, gg #endif // !defined(GGML_USE_MUSA) } else { constexpr bool use_logit_softcap = true; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -91,7 +91,7 @@ void ggml_cuda_flash_attn_ext_mma_turbo_case(ggml_backend_cuda_context & ctx, gg // the kernel receives raw quantized KV + the true byte pitch. stream_k = true. launch_fattn (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, - /*need_f16_K=*/false, /*need_f16_V=*/false, /*stream_k=*/true, warp_size_host); + /*need_f16_K=*/false, /*need_f16_V=*/false, /*stream_k=*/true, /*use_sparse=*/false, warp_size_host); } diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index 55ce2cc0bfb0..1103cff22794 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -1181,7 +1181,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1197,7 +1197,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1209,7 +1209,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1221,7 +1221,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1233,7 +1233,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1244,7 +1244,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index a5b92ebb2929..19c0eeb4402b 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -792,7 +792,7 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false, false); } template diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 8f174f10f4bf..008c08a9879d 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -6,11 +6,144 @@ #include "fattn-vec.cuh" #include "fattn.cuh" +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +__launch_bounds__(256, 1) +static __global__ void flash_attn_mask_to_sparse_indices( + const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max, + const int64_t s31, const int64_t s33) { + ggml_cuda_pdl_sync(); + + constexpr int values_per_lane = 8; + const int tid = threadIdx.x; + const int warp = tid / WARP_SIZE; + const int lane = tid % WARP_SIZE; + const int sequence = blockIdx.y; + const int query = blockIdx.x; + + const half * mask = mask_ptr + sequence*s33 + query*s31; + int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + query)*n_kv_max; + + __shared__ int warp_offsets[256/WARP_SIZE]; + __shared__ int row_count; + __shared__ int chunk_count; + + if (tid == 0) { + row_count = 0; + } + __syncthreads(); + + for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) { + uint32_t selected_warp[values_per_lane]; + int warp_count = 0; +#pragma unroll + for (int item = 0; item < values_per_lane; ++item) { + const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; + const bool selected = i < ne30 && isfinite(__half2float(mask[i])); + selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected); + warp_count += __popc(selected_warp[item]); + } + + if (lane == 0) { + warp_offsets[warp] = warp_count; + } + __syncthreads(); + + if (tid == 0) { + int offset = 0; +#pragma unroll + for (int iw = 0; iw < 256/WARP_SIZE; ++iw) { + const int count = warp_offsets[iw]; + warp_offsets[iw] = offset; + offset += count; + } + chunk_count = offset; + } + __syncthreads(); + + const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1; + int warp_item_offset = 0; +#pragma unroll + for (int item = 0; item < values_per_lane; ++item) { + const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; + const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask); + if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) { + indices[dst] = i; + } + warp_item_offset += __popc(selected_warp[item]); + } + __syncthreads(); + + if (tid == 0) { + row_count += chunk_count; + } + __syncthreads(); + } + + const int count = row_count; + for (int i = count + tid; i < n_kv_max; i += blockDim.x) { + indices[i] = -1; + } + __syncthreads(); + + // the dependent grid reads indices, signal once the row is complete + ggml_cuda_pdl_lc(); +} +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + +void ggml_cuda_flash_attn_ext_compact_mask( + const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(mask, indices, n_kv_max, stream); + GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA"); +#else + const int64_t s31 = mask->nb[1] / sizeof(half); + const int64_t s33 = mask->nb[3] / sizeof(half); + const dim3 blocks_num(mask->ne[1], mask->ne[3], 1); + const dim3 block_dim(256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params(blocks_num, block_dim, 0, stream); + ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33); + CUDA_CHECK(cudaGetLastError()); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +} + +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(ctx, dst); + return false; +#else + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * mask = dst->src[3]; + const int cc = ggml_cuda_info().devices[ctx.device].cc; + + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4); + return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && + mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f && + mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 && + K->ne[1] >= std::max(4096, 2LL*n_kv_max); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +} + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * Q = dst->src[0]; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + return; + } + } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ncols2 <= 8) { if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index b8881fa0543f..bc8e841c1012 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5532,6 +5532,15 @@ enum ggml_prec ggml_flash_attn_ext_get_prec( return (enum ggml_prec) prec_i32; } +void ggml_flash_attn_ext_set_n_kv_max( + struct ggml_tensor * a, + int32_t n_kv_max) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(n_kv_max >= 0); + + ggml_set_op_params_i32(a, 4, n_kv_max); +} + void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index e79cba6dcee2..ca2fb42d7aa6 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2549,6 +2549,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * kq_mask, ggml_tensor * sinks, ggml_tensor * v_mla, + int64_t n_kv_max, float kq_scale, int il) const { const bool v_trans = v->nb[1] > v->nb[2]; @@ -2590,6 +2591,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); + GGML_ASSERT(n_kv_max >= 0 && n_kv_max <= INT32_MAX); + ggml_flash_attn_ext_set_n_kv_max(cur, static_cast(n_kv_max)); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); // TurboQuant: inverse WHT on FA output when V values are WHT-rotated. @@ -2767,7 +2770,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = k_cur; ggml_tensor * v = v_cur; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + 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); if (wo) { @@ -2880,7 +2883,7 @@ ggml_tensor * llm_graph_context::build_attn( 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); + 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); // TurboQuant: if V was padded, the output has padded dimensions. @@ -3013,7 +3016,7 @@ ggml_tensor * llm_graph_context::build_attn( 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); + 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); // TurboQuant: if V was padded (MLA: V is view of K, may have padded dim), @@ -3124,7 +3127,7 @@ 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); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -3214,7 +3217,7 @@ ggml_tensor * llm_graph_context::build_attn( 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); + 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); // TurboQuant: if V was padded, extract original V head_dim after inverse WHT @@ -3312,7 +3315,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = k; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + 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); if (k_rot) { @@ -3371,7 +3374,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = k_cur; ggml_tensor * v = v_cur; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + 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); if (wo) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 66c0301a2a53..d7f858e35e0a 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1138,6 +1138,7 @@ struct llm_graph_context { ggml_tensor * kq_mask, ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + int64_t n_kv_max, float kq_scale, int il) const; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index e68dc49b6dff..c7658e462053 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -781,7 +781,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + const int64_t n_kv_max = std::min(raw_mask->ne[0], hparams.n_swa) + top_k->ne[0]; + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, n_kv_max, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -836,7 +837,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); cb(kq_mask, "hca_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -872,7 +873,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * k = mctx_cur->get_k(ctx0, il); - ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 5845096333dc..27eb479cebfb 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -634,7 +634,7 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr ggml_tensor * mask_b = ggml_view_2d(ctx0, kq_mask, n_kv, width, kq_mask->nb[1], (size_t) row0*kq_mask->nb[1]); - cur_b = build_attn_mha(Q_b, k_view, v_view, nullptr, mask_b, nullptr, nullptr, kq_scale, il); + cur_b = build_attn_mha(Q_b, k_view, v_view, nullptr, mask_b, nullptr, nullptr, 0, kq_scale, il); cur_b = ggml_mul(ctx0, cur_b, ggml_sigmoid(ctx0, gate_b)); cur_b = build_lora_mm(layer.wo, cur_b, layer.wo_s); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 335d139d9351..1c979b30f786 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -719,7 +719,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il); cb(cur, "kqv_out", il); // the rotation is its own inverse, so undo it on the value side of the output diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a040e82c12bc..82417b53062b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -206,6 +206,33 @@ static void init_tensor_kq_mask(ggml_tensor * tensor, float min = -1.0f, float m ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); } +static void init_tensor_kq_mask_sparse(ggml_tensor * tensor, int64_t n_kv_max) { + GGML_ASSERT(tensor->type == GGML_TYPE_F16); + GGML_ASSERT(n_kv_max > 1 && n_kv_max <= tensor->ne[0]); + + const int64_t ne0 = tensor->ne[0]; + const int64_t nrows = ggml_nrows(tensor); + std::vector data_f32(ggml_nelements(tensor), -INFINITY); + std::vector data_f16(ggml_nelements(tensor)); + std::vector order(ne0); + for (int64_t i = 0; i < ne0; ++i) { + order[i] = i; + } + + std::mt19937 gen(0x5A17); + for (int64_t row = 0; row < nrows; ++row) { + std::shuffle(order.begin(), order.end(), gen); + const int64_t count = n_kv_max - row % std::min(n_kv_max, 17); + std::sort(order.begin(), order.begin() + count); + for (int64_t i = 0; i < count; ++i) { + data_f32[row*ne0 + order[i]] = -0.03125f * (1 + (i + row) % 7); + } + } + + ggml_fp32_to_fp16_row(data_f32.data(), data_f16.data(), data_f16.size()); + ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); +} + // generate a lower triangular matrix static void init_tensor_tril(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { GGML_ASSERT(tensor->type == GGML_TYPE_F32); @@ -450,6 +477,7 @@ static std::string var_to_str(ggml_scale_mode mode) { #define VARS_TO_STR14(a, b, c, d, e, f, g, h, i, j, k, l, m, n) VAR_TO_STR(a) + "," + VARS_TO_STR13(b, c, d, e, f, g, h, i, j, k, l, m, n) #define VARS_TO_STR15(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) VAR_TO_STR(a) + "," + VARS_TO_STR14(b, c, d, e, f, g, h, i, j, k, l, m, n, o) #define VARS_TO_STR16(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) VAR_TO_STR(a) + "," + VARS_TO_STR15(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) +#define VARS_TO_STR17(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q) VAR_TO_STR(a) + "," + VARS_TO_STR16(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q) #ifdef GGML_USE_SYCL static bool inline _isinf(float f) { @@ -7226,9 +7254,10 @@ struct test_flash_attn_ext : public test_case { std::array permute; const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) const bool v_is_view_of_k; + const int64_t n_kv_max; std::string vars() override { - return VARS_TO_STR16(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k); + return VARS_TO_STR17(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k, n_kv_max); } double max_nmse_err() override { @@ -7245,9 +7274,9 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, - bool kv_view = true, bool v_is_view_of_k = false) + bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7307,6 +7336,7 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap); ggml_flash_attn_ext_add_sinks(out, s); + ggml_flash_attn_ext_set_n_kv_max(out, n_kv_max); ggml_flash_attn_ext_set_prec (out, prec); ggml_set_name(out, "out"); @@ -7319,7 +7349,11 @@ struct test_flash_attn_ext : public test_case { // make the sink values more noticeable in order to trigger a test failure when the implementation is wrong init_tensor_uniform(t, -10.0f, 10.0f); } else if (strcmp(t->name, "m") == 0) { - init_tensor_kq_mask(t); + if (n_kv_max > 0) { + init_tensor_kq_mask_sparse(t, n_kv_max); + } else { + init_tensor_kq_mask(t); + } } else { init_tensor_uniform(t); } @@ -10278,6 +10312,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + // Sparse mask hint: supported decode/prefill layouts and dense fallbacks. + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 2}, 4096, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 768)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 768)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2304)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); @@ -10683,6 +10725,27 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext_turbo4_vec(128)); test_cases.emplace_back(new test_flash_attn_ext_turbo4_vec(256)); + // sparse decode at long context + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 0)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 2048)); + + // q8_0 KV cases with long context (decode and prompt) + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 128, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 2048, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + for (int kv : { 4096, 8192, 16384,32768, 65536, }) { for (int hs : { 64, 128, 256, 576, }) { const int hsv = hs == 576 ? 512 : hs; From 209a77e554ecb7932d41dc4264d06feeed37b2be Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 15:39:32 -0600 Subject: [PATCH 05/16] qwen4exp: add NextN/MTP draft head (--spec-type draft-mtp) Ports upstream PR ggml-org/llama.cpp#27836 (open, unmerged) on top of this fork's qwen4exp support: the MTP head folds the next token's embedding into the trunk's wide hyper-connection residual, runs one trunk-shaped block (dense attention + MoE) over it, and collapses the result with its own hyper-connection mixer before reusing the trunk's LM head. Also updates the converter and gguf-py tensor mappings so conversion/qwen4exp.py --mtp can export the head. Also adds mtp_only/trunk_flags handling to load_arch_tensors, which the ported PR didn't include: without it, loading a standalone MTP-only checkpoint (all trunk tensors absent) fails outright instead of tolerating their absence, unlike the equivalent qwen35.cpp path. Detects MTP-only via the absence of blk.0.hc_attn_norm.weight, a tensor every trunk layer carries. Guards the embedding and LM-head fallbacks with a clear assert instead of a null-pointer crash when a checkpoint has neither a dedicated NextN embedding/head nor the trunk's own. Verified: full build, test-llama-archs shows no regressions. Loading a real Qwen3.8-Flash-Next MTP draft checkpoint (mtp-Qwen3.8-Flash- Next-shared-Q4_K_M.gguf) now gets past tensor loading correctly (the mtp_only path works) and fails with the new clear assertion rather than "token_embd.weight not found": that specific checkpoint has neither nextn.embed_tokens nor a trunk token_embd.weight anywhere, confirmed by inspecting its raw safetensors source directly (15 tensors total, no embedding table in any form). That's a model packaging gap in that specific file, not a code gap. --- conversion/qwen4exp.py | 68 ++++++- gguf-py/gguf/constants.py | 17 ++ gguf-py/gguf/tensor_mapping.py | 69 +++++++ src/llama-arch.cpp | 6 + src/llama-arch.h | 5 + src/llama-model.cpp | 3 +- src/llama-model.h | 6 + src/models/models.h | 12 +- src/models/qwen4exp.cpp | 356 +++++++++++++++++++++++++++++---- 9 files changed, 486 insertions(+), 56 deletions(-) diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py index 51a2fb7049b5..a36472a08589 100644 --- a/conversion/qwen4exp.py +++ b/conversion/qwen4exp.py @@ -1,7 +1,7 @@ from __future__ import annotations from pathlib import Path -from typing import Iterable +from typing import Callable, Iterable, cast import torch from torch import Tensor @@ -21,14 +21,13 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things: hyper-connections in place of every layer norm, QSA sparse attention on the full attention layers, and PLE n-gram hash embeddings on a single layer. + + The checkpoint also carries a NextN/MTP draft head under `mtp.*`, exported as a + trailing block; pass --no-nextn to leave it out. """ model_arch = gguf.MODEL_ARCH.QWEN4EXP - # the MTP block is a separate draft head; vLLM drops it too - supports_mtp_export = False - no_mtp = True - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # shards held only until the row stride is known, normally none @@ -39,6 +38,51 @@ def __init__(self, *args, **kwargs): self._ple_map: np.memmap | None = None self._ple_path: Path | None = None + # The MTP head is one trunk-shaped block (dense attention + MoE, wrapped in + # hyper-connections) plus a combiner, so once _QwenMtpMixin renames + # `mtp.layers.0.*` to the trailing block index its tensors ride the existing + # qwen4exp mappings unchanged. Only the two head-level pieces below differ. + + _MTP_MIXER_PREFIX = "mtp.hyper_connection_mixer." + + @classmethod + def filter_tensors(cls, item): + # the head carries its own copy of the trunk's hc_head_* output mixer, + # which qwen4exp has in place of a final norm; it is unindexed in the + # checkpoint and per-block in the GGUF + name, gen = item + if name.startswith("model." + cls._MTP_MIXER_PREFIX): + name = name.replace("model.", "", 1) + if name.startswith(cls._MTP_MIXER_PREFIX): + if cls.no_mtp: + return None + assert cls._original_block_count is not None + return f"model.layers.{cls._original_block_count}.{name[len('mtp.'):]}", gen + return super().filter_tensors((name, gen)) + + def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]: + # qwen4exp splits the combiner the shared NextN code calls eh_proj into + # fc_embedding and fc_hidden; W_e@e + W_h@h == [W_e|W_h] @ concat(e, h), + # so the two fuse back into the single expected matmul + tensors = super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + emb = tensors.pop("mtp.fc_embedding.weight", None) + hid = tensors.pop("mtp.fc_hidden.weight", None) + if emb is None and hid is None: + return tensors + if emb is None or hid is None: + raise ValueError( + "the qwen4exp MTP combiner needs both mtp.fc_embedding.weight and " + "mtp.fc_hidden.weight; pass --no-nextn to convert without the draft head" + ) + + assert self._original_block_count is not None + # fc_embedding first: the graph concatenates the token embedding ahead of + # the hidden state, so the fused weight has to be ordered to match + name = f"model.layers.{self._original_block_count}.eh_proj.weight" + tensors[name] = lambda: torch.cat([emb(), hid()], dim=1) + return tensors + def _read_hash_constants(self, suffix: str) -> list[int]: """Read an int64 PLE constant straight from the checkpoint. @@ -67,14 +111,18 @@ def set_gguf_parameters(self): self.gguf_writer.add_indexer_top_k(hp["indexer_budget"]) ratio = hp["indexer_compress_ratio"] layer_types = hp["layer_types"] - self.gguf_writer.add_attention_compress_ratios( - [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] - ) + ratios = [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] + # llama.cpp reads this array with length block_count, and the MTP blocks + # trailing the trunk attend densely, which is what a ratio of 0 selects + ratios += [0] * (self.block_count - n_layer) + self.gguf_writer.add_attention_compress_ratios(ratios) # ple_layer_ids is 1-based in the HF config; empty means no n-gram table, - # so emit no PLE keys rather than optional ones + # so emit no PLE keys rather than optional ones. + # a draft-only export carries no trunk tensors, so it carries no PLE table + # to describe either ple_layers = [i - 1 for i in hp["ple_layer_ids"]] - if not ple_layers: + if not ple_layers or self.mtp_only: return self.gguf_writer.add_ple_layers(ple_layers) self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 015d90fbfcc4..bfae97353bde 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1022,6 +1022,11 @@ class MODEL_TENSOR(IntEnum): NEXTN_HNORM = auto() NEXTN_SHARED_HEAD_HEAD = auto() NEXTN_SHARED_HEAD_NORM = auto() + # qwen4exp: the MTP head's own hyper-connection mixer, which stands in for the + # output norm the trunk does not have + NEXTN_HC_HEAD_NORM = auto() + NEXTN_HC_HEAD_DOWN = auto() + NEXTN_HC_HEAD_UP = auto() # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping @@ -1696,6 +1701,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm", MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: "blk.{bid}.nextn.hc_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: "blk.{bid}.nextn.hc_head_down", + MODEL_TENSOR.NEXTN_HC_HEAD_UP: "blk.{bid}.nextn.hc_head_up", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", @@ -2556,6 +2564,15 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PLE_NORM_QUERY, MODEL_TENSOR.PLE_NORM_CONV, MODEL_TENSOR.PLE_CONV1D, + # NextN/MTP draft head + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_HC_HEAD_NORM, + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN, + MODEL_TENSOR.NEXTN_HC_HEAD_UP, ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 8b8bf941d6f9..e8cd6fa90171 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2590,6 +2590,75 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", ), }, + MODEL_ARCH.QWEN4EXP: { + MODEL_TENSOR.HC_ATTN_NORM: ( + "model.layers.{bid}.attn_hyper_connection.hc_norm", + ), + MODEL_TENSOR.HC_ATTN_DOWN: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_down", + ), + MODEL_TENSOR.HC_ATTN_UP: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_up", + ), + MODEL_TENSOR.HC_ATTN_INJECT: ( + "model.layers.{bid}.attn_hyper_connection.block_inject_weight", + ), + MODEL_TENSOR.HC_FFN_NORM: ( + "model.layers.{bid}.mlp_hyper_connection.hc_norm", + ), + MODEL_TENSOR.HC_FFN_DOWN: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_down", + ), + MODEL_TENSOR.HC_FFN_UP: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_up", + ), + MODEL_TENSOR.HC_FFN_INJECT: ( + "model.layers.{bid}.mlp_hyper_connection.block_inject_weight", + ), + MODEL_TENSOR.HC_HEAD_NORM: ( + "model.hyper_connection_mixer.hc_norm", + ), + MODEL_TENSOR.HC_HEAD_DOWN: ( + "model.hyper_connection_mixer.input_mix_weight_down", + ), + MODEL_TENSOR.HC_HEAD_UP: ( + "model.hyper_connection_mixer.input_mix_weight_up", + ), + # the MTP head carries its own copy of the head mixer above + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: ( + "model.layers.{bid}.hyper_connection_mixer.hc_norm", + ), + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: ( + "model.layers.{bid}.hyper_connection_mixer.input_mix_weight_down", + ), + MODEL_TENSOR.NEXTN_HC_HEAD_UP: ( + "model.layers.{bid}.hyper_connection_mixer.input_mix_weight_up", + ), + MODEL_TENSOR.INDEXER_Q_NORM: ( + "model.layers.{bid}.self_attn.indexer.q_layernorm", + ), + MODEL_TENSOR.INDEXER_K_NORM: ( + "model.layers.{bid}.self_attn.indexer.k_layernorm", + ), + MODEL_TENSOR.PLE_KEY: ( + "model.layers.{bid}.ple.key_proj", + ), + MODEL_TENSOR.PLE_VALUE: ( + "model.layers.{bid}.ple.value_proj", + ), + MODEL_TENSOR.PLE_NORM_KEY: ( + "model.layers.{bid}.ple.norm_key", + ), + MODEL_TENSOR.PLE_NORM_QUERY: ( + "model.layers.{bid}.ple.norm_query", + ), + MODEL_TENSOR.PLE_NORM_CONV: ( + "model.layers.{bid}.ple.norm_conv", + ), + MODEL_TENSOR.PLE_CONV1D: ( + "model.layers.{bid}.ple.conv1d", + ), + }, } mapping: dict[str, tuple[MODEL_TENSOR, str]] diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 6df93bcd90e1..dae3807b4737 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -566,6 +566,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, + { LLM_TENSOR_NEXTN_HC_HEAD_NORM, "blk.%d.nextn.hc_head_norm" }, + { LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "blk.%d.nextn.hc_head_down" }, + { LLM_TENSOR_NEXTN_HC_HEAD_UP, "blk.%d.nextn.hc_head_up" }, { LLM_TENSOR_ATTN_SUB_NORM, "blk.%d.attn_sub_norm" }, { LLM_TENSOR_FFN_SUB_NORM, "blk.%d.ffn_sub_norm" }, { LLM_TENSOR_DEC_OUTPUT_NORM, "dec.output_norm" }, @@ -951,6 +954,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_NEXTN_HNORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_NEXTN_HC_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_NEXTN_HC_HEAD_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_NEXTN_HC_HEAD_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, // Nemotron 3 Super // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 4f16c4811f91..4aad0283b3b3 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -676,6 +676,11 @@ enum llm_tensor { LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + // qwen4exp: the MTP head ends in its own hyper-connection mixer rather than a + // plain RMSNorm, mirroring the trunk's hc_head_* (which is its output norm) + LLM_TENSOR_NEXTN_HC_HEAD_NORM, + LLM_TENSOR_NEXTN_HC_HEAD_DOWN, + LLM_TENSOR_NEXTN_HC_HEAD_UP, LLM_TENSOR_MASKED_EMBD_CENTROIDS, LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 915201355748..2102878587f4 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2378,7 +2378,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // attention KV cache for the MTP context instead of the hybrid wrapper. const bool mtp_on_hybrid_qwen = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && - (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE); + (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || + arch == LLM_ARCH_QWEN4EXP); if (llm_arch_is_recurrent(arch)) { res = new llama_memory_recurrent( diff --git a/src/llama-model.h b/src/llama-model.h index 5ec2e29e0d25..79affdd711b8 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -222,6 +222,12 @@ struct llama_layer_nextn { struct ggml_tensor * shared_head_head_s = nullptr; struct ggml_tensor * shared_head_head_in_s = nullptr; struct ggml_tensor * shared_head_norm = nullptr; + + // qwen4exp: the MTP head's own final hyper-connection mixer, which stands in for both + // the stream collapse and the output norm (the trunk has no separate output_norm either) + struct ggml_tensor * hc_head_norm = nullptr; + struct ggml_tensor * hc_head_down = nullptr; + struct ggml_tensor * hc_head_up = nullptr; }; struct llama_layer { diff --git a/src/models/models.h b/src/models/models.h index ddbba9fa2e3a..056a7f73ce62 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2135,7 +2135,12 @@ struct llama_model_qwen4exp : public llama_model_base { struct graph : public llm_build_delta_net_base { graph(const llama_model & model, const llm_graph_params & params); - private: + protected: + // tag-dispatched ctor for graph_mtp: binds the members without building the trunk + struct no_build_t {}; + graph(const llama_model & model, const llm_graph_params & params, no_build_t) : + llm_build_delta_net_base(params), model(model) {} + // HC replaces every layer norm: residual is [n_embd, hc, n_tokens] ggml_tensor * build_hc_mix( ggml_tensor * x, @@ -2225,6 +2230,11 @@ struct llama_model_qwen4exp : public llama_model_base { const llama_model & model; }; + // LLM_GRAPH_TYPE_DECODER_MTP draft head: one HC-wrapped dense-attention + MoE block + struct graph_mtp : public graph { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 1c979b30f786..7be9b2bb128d 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -24,6 +24,11 @@ static void qwen4exp_require_arr_len(llama_model_loader & ml, llm_kv kid, uint32 } void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { + // NextN/MTP: an extra decoder block appended past the trunk. Read this first, since + // n_layer() == n_layer_all - n_layer_nextn feeds every per-layer array below. + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count"); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -135,15 +140,22 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { const int64_t hc_dim = hc * n_embd; const int64_t hc_lr = hparams.hc_low_rank; - tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + // an MTP-only checkpoint (produced by conversion/qwen4exp.py --mtp) carries the trailing + // NextN block plus the trunk's shared embedding/head, but none of the trunk blocks + // themselves; blk.0.hc_attn_norm.weight exists on every trunk layer, so its absence + // marks this as MTP-only. + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.hc_attn_norm.weight") == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, trunk_flags); // there is no output_norm: the final hyper-connection mixer carries it - hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, 0); - hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, 0); - hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, 0); + hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, trunk_flags); + hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, trunk_flags); + hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, trunk_flags); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); - if (output == NULL) { + if (output == NULL && tok_embd != NULL) { output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); } @@ -169,9 +181,16 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { { hparams.ple_head_dim, ple_rows }, 0); } - for (int il = 0; il < n_layer; ++il) { + // MTP tensors sit in the trailing blocks; skip them entirely unless a draft head was asked for + const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + + for (int il = 0; il < (int) hparams.n_layer_all; ++il) { auto & layer = layers[il]; + // the MTP block is structurally a trunk block: is_recr()/is_ple() are both false past + // the trunk, so it takes the full-attention + MoE path below with no special casing + const int flags = il < n_layer ? trunk_flags : mtp_flags; + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; @@ -184,61 +203,86 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { const int64_t conv_dim = key_dim * 2 + value_dim; // two HC modules per layer: before the token mixer, before the MoE - layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, 0); - layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); - layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, 0); - layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0); - layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, 0); - layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); - layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, 0); - layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0); + layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, flags); + layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, flags); + layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, flags); + layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, flags); + layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, flags); + layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, flags); + layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, flags); + layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, flags); if (!hparams.is_recr(il)) { // full attention: wq holds [q|gate] interleaved per head - create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0); - layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0); + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags); - layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); - layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags); const int64_t idx_dim = hparams.indexer_head_size; - layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0); - layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0); - layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, 0); - layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, 0); + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, flags); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, flags); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, flags); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, flags); } else { - layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, 0); - layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0); - layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0); - layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, 0); - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, 0); - layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0); - layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0); - layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0); - layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0); + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, flags); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags); + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags); + layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags); } if (hparams.is_ple(il)) { - layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0); - layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0); - layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, 0); - layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0); - layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, 0); - layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); + layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, flags); + layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, flags); + layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, flags); + layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, flags); + layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, flags); + layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, flags); } - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); - create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags); + + layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags); - layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); - layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); - layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); - layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0); + if (il < n_layer) { + continue; + } + + // NextN/MTP head. enorm/hnorm gate the two inputs; eh_proj is the checkpoint's + // fc_embedding and fc_hidden fused side by side, so one matmul over + // concat(e, h) computes fc_embedding@e + fc_hidden@h. + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { hc_dim }, flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); + + // the head's own output mixer, mirroring the trunk's hc_head_*: it collapses the + // hc streams and stands in for the output norm, of which qwen4exp has none + layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { hc_dim }, flags); + layer.nextn.hc_head_down = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", il), { hc_dim, hc_lr }, flags); + layer.nextn.hc_head_up = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", il), { hc_lr, hc_dim }, flags); + + // qwen4exp sets mtp_use_dedicated_embeddings=false, so these are absent and the + // head falls back to the trunk's embedding table and LM head + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED); } } std::unique_ptr llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -378,6 +422,20 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa cur = build_layer_attn(inp->get_attn(), mctx_hyb, cur, inp_pos, sections, il); } + // an unmasked MTP export needs a hidden row for every token, so in that case the + // gather is deferred until after t_h_nextn is taken below + const bool gather_now = !cparams.embeddings_nextn || cparams.embeddings_nextn_masked; + + if (il == n_layer - 1 && inp_out_ids && gather_now) { + // everything below is per token, so drop the rows that produce no output + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inject = ggml_get_rows(ctx0, inject, inp_out_ids); + + res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + res_hc = build_hc_combine(res_hc, cur, inject, il); cur = build_hc_mix(res_hc, @@ -396,6 +454,23 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa cb(res_hc, "l_last", il); } + // The MTP head consumes the wide residual, before the head mixer collapses it. Export the + // combine result itself rather than a reshape of it: a pure view gets no backend assignment + // from the scheduler, and the readback in llama_context looks one up. It is contiguous, so + // [n_embd, hc, rows] already has the [n_embd_out, rows] layout the reader expects, and it + // carries exactly the right rows either way -- gathered above when masked, ungathered when not. + if (cparams.embeddings_nextn) { + cb(res_hc, "h_nextn", -1); + res->t_h_nextn = res_hc; + + // deferred from the last layer: collapse to the output rows now that the export is taken + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + } + // the final mixer is the output norm: there is no separate one ggml_tensor * cur = build_hc_mix(res_hc, model.hc_head_norm, model.hc_head_down, model.hc_head_up, @@ -415,6 +490,199 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, cur); } +// LLM_GRAPH_TYPE_DECODER_MTP draft head for qwen4exp. +// +// The head folds the next token's embedding into the trunk's wide hyper-connection residual, +// runs one trunk-style block over it, and collapses the result with its own mixer before +// reusing the trunk's LM head. The wide post-block residual is exported as t_h_nextn so the +// speculative driver can feed it straight back in for the next draft step. +// +// v1 simplification: the block attends densely. The trunk's QSA only prunes context past a +// 2048-token budget, so dense is a numerical superset; drafts are verified by the target +// either way. The indexer tensors are still loaded so the GGUF stays complete. +// TODO: wire up QSA here for long-context draft fidelity. +llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : + graph(model, params, no_build_t{}) { + GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN4EXP MTP requires n_layer_nextn > 0"); + GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN4EXP MTP currently only supports a single MTP block"); + GGML_ASSERT(ubatch.token && "QWEN4EXP MTP requires token input"); + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + GGML_ASSERT(hparams.n_embd_out() == (uint32_t) hc_dim && "QWEN4EXP MTP hidden width mismatch"); + + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.nextn.hc_head_norm && "MTP block missing nextn.hc_head_norm"); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + auto inp = std::make_unique(hc_dim); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_set_input(inp->embd); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + GGML_ASSERT((layer.nextn.embed_tokens || model.tok_embd) && + "QWEN4EXP MTP: checkpoint has neither nextn.embed_tokens nor a trunk token_embd.weight to draft from"); + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + cb(tok_embd, "mtp_tok_embd", il); + + ggml_tensor * h_state = ggml_reshape_3d(ctx0, inp->h, n_embd, hc, n_tokens); + cb(h_state, "mtp_h_state", il); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + auto * inp_attn = build_attn_inp_kv(); + + // grouped RMSNorm over the wide stream: normalise each hc stream, then scale the flattened + // [hc_dim] vector with the head's gamma, exactly as build_hc_mix does + ggml_tensor * h_norm = ggml_rms_norm(ctx0, h_state, hparams.f_norm_rms_eps); + h_norm = ggml_reshape_2d(ctx0, h_norm, hc_dim, n_tokens); + h_norm = ggml_mul(ctx0, h_norm, layer.nextn.hnorm); + h_norm = ggml_reshape_3d(ctx0, h_norm, n_embd, hc, n_tokens); + cb(h_norm, "mtp_hnorm", il); + + // the token embedding is shared across the streams, so broadcast it to hc copies + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + e_norm = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens), + n_embd, hc, n_tokens, 1); + cb(e_norm, "mtp_enorm", il); + + // eh_proj holds fc_embedding and fc_hidden side by side, so this one matmul is + // fc_embedding @ e_norm + fc_hidden @ h_norm, applied to each stream independently. + // Keeping the streams distinct here is the point of the hyper-connection residual: + // pooling them before the projection would throw that away. + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * res_hc = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(res_hc, "mtp_eh_proj", il); + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_mix(res_hc, + layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, + &inject, il); + cb(cur, "mtp_hc_attn_pre", il); + + // ---- dense attention, mirroring the trunk's full-attention branch ---- + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + cb(Qcur_full, "mtp_Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0); + Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "mtp_Qcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "mtp_gate", il); + + ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "mtp_Kcur_normed", il); + + ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + cb(Vcur, "mtp_Vcur", il); + + // IMRoPE, same convention and freq_base as the trunk + Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "mtp_Qcur", il); + cb(Kcur, "mtp_Kcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + cur = build_attn(inp_attn, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "mtp_attn_pregate", il); + + cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate)); + cb(cur, "mtp_attn_gated", il); + + cur = build_lora_mm(layer.wo, cur, layer.wo_s); + cb(cur, "mtp_attn_out", il); + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inject = ggml_get_rows(ctx0, inject, inp_out_ids); + + res_hc = ggml_reshape_2d(ctx0, res_hc, hc_dim, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + + res_hc = build_hc_combine(res_hc, cur, inject, il); + cb(res_hc, "mtp_hc_attn_post", il); + + // ---- MoE, identical to the trunk's build_layer_ffn ---- + cur = build_hc_mix(res_hc, + layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, + &inject, il); + cb(cur, "mtp_hc_ffn_pre", il); + + cur = build_layer_ffn(cur, il); + cb(cur, "mtp_ffn_out", il); + + res_hc = build_hc_combine(res_hc, cur, inject, il); + cb(res_hc, "mtp_hc_ffn_post", il); + + // The next draft step re-enters here, so export the wide stream before it is collapsed. + // As in the trunk, export the combine result rather than a reshape view of it. + cb(res_hc, "h_nextn", -1); + res->t_h_nextn = res_hc; + + // the head's own mixer collapses the streams and doubles as the output norm + cur = build_hc_mix(res_hc, + layer.nextn.hc_head_norm, layer.nextn.hc_head_down, layer.nextn.hc_head_up, + nullptr, nullptr, -1); + cb(cur, "mtp_hc_head", -1); + + // deliberately no res->t_embd: it would be n_embd wide while the context sizes its + // embedding buffer by n_embd_out (the wide stream). The driver reads t_h_nextn instead. + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w && "QWEN4EXP MTP: missing LM head (nextn.shared_head_head or model.output)"); + + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + std::pair llama_model_qwen4exp::graph::build_qkvz( ggml_tensor * input, int il) { From 7533986a2b38e9c3fa45cc3b4de7ef409d9838bd Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 18:03:33 -0600 Subject: [PATCH 06/16] fix: cast F16 norm weights to F32 before CUDA broadcast-mul qwen4exp MTP was crashing with GGML_ASSERT(nb10 % sizeof(src1_t) == 0) in ggml_cuda_op_bin_bcast when running real MTP-packaged checkpoints that store some norm weights (e.g. nextn.enorm) as F16. The CUDA broadcast-mul dispatch has a path for F16-activation x F32/F16-weight, but none for F32-activation x F16-weight, which build_norm's grouped RMSNorm+scale pattern hits whenever a norm weight isn't F32. Cast the weight to F32 first when this combination is detected; a no-op for the common case where norm weights are already F32. Fixes qwen4exp NextN/MTP speculative decoding startup crash against real Qwen3.8-Flash-Next MTP GGUFs. Co-Authored-By: Claude Sonnet 5 --- src/llama-graph.cpp | 9 +++++++++ src/models/qwen4exp.cpp | 13 ++++++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index ca2fb42d7aa6..df714d96740b 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1603,6 +1603,12 @@ ggml_tensor * llm_graph_context::build_norm( } if (mw) { + // the CUDA broadcast-mul kernel has no path for an F32 activation times an F16 operand + // (only the reverse, F16 activation times F32/F16 operand); most checkpoints keep norm + // weights in F32 so this is normally a no-op, but some conversions store them narrower + if (mw->type != cur->type && cur->type == GGML_TYPE_F32) { + mw = ggml_cast(ctx0, mw, GGML_TYPE_F32); + } cur = ggml_mul(ctx0, cur, mw); if (mb) { cb(cur, "norm_w", il); @@ -1610,6 +1616,9 @@ ggml_tensor * llm_graph_context::build_norm( } if (mb) { + if (mb->type != cur->type && cur->type == GGML_TYPE_F32) { + mb = ggml_cast(ctx0, mb, GGML_TYPE_F32); + } cur = ggml_add(ctx0, cur, mb); } diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 7be9b2bb128d..d1d7f0122d50 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -304,6 +304,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix( // the converter folded each gamma to (1 + w) ggml_tensor * xn = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps); xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt); + // the CUDA broadcast-mul kernel has no path for an F32 activation times an F16 operand + if (w_norm->type != xn->type && xn->type == GGML_TYPE_F32) { + w_norm = ggml_cast(ctx0, w_norm, GGML_TYPE_F32); + } xn = ggml_mul(ctx0, xn, w_norm); cb(xn, "hc_norm", il); @@ -554,7 +558,11 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ // [hc_dim] vector with the head's gamma, exactly as build_hc_mix does ggml_tensor * h_norm = ggml_rms_norm(ctx0, h_state, hparams.f_norm_rms_eps); h_norm = ggml_reshape_2d(ctx0, h_norm, hc_dim, n_tokens); - h_norm = ggml_mul(ctx0, h_norm, layer.nextn.hnorm); + ggml_tensor * hnorm_w = layer.nextn.hnorm; + if (hnorm_w->type != h_norm->type && h_norm->type == GGML_TYPE_F32) { + hnorm_w = ggml_cast(ctx0, hnorm_w, GGML_TYPE_F32); + } + h_norm = ggml_mul(ctx0, h_norm, hnorm_w); h_norm = ggml_reshape_3d(ctx0, h_norm, n_embd, hc, n_tokens); cb(h_norm, "mtp_hnorm", il); @@ -1478,6 +1486,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple( ggml_tensor * t = ggml_reshape_3d(ctx0, x, n_embd, hc, n_tokens); t = ggml_rms_norm(ctx0, t, hparams.f_norm_rms_eps); t = ggml_reshape_2d(ctx0, t, hc_dim, n_tokens); + if (w->type != t->type && t->type == GGML_TYPE_F32) { + w = ggml_cast(ctx0, w, GGML_TYPE_F32); + } t = ggml_mul(ctx0, t, w); return ggml_reshape_3d(ctx0, t, n_embd, hc, n_tokens); }; From 3888e18f09f1365aa645cca9fd5c4be27b8e5c22 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 18:35:19 -0600 Subject: [PATCH 07/16] fix: don't reuse the Gemma4-assistant same-position draft path for real NextN heads The speculative driver's is_mem_shared check (true whenever no separate -md draft model is given, e.g. qwen4exp's self-contained MTP) was being used to select Gemma4-assistant's same-llama_pos-for-every-draft-token behavior. That behavior is specific to Gemma4-assistant's early-exit self-speculation, not a property of "shares KV memory with target" in general - a real trained NextN head (qwen35, qwen4exp) still needs an incrementing position per draft step even in that mode. Add llama_model_uses_shared_position_draft(), gated on arch == LLM_ARCH_GEMMA4_ASSISTANT specifically, and use it instead of the blanket is_mem_shared check for that one branch. Co-Authored-By: Claude Sonnet 5 --- common/speculative.cpp | 6 +++++- src/llama-context.cpp | 4 ++++ src/llama-ext.h | 7 +++++++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index bfcc7f168d24..617bdeefaa6d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1325,6 +1325,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // neither (qwen35 / qwen35moe): a single trained MTP head. int32_t n_mtp_layers = 1; bool is_mem_shared = false; // gemma4 + bool same_position_draft = false; // gemma4-assistant only: every draft row in a round shares n_past bool chain_heads = false; // derived in the ctor: n_mtp_layers > 1 && !is_mem_shared // Per-sequence cross-batch carryover: pair (h_p, x_{p+1}) at MTP pos p+1. @@ -1418,6 +1419,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true); is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt; + same_position_draft = is_mem_shared && llama_model_uses_shared_position_draft(llama_get_model(ctx_dft)); chain_heads = n_mtp_layers > 1 && !is_mem_shared; chain_graph = !is_mem_shared && !chain_heads && chain_enabled && llama_model_supports_mtp_chain(llama_get_model(ctx_dft)); @@ -2107,12 +2109,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, chain_h[seq_id].data() + (size_t) t * n_embd, row_bytes); } - } else if (is_mem_shared) { + } else if (same_position_draft) { // note: with shared memory (e.g. Gemma4 assistants) we use the same position for all draft tokens // ref: https://github.com/huggingface/transformers/blob/effde20942e3f82a1b97449f60b3a48c5ff96145/docs/source/en/model_doc/gemma4_assistant.md?plain=1#L36-L37 common_batch_add(batch, id, dp.n_past, { seq_id }, true); std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes); } else { + // is_mem_shared models with a real trained NextN head (qwen35, qwen4exp, ...) + // still draft at incrementing positions, same as the non-shared-memory path common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true); std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes); } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3d8ff72478e6..b19b2683523f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4081,6 +4081,10 @@ bool llama_model_supports_mtp_chain(const llama_model * model) { return model != nullptr && model->arch == LLM_ARCH_QWEN35; } +bool llama_model_uses_shared_position_draft(const llama_model * model) { + return model != nullptr && model->arch == LLM_ARCH_GEMMA4_ASSISTANT; +} + void llama_set_mtp_chain(llama_context * ctx, bool value) { ctx->set_mtp_chain(value); } diff --git a/src/llama-ext.h b/src/llama-ext.h index dfe05f234c38..eb5a4752938b 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -118,6 +118,13 @@ LLAMA_API void llama_set_nextn_layer_offset(struct llama_context * ctx, int32_t LLAMA_API bool llama_model_supports_mtp_chain(const struct llama_model * model); +// True only for architectures whose "draft" is a same-position early-exit of the +// target's own trunk (gemma4-assistant), where every drafted token in a round shares +// llama_pos n_past. A dedicated trained NextN/MTP head (qwen35, qwen4exp, deepseek, ...) +// still needs an incrementing position per draft step even when it shares KV memory +// with the target (no separate -md model given). +LLAMA_API bool llama_model_uses_shared_position_draft(const struct llama_model * model); + // Run the DECODER_MTP graph in chained mode: the batch's first row carries the // real (token, h) inputs and each following row's inputs come from the previous // row's in-graph argmax and hidden state. One decode drafts n_tokens tokens. From 97686af29ee3dddd0cbfb8ac6ccc42da0468cc44 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 19:29:41 -0600 Subject: [PATCH 08/16] fix: qwen4exp MTP head should reuse the trunk's shared hc_head mixer Ported PR 27836 gave the MTP block its own private per-layer final mixer (layer.nextn.hc_head_norm/down/up, tensor names blk.N.nextn.hc_head_*). The actual, community-validated implementation (PR 27739, reconciled by LaurentZuijdwijk - the one real MTP-head exports like dzannotti/Qwen3.8-Flash-Next-MTP-GGUF are built against) instead trains a single hc mixer shared between the trunk's own final layer and the MTP head (exported as top-level output_hc_norm/down/up, loaded here as model.hc_head_norm/down/up). Files following that convention have no blk.N.nextn.hc_head_* tensors at all, so graph_mtp's GGML_ASSERT on it would abort at graph-build time; conversely, per PR 27836's convention the head ran through a mixer that was never actually trained as a distinct MTP-specific output projection. Confirmed against a real MTP head export off HF that model.hc_head_* is the correct, always-present tensor for this: switch graph_mtp to use it instead of the private per-layer copy. Co-Authored-By: Claude Sonnet 5 --- src/models/qwen4exp.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index d1d7f0122d50..bdcb041111a4 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -521,7 +521,10 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); - GGML_ASSERT(layer.nextn.hc_head_norm && "MTP block missing nextn.hc_head_norm"); + // the MTP head's final mixer is the trunk's own output_hc_* (model.hc_head_*), not a + // private per-layer copy: upstream trains one hc mixer, shared between the trunk's last + // layer and the draft head, same as the trunk's own final-output call below. + GGML_ASSERT(model.hc_head_norm && "QWEN4EXP MTP: model missing hc_head_norm (trunk output mixer)"); int sections[4]; std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); @@ -671,9 +674,10 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ cb(res_hc, "h_nextn", -1); res->t_h_nextn = res_hc; - // the head's own mixer collapses the streams and doubles as the output norm + // the final mixer is shared with the trunk's own output mixer (model.hc_head_*), not a + // private per-layer copy -- see the GGML_ASSERT above cur = build_hc_mix(res_hc, - layer.nextn.hc_head_norm, layer.nextn.hc_head_down, layer.nextn.hc_head_up, + model.hc_head_norm, model.hc_head_down, model.hc_head_up, nullptr, nullptr, -1); cb(cur, "mtp_hc_head", -1); From b7f556e8086f8c84bfcc733feb12554c08f4141b Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 19:34:38 -0600 Subject: [PATCH 09/16] fix: export qwen4exp MTP's t_h_nextn before the inp_out_ids gather graph_mtp gathered attention output and res_hc down to just the requested output rows immediately after attention, then exported that already-gathered res_hc as t_h_nextn. Any decode with fewer requested-output rows than input rows - notably the speculative driver's catch-up/prefill decode into the draft context, which requests logits for none of its rows - would export a zero-row (or otherwise truncated) t_h_nextn. The driver's per-token shift-by-one hidden-state handoff needs one row per input token regardless of which rows have requested logits. Matches PR 27739's (the community-validated implementation) ordering: res->t_h_nextn is assigned the full per-token combine result, and the inp_out_ids gather happens after, scoped to just the final hc_head mix / LM head projection. Co-Authored-By: Claude Sonnet 5 --- src/models/qwen4exp.cpp | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index bdcb041111a4..dc5feaed2496 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -645,15 +645,6 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ cur = build_lora_mm(layer.wo, cur, layer.wo_s); cb(cur, "mtp_attn_out", il); - if (inp_out_ids) { - cur = ggml_get_rows(ctx0, cur, inp_out_ids); - inject = ggml_get_rows(ctx0, inject, inp_out_ids); - - res_hc = ggml_reshape_2d(ctx0, res_hc, hc_dim, res_hc->ne[2]); - res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); - res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); - } - res_hc = build_hc_combine(res_hc, cur, inject, il); cb(res_hc, "mtp_hc_attn_post", il); @@ -669,11 +660,20 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ res_hc = build_hc_combine(res_hc, cur, inject, il); cb(res_hc, "mtp_hc_ffn_post", il); - // The next draft step re-enters here, so export the wide stream before it is collapsed. + // The next draft step re-enters here, so export the wide stream before it is collapsed -- + // and before the inp_out_ids gather below, so the driver's per-token shift-by-one handoff + // always gets one row per input token, not just the requested-output rows (a catch-up / + // prefill decode requests logits for none of its rows, which would otherwise export zero). // As in the trunk, export the combine result rather than a reshape view of it. cb(res_hc, "h_nextn", -1); res->t_h_nextn = res_hc; + if (inp_out_ids) { + res_hc = ggml_reshape_2d(ctx0, res_hc, hc_dim, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + // the final mixer is shared with the trunk's own output mixer (model.hc_head_*), not a // private per-layer copy -- see the GGML_ASSERT above cur = build_hc_mix(res_hc, From e5a17b5b9c8de47d88ed7a10479d601fc385acc9 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 19:58:12 -0600 Subject: [PATCH 10/16] fix: qwen4exp trunk double-gathered output rows, crashing MTP targets graph::graph gathers cur/res_hc down to the requested output rows exactly once - via the early gather right after the last layer's attention when gather_now is true, or the deferred one inside the embeddings_nextn block when it's false - the two conditions are complements, so exactly one path always fires when inp_out_ids exists. A third, unconditional gather right before the final output norm re-applied inp_out_ids to the already-reduced tensor, indexing with values sized for the original token count against a tensor that now only has output_row_count rows. This is silently harmless whenever n_outputs == n_tokens (every position requested logits, gathering by the identity permutation twice is a no-op), which is why normal generation never hit it. It reliably crashes (GGML_ASSERT(i01 >= 0 && i01 < ne01) in ggml_compute_forward_get_rows) whenever n_outputs < n_tokens on a context with embeddings_nextn set - concretely, the standard llama.cpp warmup decodes 2 tokens and requests logits for 1, and any target context paired with an MTP speculative draft sets embeddings_nextn unconditionally. So this fired on every server startup once a real -md draft-mtp config was used, independent of which draft checkpoint or graph_mtp bugs were involved. Removing the redundant gather - res_hc is already correctly sized by the time this runs - fixes it. Also drop the now-dead requirement on layer.nextn.hc_head_norm/down/up (graph_mtp reuses model.hc_head_* as of the previous commit): make them TENSOR_NOT_REQUIRED so files exporting per-block hyper-connection mixers in the old PR 27836 tensor layout don't fail to load, without requiring the tensor from files that (correctly) don't have it. Verified against a real, cleanly-exported MTP draft head (dzannotti/Qwen3.8-Flash-Next-MTP-GGUF) paired via -md: loads without crashing, and speculative decoding now gets real acceptance (26/26 on a trivial prompt, 71/88 = 80.7% on natural-language generation - in line with the community-reported 0.74-0.90 range for this same head). Co-Authored-By: Claude Sonnet 5 --- src/models/qwen4exp.cpp | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index dc5feaed2496..f97f15524a45 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -266,11 +266,12 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { hc_dim }, flags); layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); - // the head's own output mixer, mirroring the trunk's hc_head_*: it collapses the - // hc streams and stands in for the output norm, of which qwen4exp has none - layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { hc_dim }, flags); - layer.nextn.hc_head_down = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", il), { hc_dim, hc_lr }, flags); - layer.nextn.hc_head_up = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", il), { hc_lr, hc_dim }, flags); + // unused: graph_mtp reuses the trunk's own model.hc_head_* (see qwen4exp.cpp's + // graph_mtp). Kept optional here only so files that still carry this tensor + // (e.g. blk.N.nextn.hc_head_* from an older PR 27836-style export) still load. + layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { hc_dim }, flags | TENSOR_NOT_REQUIRED); + layer.nextn.hc_head_down = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", il), { hc_dim, hc_lr }, flags | TENSOR_NOT_REQUIRED); + layer.nextn.hc_head_up = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", il), { hc_lr, hc_dim }, flags | TENSOR_NOT_REQUIRED); // qwen4exp sets mtp_use_dedicated_embeddings=false, so these are absent and the // head falls back to the trunk's embedding table and LM head @@ -475,15 +476,14 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa } } + // res_hc is already reduced to the output rows by now -- via the early gather above when + // gather_now was true, or the deferred one just above when it wasn't -- inp_out_ids picks + // exactly one of those paths every time it's non-null, so no further gather belongs here. // the final mixer is the output norm: there is no separate one ggml_tensor * cur = build_hc_mix(res_hc, model.hc_head_norm, model.hc_head_down, model.hc_head_up, nullptr, nullptr, -1); - if (inp_out_ids) { - cur = ggml_get_rows(ctx0, cur, inp_out_ids); - } - cb(cur, "result_norm", -1); res->t_embd = cur; From 9f8cb2d3d5cbf2b4ed3ea5d92b6757bca7c520e0 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 20:49:12 -0600 Subject: [PATCH 11/16] CUDA: key captured graphs by shape, not first-node address The CUDA graph cache keyed captured graphs by the raw memory address of their first node (cgraph->nodes[0]). A captured graph hard-codes its shapes, but speculative decoding constantly alternates between different batch shapes (draft steps, verify batches, catch-up decodes) on the same context - when a new shape happens to reuse the same first- node address as a stale cached graph for a different shape, capture either reuses the wrong graph or thrashes, permanently resetting warmup instead of ever converging to steady-state replay. Hash node count and both endpoint tensors' shapes into the key instead (O(1) - walking all nodes would defeat the point of a CUDA graph), add LRU eviction capped at 64 graphs so the map can't grow unbounded now that distinct shapes get distinct entries. A shape this still fails to separate re-captures exactly as before, so it can't regress anything. Cherry-picked from ggml-org/llama.cpp#28243 (open, unmerged), which found this while working on qwen4exp MTP performance - the effect is generic to any speculative-decoding workload on this fork, not qwen4exp-specific, so pulling in just this piece rather than the rest of that PR (which also reworks qwen4exp trunk/draft tensor sharing and doesn't fix the mixer/export-timing bugs already fixed on this branch). Verified: qwen4exp MTP speculative decoding still produces correct, byte-identical (temp 0) output after this change. Speed effect is hard to isolate cleanly from the dominant MoE-cache warmup effect already documented on this branch, but the fix is justified on its own correctness merits regardless of measured delta. Co-Authored-By: Claude Sonnet 5 --- ggml/src/ggml-cuda/common.cuh | 21 +++++++++++++------ ggml/src/ggml-cuda/ggml-cuda.cu | 37 ++++++++++++++++++++++++++------- 2 files changed, 44 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index ba029e0651a2..92d2fdd4fe73 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1436,13 +1436,13 @@ struct ggml_backend_cuda_context { bool fa_f16_use_pool = false; #ifdef USE_CUDA_GRAPH - // Map from first_node_ptr to cuda_graph - allows multiple graphs per context - // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) - std::unordered_map> cuda_graphs; + std::unordered_map> cuda_graphs; + + static const size_t max_cuda_graphs = 64; int64_t last_graph_eviction_sweep = 0; - ggml_cuda_graph * cuda_graph(const void * first_node_ptr) { + ggml_cuda_graph * cuda_graph(uint64_t graph_key) { const int64_t time_now = ggml_time_us(); // sweep every 5s, evicting cuda graphs unused for >=10s @@ -1457,9 +1457,18 @@ struct ggml_backend_cuda_context { } } - auto it = cuda_graphs.find(first_node_ptr); + auto it = cuda_graphs.find(graph_key); if (it == cuda_graphs.end()) { - it = cuda_graphs.emplace(first_node_ptr, std::make_unique()).first; + while (cuda_graphs.size() >= max_cuda_graphs) { + auto lru = cuda_graphs.begin(); + for (auto c = cuda_graphs.begin(); c != cuda_graphs.end(); ++c) { + if (c->second->last_used_time < lru->second->last_used_time) { + lru = c; + } + } + cuda_graphs.erase(lru); + } + it = cuda_graphs.emplace(graph_key, std::make_unique()).first; } it->second->last_used_time = time_now; return it->second.get(); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f0b9829a1177..fc9393c11ae1 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2852,14 +2852,35 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { return use_cuda_graph; } -static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { - return cgraph->nodes[0]; +// a captured graph hard-codes its shapes, so with one key per split an alternating shape +// (a speculative verify batch) resets warmup forever. O(1) on purpose: walking nodes undoes the +// point of a cuda graph. A shape this fails to separate re-captures as before, so it cannot regress. +static uint64_t ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { + // unlike the previous key this dereferences nodes[0], so an empty graph is not safe here + if (cgraph->n_nodes <= 0) { + return 0; + } + + uint64_t key = (uint64_t) (uintptr_t) cgraph->nodes[0]; + + auto mix = [&key](uint64_t v) { + key = (key ^ v) * 0x100000001b3ull; + }; + + mix(cgraph->n_nodes); + + for (int d = 0; d < GGML_MAX_DIMS; d++) { + mix(cgraph->nodes[0]->ne[d]); + mix(cgraph->nodes[cgraph->n_nodes - 1]->ne[d]); + } + + return key; } static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { bool res = false; - const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const uint64_t graph_key = ggml_cuda_graph_get_key(cgraph); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (cgraph->uid != 0 && @@ -2898,7 +2919,7 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx return res; } -static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { +static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, uint64_t graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); #if CUDART_VERSION >= 12000 @@ -4197,7 +4218,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 0; } -static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { +static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, uint64_t graph_key) { bool graph_evaluated_or_captured = false; // flag used to determine whether it is an integrated_gpu @@ -4416,7 +4437,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } #ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { +static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, uint64_t graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (graph->graph == nullptr) { @@ -4439,7 +4460,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, bool use_cuda_graph = false; bool cuda_graph_update_required = false; - const void * graph_key = nullptr; + uint64_t graph_key = 0; // [TAG_FA_F16_CUDA_GRAPHS] default: no graph will be captured for this cgraph, so HIP flash- // attention keeps its raw (release-after-use) f16 temp path. Set true below only when the graph @@ -4538,7 +4559,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; #ifdef USE_CUDA_GRAPH - const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const uint64_t graph_key = ggml_cuda_graph_get_key(cgraph); const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); #else const bool use_cuda_graph = false; From 323c636ebc9c0a613ced9a2b663c56f383c488a8 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 22:23:17 -0600 Subject: [PATCH 12/16] cuda: fix turbo K/V swizzle write/read mismatch in flash-attn tile loaders The turbo2/turbo3/turbo4 SMEM loaders wrote dequantized K/V elements with plain linear indexing while the XOR-swizzled ldmatrix reads (used on Turing+ once nbatch_K2/V2 is bank-aligned) expected byte offsets from fattn-swizzle.cuh's bytes_rc(). Every read landed on the wrong element, producing garbage turbo attention output whenever swz was active. Add turbo_store_h2(), mirroring the f16 loader's existing if constexpr (swz) pattern, and route all six turbo write sites through it. Also fix fattn-mma-turbo.cuh's SMEM sizing to use tile_stride() instead of the stale nbatch_K2+4/nbatch_V2+4 formula, matching the actual swizzled stride. Fixes the turbo3/turbo4 FLASH_ATTN_EXT regression reported in PR #340. All 10955 FLASH_ATTN_EXT test-backend-ops cases pass, including the previously-failing turbo4_vec_q8_0_turbo4_d128_kv256/d256_kv256 cases. Co-Authored-By: Claude Sonnet 5 --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 42 +++++++++++++++++--------- ggml/src/ggml-cuda/fattn-mma-turbo.cuh | 6 +++- 2 files changed, 32 insertions(+), 16 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index dcf5d381df82..571ccf078f6b 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -541,7 +541,19 @@ static __constant__ float TURBO_CENTROIDS_4BIT_FATTN[16] = { // low nibble = elem 2c, high nibble = elem 2c+1. Hence one byte qs[col_offset+c] yields // the half2 for tile column c. sizeof(block_turbo4_0)-driven pointer math; never assume // 66/68 or a qs offset constant. -template +// Writes one half2 element (row, c) into a turbo tile at the same swizzled byte offset the +// f16 loader and every ldmatrix read use (fattn-swizzle.cuh::bytes_rc), so turbo K/V tiles +// stay readable once swz makes stride_tile bank-aligned. Linear layout when swz is false. +template +static __device__ __forceinline__ void turbo_store_h2(half2 * const __restrict__ tile_KV, const int row, const int c, const half2 v) { + if constexpr (swz) { + *(half2 *) ((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc(row, c)) = v; + } else { + tile_KV[row*stride_tile + c] = v; + } +} + +template static __device__ __forceinline__ void flash_attn_ext_turbo4_load_tile( const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, const int D2, const int stride_bytes, const int col_offset, const int i_sup) { @@ -551,7 +563,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo4_load_tile( for (int row = tid; row < nbatch_fa; row += nthreads) { if (oob_check && row >= i_sup) { for (int c = 0; c < D2; ++c) { - tile_KV[row*stride_tile + c] = make_half2(0.0f, 0.0f); + turbo_store_h2(tile_KV, row, c, make_half2(0.0f, 0.0f)); } continue; } @@ -586,7 +598,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo4_load_tile( #else const uint8_t byte = blk->qs[in_blk]; #endif - tile_KV[row*stride_tile + c] = __halves2half2(scaled[byte & 0xF], scaled[byte >> 4]); + turbo_store_h2(tile_KV, row, c, __halves2half2(scaled[byte & 0xF], scaled[byte >> 4])); } } } @@ -599,7 +611,7 @@ static __constant__ float TURBO_CENTROIDS_3BIT_FATTN[8] = { -0.190207f, -0.118786f, -0.066822f, -0.021663f, 0.021663f, 0.066822f, 0.118786f, 0.190207f }; -template +template static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile( const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, const int D2, const int stride_bytes, const int col_offset, const int i_sup) { @@ -608,7 +620,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile( #pragma unroll for (int row = tid; row < nbatch_fa; row += nthreads) { if (oob_check && row >= i_sup) { - for (int c = 0; c < D2; ++c) tile_KV[row*stride_tile + c] = make_half2(0.0f, 0.0f); + for (int c = 0; c < D2; ++c) turbo_store_h2(tile_KV, row, c, make_half2(0.0f, 0.0f)); continue; } const char * row_ptr = KV_raw + (int64_t)row * stride_bytes; @@ -645,7 +657,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile( const int shift = (j0 % 4) * 2; const uint8_t idx0 = ((qs_byte >> shift) & 0x3) | (((sgn_byte >> (j0 % 8)) & 0x1) << 2); const uint8_t idx1 = ((qs_byte >> (shift+2)) & 0x3) | (((sgn_byte >> (j0 % 8 + 1)) & 0x1) << 2); - tile_KV[row*stride_tile + c] = __halves2half2(scaled[idx0], scaled[idx1]); + turbo_store_h2(tile_KV, row, c, __halves2half2(scaled[idx0], scaled[idx1])); } } } @@ -655,7 +667,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile( static __constant__ float TURBO_CENTROIDS_2BIT_FATTN[4] = { -0.133462f, -0.039994f, 0.039994f, 0.133462f }; -template +template static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile( const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, const int D2, const int stride_bytes, const int col_offset, const int i_sup) { @@ -664,7 +676,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile( #pragma unroll for (int row = tid; row < nbatch_fa; row += nthreads) { if (oob_check && row >= i_sup) { - for (int c = 0; c < D2; ++c) tile_KV[row*stride_tile + c] = make_half2(0.0f, 0.0f); + for (int c = 0; c < D2; ++c) turbo_store_h2(tile_KV, row, c, make_half2(0.0f, 0.0f)); continue; } const char * row_ptr = KV_raw + (int64_t)row * stride_bytes; @@ -699,7 +711,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile( const int shift = (j0 % 4) * 2; const uint8_t idx0 = (qs_byte >> shift) & 0x3; const uint8_t idx1 = (qs_byte >> (shift+2)) & 0x3; - tile_KV[row*stride_tile + c] = __halves2half2(scaled[idx0], scaled[idx1]); + turbo_store_h2(tile_KV, row, c, __halves2half2(scaled[idx0], scaled[idx1])); } } } @@ -888,13 +900,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nthreads_turbo = nwarps * ggml_cuda_get_physical_warp_size(); const char * K_raw = (const char *) K_h2 + int64_t(k_VKQ_0) * stride_K; if constexpr (type_K == GGML_TYPE_TURBO4_0) { - flash_attn_ext_turbo4_load_tile + flash_attn_ext_turbo4_load_tile (K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup); } else if constexpr (type_K == GGML_TYPE_TURBO3_0) { - flash_attn_ext_turbo3_load_tile + flash_attn_ext_turbo3_load_tile (K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup); } else { - flash_attn_ext_turbo2_load_tile + flash_attn_ext_turbo2_load_tile (K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup); } __syncthreads(); @@ -1261,13 +1273,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nthreads_turbo = nwarps * ggml_cuda_get_physical_warp_size(); const char * V_raw = (const char *) V_h2 + int64_t(k_VKQ_0) * stride_V; if constexpr (type_V == GGML_TYPE_TURBO4_0) { - flash_attn_ext_turbo4_load_tile + flash_attn_ext_turbo4_load_tile (V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup); } else if constexpr (type_V == GGML_TYPE_TURBO3_0) { - flash_attn_ext_turbo3_load_tile + flash_attn_ext_turbo3_load_tile (V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup); } else { - flash_attn_ext_turbo2_load_tile + flash_attn_ext_turbo2_load_tile (V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup); } __syncthreads(); diff --git a/ggml/src/ggml-cuda/fattn-mma-turbo.cuh b/ggml/src/ggml-cuda/fattn-mma-turbo.cuh index 2dfbab3c39d6..6be63ced1dfa 100644 --- a/ggml/src/ggml-cuda/fattn-mma-turbo.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-turbo.cuh @@ -43,7 +43,11 @@ void ggml_cuda_flash_attn_ext_mma_turbo_case(ggml_backend_cuda_context & ctx, gg // turbo4 never aliases V onto K. constexpr bool V_is_K_view = false; - const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(nbatch_K2 + 4, nbatch_V2 + 4) * sizeof(half2); + // must match the swizzled tile stride flash_attn_ext_turbo{2,3,4}_load_tile write through + // (fattn-mma-f16.cuh's turbo_store_h2 / bytes_rc), same helper as fattn-mma-f16.cuh:2287. + const int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2, cc); + const int stride_tile_V = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2, cc); + const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(stride_tile_K, stride_tile_V) * sizeof(half2); const size_t nbytes_shared_Q = ncols * (DKQ/2 + 4) * sizeof(half2); const size_t nbytes_shared_mask = ncols1 * (nbatch_fa/2 + 4) * sizeof(half2); const size_t nbytes_shared_combine = nwarps*cols_per_warp * (nbatch_combine + 4) * sizeof(half2); From 4ce86aeb8c79a743f8991cb24f7659d7635d9d4f Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 22:41:26 -0600 Subject: [PATCH 13/16] tests: cover turbo2/3/4 K/V at hsk=256 with nb<=4 and GQA in FLASH_ATTN_EXT The general test_flash_attn_ext matrix already exercised turbo3/turbo4 K==V at hsk=128 with batched nb and GQA nr2>=2, but excluded hsk=256 entirely (only 64/72/128 were allowed for non-F16 types) and never included turbo2_0 at all. That's exactly the gap that let the turbo swizzle write/read mismatch ship: the only hsk=256 turbo coverage was the hand-rolled turbo4_vec case, which is nb=1/nr2=1/no-mask only and doesn't exercise the batched MMA path real speculative-decode batches (nb up to 4) or GQA-packed decode go through. Requested in PR #340 review: hsk 128 and 256, nb<=4, nr2>=2. Verified: 2592/2592 filtered FLASH_ATTN_EXT cases pass (turbo2/3/4, hsk=256, nb in {1,3,32,75}, nr23=[4,1], masked and unmasked). Co-Authored-By: Claude Sonnet 5 --- tests/test-backend-ops.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 82417b53062b..8f2a2ec10552 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10256,9 +10256,12 @@ static std::vector> make_test_cases_eval() { for (int nb : { 1, 3, 32, 75, }) { for (ggml_prec prec : {GGML_PREC_F32, GGML_PREC_DEFAULT}) { if (hsk != 128 && prec == GGML_PREC_DEFAULT) continue; - for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO4_0}) { - if ((type_KV == GGML_TYPE_TURBO3_0 || type_KV == GGML_TYPE_TURBO4_0) && hsk < 128) continue; - if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72 && hsk != 128) continue; + for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL, GGML_TYPE_TURBO2_0, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO4_0}) { + const bool is_turbo_kv = type_KV == GGML_TYPE_TURBO2_0 || type_KV == GGML_TYPE_TURBO3_0 || type_KV == GGML_TYPE_TURBO4_0; + if (is_turbo_kv && hsk < 128) continue; + // turbo MMA/VEC kernels are also instantiated at hsk=256 (fattn-mma-turbo.cuh + // DECL_FATTN_MMA_TURBO_ALL(256,256,...)); exercise the swizzled-write path there too. + if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72 && hsk != 128 && !(is_turbo_kv && hsk == 256)) continue; // DeepSeek MLA: the V cache is a sub-view of the K cache const bool v_is_view_of_k = hsk == 576; test_cases.emplace_back(new test_flash_attn_ext( From 283e56cd08272152fd405a845277c9079d16fa33 Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 22:41:33 -0600 Subject: [PATCH 14/16] address PR #340 review follow-ups: assert, buffer-check, comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four smaller items from TheTom's review, none of them blockers: - deepseek4.cpp: assert n_kv_max (raw SWA window + top-k count) doesn't exceed k_all's actual concat length (raw_k + csa_k), documenting an invariant that held but was previously unchecked. - llama-kv-cache.cpp: llm_graph_input_k_shift::set_input guarded k_rot on both null and ->buffer before use, but only null-checked k_shift. set_input_k_shift asserts on ggml_backend_buffer_is_host(dst->buffer), so a graph-reserve pass (tensors exist, backends not yet allocated) would fault there. Add the same buffer check used everywhere else in llama-graph.cpp's set_input overrides. - llama-graph.cpp build_norm: note why the mw/mb F32 cast is cheap enough to re-insert per graph build (n_embd-sized weight) rather than caching across builds. - fattn.cu: clarify why quantized K/V (including turbo2/3/4) are skipped in the GQA-opt alignment check — their loaders dequantize via swizzled/padded SMEM tiles rather than reading nb[] directly. Co-Authored-By: Claude Sonnet 5 --- ggml/src/ggml-cuda/fattn.cu | 5 ++++- src/llama-graph.cpp | 6 +++++- src/llama-kv-cache.cpp | 4 +++- src/models/deepseek4.cpp | 1 + 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 008c08a9879d..14b9573b6c57 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -180,7 +180,10 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); // Edge cases like no mask, ALiBi, unpadded K/V, or misaligned addresses for large data transfers - // are put into the template specialization without GQA optimizations. + // are put into the template specialization without GQA optimizations. Quantized tensors + // (incl. turbo2/3/4) are skipped here: their loaders dequantize into SMEM via the + // swizzled/padded tile helpers rather than reading nb[] directly, so the 16-byte-stride + // alignment this loop checks for doesn't apply to them. bool use_gqa_opt = mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0; for (const ggml_tensor * t : {Q, K, V, mask}) { if (t == nullptr || ggml_is_quantized(t->type)) { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index df714d96740b..8f596e24c850 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1605,7 +1605,10 @@ ggml_tensor * llm_graph_context::build_norm( if (mw) { // the CUDA broadcast-mul kernel has no path for an F32 activation times an F16 operand // (only the reverse, F16 activation times F32/F16 operand); most checkpoints keep norm - // weights in F32 so this is normally a no-op, but some conversions store them narrower + // weights in F32 so this is normally a no-op, but some conversions store them narrower. + // This re-inserts a cast node into the graph on every build (every token), but norm + // weight tensors are n_embd-sized (a few KB), so the added cost is not worth caching + // across builds versus upcasting these specific tensors once at load time. if (mw->type != cur->type && cur->type == GGML_TYPE_F32) { mw = ggml_cast(ctx0, mw, GGML_TYPE_F32); } @@ -1616,6 +1619,7 @@ ggml_tensor * llm_graph_context::build_norm( } if (mb) { + // see mw cast note above if (mb->type != cur->type && cur->type == GGML_TYPE_F32) { mb = ggml_cast(ctx0, mb, GGML_TYPE_F32); } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index d90c38013ad3..490446132fbc 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2241,7 +2241,9 @@ class llm_graph_input_k_shift : public llm_graph_input_i { void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) { GGML_UNUSED(ubatch); - if (k_shift) { + // buffer check guards the graph-reserve pass, where tensors exist but backends aren't + // allocated yet; set_input_k_shift asserts on dst->buffer, so this must not be dropped. + if (k_shift && k_shift->buffer) { kv_self->set_input_k_shift(k_shift); } diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index c7658e462053..1cd0525a83a5 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -782,6 +782,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( cb(kq_mask, "csa_lid_kq_mask", il); const int64_t n_kv_max = std::min(raw_mask->ne[0], hparams.n_swa) + top_k->ne[0]; + GGML_ASSERT(n_kv_max <= k_all->ne[2]); // must not exceed raw_k + csa_k concat length ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, n_kv_max, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); From 84192516d157e18a348f8887ac09c83bdfc6695b Mon Sep 17 00:00:00 2001 From: giveen Date: Wed, 2 Sep 2026 22:44:37 -0600 Subject: [PATCH 15/16] fix follow-up items to match what review actually asked for Corrected three mistargeted fixes from the previous commit against TheTom's actual review text (fetched via gh pr view, not re-derived from memory): - deepseek4.cpp:784: the ask was an assert catching a misconfigured n_swa == 0 model (which would silently zero out the raw SWA window and leave only csa top-k entries), not a bounds check against k_all->ne[2]. Add GGML_ASSERT(hparams.n_swa > 0); keep the bounds assert too, it documents a real, separate invariant. - llama-kv-cache.cpp:2248: the ask was about k_rot's own silent-skip semantics (is it ever legitimate, and when), not k_shift's missing buffer guard. Document that k_rot is null pre-attn_rot_k-setup for non-rotating caches and unallocated only during graph-reserve. (k_shift's missing ->buffer check from the prior commit was a real, separate bug -- set_input_k_shift asserts on it -- and stays fixed.) - fattn.cu: the stale comment was "DEFAULT OFF" above ggml_cuda_turbo_mma_fused(), whose code has defaulted ON for a while; not the alignment-loop comment at the review's line 190 (line numbers had shifted off the review's head commit by the time this landed). Verified: cmake --build --target llama compiles clean, test-llama-archs still green for qwen4exp/deepseek/deepseek2/deepseek32 (the new n_swa assert doesn't trip on any registered synthetic config). Co-Authored-By: Claude Sonnet 5 --- ggml/src/ggml-cuda/fattn.cu | 6 +++--- src/llama-kv-cache.cpp | 5 +++++ src/models/deepseek4.cpp | 5 +++++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 14b9573b6c57..c628699c5721 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -313,7 +313,7 @@ static void ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2(ggml_backend_cuda_c ggml_cuda_flash_attn_ext_mma_turbo_case(ctx, dst); // ncols2 = 1 -> (8,1) } -// Env latch for the fused turbo4 MMA decode path. DEFAULT OFF. +// Env latch for the fused turbo MMA decode path. DEFAULT ON. // // The MMA path is correctness-validated (coherent output, KLD == VEC baseline 0.008396) // and faster than VEC at every depth (beats rival "buun"), BUT it is NOT bit/token-identical @@ -321,8 +321,8 @@ static void ggml_cuda_flash_attn_ext_mma_turbo_switch_ncols2(ggml_backend_cuda_c // reduction trees (tensor-core fragment order vs per-thread VEC order), so a near-tie greedy // token can flip (~1 in ~25 tokens on a hard tie). This is the same irreducible f16-order // difference that exists between the base f16-MMA and f16-VEC kernels — not a regression — but -// it fails strict token-identity. We therefore keep VEC the default and expose the faster MMA -// path as opt-in via GGML_TURBO_MMA_FUSED=1. +// it fails strict token-identity. GGML_TURBO_MMA_FUSED=0 is the VEC kill-switch for anyone who +// needs that identity guarantee back. static bool ggml_cuda_turbo_mma_fused() { static const bool v = []{ const char * s = getenv("GGML_TURBO_MMA_FUSED"); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 490446132fbc..7bf91d3d2e22 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2247,6 +2247,11 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) { kv_self->set_input_k_shift(k_shift); } + // k_rot is null (not just unallocated) whenever attn_rot_k is false: build_input_k_rot + // only allocates a real tensor for quantized K-caches with rotation enabled, or for + // DeepSeek32/DeepSeek4's lightning-indexer cache (see attn_rot_k's setup). So a skip + // here is either "this cache doesn't rotate" (k_rot == nullptr) or "graph-reserve pass" + // (k_rot->buffer == nullptr) -- never a case that should silently drop a real input. if (k_rot && k_rot->buffer) { kv_self->set_input_k_rot(k_rot); } diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 1cd0525a83a5..91ce4134ec2b 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -781,6 +781,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); + // n_kv_max bounds the finite (non -INFINITY) mask entries per row that + // flash_attn_mask_to_sparse_indices will keep; entries past it are silently dropped. + // n_swa == 0 would zero out the raw SWA window and leave only the csa top-k entries, + // which is never a valid config for this path. + GGML_ASSERT(hparams.n_swa > 0); const int64_t n_kv_max = std::min(raw_mask->ne[0], hparams.n_swa) + top_k->ne[0]; GGML_ASSERT(n_kv_max <= k_all->ne[2]); // must not exceed raw_k + csa_k concat length ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, n_kv_max, kq_scale, il); From 69499d4d309a58f4d99bb5d006321ce3a610a9b2 Mon Sep 17 00:00:00 2001 From: giveen Date: Thu, 3 Sep 2026 15:30:47 -0600 Subject: [PATCH 16/16] memory: skip V-cache allocation for the qwen4exp indexer KV cache llama_memory_hybrid_idx's indexer cache (mem_idx) only ever reads K (the lightning-indexer top-k selection needs keys, never values), but its hparams were only narrowing n_embd_head_k_full to indexer_head_size -- n_embd_head_v_full stayed at the full model's real V head dim, and nothing marked the cache as MLA/K-only. llama_kv_cache always allocates V storage unless hparams.is_mla() is true (has_v = !is_mla), so this cache was wasting VRAM on a same-sized V-cache buffer it never touches. Fix: set n_embd_head_v_full and both n_embd_head_{k,v}_mla_impl to indexer_head_size, mirroring dsv4_make_k_only's hparams_lid setup in llama-kv-cache-dsv4.cpp (deepseek4's own lightning-indexer cache uses the exact same K-only pattern already). Upstream hit the same bug independently: ggml-org/llama.cpp#28330. Verified: test-llama-archs qwen4exp still OK (NMSE 8.94e-08, roundtrip OK) on CUDA/CPU. Co-Authored-By: Claude Sonnet 5 --- src/llama-memory-hybrid-idx.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 31e391da9a8b..b12f20274408 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -50,6 +50,13 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + hparams_idx.n_embd_head_v_full = model.hparams.indexer_head_size; + + // the indexer never reads V; mark this cache MLA/K-only (same mechanism as + // dsv4_make_k_only in llama-kv-cache-dsv4.cpp's hparams_lid) so llama_kv_cache's + // has_v = !is_mla skips allocating the unused V-cache tensors entirely. + hparams_idx.n_embd_head_k_mla_impl = model.hparams.indexer_head_size; + hparams_idx.n_embd_head_v_mla_impl = model.hparams.indexer_head_size; LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);