Tensor-parallel scaling

Four numbers per TP size, and one residual:

1. Decode-step latency at fixed batch and context, so the sweep is a clean function of $p$.
2. Speedup $S(p) = t(1)/t(p)$, and scaling efficiency $S(p)/p$.
3. Output-token throughput at a batch large enough to matter, which scales differently from latency and will disagree with it.
4. All-reduce time per step, taken from a profile, in microseconds and as a share of the step.

The residual is the exercise: predicted step time minus measured step time, with the gap attributed to a specific collective at a specific message size. A sweep that lands at, say, 62% efficiency at TP=8 and stops there has measured nothing you could act on.

The model you are predicting against

§5.1 gives two terms and they behave completely differently. Let $p$ be the TP size, $L$ the layer count, $d$ the model width, $n$ the tokens in the step, $b$ bytes per element, $\beta$ HBM bandwidth and $B_{\text{net}}$ the interconnect bandwidth. A ring all-reduce of $S$ bytes moves $2\frac{p-1}{p}S$ per rank, so the bandwidth term is

and the latency term is $2L\lambda$, where $\lambda$ is the fixed cost of one small collective — a constant that tensor parallelism does not amortise. For Llama-3-8B ($L=32$, $d=4096$, bf16) that is 64 collectives per forward pass and an 8 KiB residual slice per token. The step's floor is the weight stream: §0.4 puts Llama-3-8B at 15.01 GB streamed per decode step — 7.505 B parameters, because the input embedding table is gathered rather than streamed — which at the H100 SXM's 3.35 TB/s is 4.481 ms at TP=1, and $1/p$ of that under tensor parallelism because weights, embedding and LM head all shard.

Read that table before you run anything. At batch 1 the bandwidth term is two parts in ten thousand — it cannot possibly explain a 40% efficiency loss. This rules out that small idealized bandwidth term as the entire explanation, but does not identify the residual. Measure collective duration, launch gaps, KV traffic and achieved kernel throughput separately.

Three curves, one free parameter. Fit $\lambda$ to your measured TP=8 point, then check the fit against TP=2 and TP=4 — if one $\lambda$ fits all three, the curve is consistent with that model, not proof of a unique cause. If it does not, something in the next section is happening to you.

Part A needs no GPU:

$ python3 run.py --predict-only --model-preset llama-3-8b --tp 1 2 4 8 --lambda-us 2 5 10
$ python3 run.py --help          # every knob, including --streamed-gb and --nvlink-gb-s

Timing contract: the vLLM harness reports full-generation mean latency, while SGLang reports median decode-only latency. Dividing vLLM's mean by output length does not remove prefill or turn a mean into a median. The script labels these separately and skips decode-cost attribution for that vLLM statistic. For a cross-engine comparison, capture matched decode-only windows or use the same end-to-end harness for both; never place the two different statistics in one speedup ranking.

Part B drives the engine once per TP size. run.py shells out to each engine's own single-batch harness rather than reimplementing a timing loop, because the harnesses already handle warmup and CUDA-graph capture correctly. vLLM's takes an explicit warmup count and writes a JSON of per-iteration latencies:

        "--num-iters-warmup",
        type=int,
        default=10,
        help="Number of iterations to run for warmup.",
    )
    parser.add_argument(
        "--num-iters", type=int, default=30, help="Number of iterations to run."
    )

    if args.output_json:
        results = {
            "avg_latency": np.mean(latencies),
            "latencies": latencies.tolist(),
            "percentiles": dict(zip(percentages, percentiles.tolist())),
        }
        with open(args.output_json, "w") as f:
            json.dump(results, f, indent=4)

SGLang's single-batch harness is the closer instrument for this lab, because it separates the prefill from the decode steps and reports the median decode latency — which is the quantity the model above predicts, with the prefill and the first token excluded:

    if output_len > 1:
        med_decode_latency = np.median(decode_latencies)
        med_decode_throughput = batch_size / med_decode_latency
        rank_print(
            f"Decode.  median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s"
        )
        measurement_results["median_decode_latency"] = med_decode_latency
        measurement_results["median_decode_throughput"] = med_decode_throughput

# vLLM: --tensor-parallel-size is the knob (vllm/config/parallel.py); latency in the JSON
vllm bench latency --model meta-llama/Meta-Llama-3-8B-Instruct \
    --tensor-parallel-size 8 --batch-size 1 --input-len 128 --output-len 128 \
    --num-iters-warmup 10 --num-iters 30 --output-json tp8.json

