mini-sglang from Scratch Stage 4 · Faster and bigger

GPU attention backends: FlashInfer and FlashAttention

The reference backend computes attention one request at a time in a Python loop, which is clear but slow. On a GPU mini-sglang uses two specialized libraries: FlashInfer (very strong at decode and good at prefill) and FlashAttention 3 (the fastest prefill on Hopper). Their interfaces differ, since FlashInfer plans before it runs and FlashAttention takes a page-number table directly, but both express the metadata of chapter 5. This chapter implements both backends the way upstream does and verifies the argument semantics on a CPU through PyTorch fakes with the same interfaces.

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

1. What do FlashInfer's plan and run do? Why does plan happen once per batch?
2. mini-sglang treats the KV pool as page size 1 for FlashInfer. What goes into indices? Why is last_page_len always 1?
3. Does FlashAttention's page_table hold token locations or page numbers? How is it derived from mini-sglang's global page table?
4. How do you verify a backend that only runs on a GPU, on a CPU?

Answers (try it yourself first, then expand)

1. plan takes this batch's metadata (each request's length, the page table) and works out the schedule: it divides the work, allocates temporary buffers and copies the metadata to the GPU. run then executes that plan per layer. Every layer of a batch has the same metadata, so plan runs once, deferred until the first layer.
2. indices holds the pool location of each token's KV (with a page size of 1 the page number is the token location), that is, one row of the page table; each page holds one token, so the last page is always full and last_page_len is always 1.
3. Page numbers. Take every page_size-th location from the global page table (which stores token locations) and divide by the page size to get each page's number.
4. Write a fake with the same interface as the real library: translate both libraries' argument semantics into one reference implementation (Python plus SDPA), run the backend code on a CPU, and require the result to match the reference backend exactly.

Files you will write: attention/fi.py, attention/fa.py, plus tests/fakes/flashinfer/ and tests/fakes/sgl_kernel/ for the tests.

Files so far: 75 (2 new, 0 modified, 0 brought in early)

minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/
│   ├── __init__.py
│   ├── base.py
│   ├── utils.py
│   ├── torch_backend.py
│   ├── fi.py  ← new in this chapter
│   ├── fa.py  ← new in this chapter
├── distributed/  (3 files, unchanged)
├── engine/  (5 files, unchanged)
├── kernel/  (2 files, unchanged)
├── kvcache/  (5 files, unchanged)
├── layers/  (8 files, unchanged)
├── llm/  (2 files, unchanged)
├── message/  (5 files, unchanged)
├── models/  (7 files, unchanged)
├── moe/  (2 files, unchanged)
├── scheduler/  (9 files, unchanged)
├── server/  (4 files, unchanged)
├── tokenizer/  (4 files, unchanged)
├── utils/  (8 files, unchanged)

This step's main: examples/ch17_gpu_attention.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

FlashInfer: plan and run

FlashInfer provides a "batch prefill" and a "batch decode" wrapper for paged KV caches. Using them takes two steps:

• plan(...): pass in this batch's shape information, namely each request's query range (qo_indptr), which pages make up each request's KV (kv_indptr dividing kv_indices), and how full the last page is (last_page_len). From that FlashInfer works out the division of labour on the CPU (which thread block handles which stretch of KV) and copies it to the GPU asynchronously;
• run(q, paged_kv_cache): called once per layer, computing according to the plan.

All 28 layers have the same shapes, so plan is needed only once. mini-sglang defers it into the first layer's forward (_initialize_metadata_once) rather than doing it in prepare_metadata:

minisgl/attention/fi.py · FlashInferBackend.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, FIMetadata)
    self._initialize_metadata_once(metadata)  # plan on the first layer, then just run for the rest
    self.kvcache.store_kv(k, v, batch.out_loc, layer_id)

    def flatten(cache: torch.Tensor) -> torch.Tensor:  # [pages, page_size, H, D] -> [tokens, 1, H, D]
        return cache.view(-1, 1, cache.shape[2], cache.shape[3])

    kv = (flatten(self.kvcache.k_cache(layer_id)), flatten(self.kvcache.v_cache(layer_id)))
    return metadata.wrapper.run(q=q, paged_kv_cache=kv)

minisgl/attention/fi.py · FlashInferBackend.prepare_metadata

