Tensor Parallelism: Measure the Cost
Calculate layer communication and benchmark one, two, and four GPUs to find the smallest shard count that repays all-reduce cost.
Tensor parallelism improves LLM inference only when the compute saved on each GPU repays the collective communication inserted into every transformer layer. This guide turns that boundary into a layer calculator and a one-, two-, and four-GPU benchmark receipt.
The intended reader is deciding how to place a model that does not comfortably meet latency or memory goals on one accelerator. You will leave with an all-reduce cost model, a topology-aware sweep, and a rule for rejecting extra shards that reduce useful throughput.
The operating vocabulary connects LLM inference parallelism, all-reduce latency, GPU sharding, and Megatron-LM as parts of one placement trade-off.
- Shard A
- Shard B
- All-reduce
- Next layer
Tensor parallelism starts with one layer
tensor parallelism begins with writing the local matrix work, activation bytes, collective count, and synchronization sequence for one transformer block. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision. The Megatron model-parallelism paper describes intra-layer model parallelism that partitions transformer matrix multiplications while limiting synchronization points.
Work through four explicit moves:
- Record hidden size and intermediate width
- Partition the column- and row-parallel matrices
- Calculate communicated elements and dtype bytes
- Multiply the layer boundary across model depth
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is dividing total model FLOPs by GPU count. Its consequence is communication and synchronization vanish from the capacity estimate.
Mitigate it with a layer worksheet that keeps compute and collective time separate. The release receipt is microseconds of local work and network work per layer and generated token. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
Model all-reduce on the real topology
A useful tensor parallelism decision depends on using measured bandwidth and latency for the exact NVLink, NVSwitch, PCIe, or cross-node path instead of a vendor peak. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision. The NCCL collective operations documentation defines collective semantics, counts, and communicator behavior needed to model the all-reduce on real devices.
Work through four explicit moves:
- Map rank placement to physical links
- Benchmark relevant message sizes
- Capture warm p50 and p95 collective time
- Repeat under expected concurrent traffic
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is using theoretical link bandwidth as application throughput. Its consequence is protocol overhead and topology contention are ignored.
Mitigate it with a collective microbenchmark on the deployment communicator. The release receipt is message-size curves tagged with hosts, links, driver, and NCCL versions. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
Calculate the compute-communication boundary
The worked tensor parallelism fixture makes comparing saved matrix time per shard with the added collective time and fixed launch overhead at each candidate width. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision.
Work through four explicit moves:
- Estimate local GEMM duration by shape
- Add measured collectives per layer
- Include launch and synchronization gaps
- Find the shard count where net token time rises
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is assuming communication overlaps completely. Its consequence is the model predicts linear scaling that the dependency path cannot achieve.
Mitigate it with a pessimistic non-overlap bound beside observed overlap. The release receipt is a calculator whose predicted token latency is reconciled with measured latency. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
| Signal | Decision | Proof |
|---|---|---|
| 1 GPU | Control | 46 tok/s · 38 ms ITL |
| 2 GPUs | Candidate | 78 tok/s · 27 ms ITL |
| 4 GPUs | Reject | 92 tok/s · 31 ms ITL |
Reproduce one two and four GPU runs
tensor parallelism needs an explicit rule for replaying identical prompts, output lengths, batching, cache state, numerical settings, and arrival traces at each tensor-parallel width. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision.
Work through four explicit moves:
- Freeze model and tokenizer revisions
- Warm every rank before measurement
- Replay the same seeded request trace
- Store request samples and collective telemetry
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is raising batch size only for the wider configuration. Its consequence is batching improvement is credited to distributed matrix work.
Mitigate it with one controlled sweep plus a separate capacity sweep. The release receipt is a test that chooses the lowest-latency configuration meeting the memory requirement. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
Runnable artifact. Save this as tensor-parallelism-llm-inference.test.mjs and run node --test tensor-parallelism-llm-inference.test.mjs. Expected result: PASS: tp2 is the smallest profitable shard count. The checked-in copy lives with this batch's evidence.
import assert from "node:assert/strict";
import test from "node:test";
const runs = [
{ id: "tp1", fits: false, itl: 38, tokensPerGpuSecond: 46 },
{ id: "tp2", fits: true, itl: 27, tokensPerGpuSecond: 39 },
{ id: "tp4", fits: true, itl: 31, tokensPerGpuSecond: 23 },
];
test("chooses the smallest width that fits and clears 30 ms ITL", () => {
const choice = runs.find((run) => run.fits && run.itl <= 30);
assert.equal(choice.id, "tp2");
console.log("PASS: tp2 is the smallest profitable shard count");
});
Measure prefill and decode separately
In production, tensor parallelism turns on recognizing that large prompt matrix operations amortize synchronization differently from small sequential decode operations. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision.
Work through four explicit moves:
- Bucket by prompt and output length
- Report first-token and inter-token latency
- Measure batch composition at each phase
- Weight phase results by traffic
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is reporting aggregate tokens per second alone. Its consequence is prompt throughput can conceal slower visible token cadence.
Mitigate it with phase-specific objectives and per-request percentiles. The release receipt is TTFT, ITL, and throughput by workload cohort and shard width. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
Account for memory beyond weights
Safe tensor parallelism requires measuring sharded weights, duplicated buffers, KV cache ownership, graph captures, allocator reserve, and failover headroom. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision.
Work through four explicit moves:
- Record resident bytes after warm-up
- Separate sharded and replicated allocations
- Sweep active sequences and cache length
- Reserve space for rank recovery
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is selecting tensor parallelism from weight size alone. Its consequence is runtime buffers or cache pressure still trigger eviction and out-of-memory faults.
Mitigate it with a full peak-memory ledger at target concurrency. The release receipt is bytes per category and maximum admitted sequence count on every rank. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
- ProjectProject
Split matrix columns and compute local partial activations.
- CollectCollect
Exchange partial results over the actual interconnect.
- MergeMerge
Reduce values into the activation required by the next operation.
- RepeatRepeat
Pay the same synchronization boundary at every relevant layer.
Protect rank symmetry and failure handling
A tensor parallelism rollout should preserve ensuring every rank receives compatible work, times out together, and terminates or retries under one request owner. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision.
Work through four explicit moves:
- Validate model and configuration hashes
- Bound collective timeouts
- Fail the whole request on rank loss
- Drain communicators before replacement
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is retrying one failed rank independently. Its consequence is collectives deadlock or combine activations from different attempts.
Mitigate it with group-scoped health and request identity. The release receipt is a fault-injection trace with one terminal outcome and no stranded rank. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
Choose the smallest profitable shard count
The evidence for tensor parallelism is strongest when requiring the configuration to meet memory first and then improve user-visible latency or cost-adjusted throughput after communication. This is the narrow boundary for this section; everything outside it belongs in a separate capacity, policy, or product decision.
Work through four explicit moves:
- Reject widths that cannot hold the target workload
- Apply TTFT and ITL service limits
- Calculate useful tokens per accelerator-second
- Keep the nearest smaller width as control
In a working review, I would put the first move beside the input fixture, use the second to expose the decision boundary, and make the third observable before polishing the interface. The fourth move is the release check. This order matters because a convincing happy path can still conceal incompatible state, unfair scheduling, inaccessible fallback, or ownership ambiguity. Keeping each move named also lets another engineer reproduce the result without inheriting private context.
The named failure mode is using every available GPU because throughput increases. Its consequence is higher cost and worse token cadence become the default deployment.
Mitigate it with a lexicographic memory, latency, then cost rule. The release receipt is a decision note naming where two GPUs win and four GPUs stop paying back. Those fields connect the implementation to the article's single question and make a later update comparable instead of anecdotal.
Put the decision into practice
Choose tensor parallelism only after a layer-level compute and collective model is reconciled with one-, two-, and four-GPU measurements on the deployment topology. The smallest width that fits the target workload and clears first-token, token-cadence, and cost objectives is usually the defensible operating point.
Start with the three-row selector, replace its illustrative values with synchronized request and NCCL samples, and retain the smaller control. Reopen the decision when model shapes, interconnects, batching, or traffic cohorts change rather than treating shard count as a permanent model property.
The method connects to four existing Journal notes: disaggregated LLM inference, multi-LoRA serving, KV cache optimization, LLM admission control. Each link covers an adjacent boundary while this article stays focused on one outcome. Keep the fixture, visual evidence, command output, and release receipt together so the next review can test the claim against the same starting conditions.