GLM prefill now attends only to the keys its indexer picks

This pull request makes long prompts much faster to prefill on GLM models with DeepSeek sparse attention. The indexer already picks the top-k keys for each query. Now prefill attention reads only those keys, gathered from the cache, instead of masking all the others. It's on by default once the cache holds about 16,000 tokens.

Slide text

The change How it works Evidence

GLM prefill now attends only to the keys its indexer picks

sparse_prefill · on by default from 16,384 cached tokens

Before this change, prefill turned the indexer's selection into a dense boolean mask, one row per query and one column per key, and then ran masked attention over every key. In the author's words, the selection is computed, then the work is spent on all N keys anyway. The mask and the expanded keys and values also grow linearly with context.

Here's how the sparse path works.

Two new fields on ModelArgs control it. Sparse prefill is on by default, with a minimum context of 16,384 tokens. It only runs for a single sequence, so batches keep the masked dense path. It only runs in prefill, so decode is untouched. And below the minimum context the dense path is faster, so the code keeps it there.

class ModelArgs(BaseModelArgs):
    sparse_prefill: bool = True
    sparse_prefill_min_context: int = 16384

# GlmMoeDsaAttention.__call__
sparse = (
    self.sparse_prefill
    and B == 1
    and L > 1
    and kv_latent.shape[2] >= self.sparse_prefill_min_context
)

The attention half is a new function, sparse prefill. It joins the latent and rope caches into one 576-wide key per position, then gathers just the top-k rows each query picked. All heads share that gathered block. Attention runs in the absorbed latent space, the same form decode already uses, so each chunk costs L times top-k instead of L times N. Chunks longer than 1,024 queries split into query blocks, which keeps the gathered keys small.

The indexer gets the same treatment, in a function called blocked top-k. It scores the keys in blocks of 4,096 and folds each block into a running top-k list, so each argpartition sorts rows of top-k plus 4,096 instead of N. The full score tensor, batch by 32 heads by L by N, is no longer built. It selects the same keys as the stock indexer, up to the order of equal scores.

Almost all of it lands in two files: the model file, with about 100 added lines, and a new test file of 175 lines. The rest rides along. Pillow joins the package dependencies, the black pre-commit hook gets a version bump, and two init files turn the test folders into packages.

Now the evidence.

This is an M3 Ultra running GLM 5.3 at 4-bit, with 512-token prefill chunks. Black is the patch, and the dashed gray line is the baseline. Up to 4,000 tokens of context the two match, because the gate keeps the dense path. From 16,000 on, the baseline falls off, and the patch slows far more gently. At 128,000 tokens, prefill goes from about 12 to 96 tokens per second, roughly eight times faster. Decode speed is unchanged.

Memory follows the same curve. With 1,024-token chunks, the baseline's peak climbs fast with context, and the patch's climbs slowly. At 128,000 tokens it drops from about 480 to about 440 gigabytes. With 2,048-token chunks at 128,000 tokens, the baseline was skipped because its predicted peak was above the GPU working set. The patch ran it, peaking at about 440 gigabytes.

Output stays close to stock. On a prompt of about 12,000 tokens, mean KL divergence is 0.0046, and the top token matches stock at almost 99 of every 100 positions. Eight new tests cover the path. Sparse output matches dense within a ten-thousandth, batches and short or gated contexts take the dense path unchanged, query blocks match, and decode works after a sparse prefill. Three of them fail without this change.

Two limits. Batches still take the dense path, and DeepSeek models with sparse attention aren't covered yet. They're planned for a follow-up, with a smaller expected gain, since their indexer runs on every layer. It's on by default, and setting sparse prefill to false in the model config turns it off. GLM prefill now attends only to the keys its indexer picks.