def prepare_metadata(self, batch: Batch) -> None:
    reqs = batch.padded_reqs
    bs = len(reqs)
    seqlens_q = [r.extend_len for r in reqs]
    seqlens_k = [r.device_len for r in reqs]
    kw = dict(dtype=torch.int32, pin_memory=pin(self.device))
    seq_lens_cpu = torch.tensor(seqlens_k, **kw)
    cu_k = torch.tensor([0] + seqlens_k, **kw).cumsum_(dim=0)
    if max(seqlens_q) == 1:  # decode: 1 query per request
        cu_q = torch.arange(0, bs + 1, **kw)
    elif all(r.cached_len == 0 for r in reqs):  # a prefill with no cache hit: q and k are the same length
        cu_q = cu_k
    else:
        cu_q = torch.tensor([0] + seqlens_q, **kw).cumsum_(dim=0)
    page_table = get_global_ctx().page_table
    batch.attn_metadata = FIMetadata(
        cu_seqlens_q_cpu=cu_q,
        cu_seqlens_k_cpu=cu_k,
        cu_seqlens_q_gpu=cu_q.to(self.device, non_blocking=True),
        indices=torch.cat([page_table[r.table_idx, : r.device_len] for r in reqs]),
        last_page_len_cpu=self._get_ones_cpu(bs),
        num_qo_heads=self.qo_head_local,
        num_kv_heads=self.kv_head_local,
        head_dim=self.config.head_dim,
        seq_lens_cpu=seq_lens_cpu,
        dtype=self.kvcache.dtype,
        wrapper=self.decode_wrapper if batch.is_decode else self.prefill_wrapper,
    )

A few designs worth noticing:

• treating the KV pool as page size 1: forward views the pool as [tokens, 1, heads, dim], so one row of the global page table is kv_indices and last_page_len is always 1. Whatever the system's page size, the FlashInfer backend never has to convert anything, which is the payoff of chapter 4's "the page table stores token locations";
• two fast paths: in pure decode qo_indptr = [0, 1, …, bs]; in a prefill with no prefix hit each request's query count equals its KV length, so qo_indptr simply reuses kv_indptr;
• pinned memory and events: plan reuses a block of pinned memory for its asynchronous copy, so the previous copy has to finish before the next plan modifies it (last_event.synchronize());
• Tensor Cores: when the GQA group is 4 or more, decode is effectively a series of small matmuls and the Tensor Core decode kernel is faster.

minisgl/attention/fi.py · FlashInferBackend._initialize_metadata_once

def _initialize_metadata_once(self, metadata: FIMetadata) -> None:
    if metadata.initialized:
        return
    from flashinfer import BatchDecodeWithPagedKVCacheWrapper

    metadata.initialized = True
    # plan reuses one block of pinned memory for its asynchronous copy, so the previous copy must finish before it is overwritten
    self.last_event.synchronize()
    if isinstance(metadata.wrapper, BatchDecodeWithPagedKVCacheWrapper):
        metadata.wrapper.plan(
            indptr=metadata.cu_seqlens_k_cpu, indices=metadata.indices,
            last_page_len=metadata.last_page_len_cpu, num_qo_heads=metadata.num_qo_heads,
            num_kv_heads=metadata.num_kv_heads, head_dim=metadata.head_dim, page_size=1,
            pos_encoding_mode="NONE", seq_lens=metadata.seq_lens_cpu,
            data_type=metadata.dtype, q_data_type=metadata.dtype,
            kv_data_type=metadata.dtype, non_blocking=True,
        )
    else:
        metadata.wrapper.plan(
            qo_indptr=metadata.cu_seqlens_q_cpu, paged_kv_indptr=metadata.cu_seqlens_k_cpu,
            paged_kv_indices=metadata.indices, paged_kv_last_page_len=metadata.last_page_len_cpu,
            num_qo_heads=metadata.num_qo_heads, num_kv_heads=metadata.num_kv_heads,
            head_dim_qk=metadata.head_dim, page_size=1, pos_encoding_mode="NONE",
            seq_lens=metadata.seq_lens_cpu, q_data_type=metadata.dtype,
            kv_data_type=metadata.dtype, non_blocking=True, causal=True,
        )
    self.last_event.record()

FlashAttention: a page-number table

flash_attn_with_kvcache in sgl_kernel (FlashAttention 3) reads the paged KV pool directly and wants page_table (each request's page numbers), cache_seqlens (each request's KV length) and cu_seqlens_q (the variable-length query split). There is no plan step.

minisgl/attention/fa.py · FlashAttentionBackend.prepare_metadata

