Inference Systems · Accelerator Scheduling

Tensor-Parallel Residency and Interconnect Pressure

Sharding a model across devices splits weights and key-value heads, but every decode step becomes a collective; how that changes the memory and latency balance.

A model too large for one device must be split, and the split that keeps every device busy on every token is tensor parallelism: each weight matrix is partitioned along its output or input dimension, and every device computes a slice of every layer. The residency consequences are usually described as a benefit, since each device holds only a fraction of the weights and therefore has more room for cache. The latency consequences are the part that scheduling has to absorb, because a sharded layer is not complete until its devices exchange partial results, and that exchange happens once per layer per step.

What sharding does to the weights and the cache

Consider a linear layer with weight $W$ of shape $[d_{out}, d_{in}]$ split across $k$ devices. Column parallelism gives each device $[d_{out}/k, d_{in}]$, and each produces a slice of the output, which requires no communication before the next operation if that operation is also column-parallel. Row parallelism gives each device $[d_{out}, d_{in}/k]$, and each produces a partial sum over its slice of the input, so the outputs must be summed across devices before proceeding. The standard transformer block alternates the two so that only one all-reduce per block is needed at the point where the partial sums must be combined.

For a 32-layer model with $k = 8$, a decode step performs 32 all-reduces, one per layer, each over an activation tensor of size batch times hidden dimension. With a hidden dimension of 8,192 and a batch of 64 in two-byte precision, each all-reduce moves 1 MiB per device per layer before accounting for the reduction algorithm’s own overhead. Across 32 layers that is 32 MiB per step per device, and at a step rate of tens per second it is a steady stream.

Key-value heads shard too. If the attention uses $H_k$ key heads and $k$ devices, each device holds $H_k/k$ heads, so the per-device cache cost per token falls by $k$. This is the residency benefit: the pool on each device can hold $k$ times as many tokens as a single-device deployment of the same total memory. But the attention itself now runs over a shard of the heads, and the output must be gathered so that every device has the full attention output for the next layer, which is a second communication point per layer unless the following operation is arranged to consume a sharded input.

The latency floor per step

The important scheduling fact is that a tensor-parallel decode step has a latency floor that does not shrink with batch size. The all-reduce time depends on the message size, which grows with batch, and on the interconnect’s latency and bandwidth. For small messages the fixed latency term dominates: a collective over eight devices has a launch and synchronization cost measured in microseconds regardless of how few bytes move. Multiply that by the number of layers and the result is a per-step floor that a larger batch cannot amortise away.

That changes the meaning of batching. On a single device, a larger batch amortises the weight read and costs nothing in latency until compute saturates. On a tensor-parallel deployment, a larger batch increases the message size of every collective, so the step time rises with batch from the first sequence onward. There is still a region where the marginal cost of a sequence is small, but it is bounded above by the point at which the collectives become bandwidth-limited, and that point arrives earlier than the compute limit would suggest.

The practical consequence is that tensor parallelism raises throughput per device-hour and worsens per-token latency. A deployment that must meet an inter-token latency requirement cannot use as much parallelism as one that only maximises tokens per second, because each additional shard adds a collective whose latency is paid on every token. This is the central tension the scheduler has to resolve, and it resolves differently for prefill and for decode. Prefill has enough arithmetic per collective to hide the latency; decode does not, which is one of the reasons the two phases are sometimes given different parallelism degrees.

Residency arithmetic under sharding

The per-token cache cost on each device is $L \times 2 \times (H_k/k) \times d \times b$ for a device holding all $L$ layers. That is the same formula as how key-value bytes set the batch ceiling with the key head count divided by the shard count.

Two effects offset the apparent gain. First, the weights on each device are $N/k$ parameters, so the pool is total memory divided by $k$ minus $Nb/k$ minus workspace, and the workspace term does not shrink with $k$ because activations and collective buffers are per-device and sized by batch and hidden dimension. Second, the collective buffers are themselves device memory, and for a large hidden dimension and batch they are not negligible. A deployment with eight shards and a large batch may spend a meaningful fraction of each device’s memory on buffers that exist only to be reduced.

