qwen4_exp QSA assumes the visible set is contiguous — true today, false the moment it is batched #324

Closed
opened 2026-09-07 12:49:11 +00:00 by grenade · 1 comment
Owner

Split out of #311 so closing that issue does not bury it. Not a defect
today; a documented assumption that stops holding under a change
someone will make.

The assumption

BlockSelector derives block extents arithmetically from the block
index: block b covers positions b * block_size .. (b+1) * block_size,
and the tail is n_blocks * block_size .. n_visible. That is only
correct if the set of positions a query may attend is the contiguous
range 0..n_visible.

The reference does not assume it. From
transformers/models/qwen4_exp/modeling_qwen4_exp.py:

visible_token_indices = attention_mask if attention_mask.dtype == torch.bool else attention_mask == 0
...
local_visible_indices = torch.nonzero(visible_token_indices[batch_idx, 0, query_idx]).flatten()
num_complete_blocks = local_visible_indices.shape[-1] // self.compress_ratio
block_token_indices = local_visible_indices[: num_complete_blocks * self.compress_ratio].view(...)

Blocks are formed from whichever positions the mask allows, in order.
For a plain causal mask that is 0..n_visible and the two agree
exactly — which is why full-model parity passes (#323).

When it stops being true

Any mask with holes:

  • left padding in a batch, where short sequences are padded at the
    front;
  • ragged prompts under the batched decode engine — the lockstep
    engine multiplexes sequences of different lengths and joins new ones
    mid-stream (#98);
  • anything else that masks out positions inside a sequence.

qwen4_exp is not batched yet — supports_kv_snapshot is true since
25157ca4, but the arch has not been through the engine — so nothing
is wrong right now. The failure mode when it is: blocks composed of the
wrong positions, scored against the wrong keys, and a mask that permits
attention to positions the model should not see. It produces plausible
output on the sequence that happens to be contiguous, which is most of
them, and the wrong answer on the ones that are not.

What to do

Either follow the reference and derive blocks from the actual visible
set, or assert contiguity at the boundary so a padded or ragged batch
fails loudly instead of quietly. The second is cheap and is probably the
right first move: it converts an invisible wrong answer into an error
that names its own cause, and it can stand until someone needs the
first.

Worth doing before qwen4_exp is enabled for batched decode, not
after — the whole point is that the symptom does not look like the
cause.

Related: #311 (where this was found), #323 (parity, which passes
precisely because the causal case is contiguous), #98 (the engine whose
ragged prompts would break it).

Split out of #311 so closing that issue does not bury it. Not a defect today; a documented assumption that stops holding under a change someone will make. ## The assumption `BlockSelector` derives block extents arithmetically from the block index: block `b` covers positions `b * block_size .. (b+1) * block_size`, and the tail is `n_blocks * block_size .. n_visible`. That is only correct if the set of positions a query may attend is the contiguous range `0..n_visible`. The reference does not assume it. From `transformers/models/qwen4_exp/modeling_qwen4_exp.py`: ```python visible_token_indices = attention_mask if attention_mask.dtype == torch.bool else attention_mask == 0 ... local_visible_indices = torch.nonzero(visible_token_indices[batch_idx, 0, query_idx]).flatten() num_complete_blocks = local_visible_indices.shape[-1] // self.compress_ratio block_token_indices = local_visible_indices[: num_complete_blocks * self.compress_ratio].view(...) ``` Blocks are formed from *whichever* positions the mask allows, in order. For a plain causal mask that is `0..n_visible` and the two agree exactly — which is why full-model parity passes (#323). ## When it stops being true Any mask with holes: - **left padding** in a batch, where short sequences are padded at the front; - **ragged prompts** under the batched decode engine — the lockstep engine multiplexes sequences of different lengths and joins new ones mid-stream (#98); - anything else that masks out positions inside a sequence. `qwen4_exp` is not batched yet — `supports_kv_snapshot` is true since `25157ca4`, but the arch has not been through the engine — so nothing is wrong right now. The failure mode when it is: blocks composed of the wrong positions, scored against the wrong keys, and a mask that permits attention to positions the model should not see. It produces plausible output on the sequence that happens to be contiguous, which is most of them, and the wrong answer on the ones that are not. ## What to do Either follow the reference and derive blocks from the actual visible set, or assert contiguity at the boundary so a padded or ragged batch fails loudly instead of quietly. The second is cheap and is probably the right first move: it converts an invisible wrong answer into an error that names its own cause, and it can stand until someone needs the first. Worth doing **before** qwen4_exp is enabled for batched decode, not after — the whole point is that the symptom does not look like the cause. Related: #311 (where this was found), #323 (parity, which passes precisely because the causal case is contiguous), #98 (the engine whose ragged prompts would break it).
Author
Owner

Fixed in 0086a543, merged as 230d243

Blocks are now cut from the visible set rather than the position axis, following the reference's nonzero-then-chunk.