def prepare_metadata(self, batch: Batch) -> None:
    reqs = batch.padded_reqs
    bs = len(reqs)
    seqlens_q = [r.extend_len for r in reqs]
    seqlens_k = [r.device_len for r in reqs]
    max_seqlen_k, max_seqlen_q = max(seqlens_k), max(seqlens_q)
    kw = dict(dtype=torch.int32, pin_memory=pin(self.device))
    cache_seqlens = torch.tensor(seqlens_k, **kw).to(self.device, non_blocking=True)
    cu_k = torch.tensor([0] + seqlens_k, **kw).cumsum_(dim=0).to(self.device, non_blocking=True)
    if max_seqlen_q == 1:
        cu_q = torch.arange(0, bs + 1, dtype=torch.int32, device=self.device)
    elif all(r.cached_len == 0 for r in reqs):
        cu_q = cu_k
    else:
        cu_q = torch.tensor([0] + seqlens_q, **kw).cumsum_(dim=0).to(self.device,
                                                                      non_blocking=True)
    # the global page table stores locations per token; take every page_size-th and divide by page_size for the page number
    page_table = get_global_ctx().page_table
    table = torch.stack([page_table[r.table_idx, :max_seqlen_k : self.page_size] for r in reqs])
    if self.page_size > 1:
        table = table.div(self.page_size, rounding_mode="floor")
    batch.attn_metadata = FAMetadata(cu_k, cu_q, cache_seqlens, max_seqlen_k, max_seqlen_q, table)

The global page table stores token locations, so every page_size-th location is a page's start and dividing by the page size gives its number: page_table[r, :max_len:page_size] // page_size.

Verifying on a CPU

Both libraries only run on NVIDIA GPUs. To check that the backend code passes the right arguments, we wrote a fake with the same interface for each, under tests/fakes/, placed at the front of the import path during tests:

tests/fakes/flashinfer/__init__.py · BatchPrefillWithPagedKVCacheWrapper

class BatchPrefillWithPagedKVCacheWrapper(_Base):
    def plan(self, qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len,
             num_qo_heads, num_kv_heads, head_dim_qk, page_size, causal=False, **kwargs: Any):
        assert causal
        self._plan(qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len, head_dim_qk)

tests/fakes/sgl_kernel/flash_attn.py · flash_attn_with_kvcache

def flash_attn_with_kvcache(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                            page_table: torch.Tensor, cache_seqlens: torch.Tensor,
                            cu_seqlens_q: torch.Tensor, cu_seqlens_k_new: torch.Tensor | None = None,
                            max_seqlen_q: int = 1, softmax_scale: float | None = None,
                            causal: bool = True, **kwargs) -> torch.Tensor:
    assert causal
    page_size = k_cache.shape[1]
    lens = cache_seqlens.tolist()
    pages = [page_table[i, : (n + page_size - 1) // page_size] for i, n in enumerate(lens)]
    scale = softmax_scale if softmax_scale is not None else q.shape[-1] ** -0.5
    return paged_attention(q, k_cache, v_cache, cu_seqlens_q.tolist(), pages, lens, scale)

Both translate their arguments into "each request's query range, KV page list and KV length" and hand them to one reference implementation:

tests/fakes/_fake_ref.py · paged_attention

def paged_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                    q_starts: List[int], kv_pages: List[torch.Tensor], kv_lens: List[int],
                    scale: float) -> torch.Tensor:
    """q: [total queries, Hq, D]; k_cache: [pages, page_size, Hkv, D].

    Request i's queries are q[q_starts[i]:q_starts[i+1]], and its KV is the first kv_lens[i] tokens
    of the pages kv_pages[i] concatenated. The queries are the sequence's last n positions (the
    causal mask aligns bottom right).
    """
    out = torch.empty_like(q)
    for i, (pages, kv_len) in enumerate(zip(kv_pages, kv_lens)):
        qs, qe = q_starts[i], q_starts[i + 1]
        n = qe - qs
        if n == 0:
            continue
        k = k_cache[pages.long()].flatten(0, 1)[:kv_len].transpose(0, 1)  # [Hkv, L, D]
        v = v_cache[pages.long()].flatten(0, 1)[:kv_len].transpose(0, 1)
        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
        o = F.scaled_dot_product_attention(q[qs:qe].transpose(0, 1), k, v, attn_mask=mask,
                                           scale=scale, enable_gqa=True)
        out[qs:qe] = o.transpose(0, 1)
    return out

So any argument the backend gets wrong, such as an indptr missing its leading 0, a page number not divided by the page size, or a mistaken qo_indptr during decode, makes the fake's result differ from the reference backend and the test fail. These fakes are also the shortest path to understanding what FlashInfer's and FlashAttention's arguments mean.

examples/ch17_gpu_attention.py

"""Chapter 17: what the metadata handed to the FlashInfer and FlashAttention backends looks like.

Neither library exists on a CPU, so tests/fakes (PyTorch fakes with the same interfaces) goes on
the import path, the backend code written the official way runs unchanged, and the result is
compared against the reference backend.
"""

import sys

sys.path.insert(0, "tests/fakes")

import torch  # noqa: E402
from minisgl.core import Batch, Req, SamplingParams, reset_global_ctx  # noqa: E402
from minisgl.distributed import DistributedInfo, reset_tp_info  # noqa: E402
from minisgl.engine import Engine, EngineConfig  # noqa: E402

PROMPTS = [[785, 6722, 315, 9625, 374], [16, 488, 220, 16, 284, 220]]


def build(backend: str, page_size: int) -> Engine:
    reset_global_ctx()
    reset_tp_info()
    e = Engine(EngineConfig(model_path="models/Qwen3-0.6B", tp_info=DistributedInfo(0, 1),
                            dtype=torch.float32, attention_backend=backend, page_size=page_size,
                            max_running_req=2, num_page_override=64 // page_size, max_seq_len_override=32))
    e.page_table[0, :16] = torch.arange(16, 32)  # request 0: positions 16 to 31
    e.page_table[1, :16] = torch.arange(40, 56)  # request 1: positions 40 to 55
    return e


def batch_of(engine: Engine, phase: str, reqs):
    b = Batch(reqs=reqs, phase=phase)
    b.padded_reqs = reqs
    b.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).int()
    b.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs])
    b.out_loc = torch.cat([engine.page_table[r.table_idx, r.cached_len:r.device_len] for r in reqs])
    engine.attn_backend.prepare_metadata(b)
    return b


