Skip to content

qwen4exp: follow up fixes #27941

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggerganov merged 8 commits into from
Sep 1, 2026
Merged

qwen4exp: follow up fixes #27941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3dd5505
qwen4exp: follow up fixes
danielhanchen Aug 29, 2026
4cbf055
-kvu NaN collapse fix
danielhanchen Aug 30, 2026
9bfa91c
indexer cache ext.x/ext.y restore fix
danielhanchen Aug 30, 2026
c76ea01
kv-cells: rename seq_set to seq_get_all
danielhanchen Aug 30, 2026
4ebbc37
memory-hybrid-idx: implement set_input_qsa on the memory class
danielhanchen Aug 31, 2026
52bc9ff
tests: check that a sequence state survives a save/restore round-trip
danielhanchen Aug 31, 2026
21c13c4
tests: give the synthetic qwen4exp a PLE so the state test bites
danielhanchen Aug 31, 2026
c961cd3
llama: disable -sm tensor for qwen4exp
danielhanchen Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/llama-arch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1142,6 +1142,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_QWEN3TTS:
case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs
return false;
default:
return true;
Expand Down
16 changes: 16 additions & 0 deletions src/llama-kv-cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2387,6 +2387,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<llama_kv_cell_ext> exts;
if (has_cell_ext()) {
exts.resize(cell_count);
}

for (uint32_t i = 0; i < cell_count; ++i) {
llama_pos pos;
uint32_t n_seq_id;
Expand All @@ -2410,6 +2416,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32

// apply_ubatch() below restores ext.tok from the ubatch tokens
ubatch.token[i] = ext.tok;

exts[i] = ext;
}

// read the sequence id, but directly discard it - we will use dest_seq_id instead
Expand Down Expand Up @@ -2461,6 +2469,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]);
}
Comment on lines +2472 to +2478

@ggerganov ggerganov Aug 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we can extend the tests-save-load-state with a test that would demonstrate this issue on master. Though from the description in unslothai#143, it might be difficult to figure out such a test. Low prio, but mentioning just in case you can think of something simple.

@ggerganov ggerganov Aug 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, 02eb201 is a good idea.

However, even without this patch, the test still succeeds. So it's missing something.

@ggerganov ggerganov Aug 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@danielhanchen In case you missed this comment. The problem seems to be that the generated qwen4 dummy model by test-llama-archs does not have PLE. This makes has_cell_ext() return false:

llama.cpp/src/llama-kv-cache.cpp

Lines 1824 to 1827 in 02eb201

bool llama_kv_cache::has_cell_ext() const {
// M-RoPE needs the 2D position, the PLE n-gram hash needs the token id
return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0;
}

So the added test is not effective. I think we have to extend the dummy model to have valid PLE data.

@danielhanchen danielhanchen Aug 31, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Change the test to actually invoke PLE if that helps


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
Expand Down
11 changes: 9 additions & 2 deletions src/llama-kv-cells.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ struct llama_kv_cell_ext {
// TODO: add unit tests
class llama_kv_cells {
public:
using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;

void reset() {
for (uint32_t i = 0; i < pos.size(); ++i) {
pos[i] = -1;
Expand Down Expand Up @@ -301,6 +303,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());
Expand Down Expand Up @@ -511,8 +520,6 @@ class llama_kv_cells {
//
std::vector<llama_pos> shift;

using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;

// the bitset seq[i] tells us which sequences are currently occupying the i-th cell
std::vector<seq_set_t> seq;

Expand Down
Loading
Loading