qwen4exp: fix seq_cp, block position keying, mtmd input, cuda abort, add tests (#27941)

* qwen4exp: follow up fixes

* -kvu NaN collapse fix

Assisted-by: Claude

* indexer cache ext.x/ext.y restore fix

Assisted-by: Claude

* kv-cells: rename seq_set to seq_get_all

seq_get is already taken by the single-id getter, so the suggested name
cannot be overloaded on return type alone.

Assisted-by: Claude

* memory-hybrid-idx: implement set_input_qsa on the memory class

The context held the whole implementation, where the pattern elsewhere is a
thin context forwarding to the memory class, as llama_kv_cache_context does
for set_input_kq_mask. The body reads no context state, so it moves unchanged
and the context keeps a forwarder.

Also shortens the seq_get_all comment as suggested.

* tests: check that a sequence state survives a save/restore round-trip

Saves seq 0, erases it, restores the blob and saves again, requiring the two
blobs to match. Compares blobs rather than generated text, which cannot see a
field dropped on the way back in.

Note this passes on master for qwen4exp, so it does not demonstrate the
ext.x/ext.y drop this PR fixes; reaching that needs 2D mrope content.

* tests: give the synthetic qwen4exp a PLE so the state test bites

has_cell_ext() is n_pos_per_embd() > 1 || ple_n_heads > 0, and the indexer
cache sets rope_type = NONE, so without a PLE it serializes no cell ext at
all and the round-trip test cannot see a dropped ext.x/ext.y. With one,
removing the ext_set restore in state_read_meta fails the test: 198 of
335692 bytes differ, first at offset 282092.

Loading such a model needed two fixes:

- the row count of per_layer_token_embd came from require_weight(), which a
  model synthesised from metadata alone has no file to answer. Derive it
  from the head ranges and prefer the file's padded count where there is one.
- the PLE conv history is a row of the recurrent cache, so a PLE on a full
  attention layer dereferenced a null p_l. Reject it at load time instead.

The meta mirror is skipped for qwen4exp. It returned NaN logits before this
fixture carried a PLE, which the nmse check passes since a NaN comparison is
false, and aborts with one. -sm tensor on real devices works.

Assisted-by: Claude

* llama: disable -sm tensor for qwen4exp

test-llama-archs skipped the tensor split for this arch from inside the
test, so the arch still advertised support it does not have. Declare it in
llm_arch_supports_sm_tensor instead and drop the test-side exception; the
existing llm_arch_supports_sm_tensor branch then does the skipping.

Assisted-by: Claude
This commit is contained in:
Daniel Han
2026-09-01 13:22:04 +03:00
committed by GitHub
parent d086dbb348
commit 36b1015438
8 changed files with 529 additions and 147 deletions
+70 -19
View File
@@ -6,6 +6,23 @@
#include <algorithm>
#include <cinttypes>
// 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);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
@@ -18,21 +35,30 @@ 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);
GGML_ASSERT(hparams.ssm_d_conv > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 &&
hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0);
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 && hparams.hc_low_rank > 0);
// a count of 1 has nothing to mix: transformers configuration_qwen4_exp.py:196, vLLM
// config.py:49 and SGLang configs/qwen4_exp.py:38 all raise on hc_count <= 1
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);
GGML_ASSERT(hparams.indexer_n_head > 0
&& hparams.indexer_head_size > 0
&& hparams.indexer_top_k > 0);
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);
// PLE n-gram hash embeddings; if the key group is absent every field stays zero
@@ -44,7 +70,11 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
if (n_ple > 0) {
std::vector<uint32_t> ple_layers;
ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers);
GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer");
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) {
if (il >= hparams.n_layer_all) {
throw std::runtime_error(format("PLE layer %u is out of range", il));
@@ -59,7 +89,8 @@ 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);
GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0);
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;
@@ -70,6 +101,10 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads));
}
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);
// the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses
@@ -93,12 +128,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);
GGML_ASSERT(full_attn_interval > 0);
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;
@@ -124,18 +166,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) {
const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
const auto & ple_w = ml.require_weight(ple_name.c_str());
const int64_t ple_rows = ple_w.tensor->ne[1];
// sanity check
// 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) {
if ((int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h] > ple_rows) {
throw std::runtime_error(format("PLE head %u range exceeds the %" PRId64 " table rows", h, ple_rows));
}
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();
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 }, TENSOR_READ_LAZY);
}
@@ -556,9 +604,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);