def reqs(cached):
    return [Req(input_ids=torch.tensor(p, dtype=torch.int32), table_idx=i, cached_len=c, output_len=4,
                uid=i, sampling_params=SamplingParams(), cache_handle=None)
            for i, (p, c) in enumerate(zip(PROMPTS, cached))]


results = {}
for backend, ps in (("torch", 1), ("fi", 1), ("fa", 4)):
    engine = build(backend, ps)
    with engine.ctx.forward_batch(b := batch_of(engine, "prefill", reqs([0, 0]))):
        engine.model.forward()  # compute both requests' KV first
    rs = reqs([3, 5])  # request 0 hits its first 3 tokens, request 1 has only its last one left (like a decode)
    b = batch_of(engine, "prefill", rs)
    m = b.attn_metadata
    if backend == "fi":
        print("FlashInfer (KV pool used as page_size=1):")
        print("  qo_indptr      =", m.cu_seqlens_q_cpu.tolist())
        print("  kv_indptr      =", m.cu_seqlens_k_cpu.tolist())
        print("  kv_indices     =", m.indices.tolist())
        print("  last_page_len  =", m.last_page_len_cpu.tolist())
    if backend == "fa":
        print("FlashAttention (page_size=4, page table holds page numbers):")
        print("  cu_seqlens_q   =", m.cu_seqlens_q.tolist())
        print("  cache_seqlens  =", m.cache_seqlens.tolist())
        print("  page_table     =", m.page_table.tolist())
    with engine.ctx.forward_batch(b):
        results[backend] = engine.model.forward()
    if backend == "fi":
        print("  plan calls on the prefill wrapper:", engine.attn_backend.prefill_wrapper.plan_count,
              "(once per batch; one plan shared by that batch's 28 layers)")
    engine.shutdown()
for name in ("fi", "fa"):
    print(f"{name} logits, max error against the reference backend: {(results[name] - results['torch']).abs().max().item():.2e}")

What it prints

