SGLang dsa

profile

goal

The goals of this experiment are:

1. Draw heatmaps to observe attention patterns within a sequence.
2. Measure the intersections and unions of selected KV indices within different block sizes.

layer0

(Layer 0 has an unusual attention pattern. Deeper layers are more similar to one another, so this section is less important than the later analysis.)

I selected a long document from LongBench and ran one prefill. The first 2,048 tokens use standard MHA, meaning dense attention.

I plotted several heatmaps covering tokens 2,048–8,192. The chunk-prefill limit is 8,192; see srt/managers/scheduler.py.

The horizontal and diagonal lines are very clear. Attention exhibits strong sink and relative-position characteristics; the lightning indexer tends to select KV entries in two ways:

1. sink attention: select KV entries at particular absolute positions to preserve global information.
2. relative attention: looking backward from the current query ID, select KV entries at periodic intervals.

As the window grows, two further patterns appear:

1. The horizontal lines remain roughly stable.
2. The diagonal lines gradually weaken.

In other words, relative attention yields to sink attention as the window grows.

Intersections and unions of indices selected within a block:

This plot appears to support the observations above. The intersection seems to represent the sink pattern, while the portion of the union outside the intersection likely represents relative attention.

Based on these ratios and a table of actual selection counts—not shown here—the sizes of the intersection and union have no significant relationship with block position. Once sufficiently far from the dense region, both approach stable values.

If the intersection represents the sink, then as the sequence grows the earlier sink narrows to give later tokens a chance to appear in attention. The number of sink tokens remains relatively stable, though it begins to decline slightly around 7–8K; perhaps it falls further at greater lengths.

Other Layers

Layers after layer 0 have broadly similar patterns.

The heatmaps show that attention is primarily local, with strided attention as a secondary pattern.

The curve retains its shape up to a chunk size of 128, indicating high similarity among the KV indices selected by 128 adjacent tokens. The average overlap reaches 755, roughly one third of the total.

The following plots analyze locality.

Local selections are also highly similar within groups of 128 queries.

trade off

SGLang uses standard MHA directly below 2,048 tokens and enables MHA-mode MLA only beyond that threshold.

Reference

SGLang blog

SGLang attn backend

https://github.com/sgl-project/sglang/issues/11060

https://github.com/sgl-project/sglang/blob/main/docs/basic_usage/deepseek_v32.md

https://github.com/deepseek-ai/FlashMLA

sglang-pr