The net residency per device is therefore not simply the single-device pool divided by $k$; it is that figure minus a fixed per-device overhead that is independent of $k$. At high shard counts the overhead dominates the benefit, and the pool stops growing. This is why tensor parallelism is usually paired with pipeline parallelism beyond a certain size: pipeline stages split layers, so each device holds fewer layers’ weights and fewer layers’ caches, and the per-stage pool is larger without adding collectives to each layer’s critical path.

Pipeline parallelism has its own scheduling cost, which is bubble time: a stage cannot start its next microbatch until the previous stage’s output arrives, so a pipeline of $s$ stages processing $m$ microbatches has an idle fraction proportional to $(s-1)/(m+s-1)$. In inference with continuous batching, the pipeline is kept full by having many sequences in flight at different stages, which is exactly the kind of scheduling problem that iteration-level batching creates. The mechanism of iteration-level batching is what makes pipelined inference viable at all, because a fixed batch would leave the pipeline idle at every stage boundary.

Preemption across a shard group

Preempting a sequence in a tensor-parallel deployment is a group operation. Every device holds a shard of that sequence’s cache, so every device must free its shard, and the eviction is complete only when the slowest device has done so. A swap moves $Tc$ bytes in total but distributed across $k$ links, so the wall-clock cost is set by one link’s bandwidth, not the aggregate. A recompute re-runs the prefill across all shards, so its cost is the same as a single-device prefill over the same length but with $k$ times the collective traffic.

The asymmetry between swap and recompute shifts under sharding. Recomputation’s collectives are per layer and per chunk, so a long recompute adds a long series of collectives to a link that is already carrying the decode traffic. Swapping, by contrast, uses the same link but in one bulk transfer per layer’s cache shard, which can overlap with computation more easily. On a fast interconnect the difference is small; on a slower one, recompute’s many small collectives are the worse option.

There is also a correctness constraint that does not exist on a single device: all shards must agree on which sequences are running. If one device’s allocator fails an allocation while another’s succeeds, the group is in an inconsistent state and the step cannot proceed. The allocation decision must therefore be made once, centrally, and broadcast, or every device must be given the same deterministic allocator state. Implementations that allocate per device independently produce rare, hard-to-reproduce failures where one rank’s batch differs from another’s.

Choosing the parallelism degree

The scheduling question is how many devices to place under one model instance. The variables are the model’s weight footprint, the memory needed for the cache pool, the latency floor the collectives impose, and the batch size the workload can sustain.

A useful way to frame it: sharding is worthwhile when the model does not fit on one device, or when the pool on one device cannot hold enough sequences to reach the bandwidth-bound regime. It is counterproductive when the model fits and the batch is small, because the collectives add latency without adding throughput. The intermediate case, where the model fits but the cache does not, is where the decision is genuinely close, and the deciding quantity is whether the additional resident sequences produce enough extra tokens to pay for the added per-step latency.

Since that quantity depends on the workload’s arrival rate and length distribution, the choice cannot be made from the model alone. A deployment serving long-context requests has a cache-bound pool and gains from sharding; a deployment serving short requests with low concurrency has a compute-bound pool and loses. The scheduler sees this as the difference between a step whose time is set by weight reads and one whose time is set by collectives, and the two regimes want different batch sizes and different admission policies.

The remaining structural option is to avoid the per-layer collective altogether by giving each device a whole model and routing requests between devices. That is data parallelism, and it preserves the single-device latency profile at the cost of duplicating weights. For models that fit comfortably, it is usually the better answer, and the decision between it and tensor parallelism is the first scheduling choice a deployment makes. What comes after, in either case, is the problem of meeting latency targets with a batch assembled per step, covered in deadline-aware batch assembly.

References

  1. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al., arXiv 2019
  2. Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP 2023
  3. NCCL documentation — NVIDIA