手写 mini-sglang 阶段一 · 算得对：把骨架换成正式结构

算子层：BaseOP 与各种 Layer

mini-sglang 没有用 torch.nn.Module，而是自己写了一个 99 行的 BaseOP。这一章先弄清为什么，再实现组成 decoder 的所有层：线性层、词嵌入与输出层、RMSNorm、RoPE、激活函数和注意力层。层本身都很薄，真正的计算交给 minisgl.kernel 里的算子，它们有 PyTorch 参考实现，在 GPU 上换成 FlashInfer 或自定义 kernel。

自测：能答上来就可以跳过本章

1. 推理引擎需要 nn.Module 的哪些功能？不需要哪些？
2. 模型为什么要先在 meta 设备上构建？这时候 RoPE 的 cos/sin 表会遇到什么问题？
3. "融合的残差加 RMSNorm"返回的两个张量分别是什么？
4. prefill 时 LM head 为什么只算每个请求最后一个位置？

自测参考答案（先自己答，再展开对照）

1. 需要：按层次组织参数、按名字收集和加载权重；不需要：autograd、训练 / 评估模式、hook、parameters() 的注册机制、to() 的设备搬运等——推理只要前向，权重加载一次就不变。
2. 先在 meta 设备上建模型只分配元数据，不占内存、不做初始化，之后直接用真实的权重替换，秒建模型。但 RoPE 的 cos / sin 表不是权重，在 meta 设备上算出来也没有数据，所以要放在真实的设备上单独计算（所有层共用一张表）。
3. 一个是归一化后的输出（送进下一个子层），另一个是加上残差之后的新残差（留给下一次残差相加）。
4. 只有每个请求最后一个位置的 logits 用来采样下一个 token，其余位置的 logits 没用；LM head 是最大的矩阵乘之一（词表很大），只算最后位置省下大部分计算和显存。

本章要写的文件：layers/base.py、layers/linear.py、layers/embedding.py、layers/norm.py、layers/rotary.py、layers/activation.py、layers/attention.py，以及算子 kernel/torch_ops.py、kernel/__init__.py；把算子拼成一层的 models/utils.py、models/decoder.py，和描述形状的 models/config.py、models/base.py（下一章把层拼成整个模型并加载权重）；工具函数 utils/。

到这一章为止的文件：26 个（新建 20，修改 0，提前引入 4）

minisgl/
├── __init__.py
├── core.py
├── distributed/
│   ├── __init__.py  ← 提前引入（第 16 章讲）
│   ├── info.py  ← 提前引入（第 16 章讲）
│   ├── impl.py  ← 提前引入（第 16 章讲）
├── kernel/
│   ├── __init__.py  ← 本章新建
│   ├── torch_ops.py  ← 本章新建
├── layers/
│   ├── __init__.py  ← 本章新建
│   ├── base.py  ← 本章新建
│   ├── linear.py  ← 本章新建
│   ├── embedding.py  ← 本章新建
│   ├── norm.py  ← 本章新建
│   ├── rotary.py  ← 本章新建
│   ├── activation.py  ← 本章新建
│   ├── attention.py  ← 本章新建
├── models/
│   ├── config.py  ← 本章新建
│   ├── base.py  ← 本章新建
│   ├── utils.py  ← 本章新建
│   ├── decoder.py  ← 本章新建
│   ├── __init__.py  ← 提前引入（第 3 章讲）
├── utils/
│   ├── __init__.py  ← 本章新建
│   ├── misc.py  ← 本章新建
│   ├── hf.py  ← 本章新建
│   ├── torch_utils.py  ← 本章新建
│   ├── registry.py  ← 本章新建
│   ├── logger.py  ← 本章新建

这一步的 main：examples/ch02_layers.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

为什么不用 nn.Module