pub enum VisibleSet<'a> {
    Contiguous(usize),   // a plain causal mask, materialising nothing
    Sparse(&'a [usize]), // left padding, ragged batches, any mask with holes
}

Contiguous exists because the set reaches the context window in size and is rebuilt for every query — a Vec per query would cost more than the selection it feeds. Both variants go through the same arithmetic, which is the point: blocks are cut at set offsets and mapped back to absolute positions, and those coincide only when the set starts at zero without holes.

Two departures from the plan above

I took the first option, not the second. The issue argued the assertion was the right first move. It converts a silent wrong answer into a loud failure, which is real value — but it leaves the model unbatchable, so it defers the work rather than doing it. The selector generalisation is what makes the eventual plumbing a matter of passing Sparse instead of Contiguous.

The assertion already existed. I added one to positions_from, then wrote a test to prove it fires — and the test failed, because Indexer::advance was already enforcing total == past_len + seq_len one level up. Mine was unreachable. Deleted it and widened the existing message to name the contiguity assumption it protects, rather than only the cache divergence that is its symptom.

Tests

  • blocks_are_cut_from_the_visible_set_not_the_position_axis — visible [3,5,7,11,13,17,19,23,29], budget 1 block: selection returns [13,17,19,23,29]. Under the old arithmetic it returned 4..8 plus 8, and 4, 6, 8, 9, 10 are positions the mask forbids.
  • contiguous_and_sparse_agree_on_the_same_positions — the two representations must not drift, since only the explicit one is checked against the reference.
  • a_cache_inconsistent_with_past_len_is_refused / ..._is_served — drives the guard both ways, because an assertion that cannot be made to fire is not evidence.

Mutations verified to fail: index arithmetic in select, index arithmetic in block_starts, and removing the guard — which returns [[[0,1,2,3,4,5]]] for an inconsistent cache, the exact plausible-looking wrong answer this issue described.

Both parity fixtures still match the reference, confirming the contiguous path is untouched.

Correction to this issue's framing

This was described — including by me, elsewhere — as the blocker for batching. It is not. model.rs builds its own causal mask and no padding mask reaches this arch, so there is currently no source of a non-contiguous visible set at all.

The actual obstacle is that past_len is a single scalar for the whole batch: ragged sequences need per-sequence offsets, through the indexer, the GDN state and PLE's rolling conv context alike. Filed as #326 rather than left implied inside a closed issue.

Closing.

## Fixed in `0086a543`, merged as `230d243` Blocks are now cut from the **visible set** rather than the position axis, following the reference's `nonzero`-then-chunk. ```rust pub enum VisibleSet<'a> { Contiguous(usize), // a plain causal mask, materialising nothing Sparse(&'a [usize]), // left padding, ragged batches, any mask with holes } ``` `Contiguous` exists because the set reaches the context window in size and is rebuilt for every query — a `Vec` per query would cost more than the selection it feeds. Both variants go through the same arithmetic, which is the point: blocks are cut at *set* offsets and mapped back to absolute positions, and those coincide only when the set starts at zero without holes. ### Two departures from the plan above **I took the first option, not the second.** The issue argued the assertion was the right first move. It converts a silent wrong answer into a loud failure, which is real value — but it leaves the model unbatchable, so it defers the work rather than doing it. The selector generalisation is what makes the eventual plumbing a matter of passing `Sparse` instead of `Contiguous`. **The assertion already existed.** I added one to `positions_from`, then wrote a test to prove it fires — and the test failed, because `Indexer::advance` was already enforcing `total == past_len + seq_len` one level up. Mine was unreachable. Deleted it and widened the existing message to name the contiguity assumption it protects, rather than only the cache divergence that is its symptom. ### Tests - `blocks_are_cut_from_the_visible_set_not_the_position_axis` — visible `[3,5,7,11,13,17,19,23,29]`, budget 1 block: selection returns `[13,17,19,23,29]`. Under the old arithmetic it returned `4..8` plus `8`, and 4, 6, 8, 9, 10 are positions the mask forbids. - `contiguous_and_sparse_agree_on_the_same_positions` — the two representations must not drift, since only the explicit one is checked against the reference. - `a_cache_inconsistent_with_past_len_is_refused` / `..._is_served` — drives the guard both ways, because an assertion that cannot be made to fire is not evidence. Mutations verified to fail: index arithmetic in `select`, index arithmetic in `block_starts`, and removing the guard — which returns `[[[0,1,2,3,4,5]]]` for an inconsistent cache, the exact plausible-looking wrong answer this issue described. Both parity fixtures still match the reference, confirming the contiguous path is untouched. ### Correction to this issue's framing This was described — including by me, elsewhere — as the blocker for batching. It is not. `model.rs` builds its own causal mask and no padding mask reaches this arch, so there is currently **no source of a non-contiguous visible set at all**. The actual obstacle is that `past_len` is a single scalar for the whole batch: ragged sequences need per-sequence offsets, through the indexer, the GDN state and PLE's rolling conv context alike. Filed as #326 rather than left implied inside a closed issue. Closing.
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: helexa/helexa#324