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

The op layer: BaseOP and the layers

mini-sglang does not use torch.nn.Module. It writes its own 99-line BaseOP instead. This chapter works out why, then implements every layer a decoder is made of: linear layers, the embedding and output layers, RMSNorm, RoPE, the activation and the attention layer. The layers themselves are thin, and the real computing goes to the ops in minisgl.kernel, which have PyTorch reference implementations and switch to FlashInfer or custom kernels on a GPU.

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

1. Which parts of nn.Module does an inference engine need? Which does it not?
2. Why is the model built on the meta device first? What problem does that create for RoPE's cos/sin table?
3. What are the two tensors returned by the "fused residual add plus RMSNorm"?
4. Why does the LM head only compute each request's last position during prefill?

Answers (try it yourself first, then expand)

1. Needed: organizing parameters hierarchically, and collecting and loading weights by name. Not needed: autograd, train / eval mode, hooks, the registration machinery behind parameters(), device moves through to(). Inference only runs forward, and the weights are loaded once and never change.
2. Building on the meta device only allocates metadata, so it takes no memory and runs no initialization, and the real weights replace it afterwards, which makes building a model instant. But RoPE's cos / sin table is not a weight, and computing it on the meta device yields no data, so it has to be computed separately on a real device (one table shared by every layer).
3. One is the normalized output that goes into the next sub-layer, the other is the new residual after the addition, kept for the next residual add.
4. Only the logits at each request's last position are used to sample the next token; the logits elsewhere are useless. The LM head is one of the largest matmuls, because the vocabulary is big, so computing only the last position saves most of the compute and memory.

Files you will write: layers/base.py, layers/linear.py, layers/embedding.py, layers/norm.py, layers/rotary.py, layers/activation.py, layers/attention.py, plus the ops in kernel/torch_ops.py and kernel/__init__.py.

Files so far: 26 (20 new, 0 modified, 4 brought in early)

minisgl/
├── __init__.py
├── core.py
├── distributed/
│   ├── __init__.py  ← brought in early (explained in chapter 16)
│   ├── info.py  ← brought in early (explained in chapter 16)
│   ├── impl.py  ← brought in early (explained in chapter 16)
├── kernel/
│   ├── __init__.py  ← new in this chapter
│   ├── torch_ops.py  ← new in this chapter
├── layers/
│   ├── __init__.py  ← new in this chapter
│   ├── base.py  ← new in this chapter
│   ├── linear.py  ← new in this chapter
│   ├── embedding.py  ← new in this chapter
│   ├── norm.py  ← new in this chapter
│   ├── rotary.py  ← new in this chapter
│   ├── activation.py  ← new in this chapter
│   ├── attention.py  ← new in this chapter
├── models/
│   ├── config.py  ← new in this chapter
│   ├── base.py  ← new in this chapter
│   ├── utils.py  ← new in this chapter
│   ├── decoder.py  ← new in this chapter
│   ├── __init__.py  ← brought in early (explained in chapter 3)
├── utils/
│   ├── __init__.py  ← new in this chapter
│   ├── misc.py  ← new in this chapter
│   ├── hf.py  ← new in this chapter
│   ├── torch_utils.py  ← new in this chapter
│   ├── registry.py  ← new in this chapter
│   ├── logger.py  ← new in this chapter

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

Why not nn.Module

An inference engine uses nn.Module for exactly one thing: collecting and loading weights by name, where the names must match the keys in the checkpoint (model.layers.0.self_attn.o_proj.weight). Everything else in nn.Module, from Parameter and autograd to hooks, train()/eval() and the extra overhead of __call__, goes unused. BaseOP implements collecting and loading by name through a Python object's __dict__ and has nothing else:

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)  # replace the reference outright: the model was built on the meta device and has no real memory
            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())}")

There are only three rules:

1. a Tensor attribute that does not start with an underscore is a weight, and its name is the attribute path;
2. a BaseOP attribute that does not start with an underscore is a sub-module, handled recursively;
3. attributes starting with an underscore (RoPE's _cos_sin_cache, a linear layer's _comm) take no part.

load_state_dict does not copy data into existing tensors; it replaces the attribute with the tensor from the checkpoint (setattr). That fits "build the model on the meta device first": tensors created on torch.device("meta") have a shape but no memory, so building a 70B model is instant and costs nothing, and loading swaps in the real tensors. If any key in state_dict is left unclaimed at the end, it raises, because a leftover key usually means the model structure is wrong.

OPList names its sub-modules 0, 1, 2 and so on, matching layers.0 and layers.1 in the checkpoint; StateLessOP is for layers with no weights (RoPE, the attention layer itself), whose state_dict is always empty.

minisgl/layers/base.py · OPList

class OPList(BaseOP, Generic[T]):
    """A list of ops named 0, 1, 2 and so on, matching layers.0 and layers.1 in the checkpoint."""

    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())}")

Building one decoder layer on the meta device with the Qwen3-0.6B config, and looking at the weight names it collects:

examples/ch02_layers.py