推理引擎用到 nn.Module 的地方只有一个：按名字收集和加载权重，名字要与 checkpoint 里的键名一致（model.layers.0.self_attn.o_proj.weight）。而 nn.Module 的其余机制——Parameter、自动求导、hook、train()/eval()、__call__ 的额外开销——都用不到。BaseOP 用 Python 对象的 __dict__ 实现了按名字收集和加载，其余一概没有：

minisgl/layers/base.py · BaseOP

class BaseOP:
    @abstractmethod
    def forward(self, *args: Any, **kwargs: Any) -> Any: ...

    def state_dict(self, *, prefix: str = "", result: StateDict | None = None) -> StateDict:
        result = result if result is not None else {}
        for name, param in self.__dict__.items():
            if name.startswith("_"):
                continue
            if isinstance(param, torch.Tensor):
                result[_concat_prefix(prefix, name)] = param
            elif isinstance(param, BaseOP):
                param.state_dict(prefix=_concat_prefix(prefix, name), result=result)
        return result

    def load_state_dict(
        self, state_dict: StateDict, *, prefix: str = "", _internal: bool = False
    ) -> None:
        for name, param in self.__dict__.items():
            if name.startswith("_"):
                continue
            if isinstance(param, torch.Tensor):
                item = state_dict.pop(_concat_prefix(prefix, name))
                assert param.shape == item.shape, (name, param.shape, item.shape)
                assert param.dtype == item.dtype, (name, param.dtype, item.dtype)
                setattr(self, name, item)  # 直接替换引用：模型是在 meta 设备上建的，没有真实内存
            elif isinstance(param, BaseOP):
                param.load_state_dict(state_dict, prefix=_concat_prefix(prefix, name), _internal=True)
        if not _internal and state_dict:
            raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}")

规则只有三条：

1. 不以下划线开头的 Tensor 属性是权重，名字是"属性路径"；
2. 不以下划线开头的 BaseOP 属性是子模块，递归处理；
3. 以下划线开头的属性（比如 RoPE 的 _cos_sin_cache、线性层的 _comm）不参与。

load_state_dict 不是把数据拷进已有的张量，而是用 checkpoint 里的张量替换属性（setattr）。这配合"先在 meta 设备上建模型"：torch.device("meta") 上创建的张量只有形状、没有内存，建一个 70B 的模型也是瞬间完成、不占内存；加载时用真实张量换掉它们。加载结束后如果 state_dict 里还有没被取走的键，就报错——多余的键通常意味着模型结构写错了。

OPList 让子模块的名字是 0、1、2……，对应 checkpoint 里的 layers.0、layers.1；StateLessOP 用于没有权重的层（RoPE、注意力层本身），它们的 state_dict 永远为空。

minisgl/layers/base.py · OPList

class OPList(BaseOP, Generic[T]):
    """算子列表，名字是 0、1、2……，对应 checkpoint 里的 layers.0、layers.1。"""

    def __init__(self, ops: List[T]):
        self.op_list = ops

    def state_dict(self, *, prefix: str = "", result: StateDict | None = None) -> StateDict:
        result = result if result is not None else {}
        for i, op in enumerate(self.op_list):
            op.state_dict(prefix=_concat_prefix(prefix, str(i)), result=result)
        return result

    def load_state_dict(
        self, state_dict: StateDict, *, prefix: str = "", _internal: bool = False
    ) -> None:
        for i, op in enumerate(self.op_list):
            op.load_state_dict(state_dict, prefix=_concat_prefix(prefix, str(i)), _internal=True)
        if not _internal and state_dict:
            raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}")

用 Qwen3-0.6B 的配置在 meta 设备上建一层 decoder，看看收集到的权重名：

examples/ch02_layers.py

"""第 2 章：BaseOP 如何收集权重名；参考算子与 Hugging Face 实现的数值差异。"""

import torch
from minisgl.distributed import set_tp_info
from minisgl.kernel import torch_ops
from minisgl.layers.rotary import get_rope, set_rope_device
from minisgl.models import ModelConfig
from minisgl.models.decoder import DecoderLayer
from minisgl.utils import cached_load_hf_config, torch_dtype

