Skip to content
mini-sglang from Scratch Stage 1 · Computing it right: from skeleton to real structure

The attention backend and a reference implementation

Attention is the only computation in the engine that has to know the batch's structure: every request has a different length, the KV is scattered across pages of the pool, and some requests already have a cached prefix. mini-sglang puts all of that into the "attention backend": one prepare_metadata per batch turns the request information into whatever the kernel wants, and one forward per layer writes KV and computes attention. Upstream has three GPU backends, FlashInfer, FlashAttention and TensorRT-LLM; this chapter writes a PyTorch reference backend first to get the semantics straight, and chapter 17 wires up the GPU ones.

Self-test: if you can answer these, skip the chapter

  1. What does cu_seqlens_q = [0, 6, 9, 10] say?
  2. A request has 7 tokens of KV in total and 3 queries this step. How does the causal mask line up? Which keys can the first query see?
  3. Why is the metadata prepared once per batch rather than once per layer?
  4. Must the KV be written before or after attention is computed? Why?
Answers (try it yourself first, then expand)
  1. This batch has 3 requests with 6, 3 and 1 queries each (as a prefix sum: request i's queries are [cu[i], cu[i+1])).
  2. Aligned to the bottom right: the 3 queries are the request's last 3 positions (4, 5, 6). The first query sits at position 4 and sees keys 0 through 4, five in all.
  3. Within one batch the variable-length information, the KV lengths and the page table are identical for every layer, so preparing once serves them all; preparing per layer is pure waste.
  4. Before: this step's new tokens need their own K and V in the attention too, since causal attention has every token see itself. So the KV goes into the pool first, and then all of it is read back for the computation.

Files you will write: attention/base.py, attention/utils.py, attention/torch_backend.py, attention/__init__.py.

Files so far: 39 (4 new, 0 modified, 1 brought in early)
minisgl/
├── __init__.py
├── core.py
├── attention/
│   ├── __init__.py  ← new in this chapter
│   ├── base.py  ← new in this chapter
│   ├── utils.py  ← new in this chapter
│   ├── torch_backend.py  ← new in this chapter
├── distributed/  (3 files, unchanged)
├── kernel/  (2 files, unchanged)
├── kvcache/  (4 files, unchanged)
├── layers/  (8 files, unchanged)
├── models/  (7 files, unchanged)
├── scheduler/  (2 files, unchanged)
├── utils/
│   ├── __init__.py
│   ├── misc.py
│   ├── hf.py
│   ├── torch_utils.py
│   ├── registry.py
│   ├── logger.py
│   ├── device.py  ← brought in early (explained in chapter 6)
This step's main: examples/ch05_attention.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

The interface

minisgl/attention/base.py · BaseAttnBackend
class BaseAttnBackend(ABC):
    @abstractmethod
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int,
                batch: Batch) -> torch.Tensor:
        """Write this step's k and v into the KV pool first, then compute attention. q: [T, Hq, D], returns [T, Hq, D]."""

    @abstractmethod
    def prepare_metadata(self, batch: Batch) -> None:
        """Called once per batch (not once per layer): turn the request information into what the kernel wants."""

    # the three methods below serve CUDA Graph (chapter 18)
    @abstractmethod
    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: ...

    @abstractmethod
    def prepare_for_capture(self, batch: Batch) -> None: ...

    @abstractmethod
    def prepare_for_replay(self, batch: Batch) -> None: ...
- prepare_metadata(batch): the scheduler calls it once before each batch, and the result goes into batch.attn_metadata. The model has 28 layers whose attention shapes are identical, so the metadata only has to be prepared once; a GPU backend also copies data from the CPU to the GPU here, or runs FlashInfer's plan. - forward(q, k, v, layer_id, batch): called per layer. Write first, read second: put this step's new k and v into the KV pool, then read every request's full KV back out and compute attention. The order cannot be reversed, because this step's queries have to see this step's keys (the diagonal of the causal mask). - The last three methods serve CUDA Graph and wait until chapter 18.

The metadata also provides get_last_indices(bs): during prefill, each request's last token's index into the flattened q, for the LM head (chapter 2).

Variable-length, paged attention with prefixes

