Skip to content
mini-sglang from Scratch Stage 4 · Faster and bigger

Stage 4 milestone: the same code on more GPUs, with faster kernels

The engine built in the first three stages is complete, but it only uses one GPU and only PyTorch's operators. The six chapters of stage 4 let the same code split across several GPUs (tensor parallelism), switch to FlashInfer / FlashAttention on a GPU, record decode as a CUDA Graph and replay it, use custom kernels, run an MoE model, and finally measure what every optimisation is worth in one benchmark table. This page verifies "correct" on a CPU first: split into two ranks, or recorded and replayed as a graph, not one output token changes. The "fast" numbers need a real card.

What you hold at the end of this stage

  • The same Engine: with tp_size=2 each rank holds half the weights and half the KV heads, the forward pass merges with all-reduce, and the output matches a single GPU;
  • On a GPU, attention switches to FlashInfer / FlashAttention by itself, decode goes through CUDA Graph by itself, and KV writes and indexing use custom kernels;
  • An MoE model (Qwen3-MoE) runs too, with a fused-MoE Triton kernel replacing the per-expert loop on a GPU;
  • Chapter 21's benchmark table: each optimisation toggled on its own, with its throughput and latency.

Run it first

cd minisgl && python examples/stage4_milestone.py

examples/stage4_milestone.py
"""Stage 4 milestone: the same code split across two ranks, or recorded as a graph and replayed, with not one output token changed.

On a CPU this stage can only verify "correct": tensor parallelism goes through gloo and CUDA Graph through
emulation, neither any faster. The "fast" numbers need a real card: gpu_check.py minisgl in the repository
root, and chapter 21's benchmark table.
"""

import multiprocessing as mp

import torch

PATH = "models/Qwen3-0.6B"
PROMPT_IDS = [785, 6722, 315, 9625, 374]       # "The capital of France is"
STEPS = 4


def run_tp(rank: int, size: int, out: mp.Queue) -> None:
    """One rank: load its own share of the weights, drive prefill + 3 decode steps by hand (like stage 1's step())."""
    from minisgl.core import Batch, Req, SamplingParams
    from minisgl.distributed import DistributedInfo
    from minisgl.engine import Engine, EngineConfig

    engine = Engine(EngineConfig(model_path=PATH, tp_info=DistributedInfo(rank, size), dtype=torch.float32,
                                 max_running_req=2, num_page_override=256, max_seq_len_override=128,
                                 distributed_port=29610))
    engine.page_table[0, :128] = torch.arange(128)
    req = Req(input_ids=torch.tensor(PROMPT_IDS, dtype=torch.int32), table_idx=0, cached_len=0,
              output_len=STEPS, uid=0, sampling_params=SamplingParams(max_tokens=STEPS), cache_handle=None)
    tokens = []
    for phase in ["prefill"] + ["decode"] * (STEPS - 1):
        batch = Batch(reqs=[req], phase=phase)
        engine.graph_runner.pad_batch(batch)
        batch.positions = torch.arange(req.cached_len, req.device_len).int()
        batch.input_ids = req.input_ids[req.cached_len:req.device_len]
        batch.out_loc = engine.page_table[0, req.cached_len:req.device_len]
        engine.attn_backend.prepare_metadata(batch)
        o = engine.forward_batch(batch, engine.sampler.prepare(batch))
        req.append_host(o.next_tokens_cpu[:1])
        tokens.append(int(o.next_tokens_cpu[0]))
    qkv = engine.model.model.layers.op_list[0].self_attn.qkv_proj.weight
    out.put((rank, tokens, tuple(qkv.shape), tuple(engine.kv_cache.k_cache(0).shape)))
    engine.shutdown()


def run_graph(cuda_graph_max_bs: int):
    """CUDA Graph emulation on a CPU: record graphs for fixed batch sizes, and decode goes through replay."""
    from minisgl.core import SamplingParams
    from minisgl.env import ENV
    from minisgl.llm import LLM

    ENV.DISABLE_OVERLAP_SCHEDULING.value = True
    llm = LLM(PATH, dtype=torch.float32, cuda_graph_max_bs=cuda_graph_max_bs, max_running_req=8,
              num_page_override=512, max_seq_len_override=128)
    replays = []
    orig = llm.engine.graph_runner.replay
    llm.engine.graph_runner.replay = lambda b: (replays.append(f"{b.size}->{b.padded_size}"), orig(b))[1]
    out = llm.generate(["The capital of France is", "List three prime numbers:", "def fibonacci(n):"],
                       SamplingParams(max_tokens=6, ignore_eos=True))
    llm.shutdown()
    return [o["token_ids"] for o in out], replays