FlashInfer (KV pool used as page_size=1):
  qo_indptr      = [0, 2, 3]
  kv_indptr      = [0, 5, 11]
  kv_indices     = [16, 17, 18, 19, 20, 40, 41, 42, 43, 44, 45]
  last_page_len  = [1, 1]
  plan calls on the prefill wrapper: 2 (once per batch; one plan shared by that batch's 28 layers)
FlashAttention (page_size=4, page table holds page numbers):
  cu_seqlens_q   = [0, 2, 3]
  cache_seqlens  = [5, 6]
  page_table     = [[4, 5], [10, 11]]
fi logits, max error against the reference backend: 0.00e+00
fa logits, max error against the reference backend: 0.00e+00

Two requests: request 0 hits its first 3 tokens and computes 2 this step; request 1 has 5 in the cache and computes 1. FlashInfer receives token-level kv_indices (16-20 and 40-45) with last_page_len all 1; FlashAttention, with a page size of 4, receives page numbers (location 16 is page 4, location 40 is page 10). Both backends' logits are identical to the reference backend's.

The official implementation

• attention/fi.py · FlashInferBackend (lines 80-271)
• attention/fa.py · FlashAttentionBackend (lines 36-136)
• Upstream also has a TensorRT-LLM backend, attention/trtllm.py (the default on Blackwell, supporting only page sizes 16, 32 and 64), which we do not implement.

Our two backends follow upstream's logic exactly and only replace pin_memory=True and torch.cuda.Event with the device abstraction's functions.

Tests

tests/test_ch17_gpu_backends.py · test_backend_logits_match_torch_backend

@pytest.mark.parametrize("backend", ["fi", "fa"])
def test_backend_logits_match_torch_backend(fake_kernels, backend):
    """On one batch (a prefill with a prefix hit plus a decode), FlashInfer/FlashAttention agree with the reference backend."""
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(QWEN3)
    ids = [tok(p).input_ids for p in PROMPTS[:2]]
    logits = {}
    for name in ("torch", backend):
        engine = build_engine(QWEN3, attention_backend=name)
        identity_page_table(engine, 2)
        forward(engine, [new_req(ids[0][:3], 0), new_req(ids[1], 1)], "prefill")  # fill in the KV first
        a = new_req(ids[0], 0, cached_len=3)  # the first 3 tokens hit, and the rest are computed
        b = new_req(ids[1] + [100], 1, cached_len=len(ids[1]))  # one decode step
        logits[name] = forward(engine, [a, b], "prefill")
        engine.shutdown()
    assert (logits["torch"] - logits[backend]).abs().max().item() < 1e-4

tests/test_ch17_gpu_backends.py also checks that plan runs once per batch, and compares end to end against Hugging Face under overlap scheduling and chunked prefill in three configurations: FlashInfer (page size 1), FlashAttention (page size 4) and the fa,fi pair.

How to explain it

On GPU attention backends: FlashInfer splits into plan and run, where plan works out the task division and temporary buffers from this batch's length distribution (CPU work, once per batch, deferred to the first layer) and run is called per layer; mini-sglang treats the KV pool as page size 1, so a page-table row is kv_indices and last_page_len is always 1. FlashAttention takes a page-number table, the KV lengths and cu_seqlens_q directly, with page numbers obtained by taking every "page size"-th location from the global page table and dividing. Without a GPU, fakes with the same interfaces translate both libraries' argument semantics into a reference implementation and verify that the backends build their metadata correctly.

Exercises

1. To have FlashInfer use the real page size (16, say) instead of flattening to 1, what changes in prepare_metadata? How is last_page_len computed?
2. Deliberately break the fake FlashInfer's interpretation of qo_indptr (ignore the leading 0, say). Which test fails?
3. Why does CUDA Graph only ever deal with the decode backend, and why does FlashAttention's prefill not need it?

Answers

1. kv_indices becomes page numbers: page_table[r, :device_len:ps] // ps; kv_indptr accumulates page counts (ceil(device_len / ps)); last_page_len = device_len - (pages - 1) * ps; and forward stops viewing the pool as page size 1 and passes the [pages, ps, heads, dim] pool directly.
2. test_backend_logits_match_torch_backend[fi] and every end-to-end test using fi fail, which is exactly what the fakes are for.
3. A prefill batch's shape (token count, per-request lengths) varies endlessly and no graph can be captured for every shape; prefill also does a lot of compute, so kernel-launch overhead is a small share and graphs are unnecessary. A decode batch's shape depends only on the batch size and each step does very little compute, so launch overhead is a large share.

Summary

• FlashInfer: plan (once per batch, deferred to the first layer) plus run (per layer); mini-sglang treats the KV pool as page size 1, so a page-table row is kv_indices.
• FlashAttention: takes a page-number table, the KV lengths and cu_seqlens_q directly, with page numbers taken every page_size-th location from the global page table and divided.
• Fakes with the same interfaces translate both libraries' argument semantics into one reference implementation, verifying the backend code on a CPU against results identical to the reference backend.

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.

• 200. FlashInfer 的分页元数据 (medium)

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