set_tp_info(0, 1)
set_rope_device(torch.device("cpu"))  # RoPE 表是真实数据，不能放在 meta 设备上
config = ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B"))
with torch.device("meta"), torch_dtype(torch.bfloat16):  # 只建结构，不分配内存
    layer = DecoderLayer(config, layer_id=0, has_qk_norm=True, use_moe=False)
for name, t in layer.state_dict(prefix="model.layers.0").items():
    print(f"{name:<50} {str(tuple(t.shape)):<14} {t.dtype} {t.device}")

# 与 HF 的 RMSNorm、RoPE 比较
from transformers.models.qwen3.modeling_qwen3 import Qwen3RMSNorm, apply_rotary_pos_emb

torch.manual_seed(0)
x, w = torch.randn(4, 1024), torch.randn(1024)
ref = Qwen3RMSNorm(1024, eps=1e-6)
ref.weight.data.copy_(w)
print("RMSNorm 最大误差:", (torch_ops.rmsnorm(x, w, 1e-6) - ref(x)).abs().max().item())

rope = get_rope(128, 128, 4096, 1e6)
pos = torch.tensor([0, 7, 4095])
q, k = torch.randn(3, 16 * 128), torch.randn(3, 8 * 128)
inv = 1.0 / 1e6 ** (torch.arange(0, 128, 2).float() / 128)
emb = torch.cat([pos.float()[:, None] * inv] * 2, dim=-1)
q_ref, k_ref = apply_rotary_pos_emb(q.view(3, 16, 128).transpose(0, 1)[None],
                                    k.view(3, 8, 128).transpose(0, 1)[None],
                                    emb.cos()[None], emb.sin()[None])
rope.forward(pos, q, k)
print("RoPE 最大误差:", (q.view(3, 16, 128).transpose(0, 1) - q_ref[0]).abs().max().item())

运行结果

model.layers.0.self_attn.qkv_proj.weight           (4096, 1024)   torch.bfloat16 meta
model.layers.0.self_attn.q_norm.weight             (128,)         torch.bfloat16 meta
model.layers.0.self_attn.k_norm.weight             (128,)         torch.bfloat16 meta
model.layers.0.self_attn.o_proj.weight             (1024, 2048)   torch.bfloat16 meta
model.layers.0.mlp.gate_up_proj.weight             (6144, 1024)   torch.bfloat16 meta
model.layers.0.mlp.down_proj.weight                (1024, 3072)   torch.bfloat16 meta
model.layers.0.input_layernorm.weight              (1024,)        torch.bfloat16 meta
model.layers.0.post_attention_layernorm.weight     (1024,)        torch.bfloat16 meta
RMSNorm 最大误差: 0.0
RoPE 最大误差: 0.0

名字与 Hugging Face 的 checkpoint 一致，只有两处不同：q_proj、k_proj、v_proj 合并成了 qkv_proj，gate_proj、up_proj 合并成了 gate_up_proj。合并后三个矩阵乘变成一个，输入只读一次。合并由下一章的权重加载器在读取时完成。

注意 q_norm 同时是 RopeAttn 和它内部 AttentionLayer 的属性，但只出现了一次：AttentionLayer 是 StateLessOP，不收集任何权重，它只是引用 RopeAttn 持有的那一份。

算子：参考实现与分派

层只负责组织参数，计算交给 minisgl.kernel。每个算子都有一个 PyTorch 参考实现，计算顺序刻意与 Hugging Face 保持一致，这样在 float32 下可以逐位对齐：

minisgl/kernel/torch_ops.py · rmsnorm

def rmsnorm(
    x: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor | None = None
) -> torch.Tensor:
    xf = x.float()
    y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
    y = weight * y.to(x.dtype)
    if out is not None:
        out.copy_(y)
        return out
    return y

minisgl/kernel/torch_ops.py · fused_add_rmsnorm

