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

Tensor parallelism

When a model will not fit on one GPU, or one GPU's bandwidth cannot reach the latency you want, each layer's weights have to be split across several GPUs: that is tensor parallelism (TP). All those get_tp_info().size calls in the linear layers, the embeddings and the KV pool of earlier chapters finally come into play here. We communicate over gloo on a CPU, and TP=2 and TP=4 produce output identical to single-GPU Hugging Face.

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

  1. How many all-reduces does one decoder layer need under tensor parallelism? Where are they?
  2. Which dimension do column parallel and row parallel split? Why is the MLP "column then row"?
  3. Qwen3-0.6B has 8 KV heads. How are they divided at TP=16?
  4. Vocabulary-parallel embeddings use an all-reduce and the output layer an all-gather. Why the difference?
Answers (try it yourself first, then expand)
  1. Two: one after attention's o_proj (row parallel) and one after the MLP's down_proj (row parallel).
  2. Column parallel splits along the output dimension (each rank computes part of the output); row parallel splits along the input dimension (each rank gets a partial sum that needs an all-reduce). Column then row: what column parallel outputs is exactly the slice of input row parallel needs, and the activation in between is elementwise and needs no communication, so the whole MLP only needs the one all-reduce at the end.
  3. 8 KV heads do not go round 16 ranks: each KV head is replicated onto 2 ranks, so each rank holds 1 KV head (along with its query heads).
  4. The embedding splits the vocabulary: each rank looks up only the tokens in its own slice and outputs zeros for the rest, and an all-reduce sums them into the complete embedding. The output layer, also split by vocabulary, computes logits for part of the vocabulary on each rank, and they have to be all-gathered into the complete distribution before sampling.

Files to revisit: layers/linear.py, layers/embedding.py, shard_tensor in models/weight.py, distributed/, and _init_communication in engine/engine.py.

Files so far: 73 (3 new, 0 modified, 0 brought in early)
minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/  (4 files, unchanged)
├── distributed/
│   ├── __init__.py  ← new in this chapter
│   ├── info.py  ← new in this chapter
│   ├── impl.py  ← new in this chapter
├── 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/ch16_tp.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Animation: how tensor parallelism splits a layer and how often it communicates (about 1.5 minutes, Chinese narration and subtitles)

Megatron-style sharding

A decoder layer has two "matmul, elementwise, matmul" structures: attention (qkv_proj, attention, o_proj) and the MLP (gate_up_proj, SiLU × up, down_proj). Megatron-LM's way of splitting them (the idea is in the tensor-parallelism chapter of the Inference Systems handbook):

  • split the first matmul along the output dimension (column parallel): each rank computes some of the output channels. Attention splits by head: each rank owns a few heads, and heads are independent of each other, so attention itself needs no communication at all. The MLP splits along the intermediate dimension, and SiLU × up is elementwise, which needs no communication either.
  • split the second matmul along the input dimension (row parallel): each rank multiplies its slice of the intermediate result by its slice of the weight and gets a full-shaped partial sum, and one all-reduce adds them at the end.

So each layer has two all-reduces, after o_proj and after down_proj, and everything else is local.

one decoder layer at TP=2 (color = rank) x full qkv_proj column parallel split by head qkv_proj column parallel split by head attention own heads per rank no communication attention own heads per rank no communication o_proj row parallel partial sum o_proj row parallel partial sum all-reduce add the partial sums same for the MLP gate_up col → down row another all-reduce Two all-reduces per layer; the embedding is split by vocabulary and all-reduced, the LM head split by vocabulary and all-gathered. With fewer KV heads than ranks, several ranks replicate one KV head (Qwen3-0.6B has 8 KV heads, so at TP=16 every two ranks share one). A fused qkv must be split per q, k and v and then concatenated, so each rank holds [q_local; k_local; v_local].
the data flow of one decoder layer under tensor parallelism

minisgl/layers/linear.py · LinearQKVMerged
class LinearQKVMerged(_LinearTPImpl):
    """The merged QKV projection, split by attention head; with fewer KV heads than ranks, a KV head is replicated across ranks."""

    def __init__(self, hidden_size: int, head_dim: int, num_qo_heads: int, num_kv_heads: int,
                 has_bias: bool):
        tp_size = get_tp_info().size
        local_qo = div_even(num_qo_heads, tp_size)
        local_kv = div_even(num_kv_heads, tp_size, allow_replicate=True)
        full_osize = (num_qo_heads + 2 * num_kv_heads) * head_dim
        local_osize = (local_qo + 2 * local_kv) * head_dim
        super().__init__(hidden_size, full_osize, hidden_size, local_osize, has_bias)
