transformers
d75219d9 - Run Qwen4-Exp under continuous batching: indexed attention and per-layer n-gram embeddings

Commit
1 day ago
Run Qwen4-Exp under continuous batching: indexed attention and per-layer n-gram embeddings Indexed attention gets its allocator: full attention rows plus the indexer's keys, which the layer declares as a third row and writes before the attention writes its keys and values. Under the engine, the indexer selects for every packed query among the keys of its own sequence, one sequence at a time, and rotates the pooled block keys with key_position_embeddings, the rotary embeddings of the positions the batch reads, which the model computes once per forward. The selection goes on top of the engine's mask, so these layers run with sdpa and without the decode fast path. The per-layer n-gram embeddings keep two states per request: the last tokens of the sequence, which start as end-of-sequence tokens, and the last inputs of their dilated convolution. The packed helpers gain a dilation and token windows. The engine gives each attention type reading keys one of the two key-length rows, instead of fixed rows for full and sliding attention. Tests: dilated packed convolution, packed token windows, and for Qwen3.5 and Qwen4-Exp the engine's logprobs against a plain forward over each sequence (4.8e-6 in float32 on Qwen4-Exp). That test fails when the indexer selects every key, when the n-gram context starts at 0, when the convolution state is dropped, and, on both models, with the mask collapse fixed in the base branch.
Author
Parents
Loading