The tensors in a batch are one-dimensional arrays with the requests packed end to end. An example makes clear what attention needs to know: request A prefills 6 tokens from scratch; request B hits the prefix cache on its first 4 tokens and computes 3 this step; request C already has 8 tokens and decodes 1.

  • Where this step's queries are for each request: cu_seqlens_q = [0, 6, 9, 10] (cumulative sequence lengths), so request i is q[cu[i]:cu[i+1]];
  • How long each request's KV is: cache_seqlens = [6, 7, 9], each request's device_len;
  • Where each request's KV sits in the pool: the matching row of the page table, taking its first cache_seqlens[i] locations;
this step's q (requests packed end to end, no padding) A0 A1 A2 A3 A4 A5 B0 B1 B2 C0 0 6 9 10 cu_seqlens_q = [0, 6, 9, 10] get_last_indices → [5, 8, 9] a request's whole KV, gathered from the pool through the page table A 0 1 2 3 4 5 cache_seqlens = 6 B 40 41 42 43 20 21 22 cache_seqlens = 7 (first 4 cached) C 50 51 52 53 54 55 56 57 58 cache_seqlens = 9 (first 8 cached) already cached written this step (out_loc) request B's mask (bottom-right aligned) q0 q1 q2 0 1 2 3 4 5 6 q0 sits at position 4: it sees keys 0-4 this step's n queries are the last n positions
how a mixed batch splits its queries, where the KV comes from, and the causal mask
  • How the causal mask lines up: this step's n queries are the last n positions of the sequence. With 7 tokens of KV and 3 queries, the first query sits at absolute position 4 and sees keys 0 through 4. This is "bottom-right alignment": draw the [n, kv_len] mask matrix and its diagonal lands in the bottom-right corner.

The reference backend is that description turned straight into code:

minisgl/attention/torch_backend.py · TorchAttnBackend.forward
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int,
            batch: Batch) -> torch.Tensor:
    metadata = batch.attn_metadata
    assert isinstance(metadata, TorchAttnMetadata)
    self.kvcache.store_kv(k, v, batch.out_loc, layer_id)  # write first, read second: this step's tokens take part in the attention too
    k_pool = self.kvcache.k_cache(layer_id).flatten(0, 1)  # [num_tokens, Hkv, D]
    v_pool = self.kvcache.v_cache(layer_id).flatten(0, 1)
    cu_q = metadata.cu_seqlens_q.tolist()
    seq_lens = metadata.cache_seqlens.tolist()
    out = torch.empty_like(q)
    for i, kv_len in enumerate(seq_lens):
        qs, qe = cu_q[i], cu_q[i + 1]
        n = qe - qs  # this request has n queries this step, which are the sequence's last n positions
        if n == 0:
            continue
        loc = metadata.page_table[i, :kv_len].long()
        qi = q[qs:qe].transpose(0, 1)  # [Hq, n, D]
        ki = k_pool[loc].transpose(0, 1)  # [Hkv, kv_len, D]
        vi = v_pool[loc].transpose(0, 1)
        # query r sits at absolute position kv_len - n + r and sees every key at or before it
        q_pos = torch.arange(kv_len - n, kv_len, device=q.device).unsqueeze(1)
        mask = torch.arange(kv_len, device=q.device).unsqueeze(0) <= q_pos
        oi = F.scaled_dot_product_attention(qi, ki, vi, attn_mask=mask, scale=self.scale,
                                            enable_gqa=True)
        out[qs:qe] = oi.transpose(0, 1)
    return out
For each request: gather its full K and V from the pool through the page table, build the bottom-right-aligned causal mask, and call scaled_dot_product_attention (with enable_gqa=True so 16 q heads share 8 KV heads). It is a Python loop and it is not fast, but it is completely unambiguous and it runs on a CPU.

Preparing the metadata:

minisgl/attention/torch_backend.py · TorchAttnBackend.prepare_metadata
def prepare_metadata(self, batch: Batch) -> None:
    reqs = batch.padded_reqs
    seqlens_q = [req.extend_len for req in reqs]
    seqlens_k = [req.device_len for req in reqs]
    max_seqlen_k = max(seqlens_k)
    page_table = get_global_ctx().page_table
    use_pin = pin(self.device)
    batch.attn_metadata = TorchAttnMetadata(
        cu_seqlens_q=make_cu_seqlens(seqlens_q, self.device, use_pin),
        cache_seqlens=torch.tensor(seqlens_k, dtype=torch.int32, pin_memory=use_pin)
        .to(self.device, non_blocking=True),
        page_table=torch.stack([page_table[r.table_idx, :max_seqlen_k] for r in reqs]),
        max_seqlen_q=max(seqlens_q),
    )
The page table takes the first max_seqlen_k columns of each request's row, stacked into [bs, max_seqlen_k], which is what the official FlashAttention backend does. It uses batch.padded_reqs rather than batch.reqs, because the dummy requests CUDA Graph pads in need metadata too.

Running the example above and comparing against "complete causal attention with no paging and no packing":

examples/ch05_attention.py
"""Chapter 5: a mixed batch's attention metadata, and the reference backend against unpaged complete attention.

Request A prefills 6 tokens from scratch; request B hits the cache on its first 4 and computes 3
more; request C decodes one step.
"""

import torch
import torch.nn.functional as F
from minisgl.attention.torch_backend import TorchAttnBackend
from minisgl.core import Batch, Context, Req, SamplingParams, set_global_ctx
from minisgl.distributed import set_tp_info
from minisgl.kvcache.mha_pool import MHAKVCache
from minisgl.models import ModelConfig
from minisgl.utils import cached_load_hf_config

set_tp_info(0, 1)
H, HKV, D = 16, 8, 128
ctx = Context(page_size=1)
set_global_ctx(ctx)
ctx.kv_cache = MHAKVCache(HKV, num_layers=1, head_dim=D, num_pages=64, page_size=1,
                          dtype=torch.float32, device=torch.device("cpu"))
