9 posts tagged with "Inference"

Running multi-trillion-parameter Mixture-of-Experts (MoE) models in production shifts the primary engineering bottleneck from model weights to distributed systems orchestration. With Moonshot AI's Kimi K3.1 Pro (moonshotai/Kimi-K3.1-Pro), infrastructure engineers must accommodate a ~2.8-trillion-parameter footprint (~104B active parameters per token across 896 routed experts), a 1M-token context window, and a hybrid 93-layer attention stack (69 Kimi Delta Attention layers paired with 24 Gated Multi-head Latent Attention layers).

At this footprint, the model weights alone surpass the High-Bandwidth Memory (HBM) capacity of any single machine, requiring multi-node Tensor and Expert Parallelism across high-bandwidth interconnects. Compounding the hardware footprint, long-horizon agentic sessions expose the limits of stateless Kubernetes load balancing: 1M-token KV caches rapidly exhaust accelerator HBM, uncoordinated routing forces replicas to recompute identical multi-turn prefixes, and compute-heavy prefill bursts stall ongoing token generation.

To eliminate these bottlenecks at launch, we worked with Inferact (vLLM) and RadixArk (SGLang) to deliver Day-0 distributed serving recipes on llm-d. Running llm-d as a unified, Kubernetes-native inference control plane, the same stack serves Kimi K3.1 Pro across both TPUs (Ironwood TPU v7x) and NVIDIA GPUs (GB300 NVL72), here on Google Cloud (A4X-Max). Because the control plane is accelerator- and cloud-independent, platform teams gain a portable serving stack that handles state-aware KV routing, hierarchical memory offloading, and prefill/decode disaggregation out of the box.

Lessons from bringing up RDMA-backed Prefill/Decode serving across H200, B200, GB200, and GB300 on GKE

Prefill/Decode (P/D) disaggregation splits LLM serving into two pools of workers: prefill workers process the prompt, and decode workers generate the output tokens. Between the two, every request hands its KV cache over the network. The path those bytes take (which NIC, which switches, how many hops) is set by three layers that don't coordinate by default: the platform decides which network paths exist at all, the resource allocator picks each worker's GPU and NIC, and the router picks which prefill worker pairs with which decode worker.

We learned how much that matters while bringing up llm-d's P/D path on four GCP GPU platforms: A3 Ultra (H200), A4 (B200), A4X (GB200 NVL72), and A4X Max (GB300 NVL72). The bring-up itself was a first: llm-d's P/D guide had no high-performance GKE RDMA recipe before this work built one on A3 Ultra, then carried the same allocation-and-transport contract through A4, A4X, and A4X Max (allocation). The problems we hit shared one shape: one layer lost what another knew, and the symptom appeared far from the cause.

The lesson is that topology is not only a transport optimization; it is an end-to-end contract spanning allocator, kernel, NIC, transport, engine, and router. Keeping RDMA fast depends on every layer holding its part: a mismatched GPU and NIC, a missing route, a silent fallback to TCP, or a topology-blind pairing anywhere along the way degrades the whole path. Concretely, the platform must discover topology, publishing each device's PCIe root, NIC, and rail as attributes Kubernetes can see. Allocation must preserve it, giving each worker a GPU and a NIC that belong together. The runtime must verify it, confirming that KV really moves over RDMA. And the router must consume it, pairing prefill and decode workers that are close on the network. Two results from our clusters show what that is worth:

• A silent fallback from RDMA to TCP is not a slowdown but an outage. Forcing one on purpose under benchmark load made 436 of 5,400 requests time out; eliminating one on A4X cut mean KV-transfer time by more than 93%.
• Mean time-to-first-token (TTFT) fell from 2.15 s to 1.157 s once each worker had a matched GPU and NIC and a verified RDMA channel, and then to 0.33 s, a further 3.5× reduction, once the router also paired workers by topology (evaluation).

