mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-19 17:24:57 +02:00
llama: carry the qwen4exp PLE conv state across ubatches
The PLE depthwise conv was zero-padding on the left, which is only right for a prefill that starts at position 0. Decode and chunked prefill saw a truncated history for the first (kernel-1)*ngram_size positions of every ubatch. The PLE module sits on a layer that is also a delta-net layer, so both need a conv history in the same recurrent row. Rather than plumb a per-layer state size through build_rs and build_conv_state, the row is widened once and each convolution addresses its own slice through a local helper. n_embd_r() gains the extra span, which is zero for every other architecture because it is derived from ple_n_heads. Verified by feeding the same 1024 token sequence in chunks instead of one shot: at 64 tokens per decode the logits are bit-identical to the single-shot run, 1023 of 1023 top-1 and a maximum logprob deviation of exactly zero. At one token per decode they differ slightly, but the no-PLE model differs more under the same test (94.6% against 97.1%), so that is the usual gemv-versus- gemm accumulation difference and not the state. The conv branch is also no longer unverified. With non-zero conv weights the port sits 6.3 points of top-1 below the numerical floor, the same distance as with the weights zeroed and as the model with no PLE at all, so the branch adds no error of its own. test-llama-archs passes every existing architecture at 0.00e+00, including the delta-net models that share this code path.
This commit is contained in:
+16
-1
@@ -201,7 +201,13 @@ uint32_t llama_hparams::n_embd_r() const {
|
||||
// TODO: maybe support other convolution strides than 1
|
||||
// NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed
|
||||
// Corresponds to Mamba's conv_states size
|
||||
return (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state);
|
||||
const uint32_t n_conv = (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state);
|
||||
|
||||
// qwen4exp hosts a PLE module on a layer that is also a delta-net layer, so
|
||||
// that row has to carry a second, dilated conv state after the first. The
|
||||
// rows are uniform across layers, so every recurrent layer reserves it.
|
||||
// ple_n_heads is zero for every other architecture, leaving n_conv alone.
|
||||
return n_conv + ple_conv_state();
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_s() const {
|
||||
@@ -236,6 +242,15 @@ bool llama_hparams::is_recr(uint32_t il) const {
|
||||
GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all);
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::ple_conv_state() const {
|
||||
if (ple_n_heads == 0 || ple_conv_kernel == 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// dilation equals the n-gram size, matching the reference module
|
||||
return (ple_conv_kernel - 1) * ple_ngram_size * dsv4_hc_mult * n_embd;
|
||||
}
|
||||
|
||||
bool llama_hparams::is_ple(uint32_t il) const {
|
||||
if (il < n_layer_all) {
|
||||
return is_ple_impl[il];
|
||||
|
||||
@@ -288,6 +288,10 @@ struct llama_hparams {
|
||||
|
||||
bool is_ple(uint32_t il) const;
|
||||
|
||||
// rows of the PLE depthwise conv history: (kernel - 1) * dilation, where
|
||||
// the dilation is the n-gram size. zero unless the model has a PLE module.
|
||||
uint32_t ple_conv_state() const;
|
||||
|
||||
// qwen3vl deepstack
|
||||
// When parsed from GGUF, this implies the first N layers consume the first
|
||||
// N deepstack embeddings. Use deepstack_mapping_arr if you need a more
|
||||
|
||||
@@ -2330,6 +2330,17 @@ struct llama_model_qwen4exp : public llama_model_base {
|
||||
ggml_tensor * gate,
|
||||
int layer);
|
||||
|
||||
// conv history at an explicit offset in the recurrent row: this arch
|
||||
// packs the delta-net and PLE conv states into the same row
|
||||
ggml_tensor * build_conv_state_at(
|
||||
llm_graph_input_rs * inp,
|
||||
ggml_tensor * conv_states_all,
|
||||
ggml_tensor * x,
|
||||
int64_t state_cols,
|
||||
int64_t channels,
|
||||
int64_t row_offset,
|
||||
int il);
|
||||
|
||||
ggml_tensor * build_ple(
|
||||
llm_graph_input_rs * inp,
|
||||
ggml_tensor * hidden,
|
||||
|
||||
+75
-16
@@ -503,7 +503,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
|
||||
const int64_t conv_kernel_size = conv_kernel->ne[0];
|
||||
const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state;
|
||||
|
||||
ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il);
|
||||
// offset 0: the delta-net history sits at the front of the row, with the
|
||||
// PLE history (if this model has one) after it
|
||||
ggml_tensor * conv_input = build_conv_state_at(inp, conv_states_all, qkv_mixed,
|
||||
conv_kernel_size - 1, conv_channels, 0, il);
|
||||
|
||||
ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs);
|
||||
state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs);
|
||||
@@ -770,6 +773,60 @@ void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) {
|
||||
ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows));
|
||||
}
|
||||
|
||||
// Fetch one conv history out of the recurrent row and write the updated tail
|
||||
// back, at an explicit offset within that row.
|
||||
//
|
||||
// The shared build_conv_state assumes the whole row belongs to one convolution.
|
||||
// Here the row carries the delta-net conv history followed by the PLE one, so
|
||||
// each caller addresses its own slice. Same structure as the shared helper,
|
||||
// only with an offset and an explicit dilation.
|
||||
ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
|
||||
llm_graph_input_rs * inp,
|
||||
ggml_tensor * conv_states_all,
|
||||
ggml_tensor * x,
|
||||
int64_t state_cols,
|
||||
int64_t channels,
|
||||
int64_t row_offset,
|
||||
int il) {
|
||||
const auto * mctx_cur = inp->mctx;
|
||||
|
||||
const auto kv_head = mctx_cur->get_head();
|
||||
const auto mem_size = mctx_cur->get_size();
|
||||
|
||||
const int64_t n_seqs = ubatch.n_seqs;
|
||||
const int64_t row_total = hparams.n_embd_r();
|
||||
|
||||
// the gather needs the whole row, then this convolution takes its slice
|
||||
ggml_tensor * rows = build_rs(inp, conv_states_all, row_total, n_seqs);
|
||||
|
||||
const size_t esz = ggml_element_size(rows);
|
||||
|
||||
ggml_tensor * state = ggml_cont(ctx0,
|
||||
ggml_view_2d(ctx0, rows, state_cols * channels, n_seqs,
|
||||
rows->nb[1], row_offset * esz));
|
||||
state = ggml_reshape_3d(ctx0, state, state_cols, channels, n_seqs);
|
||||
cb(state, "conv_state_at", il);
|
||||
|
||||
ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);
|
||||
|
||||
// keep the last state_cols columns for the next ubatch
|
||||
const size_t row_size = ggml_row_size(conv_states_all->type, row_total);
|
||||
|
||||
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));
|
||||
|
||||
ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
|
||||
state_cols * channels, n_seqs,
|
||||
conv_states_all->nb[1],
|
||||
kv_head * row_size + row_offset * ggml_element_size(conv_states_all));
|
||||
|
||||
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,
|
||||
ggml_tensor * hidden,
|
||||
@@ -837,27 +894,29 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple(
|
||||
//
|
||||
// out[c, t] = sum_k w[k, c] * x[c, t - (K-1-k)*dilation]
|
||||
//
|
||||
// Positions before the sequence start read as zero, which is what the
|
||||
// reference's zero-padded nn.Conv1d does.
|
||||
// History from earlier ubatches is prepended, so decode and chunked prefill
|
||||
// see the same context a single-shot prefill would. A fresh sequence starts
|
||||
// with a zeroed state, which is what the reference's zero-padded nn.Conv1d
|
||||
// gives at a sequence start.
|
||||
const int64_t kern = hparams.ple_conv_kernel;
|
||||
const int64_t dil = hparams.ple_ngram_size;
|
||||
const int64_t hist = (kern - 1) * dil;
|
||||
|
||||
// [hc_dim, hist + n_tokens], transposed to put tokens on ne[0]
|
||||
ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_r_l(il),
|
||||
normalized, hist, hc_dim,
|
||||
hparams.n_embd_r() - hparams.ple_conv_state(), il);
|
||||
|
||||
ggml_tensor * conv_out = nullptr;
|
||||
for (int64_t k = 0; k < kern; ++k) {
|
||||
const int64_t shift = (kern - 1 - k) * dil;
|
||||
// tap k reads (kern-1-k)*dilation positions back
|
||||
const int64_t start = hist - (kern - 1 - k) * dil;
|
||||
|
||||
ggml_tensor * shifted;
|
||||
if (shift == 0) {
|
||||
shifted = normalized;
|
||||
} else if (shift >= n_tokens) {
|
||||
continue; // entirely out of range, contributes nothing
|
||||
} else {
|
||||
ggml_tensor * keep = ggml_view_2d(ctx0, normalized, hc_dim, n_tokens - shift,
|
||||
normalized->nb[1], 0);
|
||||
ggml_tensor * zeros = ggml_scale(ctx0,
|
||||
ggml_view_2d(ctx0, normalized, hc_dim, shift, normalized->nb[1], 0), 0.0f);
|
||||
shifted = ggml_concat(ctx0, zeros, keep, 1);
|
||||
}
|
||||
ggml_tensor * shifted = ggml_cont(ctx0,
|
||||
ggml_transpose(ctx0,
|
||||
ggml_view_2d(ctx0, padded, n_tokens, hc_dim,
|
||||
padded->nb[1],
|
||||
ggml_row_size(padded->type, start))));
|
||||
|
||||
// column k of the [kern, hc_dim] kernel is one weight per channel
|
||||
ggml_tensor * wk = ggml_cont(ctx0,
|
||||
|
||||
Reference in New Issue
Block a user