if __name__ == "__main__":
    mp.set_start_method("spawn")
    results = {}
    for size in (1, 2):                                       # TP=1 first as the baseline, then TP=2
        q: mp.Queue = mp.Queue()
        procs = [mp.Process(target=run_tp, args=(r, size, q)) for r in range(size)]
        for p in procs:
            p.start()
        results[size] = sorted((q.get() for _ in range(size)), key=lambda x: x[0])
        for p in procs:
            p.join()
    base = results[1][0]
    print(f"TP=1: qkv_proj {base[2]}, KV pool per layer {base[3]}, generated {base[1]}")
    for rank, tokens, qkv_shape, k_shape in results[2]:
        print(f"TP=2 rank {rank}: qkv_proj {qkv_shape} (half), KV pool per layer {k_shape} (KV heads halved too), "
              f"generated {tokens}, same as TP=1: {tokens == base[1]}")

    eager, _ = run_graph(0)
    graph, replays = run_graph(4)
    print(f"CUDA Graph emulation: 3 requests, decode padded and replayed every round ({replays[0]} x {len(replays)} rounds), same output as without the graph: {graph == eager}")
    print("This is as far as a CPU goes; for speed, use a real card: .venv-gpu/bin/python gpu_check.py minisgl, and chapter 21's benchmark table")
What it prints
TP=1: qkv_proj (4096, 1024), KV pool per layer (257, 1, 8, 128), generated [12095, 13, 576, 6722]
TP=2 rank 0: qkv_proj (2048, 1024) (half), KV pool per layer (257, 1, 4, 128) (KV heads halved too), generated [12095, 13, 576, 6722], same as TP=1: True
TP=2 rank 1: qkv_proj (2048, 1024) (half), KV pool per layer (257, 1, 4, 128) (KV heads halved too), generated [12095, 13, 576, 6722], same as TP=1: True
CUDA Graph emulation: 3 requests, decode padded and replayed every round (3->4 x 5 rounds), same output as without the graph: True
This is as far as a CPU goes; for speed, use a real card: .venv-gpu/bin/python gpu_check.py minisgl, and chapter 21's benchmark table

Reading these lines

With TP=2, qkv_proj goes from (4096, 1024) to (2048, 1024), and the KV pool's heads per layer from 8 to 4. Each rank loads and stores only its half — which is exactly the problem tensor parallelism solves: when the model does not fit on one card, split every linear layer by rows or columns, let two cards compute half each, and merge with one all-reduce in between. The 4 tokens are identical to TP=1: with the right split, it is mathematically the same matrix multiply. On a CPU the two processes go through gloo and are no faster; on real cards they go through NCCL, the weight-reading bandwidth doubles, and decode speeds up.

CUDA Graph emulation: 3 requests padded to 4, replayed 5 times, same output. CUDA Graph requires every replay to use fixed tensor shapes and addresses, so the batch is padded to a recorded size and every input is first copied into fixed buffers. The CPU emulation does one thing: enforce that "fixed buffers" discipline — skip copying any one input and the output is wrong at once (chapter 18 demonstrates this on purpose). On a real card that discipline buys the removal of dozens of kernel-launch overheads from every decode step.

On a CPU, this is as far as it goes. The other three chapters of this stage — FlashInfer / FlashAttention, custom CUDA kernels, fused MoE — only have same-interface stand-ins on a CPU: fakes, the CPU simulator, the Triton interpreter. They verify logic and interfaces, not speed. With a card, gpu_check.py minisgl in the repository root runs these chapters on real hardware; chapter 21's benchmark table is the real milestone of this stage.

What each chapter adds

Optimisation What it solves How it is verified on a CPU Where
tensor parallelism the model does not fit on one card; several cards read the weights together two processes over gloo, output identical to one card Tensor parallelism
FlashInfer / FlashAttention never materialise the attention score matrix; read paged KV directly same-interface PyTorch fakes GPU attention backends
CUDA Graph the CPU overhead of dozens of kernel launches per decode step EmulatedGraph, enforcing fixed buffers CUDA Graph
custom kernels launch overhead and memory traffic of small operations such as KV writes and vocabulary indexing the CUDA handbook's CPU simulator Custom CUDA kernels
fused MoE the per-expert loop becomes one grouped matrix multiply Triton interpreter mode MoE and fused MoE
benchmarks what each optimisation is actually worth real hardware only Benchmarks and ablations

How to read these six chapters

Tensor parallelism and CUDA Graph can be verified completely on a CPU; read them first. Read GPU attention backends and custom kernels alongside the matching kernel chapters of the CUDA handbook; MoE can wait until last. After the six chapters, if you have a card, take the engine to 60% of the official build by the bar in the project — that is where this book really ends.

Checkpoint: after this stage

  • Say which of the four linear-layer splits applies to which layer, and why o_proj and down_proj are followed by an all-reduce;
  • Explain why CUDA Graph pads the batch, and why skipping one input copy produces a wrong answer rather than an error;
  • Explain what FlashInfer's plan and run steps each do;
  • Run gpu_check.py minisgl on a real card, and explain where every difference in chapter 21's benchmark table comes from.

Summary

  • Tensor parallelism: half the weights and half the KV heads per rank, merged with all-reduce, output identical to one card.
  • CUDA Graph: fixed shapes and fixed buffers; the CPU emulation verifies only that discipline.
  • Attention backends, custom kernels and fused MoE can only be verified for interface and logic on a CPU; speed needs a real card.
  • Chapter 21's benchmark table is the real milestone of this stage.

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