ctx.page_table = torch.zeros((3, 16), dtype=torch.int32)
ctx.page_table[0, :6] = torch.arange(0, 6)
ctx.page_table[1, :7] = torch.tensor([40, 41, 42, 43, 20, 21, 22])  # the first 4 positions come from the prefix cache
ctx.page_table[2, :9] = torch.arange(50, 59)
backend = TorchAttnBackend(ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B")))


def req(n, cached, row):
    return Req(input_ids=torch.zeros(n, dtype=torch.int32), table_idx=row, cached_len=cached,
               output_len=4, uid=row, sampling_params=SamplingParams(), cache_handle=None)


torch.manual_seed(0)
lens, cached = [6, 7, 9], [0, 4, 8]
full = [(torch.randn(n, H, D), torch.randn(n, HKV, D), torch.randn(n, HKV, D)) for n in lens]
for i, c in enumerate(cached):  # the cached part is written into the KV pool beforehand
    if c:
        ctx.kv_cache.store_kv(full[i][1][:c], full[i][2][:c], ctx.page_table[i, :c], 0)

batch = Batch(reqs=[req(n, c, i) for i, (n, c) in enumerate(zip(lens, cached))], phase="prefill")
batch.padded_reqs = batch.reqs
batch.out_loc = torch.cat([ctx.page_table[i, c:n] for i, (n, c) in enumerate(zip(lens, cached))])
backend.prepare_metadata(batch)
m = batch.attn_metadata
print("cu_seqlens_q :", m.cu_seqlens_q.tolist(), " (each request's query range this step)")
print("cache_seqlens:", m.cache_seqlens.tolist(), "  (each request's total KV length)")
print("page_table   :", m.page_table.tolist())
print("out_loc      :", batch.out_loc.tolist())
print("last-position indices:", m.get_last_indices(3).tolist())

q = torch.cat([f[0][c:] for f, c in zip(full, cached)])
k = torch.cat([f[1][c:] for f, c in zip(full, cached)])
v = torch.cat([f[2][c:] for f, c in zip(full, cached)])
out = backend.forward(q, k, v, layer_id=0, batch=batch)
ref = torch.cat([
    F.scaled_dot_product_attention(f[0].transpose(0, 1), f[1].transpose(0, 1), f[2].transpose(0, 1),
                                   is_causal=True, scale=D**-0.5, enable_gqa=True).transpose(0, 1)[c:]
    for f, c in zip(full, cached)])
print("max error against complete causal attention:", (out - ref).abs().max().item())
What it prints
cu_seqlens_q : [0, 6, 9, 10]  (each request's query range this step)
cache_seqlens: [6, 7, 9]   (each request's total KV length)
page_table   : [[0, 1, 2, 3, 4, 5, 0, 0, 0], [40, 41, 42, 43, 20, 21, 22, 0, 0], [50, 51, 52, 53, 54, 55, 56, 57, 58]]
out_loc      : [0, 1, 2, 3, 4, 5, 20, 21, 22, 58]
last-position indices: [5, 8, 9]
max error against complete causal attention: 5.960464477539062e-07
A few observations:

  • Request B's KV is not contiguous in the pool: the 4 tokens that hit the prefix cache are at 40-43 and the 3 new ones at 20-22. out_loc only covers this step's new tokens.
  • The page table is cut at max_seqlen_k = 9, so shorter requests have irrelevant zeros after their data, which cache_seqlens keeps out of reach.
  • get_last_indices gives [5, 8, 9], the indices of the three requests' last queries in the flattened q.
  • Against per-request complete causal attention the difference is at the level of float32 rounding.

Creating a backend: one or a pair

minisgl/attention/__init__.py · create_attention_backend
def create_attention_backend(backend: str, config: ModelConfig) -> BaseAttnBackend:
    """backend is either one name or a "prefill backend,decode backend" pair."""
    validate_attn_backend(backend, allow_auto=False)
    if "," in backend:
        p_name, d_name = backend.split(",", 1)
        if p_name != d_name:
            logger.info(f"Hybrid attention backend: prefill={p_name}, decode={d_name}")
            return HybridBackend(create_attention_backend(p_name, config),
                                 create_attention_backend(d_name, config))
        backend = p_name
    return SUPPORTED_ATTENTION_BACKENDS[backend](config)
Backends are created by name through the registry. "fa,fi" means FlashAttention for prefill and FlashInfer for decode, dispatched by HybridBackend on batch.is_prefill:

minisgl/attention/base.py · HybridBackend
class HybridBackend(BaseAttnBackend):
    """Different backends for prefill and decode, such as FA3 for prefill and FlashInfer for decode on Hopper."""

    def __init__(self, prefill_backend: BaseAttnBackend, decode_backend: BaseAttnBackend) -> None:
        self.prefill_backend = prefill_backend
        self.decode_backend = decode_backend

    def _pick(self, batch: Batch) -> BaseAttnBackend:
        return self.prefill_backend if batch.is_prefill else self.decode_backend

    def forward(self, q, k, v, layer_id: int, batch: Batch):
        return self._pick(batch).forward(q, k, v, layer_id, batch)

    def prepare_metadata(self, batch: Batch) -> None:
        self._pick(batch).prepare_metadata(batch)

    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
        self.decode_backend.init_capture_graph(max_seq_len, bs_list)  # only decode is ever captured

    def prepare_for_capture(self, batch: Batch) -> None:
        self.decode_backend.prepare_for_capture(batch)

    def prepare_for_replay(self, batch: Batch) -> None:
        self.decode_backend.prepare_for_replay(batch)
CUDA Graph only captures decode, so the three graph-related methods only forward to the decode backend. With --attn auto the engine picks torch on a CPU, fa,fi on Hopper (SM90) and fi on other GPUs. There is one more constraint, on dtype: FlashInfer's and FlashAttention's kernels are compiled for fp16 / bf16 only, so --dtype float32 — the precision this book uses when comparing token by token against Hugging Face — falls back to the torch backend even on a GPU.

Difference from upstream: the reference backend

Upstream has no PyTorch backend (production SGLang has a torch_native backend that plays a similar role). Our torch backend uses the same metadata format as the official FlashAttention backend and implements the three CUDA Graph hooks the way that backend does, which is what lets chapter 18 verify CUDA Graph's data flow on a CPU. Upstream defaults to the trtllm backend on Blackwell (SM100), which we do not implement.

The official implementation

Tests

tests/test_ch05_attention.py · test_torch_backend_mixed_batch_with_prefix_hits
def test_torch_backend_mixed_batch_with_prefix_hits():
    """Two requests: A prefills 6 tokens from scratch; B has its first 4 tokens cached and computes only the last 3."""
    engine = build_engine(QWEN3)
    identity_page_table(engine, 2)
    backend, layer, H, Hkv, D = engine.attn_backend, 5, 16, 8, 128
    torch.manual_seed(0)
    full = {name: (torch.randn(n, H, D), torch.randn(n, Hkv, D), torch.randn(n, Hkv, D))
            for name, n in (("A", 6), ("B", 7))}
    # B's first 4 tokens are written into the KV pool beforehand (imitating a prefix-cache hit)
    loc_b = engine.page_table[1, :4]
    engine.kv_cache.store_kv(full["B"][1][:4], full["B"][2][:4], loc_b, layer)
    a, b = new_req(list(range(6)), 0), new_req(list(range(7)), 1, cached_len=4)
    from minisgl.core import Batch

    batch = Batch(reqs=[a, b], phase="prefill")
    batch.padded_reqs = batch.reqs
    batch.out_loc = torch.cat([engine.page_table[0, :6], engine.page_table[1, 4:7]])
    backend.prepare_metadata(batch)
    q = torch.cat([full["A"][0], full["B"][0][4:]])
    k = torch.cat([full["A"][1], full["B"][1][4:]])
    v = torch.cat([full["A"][2], full["B"][2][4:]])
    out = backend.forward(q, k, v, layer, batch)
    scale = D**-0.5
    ref_a = dense_causal(*full["A"], scale)
    ref_b = dense_causal(*full["B"], scale)[4:]
    assert torch.allclose(out[:6], ref_a, atol=1e-5) and torch.allclose(out[6:], ref_b, atol=1e-5)
    # the LM head only needs each request's last position
    assert batch.attn_metadata.get_last_indices(2).tolist() == [5, 8]
    engine.shutdown()

How to explain it

On the attention backend: keep the two jobs apart. One prepare_metadata per batch describes the variable lengths and the paging (cu_seqlens_q splits the queries per request, cache_seqlens gives the KV lengths, the page table gives the KV locations), and one forward per layer writes this step's K and V into the cache and then computes attention, because otherwise this step's tokens cannot see themselves. In a prefill with a prefix the causal mask aligns bottom right: with 7 tokens of KV and 3 queries, the first query sees the first 5 keys. The metadata is prepared once per batch because every layer sees the same batch structure. The reference backend is a Python loop plus SDPA and serves as the baseline for checking the FlashInfer and FlashAttention backends; prefill and decode may use different backends.

Exercises

  1. The reference backend builds an [n, kv_len] mask per request. During decode n = 1 and the mask is all true. Write a fast path for decode: pad every request's KV to the maximum length as [bs, max_len, H, D] and do it in one batched SDPA.
  2. What happens to a prefill if you forget to call store_kv first in forward? What about a decode?
  3. What does cu_seqlens_q equal for "a prefill with no prefix hit"? And for "pure decode"? (Hint: look at the two fast paths in the FlashInfer backend in chapter 17.)
Answers
  1. Build an [bs, 1, max_len] length mask from cache_seqlens (arange(max_len) < len) and gather the KV in one go with page_table[:, :max_len]; make sure the padding is masked out.
  2. During prefill this step's keys read stale data from the pool (uninitialized, or another request's), so the result is wrong; during decode the KV at the last position, the current token, is stale, which is wrong in the same way. Only "write first, read second" lets this step's token see itself.
  3. With no prefix hit each request's query count equals its KV length, so cu_seqlens_q == cu_seqlens_k; in pure decode each request has 1 query, so cu_seqlens_q = [0, 1, 2, ..., bs].

Summary

  • The attention backend does two things: one prepare_metadata per batch, and one forward per layer that writes KV before computing attention.
  • The metadata describes variable lengths and paging: cu_seqlens_q splits the queries, cache_seqlens gives the KV lengths, the page table gives the KV locations, and the causal mask aligns bottom right.
  • The reference backend is a Python loop plus SDPA: clear semantics, runs on a CPU, and serves as the baseline for the GPU backends.
  • Backends are created by name, and prefill and decode may use different ones.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue