A symphony orchestra performing a score too dense for any single musician to sight-read alone splits the parts across sections: strings carry one set of measures, brass another, percussion a third, and every section has to land on the same beat within a few milliseconds or the piece falls apart. Serving a 70 billion parameter language model has the same shape of problem. The model’s weights in FP16 come to roughly 140GB, and the largest single GPU commonly available for inference, an A100 80GB, holds less than half of that. There is no version of this problem where one GPU does the job. Eight have to act like one, and every one of them has to agree, layer by layer, on what just happened.

The naive fix looks simple: cut the model into eight pieces, put one piece on each GPU, and pass activations down the line like a relay baton. That is pipeline parallelism, and it works, but only if every runner is handed the baton at a steady cadence. An 80-layer, 70B-parameter model split eight ways gives each GPU ten layers to run before handing off. While GPU 0 computes its ten layers for the next token, GPUs 1 through 7 are either waiting for GPU 0’s output or already idle because they finished their slice of the previous token and have nothing queued behind it. That idle time is a pipeline bubble, and at low concurrency it can eat well over half of every GPU’s cycles. Sharding fixes the memory problem and creates a new scheduling problem in the same motion.

The alternative, tensor parallelism, cuts every matrix multiply inside every layer into slices instead of cutting layers between GPUs. Now all eight GPUs work on every layer simultaneously, each holding a slice of every weight matrix, and there is no idle relay handoff. The cost shows up somewhere else: after every attention block and every MLP block, the GPUs have to exchange partial results and sum them, an operation called all-reduce, and that exchange has to complete fast enough that it never becomes the bottleneck it was supposed to avoid. Across 8 GPUs connected by NVLink at roughly 600GB/s inside one node, that exchange is nearly free. Cross a node boundary onto InfiniBand at around 25GB/s effective bandwidth with meaningfully higher per-hop latency, and the same exchange can cost more than the compute it is protecting.

Quantization changes the equation from a completely different angle: cut every weight from 16 bits down to 8 or fewer, and the 140GB footprint shrinks to 70GB or less. That can mean fewer GPUs per replica and more replicas per node, trading a small amount of numerical precision for a large amount of shard count. We need to solve for three things simultaneously: how to split the model across GPUs so neither all-reduce overhead nor pipeline idle time dominates, how much of each GPU’s memory to reserve for weights versus the KV cache that holds every active conversation’s attention history, and how to route requests across many replica groups so no single shard group turns into a hot spot while another sits half empty.

Requirements and Constraints

Functional Requirements

• Load a checkpoint for a 70B+ parameter dense transformer and partition its weights across a fixed GPU topology (8x A100 80GB per node) according to a chosen tensor parallelism degree and pipeline parallelism degree
• Execute a full forward pass, both prefill and decode, across the sharded layout, synchronizing partial results via NCCL collectives at every tensor-parallel boundary
• Stream tokens back to the client as they are generated, producing the same output distribution the unsharded model would (numerically exact for FP16, bounded-delta for quantized modes)
• Run multiple replica groups concurrently, each an independent shard of GPUs, and route new requests to whichever replica group has KV cache headroom
• Support re-sharding to a new TP/PP degree or quantization mode through rolling replica replacement, without a full-fleet outage
• Detect a shard group whose GPUs report a communication failure (NCCL timeout) and evict it before it silently serves degraded or hung requests

Non-Functional Requirements

• Latency: time-per-output-token (TPOT) p99 under 25ms, sustaining at least 40 tokens/second per active sequence during decode
• Throughput: aggregate decode tokens/second per replica, and total node throughput as replica density scales with quantization mode
• Quality: sharded execution must match a trusted reference shard plan within a bounded perplexity delta on a golden eval set; FP16 plans should be near bit-identical modulo floating-point summation order in the all-reduce
• Cost: dollars per million output tokens must improve, or at minimum not regress, as the fleet moves from a low-shard-count plan to a higher-density plan enabled by quantization
• Capacity: GPU memory footprint per shard is weights plus activation buffers plus KV cache, with KV cache capped at 20GB per GPU shard to leave headroom for in-flight micro-batches and allocator fragmentation

Constraints

• Assume a 70B-class dense transformer matching the Llama-3 70B architecture shape: 80 transformer layers, hidden size 8192, 64 attention heads, 8 KV heads via grouped-query attention (GQA), head dim 128, FP16 weights
• Assume 8x A100 80GB GPUs per node, fully connected by NVLink/NVSwitch inside the node at roughly 600GB/s, and InfiniBand (200Gb/s, ~25GB/s effective) between nodes
• Assume checkpoint shards are pre-sharded and stored in an object store; re-sharding a live checkpoint from one layout to another on the fly is out of scope
• Out of scope: training or fine-tuning workloads. There is no gradient all-reduce and no optimizer state sharding here, this is serving-only
• Out of scope: mixture-of-experts routing. This design assumes a dense architecture where every token visits every layer