An earlier post, Networking for Distributed Inference in llm-d, covered the transport stack itself: NIXL (the KV-transfer library), UCX (the communication library underneath it), and backend performance. This post follows topology from the platform through allocation and transport into routing.

Text inference has a comfortable shape in llm-d. A request arrives as a few thousand tokens, prefill and decode run back to back on one replica, and a prefix-aware router keeps the KV cache warm. Every layer of the stack — scheduler, load balancer, autoscaler — is calibrated against that shape.

Then someone attaches an image, and it is tempting to think of the result as the same request with a picture bolted on. It is not. A multimodal request differs from a text-only one in almost every property a serving system actually reads: how big it is, how much it costs, how long it takes before the first output token, what identity its cache entries have, and how many stages it has to move through. The parts of the stack that were calibrated on text keep working, silently, on numbers that no longer mean what they used to.

This post walks those differences deliberately. First what makes media content different from text content. Then what that does to an inference request, difference by difference. Then, for each one, what llm-d actually does about it — the token estimator, the content-hash routing path, the retuned affinity gate, and the encode tier. Model server like vllm is the engine throughout; its per-replica mechanisms are the floor, not the answer.

How the llm-d router balances KV-cache affinity against token load

The llm-d router's default configuration has changed to something an operator can reason about, built on a single methodology: token-aware routing. The scheduler prioritizes KV-cache affinity, keeping each request on the endpoint that already holds its prefix, unless the load on that endpoint exceeds a calibrated limit; past the limit it ignores affinity and picks the endpoint by load alone. Load is measured in tokens matched to the workload's bottleneck, giving two configurations: for prefill-bound traffic (long prompts), prefix-cache affinity + token load (prefix-cache-affinity-filter + token-load-scorer), routing on uncached prefill tokens in flight; for decode-bound traffic (long outputs), prefix-cache affinity + active requests (prefix-cache-affinity-filter + active-request-scorer), routing on active streams. We call each pairing of the affinity filter with a load signal a bottleneck-matched configuration (or matched configuration for short).

Two failure modes motivated the change. The first is the previous default itself: a four-signal weighted blend (prefix-cache match, queue depth, KV utilization, LRU) whose emergent behavior was hard to predict and harder to tune. The second is the hot spotting most routers suffer from over-indexing on KV-cache affinity: affinity concentrates traffic on cache-warm endpoints, and without an explicit saturation release the warm endpoint keeps absorbing load past the point where a cold endpoint would serve the request faster. Token-aware routing pairs one signal with one bottleneck and one calibrated limit, so the scheduler is legible: an operator can predict what it will do, why it will do it, and how it will degrade under load.

The results back the change. On the benchmarks below, the matched configurations sustain 2–3× the throughput of Kubernetes Service round-robin on prefill-bound workloads while holding TTFT through the operating range, and hold parity or better everywhere else. The stack now ships as the default across llm-d's optimized-baseline, agentic-serving, multimodal-serving, and P/D-disaggregation guides, and it is running on production serving fleets, including Google Cloud Vertex AI, Red Hat, and Mistral.

Vision-language model (VLM) inference involves several stages: vision encoding, language model prefill, and language model decode. Each stage has different compute and memory characteristics, making heterogeneous disaggregation a natural fit. This gives us the flexibility to place different stages to different accelerator resources based on their compute and memory characteristics.

In this post, we evaluate a vision-heavy Kimi-VL serving workload on llm-d and SGLang. We place vision encoding stage in 4 Intel Arc Pro B60 (Intel B60) GPUs while keeping prefill and decode stage on 1 NVIDIA H200 GPU. Compared with the collocated baseline, the measured heterogeneous E/PD configuration achieved 2.4x-2.8x higher throughput and reduced mean Time To First Token (TTFT) by roughly 69%-80% under load.

For most of the transformer era, the KV cache rested on a quiet assumption: one model, one uniform cache. Every layer attended the same way, every block was the same size, and everything built on top of the cache (allocators, offload connectors, schedulers) could treat it as a single pool.

