Note: attention sinks, why the first few tokens matter disproportionately

The problem, from first principles

[05-sliding-window-attention]'s fixed window solves the compute problem cleanly, but Xiao et al. (2023, "StreamingLLM") discovered an unexpected practical failure when actually deploying it: once the FIRST few tokens of a sequence slide OUTSIDE the window (as later positions push them out of range), model quality often collapses sharply, even though those tokens' actual CONTENT rarely mattered semantically. Digging into WHY revealed something specific to how softmax works: softmax weights are required to sum to exactly 1 even when there's genuinely nowhere useful to "put" the leftover weight (a query with no strongly relevant key still has to distribute its probability mass somewhere), and empirically, models learn to dump this leftover weight onto whatever's EARLIEST and always-available, the first few tokens, turning them into a kind of attention "sink" regardless of what they actually contain. Removing those tokens from view doesn't just lose a little context, it breaks the mechanism the model was implicitly relying on to keep its softmax outputs well-behaved.

The practical fix requires no retraining: keep a small, fixed number of "sink" tokens (often just 1-4) PERMANENTLY visible, in addition to [05-sliding-window-attention]'s ordinary sliding window, regardless of how far the sequence has moved on since. This lets a long-running generation stream stay bounded in memory (like ordinary sliding-window attention) while avoiding the quality collapse that dropping the very first tokens would otherwise cause.

From theory to code

Implement attention_sink_mask(seq_len, window_size, num_sink_tokens): [05-sliding-window-attention]'s bounded-window mask, but with the first num_sink_tokens positions marked as ALWAYS visible (subject only to ordinary causality), regardless of whether they'd otherwise fall outside the window.

Constraints

• A position within num_sink_tokens of the start is visible to any LATER query (j < num_sink_tokens and j <= i implies visible), no matter how far outside the sliding window it would otherwise be.
• Causality still applies EVEN to sink tokens: a sink token in the FUTURE relative to the query (j > i) remains masked, exactly like every other position.
• Every position OUTSIDE both the sink range and the sliding window remains masked, unchanged from [05-sliding-window-attention]'s behavior.
• num_sink_tokens=0 must reduce EXACTLY to [05-sliding-window-attention]'s ordinary sliding-window mask.

Hints

positions = np.arange(seq_len)
is_causally_visible = positions[None, :] <= positions[:, None]  # j <= i
is_sink_token = positions[None, :] < num_sink_tokens             # j < num_sink_tokens
always_visible = is_sink_token & is_causally_visible