# SGLang: --tp-size is the same knob; median_decode_latency lands in the jsonl
python3 -m sglang.benchmark.one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct \
    --tp-size 8 --batch-size 1 --input-len 128 --output-len 128 \
    --result-filename tp8.jsonl

# or let run.py do the whole sweep and print the table
$ python3 run.py --engine vllm --model meta-llama/Meta-Llama-3-8B-Instruct \
    --tp 1 2 4 8 --batch-size 1 --workdir ./results

Part C is the profile. Do not guess which kernels are the collectives — both projects ship a kernel-name-to-bucket map for exactly this, and they agree on the regex:

    "llama": {
      "fused_moe_kernel|GroupProblemShape|group_gemm_starts|bmm_|GemmUniversal": "moe_gemm",
      "gemm|nvjet": "gemm",
      "moe|sigmoid": "moe",
      "CatArrayBatched|prepare_inputs": "prepare_next",
      "ncclDevKernel|cross_device_reduce": "nccl_and_custom_ar",
      "_norm_|Norm": "norm",
      "act_and_mul_": "activation",
      "Rotary": "rope",
      "SoftMax": "softmax",
      "flash|fmha": "attn",
      "elementwise": "elementwise",
      "fp8_quant|cvt_": "quantize",
      "reduce_kernel": "reduce",
      "triton": "triton_kernel",
      "CUDA mem": "non-gpu-H_D_memops",
      ".*": "misc"

    "llama": {
      "gemm|nvjet": "gemm",
      "fused_moe_kernel|GroupProblemShape|group_gemm_starts|bmm_|GemmUniversal": "moe_gemm",
      "moe|sigmoid": "moe",
      "CatArrayBatched|prepare_inputs": "prepare_next",
      "ncclDevKernel|cross_device_reduce": "nccl_and_custom_ar",
      "_norm_|Norm": "norm",
      "topk": "topk",
      "act_and_mul_": "activation",
      "Rotary": "rope",
      "SoftMax": "softmax",
      "flash|fmha": "attn",
      "elementwise": "elementwise",
      "fp8_quant|cvt_|quantize": "quantize",
      "reduce_kernel": "reduce",
      "triton": "triton_kernel",
      "CUDA mem": "non-gpu-H_D_memops",
      ".*": "misc"

The nccl_and_custom_ar bucket is your answer. It matches two families of symbol: NCCL's device kernels, whose names begin ncclDevKernel, and the engines' own one-shot and two-shot kernels, whose names begin cross_device_reduce. vLLM's are readable in tree:

template <typename T, int ngpus>
__global__ void __launch_bounds__(512, 1)
    cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg,
                               T* __restrict__ result, int rank, int size) {
  using P = typename packed_t<T>::P;
  using A = typename packed_t<T>::A;
  // note: we don't reorder the address so the accumulation order is the same

Capture with nsys, then run the shipped script — it produces a per-bucket CSV and an HTML chart, and the bucket you want is one row of it. The invocation is documented in the repo's own README (tools/profiler/nsys_profile_tools/README.md):

nsys profile -t cuda -o tp8 -f true --trace-fork-before-exec=true \
    --cuda-graph-trace=node --delay 60 --duration 60 \
    vllm serve meta-llama/Meta-Llama-3-8B-Instruct --tensor-parallel-size 8

python3 $VLLM/tools/profiler/nsys_profile_tools/gputrc2graph.py \
    --in_file tp8.nsys-rep,vllm,llama,0 --title "TP=8 kernel budget"

If you would rather stay inside torch.profiler, both engines expose start/stop over HTTP; §10.5 owns the flags (vLLM's --profiler-config.profiler=torch plus --profiler-config.torch_profiler_dir; SGLang's SGLANG_TORCH_PROFILER_DIR and a POST body). Nsight Systems is the better tool here for one specific reason: it is the only one of the three that shows rank skew, and a TP sweep that decays faster than the model predicts is usually one slow rank making seven others wait.

The simple model predicts decreasing efficiency; real shape/backend transitions can be nonmonotonic. Five things make it fall faster than $t(p) = t(1)/p + 2L\lambda$ predicts, and they are worth checking in this order.

1. The all-reduce implementation changes under you, by message size

This is the one that ambushes people, because nothing logs it. vLLM tries a ladder of all-reduce implementations per call and falls through to NCCL when none of the fast paths accepts the tensor:

        # always try quick reduce first, then flashinfer, then the AITER or vLLM
        # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*)
        qr_comm = self.qr_comm
        # ... flashinfer and AITER branches elided ...
        ca_comm = self.ca_comm
        if (
            ca_comm is not None
            and not ca_comm.disabled
            and ca_comm.should_custom_ar(input_)
        ):
            out = ca_comm.custom_all_reduce(input_)
            assert out is not None
            return out
        symm_mem_comm = self.symm_mem_comm
        if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_):
            out = symm_mem_comm.all_reduce(input_)
            assert out is not None
            return out
        pynccl_comm = self.pynccl_comm
        if pynccl_comm is None or pynccl_comm.disabled:
            out = input_.clone()
            torch.distributed.all_reduce(out, group=self.device_group)
            return out
        assert pynccl_comm is not None
        out = pynccl_comm.all_reduce(input_)

The gate on the custom kernel is a byte count:

    def should_custom_ar(self, inp: torch.Tensor):
        if self.disabled or self.world_size > 8:
            return False
        inp_size = inp.numel() * inp.element_size()
        # custom allreduce requires input byte size to be multiples of 16
        if inp_size % 16 != 0:
            return False
        if not is_weak_contiguous(inp):
            return False
        # for 4 or more non NVLink-capable GPUs, custom allreduce provides
        # little performance improvement over NCCL.
        if self.world_size == 2 or self.fully_connected:
            return inp_size < self.max_size
        return False

and max_size is a per-architecture, per-world-size table:

CUSTOM_ALL_REDUCE_MAX_SIZES = {
    "9.0": {
        2: 64 * MiB,  # 64 MB
        4: 32 * MiB,  # 32 MB
        6: MiB // 2,  # 512 KB
        8: MiB // 4,  # 256 KB
    },

On an H100 (compute capability 9.0) at TP=8 the ceiling is 256 KB. A Llama-3-8B residual slice is 8 KiB per token, so the custom kernel handles up to 31 tokens per step and batch 32 and above runs on NCCL instead. That is a discontinuity in the middle of any batch-size sweep, invisible from outside, and visible in a profile as the kernel name changing from cross_device_reduce_* to ncclDevKernel*.

SGLang gates the same kernel on a single flat constant — _MAX_CAR_SIZE = 8192 * 1024, 8 MiB, at every world size (python/sglang/srt/distributed/device_communicators/custom_all_reduce.py:L41-L42) — so at TP=8 it keeps its custom kernel out to 1,024 tokens where vLLM has long since switched. If your two engines diverge sharply at TP=8 and moderate batch, this is why, and it is a threshold choice rather than a kernel-quality difference.

Both engines let you delete the variable entirely, under the same flag spelling — vLLM's field is disable_custom_all_reduce at vllm/config/parallel.py:L205, SGLang's is declared here. Force NCCL on both sides and re-run: the difference between the two sweeps is the custom kernel's whole contribution.

    disable_custom_all_reduce: A[
        bool,
        Arg(
            help="Disable the custom all-reduce kernel and fall back to NCCL.",
            resolvable=True,
        ),
        NS("exec.comm"),
    ] = False

2. Fixed collective latency does not amortise

Sixty-four collectives per forward pass, and $\lambda$ is a property of the fabric and the launch path, not of $p$. At TP=8 the step floor is 560 µs, so the entire budget is 560/64 = 8.75 µs per collective before communication consumes the whole step. This is why efficiency decays even at batch 1, where the bandwidth term is 0.18% of the floor.

3. KV heads stop dividing, so throughput and latency diverge

Llama-3-8B has $h_{kv} = 8$, so at TP=8 every rank holds exactly one KV head and past TP=8 they would replicate. Latency keeps improving with $p$; per-card KV capacity stops. Measure both and you will see the sweep's two halves disagree — §5.1 derives the capacity law, and lab 02 is where you read the engine's reported pool size at each TP.

4. CUDA-graph bucketing flattens the batch axis

If you sweep batch size as well as TP, remember the capture ladder: batches 9 through 16 replay the same graph. A plateau in a batch sweep is bucketing, not scaling. Sweep on the rungs, or report the padded size — §10.3 has the full confounder table, and every row of it applies to this lab.

5. One slow rank

All-reduce is a barrier. If one GPU is thermally throttled, on a different NUMA node, or sharing a PCIe switch with the NIC, all $p$ ranks run at its speed and the profile shows the other ranks parked inside the collective. This is the failure that looks exactly like "collectives are expensive" and is not. It is the reason Part C uses Nsight Systems rather than a single-rank torch.profiler trace.