minisgl/layers/linear.py · LinearRowParallel
class LinearRowParallel(_LinearTPImpl):
    """Row parallel, split along the input dimension: each rank computes a partial sum, then an all-reduce."""

    def __init__(self, input_size: int, output_size: int, has_bias: bool):
        tp_info = get_tp_info()
        self._comm = DistributedCommunicator()
        self._tp_size = tp_info.size
        local_isize = div_even(input_size, tp_info.size)
        super().__init__(input_size, output_size, local_isize, output_size, has_bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.linear(x, self.weight, self.bias)
        return self._comm.all_reduce(y) if self._tp_size > 1 else y
A merged qkv_proj needs care when sharding: cutting the merged matrix straight down the middle would give all of q to rank 0 and k and v to rank 1. Instead q, k and v are split separately and concatenated on each rank. That is exactly what the streaming loader does, shard_tensor first and merge second (chapter 3).

Replicating KV heads when there are not enough

minisgl/models/weight.py · shard_tensor
def shard_tensor(key: str, value: torch.Tensor, rank: int, size: int,
                 num_kv_heads: int) -> torch.Tensor:
    """Take shard number rank."""
    if size == 1:
        return value
    if any(s in key for s in _SPLIT_DIM_0):
        is_kv = ".k_proj" in key or ".v_proj" in key
        if is_kv and num_kv_heads < size:  # fewer KV heads than ranks: several ranks share (replicate) one KV head
            head_dim = value.shape[0] // num_kv_heads
            head = rank * num_kv_heads // size
            return value[head * head_dim:(head + 1) * head_dim].clone()
        return value.chunk(size, dim=0)[rank].clone()
    if any(s in key for s in _SPLIT_DIM_1):
        return value.chunk(size, dim=1)[rank].clone()
    if "lm_head" in key or "embed_tokens" in key:  # vocabulary parallel: split by row, with the last slice possibly shorter
        per = div_ceil(value.shape[0], size)
        return value[rank * per:min((rank + 1) * per, value.shape[0])].clone()
    return value  # everything else (the norms, the MoE router) is a full copy per rank
A GQA model has fewer KV heads than q heads. Qwen3-0.6B has 16 q heads and 8 KV heads: at TP=2 each rank gets 8 q heads and 4 KV heads; at TP=16 each rank gets 1 q head, but there are only 8 KV heads, so every two ranks share one KV head (div_even(..., allow_replicate=True) returns 1) and each keeps its own copy. That wastes a little memory but keeps attention free of communication.

Vocabulary parallelism

The embedding's vocabulary is split by row, one slice per rank; a lookup outputs zeros for tokens outside the slice, and an all-reduce sums them, since each token finds its real vector on exactly one rank. The output layer is split by vocabulary too, with each rank computing the logits for its slice; sampling needs the complete distribution, so an all-gather stitches the slices together (ParallelLMHead.forward in chapter 2). The first adds the ranks' results, the second concatenates them.

Running it

examples/ch16_tp.py
"""Chapter 16: the weight shapes each rank gets under tensor parallelism; a two-process (gloo) forward pass against single-GPU, with the communication counted."""

import multiprocessing as mp

import torch

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


def show_shapes() -> None:
    import safetensors
    from minisgl.models import shard_tensor

    names = ["model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight",
             "model.layers.0.self_attn.o_proj.weight", "model.layers.0.mlp.gate_proj.weight",
             "model.layers.0.mlp.down_proj.weight", "model.embed_tokens.weight",
             "model.layers.0.input_layernorm.weight"]
    with safetensors.safe_open(f"{PATH}/model.safetensors", "pt") as f:
        print(f"{'weight':<42} {'full':<14} {'TP=2 per rank':<16} TP=16 per rank")
        for n in names:
            w = f.get_tensor(n)
            s2 = tuple(shard_tensor(n, w, 0, 2, num_kv_heads=8).shape)
            s16 = tuple(shard_tensor(n, w, 0, 16, num_kv_heads=8).shape)
            print(f"{n.removeprefix('model.'):<42} {str(tuple(w.shape)):<14} {str(s2):<16} {s16}")


def run(rank: int, size: int, out: mp.Queue) -> None:
    from minisgl.core import Batch, Req, SamplingParams
    from minisgl.distributed import DistributedCommunicator, 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=29600))
    counts = {"all_reduce": 0, "all_gather": 0}
    impl = DistributedCommunicator.plugins[-1]
    for op in counts:
        fn = getattr(impl, op)

        def wrapped(x, fn=fn, op=op):
            counts[op] += 1
            return fn(x)

        setattr(impl, op, wrapped)
    n = len(PROMPT_IDS)
    engine.page_table[0, :n] = torch.arange(n)
    req = Req(input_ids=torch.tensor(PROMPT_IDS, dtype=torch.int32), table_idx=0, cached_len=0,
              output_len=4, uid=0, sampling_params=SamplingParams(), cache_handle=None)
    batch = Batch(reqs=[req], phase="prefill")
    batch.padded_reqs = batch.reqs
    batch.input_ids, batch.positions = req.input_ids.clone(), torch.arange(n, dtype=torch.int32)
    batch.out_loc = engine.page_table[0, :n]
    engine.attn_backend.prepare_metadata(batch)
    with engine.ctx.forward_batch(batch):
        logits = engine.model.forward()
    qkv = engine.model.model.layers.op_list[0].self_attn.qkv_proj.weight
    out.put((rank, logits[0].numpy().copy(), tuple(qkv.shape), counts, tuple(engine.kv_cache.k_cache(0).shape)))
    engine.shutdown()