High-Level Architecture

The system has six major components. The Shard Planner decides, at fleet provisioning time, the tensor parallelism degree, pipeline parallelism degree, and quantization mode for a replica group given the GPU topology and memory budget. The Checkpoint Loader reads pre-sharded weight files from the object store into exactly the GPU shard that owns each partition. The Tensor-Parallel Shard Group is the set of GPUs inside one pipeline stage that jointly hold every layer in that stage, split by row and column. The Pipeline Stage Runtime groups layers into stages and schedules micro-batches so stage handoffs stay full instead of idle. The NCCL Collective Coordinator issues all-reduce inside a TP group and point-to-point send/recv between pipeline stages, over NVLink or InfiniBand depending on placement. The Replica Router is the front door that assigns incoming requests to whichever replica group has KV cache headroom and rebalances as load shifts.

A request lands at the Replica Router, which checks KV cache headroom across replica groups and assigns the request to the least-loaded one. Inside that replica group, the request enters Stage 0, a 4-GPU tensor-parallel shard holding the first half of the model’s layers; every attention and MLP block in that stage runs on all four GPUs simultaneously, with the NCCL Collective Coordinator firing an all-reduce after each block to sum partial results. When Stage 0 finishes its last layer, the coordinator sends the resulting activations to Stage 1, a second 4-GPU tensor-parallel shard holding the remaining layers, which repeats the same pattern before streaming tokens back to the client over SSE.

Before any of that runs, the Shard Planner decides this replica group’s TP=4, PP=2 layout based on the fixed 8-GPU node and the model’s memory footprint, and the Checkpoint Loader pulls the matching weight shard into each of the eight GPUs from the object store, verifying a checksum against the manifest before the replica accepts traffic.

The Shard Planner

The planner’s job is to pick a tensor parallelism degree, a pipeline parallelism degree, and a quantization mode that fit the model into the fixed GPU topology without violating the KV cache budget.

The non-obvious part: a smart engineer’s first instinct is to maximize tensor parallelism degree, since TP keeps every GPU active on every layer and avoids pipeline bubbles entirely. That is correct for latency but assumes the NVLink domain is large enough to hold the whole TP group, and it ignores that a bigger TP group means a bigger all-reduce every single layer. The planner has to weigh “no bubbles” against “cheaper, smaller collectives” rather than defaulting to the largest TP degree the hardware allows.

# Shard planner: picks TP degree, PP degree, and quantization mode for a fixed GPU budget
from dataclasses import dataclass

BYTES_PER_PARAM = {"fp16": 2, "int8": 1, "fp8": 1}
NVLINK_DOMAIN_SIZE = 8  # GPUs fully connected by NVSwitch in one node

@dataclass
class ShardPlan:
    tp_degree: int
    pp_degree: int
    dtype: str
    gpus_per_replica: int
    weight_bytes_per_gpu: float
    kv_cache_budget_bytes: int

def plan_shard_layout(
    total_params: int,
    num_layers: int,
    num_kv_heads: int,
    gpu_memory_bytes: int,
    kv_cache_cap_bytes: int,
    dtype: str = "fp16",
) -> ShardPlan:
    bytes_per_param = BYTES_PER_PARAM[dtype]
    total_weight_bytes = total_params * bytes_per_param

    best = None
    for tp in (1, 2, 4, 8):
        if tp > NVLINK_DOMAIN_SIZE or num_kv_heads % tp != 0:
            continue  # TP must stay inside one NVLink domain and divide KV heads evenly
        for pp in (1, 2, 4, 8):
            gpus = tp * pp
            if gpus > 8 or num_layers % pp != 0:
                continue
            weight_per_gpu = total_weight_bytes / gpus
            activation_overhead = 3 * (1024 ** 3)  # ~3GB reserved for buffers, NCCL, CUDA context
            headroom = gpu_memory_bytes - weight_per_gpu - activation_overhead
            if headroom < kv_cache_cap_bytes:
                continue  # doesn't leave the required KV cache budget
            replicas_per_node = 8 // gpus
            score = replicas_per_node  # prefer plans that pack more replicas per node
            if best is None or score > best[0]:
                best = (score, ShardPlan(tp, pp, dtype, gpus, weight_per_gpu, kv_cache_cap_bytes))
    if best is None:
        raise ValueError("no shard layout fits the memory budget at this dtype")
    return best[1]