def fused_add_rmsnorm(
    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
) -> None:
    """原地计算：residual <- x + residual；x <- rmsnorm(residual)。"""
    residual.add_(x)
    rmsnorm(residual, weight, eps, out=x)

kernel/__init__.py 负责分派：CUDA 张量、dtype 是 fp16 / bf16、且装了 FlashInfer 时用 FlashInfer，否则用参考实现。dtype 这一条容易漏：FlashInfer 的 kernel 只为半精度编译，fp32 传进去会在 C++ 层抛 failed to dispatch data type，而本书验证正确性时用的恰恰是 fp32。

minisgl/kernel/__init__.py · rmsnorm

def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor | None = None):
    if _use_fi(x):
        return _flashinfer().rmsnorm(x, weight, eps, out=out)
    return torch_ops.rmsnorm(x, weight, eps, out=out)

与官方的差异：算子分派

官方的各层直接 from flashinfer import rmsnorm 等，只能在 GPU 上运行。我们让所有层都通过 minisgl.kernel 调用算子，CPU 上用参考实现。上面例子的最后两行显示，参考实现的 RMSNorm、RoPE 与 Hugging Face 的误差是 0。

线性层

minisgl/layers/linear.py · _LinearTPImpl

class _LinearTPImpl(BaseOP):
    def __init__(self, full_isize: int, full_osize: int, local_isize: int, local_osize: int,
                 has_bias: bool):
        self.full_input_size = full_isize
        self.full_output_size = full_osize
        self.local_input_size = local_isize
        self.local_output_size = local_osize
        self.weight = torch.empty(local_osize, local_isize)
        self.bias = torch.empty(local_osize) if has_bias else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(x, self.weight, self.bias)

_LinearTPImpl 记录了完整的和本 rank 的输入、输出维度，forward 就是一个 F.linear。四个子类对应张量并行下的四种切法：

类	用在哪	TP 下怎么切
LinearReplicated	MoE 的路由层	不切，每个 rank 一份完整的
LinearColParallelMerged	gate_up_proj	按输出维切，多个投影合并
LinearQKVMerged	qkv_proj	按注意力头切，KV 头不够分时复制
LinearRowParallel / LinearOProj	down_proj、o_proj	按输入维切，结果 all-reduce

TP=1 时它们都退化成普通的线性层。切分的细节和通信放到张量并行一章；现在只需要记住：构造函数里出现的 get_tp_info().size 在单卡时是 1。

词嵌入与输出层

minisgl/layers/embedding.py · VocabParallelEmbedding

class VocabParallelEmbedding(BaseOP):
    """词表并行的嵌入层：每个 rank 保存词表的一段，查不到的 token 填 0，最后 all-reduce 相加。"""

    def __init__(self, num_embeddings: int, embedding_dim: int):
        tp_info = get_tp_info()
        self.tp_size = tp_info.size
        self.num_embeddings = num_embeddings
        self.num_embeddings_tp = div_ceil(num_embeddings, self.tp_size)
        start = self.num_embeddings_tp * tp_info.rank
        finish = min(start + self.num_embeddings_tp, num_embeddings)
        self.vocab_range = (start, finish - start)
        self.weight = torch.empty(self.num_embeddings_tp, embedding_dim)
        self._comm = DistributedCommunicator()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = indexing(self.weight, x, self.vocab_range if self.tp_size > 1 else None)
        return self._comm.all_reduce(y) if self.tp_size > 1 else y

查表交给 indexing 算子。张量并行时每个 rank 只保存词表的一段 vocab_range = (起点, 长度)，查不到的 token 填 0，最后 all-reduce 相加——每个 token 只在一个 rank 上查到非零向量。

minisgl/layers/embedding.py · ParallelLMHead.forward

