手写 mini-sglang 阶段四 · 更快、更大

阶段四的里程碑：同一份代码，更多的卡、更快的 kernel

前三个阶段做出来的引擎是完整的，但只会用一张卡、只会用 PyTorch 的算子。阶段四的六章让同一份代码在多卡上切开跑（张量并行）、在 GPU 上换成 FlashInfer / FlashAttention、把 decode 录成 CUDA Graph 回放、换上自定义 kernel、接入 MoE 模型，最后用一张基准表把每项优化的收益量出来。这一页在 CPU 上先把"对"验证掉：切成两个 rank，或者录成 graph 回放，输出一个 token 都不变。"快"的数字要上真卡。

这一阶段结束时你手里有什么

• 同一个 Engine，tp_size=2 时每个 rank 只拿一半权重、一半 KV 头，前向中间用 all-reduce 合并，输出和单卡一致；
• GPU 上注意力自动切到 FlashInfer / FlashAttention，decode 自动走 CUDA Graph，KV 写入和索引用自定义 kernel；
• 一个 MoE 模型（Qwen3-MoE）也能跑，fused MoE 的 Triton kernel 在 GPU 上替换掉逐专家的循环；
• 第 21 章一张基准表：每项优化单独开关，吞吐和延迟各多少。

先跑起来

cd minisgl && python examples/stage4_milestone.py

examples/stage4_milestone.py

"""阶段四的里程碑：同一份代码，切成两个 rank 跑，或者录成 graph 回放，输出一个 token 都不变。

CPU 上这一阶段只能验证"对"：张量并行走 gloo、CUDA Graph 走仿真，都快不起来。
"快"的数字要上真卡：根目录的 gpu_check.py minisgl，以及第 21 章的基准表。
"""

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:
    """一个 rank：加载自己那一份权重，手工驱动 prefill + 3 步 decode（和阶段一的 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):
    """CPU 上的 CUDA Graph 仿真：录制固定批大小的 graph，decode 走 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 当基准，再 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 池每层 {base[3]}，生成 {base[1]}")
    for rank, tokens, qkv_shape, k_shape in results[2]:
        print(f"TP=2 rank {rank}：qkv_proj {qkv_shape}（一半），KV 池每层 {k_shape}（KV 头也减半），"
              f"生成 {tokens}，与 TP=1 一致：{tokens == base[1]}")

    eager, _ = run_graph(0)
    graph, replays = run_graph(4)
    print(f"CUDA Graph 仿真：3 个请求的 decode 每轮补齐后 replay（{replays[0]} x {len(replays)} 轮），与不用 graph 的输出一致：{graph == eager}")
    print("CPU 上到此为止——快不快要上真卡：.venv-gpu/bin/python gpu_check.py minisgl，以及第 21 章的基准表")

运行结果

TP=1：qkv_proj (4096, 1024)，KV 池每层 (257, 1, 8, 128)，生成 [12095, 13, 576, 6722]
TP=2 rank 0：qkv_proj (2048, 1024)（一半），KV 池每层 (257, 1, 4, 128)（KV 头也减半），生成 [12095, 13, 576, 6722]，与 TP=1 一致：True
TP=2 rank 1：qkv_proj (2048, 1024)（一半），KV 池每层 (257, 1, 4, 128)（KV 头也减半），生成 [12095, 13, 576, 6722]，与 TP=1 一致：True
CUDA Graph 仿真：3 个请求的 decode 每轮补齐后 replay（3->4 x 5 轮），与不用 graph 的输出一致：True
CPU 上到此为止——快不快要上真卡：.venv-gpu/bin/python gpu_check.py minisgl，以及第 21 章的基准表

读这几行

TP=2 时 qkv_proj 从 (4096, 1024) 变成 (2048, 1024)，KV 池每层的头数从 8 变成 4。 每个 rank 只加载、只保存自己那一半——这就是张量并行解决的问题：模型放不进一张卡时，把每个线性层按行或按列切开，两张卡各算一半，中间用一次 all-reduce 把结果合起来。4 个 token 和 TP=1 完全一样：切法正确的话，数学上就是同一个矩阵乘。CPU 上两个进程走 gloo，不会更快；真卡上两张卡走 NCCL，权重读取的带宽翻倍，decode 就快了。

CUDA Graph 仿真：3 个请求补齐成 4，replay 了 5 轮，输出一致。 CUDA Graph 要求每次回放的张量形状和地址都固定，所以 batch 要补齐到录制过的大小、所有输入都要先拷进固定的缓冲区。CPU 上的仿真只做一件事：强制这套"固定缓冲区"的纪律——漏拷任何一个输入，输出立刻出错（第 18 章故意演示了这一点）。真卡上这套纪律换来的是 decode 每步少掉几十次 kernel 启动的开销。

CPU 上到此为止。 这一阶段的其他三章——FlashInfer / FlashAttention、自定义 CUDA kernel、fused MoE——在 CPU 上都只有"同接口的替身"：假实现、CPU 模拟器、Triton 解释器。它们能验证逻辑和接口，不能给出速度。有卡的话，仓库根目录的 gpu_check.py minisgl 会把这几章在真卡上跑一遍；第 21 章的基准表是这一阶段真正的里程碑。

每一章加了什么

优化	解决什么	CPU 上怎么验证	在哪一章
张量并行	模型放不进一张卡；多卡一起读权重	gloo 上两个进程，输出与单卡一致	张量并行
FlashInfer / FlashAttention	不物化注意力分数矩阵；分页 KV 直接读	同接口的 PyTorch 假实现	GPU 注意力后端
CUDA Graph	decode 每步几十次 kernel 启动的 CPU 开销	EmulatedGraph，强制固定缓冲区	CUDA Graph
自定义 kernel	KV 写入、词表索引这类小操作的启动开销和访存	CUDA 手册的 CPU 模拟器	自定义 CUDA kernel
fused MoE	逐专家循环变成一个分组的矩阵乘	Triton 解释器模式	MoE 与 fused MoE
基准测试	每项优化到底值多少	只能上真卡	基准测试与消融

怎么读这六章

张量并行和 CUDA Graph 在 CPU 上就能完整验证，先读；GPU 注意力后端和自定义 kernel 读的时候对照 CUDA 手册里对应的 kernel 章节；MoE 可以放到最后。读完六章如果手边有一张卡，按大作业的门槛把引擎跑到正式版的 60%——那是这本书真正的终点。

验收：读完这一阶段

• 说出四种线性层切法各适用于哪一层，为什么 o_proj 和 down_proj 后面要 all-reduce；
• 解释 CUDA Graph 为什么要补齐 batch、为什么漏拷一个输入不会报错只会算错；
• 讲清 FlashInfer 的 plan / run 两步各做什么；
• 在真卡上跑过 gpu_check.py minisgl，并能解释第 21 章基准表里每一行的差异来自哪里。

小结

• 张量并行：每个 rank 一半权重、一半 KV 头，all-reduce 合并，输出与单卡一致。
• CUDA Graph：固定形状、固定缓冲区，CPU 仿真只验证这套纪律。
• 注意力后端、自定义 kernel、fused MoE 在 CPU 上只能验证接口与逻辑；速度要上真卡。
• 第 21 章的基准表是这一阶段的真正里程碑。

发现错误或有更好的写法？在 GitHub 上编辑此页 · 提交问题