if __name__ == "__main__":
    show_shapes()
    mp.set_start_method("spawn")
    results = {}
    for size in (1, 2):
        q: mp.Queue = mp.Queue()
        procs = [mp.Process(target=run, 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 = torch.from_numpy(results[1][0][1])
    for rank, logits, qkv_shape, counts, k_shape in results[2]:
        print(f"TP=2 rank {rank}: qkv_proj {qkv_shape}, KV pool per layer {tuple(k_shape)}, "
              f"one forward pass: {counts['all_reduce']} all_reduce, {counts['all_gather']} all_gather, "
              f"logits max error against TP=1: {(torch.from_numpy(logits) - base).abs().max().item():.2e}")
What it prints
weight                                     full           TP=2 per rank    TP=16 per rank
layers.0.self_attn.q_proj.weight           (2048, 1024)   (1024, 1024)     (128, 1024)
layers.0.self_attn.k_proj.weight           (1024, 1024)   (512, 1024)      (128, 1024)
layers.0.self_attn.o_proj.weight           (1024, 2048)   (1024, 1024)     (1024, 128)
layers.0.mlp.gate_proj.weight              (3072, 1024)   (1536, 1024)     (192, 1024)
layers.0.mlp.down_proj.weight              (1024, 3072)   (1024, 1536)     (1024, 192)
embed_tokens.weight                        (151936, 1024) (75968, 1024)    (9496, 1024)
layers.0.input_layernorm.weight            (1024,)        (1024,)          (1024,)
TP=2 rank 0: qkv_proj (2048, 1024), KV pool per layer (257, 1, 4, 128), one forward pass: 57 all_reduce, 1 all_gather, logits max error against TP=1: 1.14e-05
TP=2 rank 1: qkv_proj (2048, 1024), KV pool per layer (257, 1, 4, 128), one forward pass: 57 all_reduce, 1 all_gather, logits max error against TP=1: 1.14e-05
- Weight shapes: at TP=2, q_proj, k_proj and gate_proj halve along the output dimension, o_proj and down_proj halve along the input dimension, the vocabulary halves, and the norms are not split; at TP=16 k_proj still has 128 rows, one complete KV head, replicated. - Both ranks' qkv_proj has (8 + 2 × 4) × 128 = 2048 rows, and each rank's KV pool holds only 4 KV heads. - One forward pass does 57 all-reduces (28 layers × 2, plus 1 in the embedding) and 1 all-gather (the output layer), exactly as the analysis says. - The logits differ from single-GPU by about 1e-5 at most: an all-reduce changes the order of floating-point addition, so the results are no longer bit-identical, but greedy decoding picks exactly the same tokens (this chapter's tests verify TP=2 and TP=4 at the service level).

How communication is implemented

minisgl/distributed/impl.py · TorchDistributedImpl
class TorchDistributedImpl:
    """Collectives through torch.distributed: NCCL on a GPU, gloo on a CPU."""

    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
        if dist.get_world_size() == 1:
            return x
        dist.all_reduce(x, op=dist.ReduceOp.SUM)
        return x

    def all_gather(self, x: torch.Tensor) -> torch.Tensor:
        """Concatenate the ranks' tensors along dimension 0: [n, ...] -> [world_size * n, ...]."""
        world_size = dist.get_world_size()
        if world_size == 1:
            return x
        x = x.contiguous()
        out = x.new_empty((world_size * x.shape[0],) + tuple(x.shape[1:]))
        if dist.get_backend() == "gloo":  # gloo does not support all_gather_into_tensor
            dist.all_gather(list(out.chunk(world_size, dim=0)), x)
        else:
            dist.all_gather_into_tensor(out, x)
        return out
DistributedCommunicator picks an implementation from a plugin list, defaulting to torch.distributed: NCCL on a GPU, gloo on a CPU. gloo does not support all_gather_into_tensor, so on a CPU we use all_gather writing into slices of the output tensor.

The process group is created as the engine initializes:

minisgl/engine/engine.py · Engine._init_communication
def _init_communication(self, config: EngineConfig) -> Any:
    """Create the process group when TP>1. On a GPU tensors go over NCCL with a separate gloo group for CPU-side syncs; on a CPU everything is gloo."""
    if config.tp_info.size == 1:
        return None
    backend = "nccl" if is_cuda(self.device) else "gloo"
    dist.init_process_group(
        backend=backend,
        rank=config.tp_info.rank,
        world_size=config.tp_info.size,
        timeout=timedelta(seconds=config.distributed_timeout),
        init_method=config.distributed_addr,
    )
    return dist.new_group(backend="gloo") if backend == "nccl" else dist.group.WORLD
On a GPU a separate gloo group is created for control information on the CPU (the broadcast count of chapter 14, the memory sync at startup).

Difference from upstream: PyNCCL

By default (use_pynccl=True) upstream does not use torch.distributed for GPU communication but calls NCCL directly through tvm-ffi (kernel/csrc/src/pynccl.cu), with a preallocated communication buffer. That costs less overhead and gives more control over how it interacts with CUDA Graph. We implement only torch.distributed and leave PyNCCL as an exercise (it needs a GPU).

The official implementation

Tests

tests/test_ch16_tp.py · test_tensor_parallel_matches_hf
@pytest.mark.parametrize("tp", [2, 4])
def test_tensor_parallel_matches_hf(tp):
    """Tensor parallelism over gloo on a CPU: greedy output at TP=2 and TP=4 matches single-GPU HF."""
    proc, port = start_server("--tp", str(tp), "--num-pages", "2048")
    try:
        tok = AutoTokenizer.from_pretrained(QWEN3)
        refs = hf_greedy(QWEN3, PROMPTS[:3], 10)
        for prompt, ref in zip(PROMPTS[:3], refs):
            r = httpx.post(f"http://127.0.0.1:{port}/v1/chat/completions",
                           json={"model": "m", "prompt": prompt, "max_tokens": 10, "temperature": 0,
                                 "ignore_eos": True}, timeout=120).json()
            assert r["choices"][0]["message"]["content"] == tok.decode(ref)
    finally:
        os.killpg(proc.pid, signal.SIGKILL)
Each TP configuration brings up a complete service (several scheduler processes, the tokenizer, the API server), and the greedy output for 3 prompts matches single-GPU Hugging Face word for word. That also verifies chapter 14's multi-rank message synchronization.

How to explain it

On tensor parallelism: the Megatron split shards attention by head (qkv column parallel, o_proj row parallel) and the MLP column parallel first (gate/up along the intermediate dimension) then row parallel (down), so each decoder layer needs only two all-reduces, one after attention and one after the MLP. A merged qkv has to be split per q, k and v before being concatenated, or the ranks get the wrong heads; when there are fewer KV heads than ranks (Qwen3-0.6B's 8 KV heads at TP=16) each KV head is replicated onto two ranks. Vocabulary parallelism: the embedding looks up only its own slice and zeroes the rest before an all-reduce, while the output layer computes its slice's logits and all-gathers them. Each rank differs from single-GPU only by the tiny error from the order of floating-point summation.

Exercises

  1. Compute how much data the two all-reduces per layer move for Llama-3.1-70B (80 layers, hidden 8192, bf16) at TP=8 with a decode batch of 64. At 450 GB/s of NVLink bandwidth one way, roughly what fraction of the time is communication?
  2. If the TP size does not divide the q head count (Qwen2.5-0.5B has 14 q heads at TP=4), where does it fail? How would you support it?
  3. Implement a PyNCCL plugin: load libnccl.so with ctypes and implement all_reduce; compare it against torch.distributed on a GPU.
Answers
  1. Each all-reduce moves [64, 8192] bf16 values, 1 MiB; a ring all-reduce has each rank send about 2 × (8−1)/8 × 1 MiB ≈ 1.75 MiB, roughly 4 µs of bandwidth time, plus a few microseconds of fixed latency per all-reduce. With 80 layers × 2 ≈ 160 of them, that is about 1 to 2 ms in total. At small batch sizes the fixed latency dominates, which is why the TP size should not get too large during decode.
  2. The div_even(num_qo_heads, tp_size) assertion in LinearQKVMerged fails. You can pad the head count (adding empty heads) or support only divisible TP sizes; production engines usually just require divisibility.
  3. Follow the official kernel/pynccl.py and SGLang's pynccl_wrapper.py: rank 0 creates an ID with ncclGetUniqueId and broadcasts it over the CPU process group; each rank calls ncclCommInitRank; and all_reduce calls ncclAllReduce on the current CUDA stream.

Summary

  • The Megatron split: column parallel first (by head or by intermediate dimension), row parallel second, two all-reduces per layer.
  • A merged qkv is split per projection and then concatenated; KV heads are replicated across ranks when there are not enough.
  • Vocabulary parallelism: the embedding all-reduces a sum, the output layer all-gathers a concatenation.
  • Each rank's output differs from single-GPU only by the tiny error from the order of floating-point summation, and greedy decoding agrees.

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.

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