Hybrid models broke this assumption. Many recent frontier and open-weight models increasingly mix attention types within a single model (full attention next to sliding-window, linear, or Mamba layers), making the cache heterogeneous: different layers now hold different amounts of state, in different shapes, with different reuse rules. A cache block that used to be allocated as one uniform unit is now constituted of several distinct parts.

To serve a hybrid model efficiently, an AI inference platform has to handle that heterogeneity in at least three aspects of the stack:

• GPU Memory Allocation: How the cache is laid out and allocated on the GPU. vLLM solved this with its Hybrid Memory Allocator (HMA), rebuilt around a unified allocator (see Hybrid Models as First-Class Citizens in vLLM).
• KV Offloading: Extending the KV cache to CPU and storage. Without HMA awareness, an offloading connector turns the HMA off and therefore discards the GPU memory improvements or potential data movement savings.
• KV-Aware Routing: Sending each request to the right model-server replica. Ignoring hybrid memory structure may erroneously list nodes as having or not having the required KV data based on information stemming from just part of the layers.

vLLM's HMA solved hybrid GPU memory allocation when handling a single vLLM instance. This post shows how llm-d extends that to tiered KV cache management - including KV offloading to CPU and storage, and KV-aware request routing - significantly improving throughput and latency at scale for hybrid models.

Most production inference clusters today are single-vendor because that is often the simplest way to configure and operate a cluster.

That is starting to change. Procurement cycles bring new generations alongside older ones, supply planning spans multiple accelerator options, and cost/performance profiles differ by workload. Real production fleets are accumulating heterogeneity whether or not the architecture planned for it.

This is an opportunity to unlock real value: different accelerator classes can be matched to workload requirements, stranded capacity gets reclaimed, and operators gain more flexibility in capacity planning. The case is stronger still for sovereign and on-premise deployments, where data residency, regulatory alignment, and the long-term economics of high-volume inference make local fleet optimization especially important.

Making that work in practice is a non-trivial systems problem. Each accelerator stack brings its own optimized drivers, firmware, container images, runtime settings, and attention kernels. A coherent serving layer needs to preserve those platform-specific optimizations while still giving operators one control plane for routing, observability, and policy.

llm-d was born Kubernetes-native. Its workers are Deployments, its endpoints live in an InferencePool, and its guides assume a cluster is one kubectl away. That made sense: Kubernetes is where most production inference runs, and building on it gave llm-d a head start on networking, lifecycle, and scale.

But the thing that makes llm-d llm-d - KV-cache-aware scoring, prefix-cache affinity, prefill/decode disaggregation, flow control - was never fundamentally about Kubernetes. It is routing intelligence. It reasons about the state of a fleet of model servers and decides where each request should go. Nothing about that logic needs an API server. The dependency on Kubernetes was incidental, inherited from how endpoints happened to be discovered, not essential to what the router actually does.

This post is about pulling those two things apart. We introduce the EndpointDiscovery abstraction in the llm-d router that separates what endpoints exist from how to route across them, and the first plugin built on it - file discovery - which lets the full routing stack run as a plain process or container with no Kubernetes anywhere in sight: on an HPC cluster, inside a Ray job, on a bare-metal rack, or on your laptop.

Figure 1: The big picture - one routing stack under every platform. llm-d discovers endpoints through its EndpointDiscovery module (Kube Discovery against an InferencePool, File Discovery against everything else) and serves requests the same way on Kubernetes, Slurm, Ray, or bare metal (inference, HPC, and RL rollout workloads: veRL, SkyRL, prime-rl). The rest of this post explains how.

Not all LLM requests cost the same. A short prompt might complete in milliseconds, while a long one can occupy a GPU for seconds. If we can predict how long a request will take on each candidate server before dispatching it, we can make substantially better routing decisions. This post describes a system that does exactly that: a lightweight ML model trained online from live traffic that replaces manually tuned heuristic weights with direct latency predictions.