"""Chapter 2: how BaseOP collects weight names, and how the reference ops differ numerically from 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"))  # the RoPE table is real data and cannot live on the meta device
config = ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B"))
with torch.device("meta"), torch_dtype(torch.bfloat16):  # build the structure only, allocating no memory
    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}")

# against HF's RMSNorm and 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 max error:", (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 max error:", (q.view(3, 16, 128).transpose(0, 1) - q_ref[0]).abs().max().item())

What it prints

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 max error: 0.0
RoPE max error: 0.0

The names match the Hugging Face checkpoint except in two places: q_proj, k_proj and v_proj are merged into qkv_proj, and gate_proj and up_proj into gate_up_proj. Merging turns three matmuls into one and reads the input only once. The weight loader of the next chapter does the merging as it reads.

Note that q_norm is an attribute of both RopeAttn and the AttentionLayer inside it, yet appears only once: AttentionLayer is a StateLessOP and collects no weights, so it merely references the one RopeAttn holds.

Ops: reference implementations and dispatch

Layers only organize parameters; the computing goes to minisgl.kernel. Every op has a PyTorch reference implementation whose order of operations deliberately matches Hugging Face, so that in float32 they agree bit for bit:

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:
    """In place: residual <- x + residual; x <- rmsnorm(residual)."""
    residual.add_(x)
    rmsnorm(residual, weight, eps, out=x)

kernel/__init__.py does the dispatch: FlashInfer for a CUDA tensor whose dtype is fp16 or bf16 when FlashInfer is installed, the reference implementation otherwise. The dtype condition is easy to miss: FlashInfer's kernels are compiled for half precision only, and fp32 raises failed to dispatch data type down in C++ — and fp32 is exactly what this book uses to verify correctness.

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)

Difference from upstream: op dispatch

The official layers do from flashinfer import rmsnorm and friends directly, so they only run on a GPU. We have every layer call its ops through minisgl.kernel, with the reference implementation on a CPU. The last two lines of the example above show that the reference RMSNorm and RoPE differ from Hugging Face by exactly 0.

Linear layers

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 records both the full and this rank's input and output dimensions, and forward is one F.linear. Four subclasses cover the four ways to shard under tensor parallelism:

Class	Used for	How TP shards it
LinearReplicated	the MoE router	not sharded, a full copy per rank
LinearColParallelMerged	gate_up_proj	split along the output dimension, several projections merged
LinearQKVMerged	qkv_proj	split by attention head, with KV heads replicated when there are not enough
LinearRowParallel / LinearOProj	down_proj, o_proj	split along the input dimension, result all-reduced

At TP=1 they all degenerate into ordinary linear layers. The sharding details and the communication are in the tensor parallelism chapter; for now just remember that the get_tp_info().size appearing in the constructors is 1 on a single GPU.

Embedding and output layers

minisgl/layers/embedding.py · VocabParallelEmbedding

class VocabParallelEmbedding(BaseOP):
    """The vocabulary-parallel embedding: each rank keeps a slice of the vocabulary, fills zeros for tokens it does not hold, and all-reduces at the end."""

    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

The lookup goes to the indexing op. Under tensor parallelism each rank keeps only a slice of the vocabulary, vocab_range = (start, length), fills zeros for tokens it does not hold, and all-reduces at the end, since every token finds its non-zero vector on exactly one 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 concatenates along dimension 0: [tp, bs, V/tp] -> transposed into [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]

The output layer has one important optimization: during prefill it takes only the hidden state at each request's last position to compute logits. For a 1000-token prompt, prefill only needs the next-token distribution at position 1000; computing every position would waste a thousandfold on the LM head's compute and memory ([1000, 1024] × [1024, 151936]). The index of the last position comes from get_last_indices in the attention metadata (chapter 5). During decode each request has only one position anyway, so there is nothing to pick.

Qwen3-0.6B ties its embeddings (tie_word_embeddings), so ParallelLMHead uses embed_tokens's weight directly and its own state_dict is empty; if the checkpoint still carries lm_head.weight during loading, as Qwen3-0.6B does, it is dropped.

RMSNorm and the residual

minisgl/layers/norm.py · RMSNormFused

class RMSNormFused(BaseOP):
    """Fuse "add the residual" and "normalize" into one op that returns both the normalized result and the new residual."""

    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:  # the first layer: there is no residual yet, so x is the residual
            return rmsnorm(x, self.weight, self.eps), x
        fused_add_rmsnorm(x, residual, self.weight, self.eps)
        return x, residual

In a decoder layer the residual add is always followed immediately by a normalization: residual = x + residual; x = norm(residual). RMSNormFused merges the two into one op that returns both the normalized result and the new residual; on a GPU it is a single FlashInfer kernel, saving one read and write of the hidden state. In the first layer, where there is no residual yet, x is the residual. So a whole decoder layer reads like this:

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]
        # starts with an underscore: not a weight and not in state_dict; computed once at startup and looked up by position
        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 is computed once for every position at construction time, and the forward pass looks it up by positions. The op rotates in the NeoX style, pairing the two halves, which matches the Hugging Face implementation for the Llama and Qwen families:

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-style rotary position embedding (pairing the two halves), modifying query and key in place.

    query: [T, Hq * D] or [T, Hq, D]; key likewise; cos_sin_cache: [max_pos, D], cos in the first
    half and sin in the second.
    """
    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 caches with functools.cache: every layer has the same RoPE parameters and shares one table. One more detail: the model is built on torch.device("meta"), while the cos/sin table is real data and cannot live there. So when get_rope sees that the current default device is meta, it uses the real device named earlier through set_rope_device:

minisgl/layers/rotary.py · get_rope

@functools.cache  # every layer shares one RoPE table
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":
        # the model is built on the meta device, but the RoPE table is real data and has to live on a real one
        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)

Forget to call set_rope_device before building on meta and you get RuntimeError: Call set_rope_device() before building a model on meta device.

Llama 3.1's long-context extension (rope_type: llama3) post-processes the frequencies: high-frequency components are left alone, low-frequency ones are divided by factor, and the middle transitions smoothly. Its correctness is verified in chapter 20 against Hugging Face using a small Llama 3 model built from random weights.

Difference from upstream: YaRN

The official _get_rope also supports yarn, but it only changes the frequencies and never multiplies by the attention scaling factor from the Hugging Face implementation (0.1 * ln(factor) + 1). We implement only default and llama3 and leave YaRN as an exercise.

The attention layer

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)

The attention layer holds no weights either: it splits the merged qkv into three, applies an RMSNorm per head to q and k for Qwen3 (QK-Norm), adds RoPE, and hands q, k and v to the attention backend in the global context. Writing the KV cache and computing attention are the backend's job, implemented in chapter 5.

The official implementation

• BaseOP: layers/base.py · BaseOP (lines 15-53)
• the output layer taking only the last position: layers/embedding.py · ParallelLMHead.forward (lines 87-110)
• RoPE's meta-device handling: layers/rotary.py · get_rope (lines 125-143)
• the attention layer: layers/attention.py · AttentionLayer.forward (lines 47-57)

Upstream applies QK-Norm and RoPE in place on the views from qkv.split(...), since FlashInfer's ops accept strided inputs. Our reference ops make a .contiguous() copy first, which is more direct to read at the price of one extra copy.

Tests

tests/test_ch02_layers.py checks BaseOP's naming rules, the error on leftover keys, RMSNorm and RoPE matching Hugging Face, and the vocabulary-parallel lookup:

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)  # the vocabulary slice [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]]

How to explain it

On the op layer: an inference engine needs only a small part of nn.Module, namely collecting and loading weights by name and organizing the hierarchy; it needs no autograd, no hooks and no training state, which is why mini-sglang uses the lightweight BaseOP, where layers only organize parameters and the computing goes to the kernels (a reference implementation that matches HF numerically on a CPU, FlashInfer on a GPU). The model is built as an empty shell on the meta device, which allocates no memory and takes no time, and loading swaps in the real weights; RoPE's cos/sin table, however, is computed rather than loaded and must live on a real device. The fused residual add plus RMSNorm returns the normalized output and the new residual. During prefill the LM head computes only each request's last position, which saves an enormous amount of useless vocabulary projection.

Exercises

1. Add a num_params() method to BaseOP that returns the total number of weight elements. Use it to compute Qwen3-0.6B's parameter count (careful not to double-count the tied embeddings).
2. Remove the "prefill takes only the last position" logic from ParallelLMHead.forward and prefill a 1000-token request. How large is the LM head's output tensor in float32?
3. Implement YaRN: add a branch to _get_rope and multiply cos/sin by the attention scaling factor. Verify it against a small model built with transformers's Qwen3Config(rope_scaling={"rope_type": "yarn", ...}).

Answers

1. sum(t.numel() for t in self.state_dict().values()). With tied embeddings ParallelLMHead.state_dict() is empty, so nothing is double-counted: Qwen3-0.6B comes to about 596 million.
2. [1000, 151936] float32 values, about 580 MiB; taking only the last position makes it 0.6 MB.
3. Follow Hugging Face's _compute_yarn_parameters: the frequencies interpolate between inv_freq / factor and inv_freq along a ramp, and attention_factor = 0.1 * ln(factor) + 1 multiplies the cos/sin table as it is built.

Summary

• BaseOP collects and loads weights by attribute path through __dict__; attributes starting with an underscore are not weights; loading replaces rather than copies, which is what makes meta-device construction instant.
• Layers only organize parameters and hand the computing to minisgl.kernel: a reference implementation that matches HF numerically on a CPU, FlashInfer on a GPU.
• During prefill the output layer computes logits only at each request's last position.
• The RoPE table must live on a real device, and every layer shares one table.
• The attention layer only splits, applies QK-Norm and RoPE, and leaves the cache writes and the attention itself to the attention 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.

• 186. BaseOP：不用 nn.Module 收集和加载权重 (medium)

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