def forward(self, x: torch.Tensor) -> torch.Tensor:
    batch = get_global_ctx().batch
    bs = batch.size
    if batch.is_prefill:
        x = x[batch.attn_metadata.get_last_indices(bs)].contiguous()
    module = self.tied_embedding or self
    logits = F.linear(x, module.weight, self.bias)
    if self.tp_size == 1:
        return logits
    # all_gather 沿第 0 维拼接：[tp, bs, V/tp] -> 转置成 [bs, tp * V/tp]
    gathered = self._comm.all_gather(logits).view(self.tp_size, *logits.shape)
    gathered = gathered.permute(1, 0, 2).reshape(logits.shape[0], -1)
    return gathered[:, : self.num_embeddings]

输出层有一个重要的优化：prefill 时只取每个请求最后一个位置的隐状态去算 logits。一个 1000 个 token 的提示词，prefill 需要的只是第 1000 个位置的下一个 token 分布；如果对所有位置都算，LM head（[1000, 1024] × [1024, 151936]）的计算量和显存都会浪费上千倍。最后位置的下标由注意力元数据的 get_last_indices 给出（第 5 章）。decode 时每个请求本来就只有一个位置，不需要挑。

Qwen3-0.6B 共享词嵌入（tie_word_embeddings），ParallelLMHead 直接用 embed_tokens 的权重，自己的 state_dict 为空；加载时如果 checkpoint 里仍然带着 lm_head.weight（Qwen3-0.6B 就带着），就丢掉。

RMSNorm 与残差

minisgl/layers/norm.py · RMSNormFused

class RMSNormFused(BaseOP):
    """把"残差相加"和"归一化"融合成一个算子，同时返回归一化结果和新的残差。"""

    def __init__(self, size: int, eps: float) -> None:
        self.eps = eps
        self.weight = torch.empty(size)

    def forward(
        self, x: torch.Tensor, residual: torch.Tensor | None = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if residual is None:  # 第一层：还没有残差，x 本身就是残差
            return rmsnorm(x, self.weight, self.eps), x
        fused_add_rmsnorm(x, residual, self.weight, self.eps)
        return x, residual

decoder 层里残差相加总是紧跟着归一化：residual = x + residual; x = norm(residual)。RMSNormFused 把两步合成一个算子，同时返回归一化结果和新的残差；在 GPU 上它是 FlashInfer 的一个 kernel，少读写一遍隐状态。第一层还没有残差时，x 本身就是残差。于是整个 decoder 层写成：

minisgl/models/decoder.py · DecoderLayer.forward

def forward(self, x: torch.Tensor,
            residual: torch.Tensor | None) -> Tuple[torch.Tensor, torch.Tensor]:
    x, residual = self.input_layernorm.forward(x, residual)
    x = self.self_attn.forward(x)
    x, residual = self.post_attention_layernorm.forward(x, residual)
    x = self.mlp.forward(x)
    return x, residual

RoPE

minisgl/layers/rotary.py · RotaryEmbedding

class RotaryEmbedding(StateLessOP):
    def __init__(self, head_size: int, rotary_dim: int, max_position: int, base: float,
                 post_process: Callable[[torch.Tensor], torch.Tensor] | None = None) -> None:
        assert rotary_dim == head_size
        self.head_size = head_size
        inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
        if post_process is not None:
            inv_freq = post_process(inv_freq)
        t = torch.arange(max_position, dtype=torch.float)
        freqs = torch.outer(t, inv_freq)  # [max_pos, D/2]
        # 以下划线开头：不是权重，不参与 state_dict；启动时一次性算好，前向时按位置查表
        self._cos_sin_cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1)

    def forward(self, positions: torch.Tensor, query: torch.Tensor,
                key: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        apply_rope_inplace(positions, query, key, self.head_size, self._cos_sin_cache)
        return query, key

cos_sin_cache 在构造时对所有位置一次性算好，前向时按 positions 查表。算子按 NeoX 风格（前后两半配对）旋转，与 Llama、Qwen 系列的 Hugging Face 实现相同：

minisgl/kernel/torch_ops.py · apply_rope_inplace

def apply_rope_inplace(
    positions: torch.Tensor,
    query: torch.Tensor,
    key: torch.Tensor,
    head_size: int,
    cos_sin_cache: torch.Tensor,
) -> None:
    """NeoX 风格（前后两半配对）的旋转位置编码，原地修改 query 和 key。

    query: [T, Hq * D] 或 [T, Hq, D]；key 同理；cos_sin_cache: [max_pos, D]，前一半是 cos、后一半是 sin。
    """
    cos, sin = cos_sin_cache[positions.long()].chunk(2, dim=-1)  # [T, D/2]
    cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
    for t in (query, key):
        x = t.view(t.shape[0], -1, head_size)
        x1, x2 = x[..., : head_size // 2], x[..., head_size // 2 :]
        c, s = cos.to(x.dtype), sin.to(x.dtype)
        o1 = x1 * c - x2 * s
        o2 = x2 * c + x1 * s
        x.copy_(torch.cat([o1, o2], dim=-1))

get_rope 用 functools.cache 缓存：所有层的 RoPE 参数相同，共用一张表。还有一个细节——模型是在 torch.device("meta") 上构建的，而 cos/sin 表是真实数据，不能放在 meta 上。所以 get_rope 检测到当前默认设备是 meta 时，改用 set_rope_device 事先指定的真实设备：

minisgl/layers/rotary.py · get_rope

@functools.cache  # 所有层共享同一个 RoPE 表
def get_rope(head_dim: int, rotary_dim: int, max_position: int, base: float,
             rope_scaling: Tuple[Tuple[str, Any], ...] | None = None) -> RotaryEmbedding:
    scaling = dict(rope_scaling) if rope_scaling is not None else None
    if torch.empty(0).device.type == "meta":
        # 模型在 meta 设备上构建，但 RoPE 表是真实数据，必须放到真实设备上
        if _ROPE_DEVICE is None:
            raise RuntimeError("Call set_rope_device() before building a model on meta device")
        with torch.device(_ROPE_DEVICE):
            return _get_rope(head_dim, rotary_dim, max_position, base, scaling)
    return _get_rope(head_dim, rotary_dim, max_position, base, scaling)

忘了调用 set_rope_device 就在 meta 上建模型，会得到 RuntimeError: Call set_rope_device() before building a model on meta device。

Llama 3.1 的长上下文扩展（rope_type: llama3）对频率做后处理：高频分量不变，低频分量除以 factor，中间平滑过渡。它的正确性由第 20 章里一个用随机权重构造的小 Llama 3 模型与 Hugging Face 对比验证。

与官方的差异：YaRN

官方的 _get_rope 还支持 yarn，但只修改了频率，没有乘 Hugging Face 实现里的注意力缩放系数（0.1 * ln(factor) + 1）。我们只实现了 default 和 llama3 两种，YaRN 留作练习。

注意力层

minisgl/layers/attention.py · AttentionLayer.forward

def forward(self, qkv: torch.Tensor) -> torch.Tensor:
    ctx = get_global_ctx()
    q, k, v = qkv.split([self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim], dim=-1)
    q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
    if self.q_norm is not None:
        self.q_norm.forward_inplace(q.view(-1, self.num_qo_heads, self.head_dim))
    if self.k_norm is not None:
        self.k_norm.forward_inplace(k.view(-1, self.num_kv_heads, self.head_dim))
    q, k = self.rotary.forward(ctx.batch.positions, q, k)
    q = q.view(-1, self.num_qo_heads, self.head_dim)
    o = ctx.attn_backend.forward(q, k, v, self.layer_id, ctx.batch)
    return o.reshape(-1, self.qo_attn_dim)

注意力层也不持有权重：它把合并的 qkv 拆成三份，Qwen3 还要对每个头的 q、k 做一次 RMSNorm（QK-Norm），然后加上 RoPE，最后把 q、k、v 交给全局上下文里的注意力后端。写 KV 缓存和计算注意力都是后端的事，第 5 章实现。

官方实现

• BaseOP：layers/base.py · BaseOP（第 15～53 行）
• 输出层只取最后位置：layers/embedding.py · ParallelLMHead.forward（第 87～110 行）
• RoPE 的 meta 设备处理：layers/rotary.py · get_rope（第 125～143 行）
• 注意力层：layers/attention.py · AttentionLayer.forward（第 47～57 行）

官方在注意力层里直接对 qkv.split(...) 得到的视图原地做 QK-Norm 和 RoPE（FlashInfer 的算子支持带步长的输入）。我们的参考算子先 .contiguous() 拷一份，写法更直白，代价是多一次拷贝。

测试

tests/test_ch02_layers.py 检查 BaseOP 的命名规则、多余键报错、RMSNorm 与 RoPE 与 Hugging Face 一致，以及词表并行查表：

tests/test_ch02_layers.py · test_vocab_parallel_indexing_masks_other_shards

def test_vocab_parallel_indexing_masks_other_shards():
    w = torch.arange(20.0).view(10, 2)  # 词表 [5, 15) 这一段
    ids = torch.tensor([3, 5, 14, 15, 9])
    out = torch_ops.indexing(w, ids, vocab_range=(5, 10))
    assert out.tolist() == [[0, 0], [0, 1], [18, 19], [0, 0], [8, 9]]

怎么讲清楚

讲算子层：推理引擎只需要 nn.Module 的一小部分——按名字收集和加载权重、组织层级；不需要自动求导、钩子、训练状态，所以 mini-sglang 用轻量的 BaseOP，层只组织参数、计算交给 kernel（CPU 上是与 HF 数值一致的参考实现，GPU 上是 FlashInfer）。模型先在 meta 设备上建空壳（不分配内存，秒建），加载时直接替换成真实权重；但 RoPE 的 cos/sin 表是计算出来的，必须放在真实设备上。融合的残差加 RMSNorm 返回归一化后的输出和新的残差；prefill 时 LM head 只算每个请求最后一个位置，省掉大量无用的词表投影。

练习

1. 给 BaseOP 加一个 num_params() 方法，返回权重元素总数。用它算出 Qwen3-0.6B 的参数量（注意共享词嵌入不要重复计算）。
2. 把 ParallelLMHead.forward 里"prefill 只取最后位置"的逻辑去掉，prefill 一个 1000 token 的请求，LM head 的输出张量有多大（float32）？
3. 实现 YaRN：在 _get_rope 里加一个分支，并让 cos/sin 乘上注意力缩放系数。用 transformers 的 Qwen3Config(rope_scaling={"rope_type": "yarn", ...}) 构造一个小模型验证。

参考答案

1. sum(t.numel() for t in self.state_dict().values())。因为共享词嵌入时 ParallelLMHead.state_dict() 为空，结果自然不会重复：Qwen3-0.6B 约 5.96 亿。
2. [1000, 151936] 个 float32，约 580 MiB；只取最后位置时是 0.6 MB。
3. 参考 Hugging Face _compute_yarn_parameters：频率按 ramp 在 inv_freq / factor 与 inv_freq 之间插值，attention_factor = 0.1 * ln(factor) + 1，构造 cos/sin 表时乘上它。

小结

• BaseOP 用 __dict__ 按属性路径收集和加载权重；下划线开头的属性不是权重；加载是替换而不是拷贝，配合 meta 设备秒建模型。
• 层只组织参数，计算交给 minisgl.kernel：CPU 上是与 HF 数值一致的参考实现，GPU 上是 FlashInfer。
• 输出层在 prefill 时只算每个请求最后一个位置的 logits。
• RoPE 表必须放在真实设备上；所有层共用一张表。
• 注意力层只做拆分、QK-Norm 和 RoPE，写缓存和算注意力交给注意力后端。

本章练习题（1 道）

在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

• 186. BaseOP：不用 nn.Module 收集和加载权重（中等）

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