This document is relevant for: Trn2, Trn3

Disaggregated Inference — Design Document#

Overview#

Disaggregated inference (DI) separates the prefill phase (prompt processing) from the decode phase (token generation) into independent server instances. This enables:

  • Independent scaling — prefill and decode have different compute profiles; scaling them separately matches hardware to workload.

  • Latency isolation — decode instances run exclusively decode batches, eliminating inter-token latency (ITL) spikes caused by prefill interference.

  • Asymmetric parallelism — each role can use different TP/DP/EP configurations optimized for its workload.

This document covers the DI architecture as implemented in upstream vLLM and the Neuron-specific extensions in vllm-neuron.

Architecture#

The diagram below illustrates the default read mode (decode pulls KV), which is what the toy proxy server drives.

  +----------+   (1) request    +-----------------------------+
  |  Client  | ---------------> |        Proxy Server         |
  |          | <--------------- | (load balancer; CPU-only,   |
  +----------+   (5) SSE        |  no Neuron cores)           |
                                +-----------------------------+
                                   |                       |
                  (2) dispatch     |                       |  (3) dispatch
                  (max_tokens=1)   |                       |  (w/ kv_transfer_params)
                                   v                       v
                +----------------------+      +----------------------+
                |    Prefill Server    |      |    Decode Server     |
                |    (kv_producer)     |      |    (kv_consumer)     |
                |  KV cache in device  |      |  generates ALL out-  |
                |  memory              |      |  put tokens, incl.   |
                |                      |      |  the first           |
                +----------------------+      +----------------------+
                           ^                              |
                           |   (3') NIXL RDMA READ:       |  (4) stream
                           |        decode pulls KV       |      tokens
                           +------------------------------+

Flow (read mode): (1) client → proxy. (2) proxy dispatches to prefill (max_tokens=1) and blocks; prefill produces KV and returns kv_transfer_params. (3) proxy dispatches to decode with those params; (3’) the decode worker pulls the KV from the prefill server via a NIXL RDMA READ. (4)/(5) decode generates and streams tokens back through the proxy to the client.

Components#

Component

Role

vLLM Configuration

Prefill server

Processes prompt, produces KV cache (the one token it samples is discarded)

kv_role: "kv_producer"

Decode server

Pulls KV, generates all output tokens including the first

kv_role: "kv_consumer"

Proxy server

Load balancer / router — relays SSE responses. Runs CPU-only (no Neuron cores)

Application-level (not vLLM)

Request Flow (read mode)#

  1. Client sends request to the proxy server.

  2. Proxy dispatches the request to a prefill server with max_tokens=1, stream=false, and do_remote_decode=true, and blocks on the response.

  3. Prefill server computes the forward pass, producing KV cache in its device memory, and returns kv_transfer_params (engine id, block ids, host/port). The single token it samples is discarded — it exists only to drive KV production.

  4. Proxy dispatches the request to a decode server, attaching the prefill’s kv_transfer_params. When the decode worker starts the request (start_load_kv), it issues a NIXL RDMA READ to pull the KV blocks from the prefill server’s device memory; the request waits in WAITING_FOR_REMOTE_KVS until the read completes.

  5. Decode server runs autoregressive generation on the transferred KV — the first output token is generated by the decode server — and streams tokens back through the proxy to the client.

In write mode, the prefill server pushes KV to decode and the proxy can dispatch both concurrently (lower TTFT); see the Transfer Modes section. The diagram shows read mode for clarity.

Upstream vLLM Implementation#

KV Connector Framework#

Located in vllm/distributed/kv_transfer/, the connector framework provides:

  • KVConnectorBase_V1 (kv_connector/v1/base.py) — abstract interface with scheduler-side and worker-side methods.

  • Scheduler-side: get_num_new_matched_tokens(), build_connector_meta(), update_state_after_alloc(), request_finished()

  • Worker-side: start_load_kv(), save_kv_layer(), wait_for_save(), register_kv_caches()

Connector Implementation#

vLLM ships several KV connectors, but Neuron DI uses NixlConnector:

Connector

Transport

Use Case

NixlConnector

RDMA (NIXL library)

Production — GPU & Neuron

Transfer Modes#

Mode

Direction

Proxy Behavior

Latency Impact

Read (decode pulls)

D ← P

Serial: await prefill, then dispatch decode

Higher TTFT

Write (prefill pushes)

P → D

Concurrent: both dispatched simultaneously

Lower TTFT

V1 Engine Integration Points#

# Scheduler integration
connector.get_num_new_matched_tokens(request, num_computed_tokens)
connector.update_state_after_alloc(request, blocks)
meta = connector.build_connector_meta(scheduler_output)

# Worker integration (context manager in model runner)
with kv_connector_context(scheduler_output):
    connector.start_load_kv(forward_context)   # Decode: load KV
    model.forward(...)                          # Compute
    connector.save_kv_layer(layer, kv, meta)   # Prefill: save KV
    connector.wait_for_save()                  # Prefill: block until done

Request States#

WAITING → WAITING_FOR_REMOTE_KVS → WAITING → RUNNING → FINISHED
                                    ↑
                          (KV transfer complete)

vllm-neuron Implementation#

Integration Strategy#

Neuron DI uses the upstream NixlConnector for all topologies, including hybrid-TP. EFA/LIBFABRIC is also upstream: it is selected via the kv_connector_extra_config {"backends": ["LIBFABRIC"]} option.

Neuron-Specific DI Features#

Hybrid TP (Asymmetric Parallelism)#

Different tensor parallelism on prefill vs decode:

Prefill: TP=4 (fewer cores, optimized for compute-heavy prompt processing)
Decode:  TP=8 (more cores, optimized for memory-bandwidth-bound generation)

The connector handles the head-count mismatch by computing head_ratio and inverse_head_ratio during the NIXL handshake.

Component DP (Decode-Side Module Sharding)#

Component DP shards individual model components (attention, embedding, MLP, lm_head) across DP replicas differently:

{
  "neuron_config": {
    "attention_dp_size": 4,
    "embedding_dp_size": 2,
    "mlp_dp_size": 4,
    "lm_head_dp_size": 2
  }
}

  • Only supported on the decode server (kv_role: "kv_consumer")

  • Requires --data-parallel-size >= 2

  • Forces the MoE engine path (DPEngineCoreProc) to maintain cross-DP process groups

Speculative Decoding with DI (Eagle3)#

Eagle3 speculative decoding works with DI:

  • Prefill server produces the target KV cache and runs spec decoding to generate the draft KV cache, so both are available to the decode side.

  • Decode server runs the Eagle3 draft model + verification with rejection sampling.

  • Tested with DP+EP topologies (e.g., dp4tp8_dp4tp8ep32).

Bi-directional KV Transfer (D→P)#

For multi-turn conversations, the decode server sends KV cache back to the prefill server so the next turn avoids recomputation:

export VLLM_NEURON_BIDIRECTIONAL_KV=1

The connector/scheduler handles the D→P direction (is_bidirectional_kv_xfer_enabled):

  • get_num_new_matched_tokens: Reports remote_block_ids from D as externally cached tokens.

  • update_state_after_alloc: Populates _reqs_need_recv for the prefill worker to pull KV from decode.

  • request_finished: Returns kv_transfer_params with remote_block_ids for the proxy to forward to the next prefill request.

Supports a remote_pull_threshold — minimum token count before pulling from remote vs recomputing locally.

Streaming KV Transfer Params#

When the decode server streams its response (SSE), kv_transfer_params are emitted in the final streaming chunk; the proxy reads them to relay decode block IDs back to the prefill server (used by bidirectional KV).

This is part of upstream vLLM (≥ 0.21.0).

Role-Specific Warmup#

Environment variables to skip compilation for the unused phase:

VLLM_NEURON_SKIP_PREFILL_WARMUP=1  # Decode-only server
VLLM_NEURON_SKIP_DECODE_WARMUP=1   # Prefill-only server

Saves compilation time — no need to compile prefill NEFFs on a decode-only server.

Configuration Reference#

Prefill Server#

vllm serve $MODEL \
    --tensor-parallel-size 8 \
    --kv-transfer-config '{
        "kv_connector": "NixlConnector",
        "kv_role": "kv_producer",
        "kv_buffer_device": "cuda",
        "kv_connector_extra_config": {"backends": ["LIBFABRIC"]}
    }'

Decode Server#

--data-parallel-size / --enable-expert-parallel are only needed for DP+EP (MoE) decode; a basic decode server omits them.

vllm serve $MODEL \
    --tensor-parallel-size 8 \
    --kv-transfer-config '{
        "kv_connector": "NixlConnector",
        "kv_role": "kv_consumer",
        "kv_buffer_device": "cuda",
        "kv_connector_extra_config": {"backends": ["LIBFABRIC"]}
    }'

# DP+EP (MoE) decode adds:
#   --data-parallel-size 4 --enable-expert-parallel

Environment Variables#

Variable

Default

Description

VLLM_NIXL_SIDE_CHANNEL_HOST

127.0.0.1

NIXL metadata exchange bind address

VLLM_NIXL_SIDE_CHANNEL_PORT

5557

NIXL metadata exchange port

VLLM_NEURON_SKIP_PREFILL_WARMUP

0

Skip prefill NEFF compilation (decode-only)

VLLM_NEURON_SKIP_DECODE_WARMUP

0

Skip decode NEFF compilation (prefill-only)

VLLM_NEURON_BIDIRECTIONAL_KV

0

Enable D→P KV transfer for multi-turn

NEURON_VISIBLE_DEVICES

all

Core partitioning for single-node DI

Proxy Server#

The proxy is a CPU-only load balancer (no Neuron cores). A reference implementation and launch script are in examples/vllm_neuron/vllm/disaggregated_inference/ — toy_proxy_server.py (round-robin over prefill/decode instances) and server.sh:

python examples/vllm_neuron/vllm/disaggregated_inference/toy_proxy_server.py \
    --prefiller-hosts <P_HOST> --prefiller-ports 8100 \
    --decoder-hosts   <D_HOST> --decoder-ports   8200 \
    --port 8000

Supported Topologies#

Topology

Description

Single-Node

Multi-Node

xPyD

General multi-prefill multi-decode (covers 1P1D, 2P1D, 1P2D, …)

✓¹

✓

Hybrid TP

Different TP per role

✓

✓

DP+EP

Data parallel with expert parallel on decode

✓

✓

Spec decode (Eagle3) + DI

Draft model on decode side

✓

✓

¹ Single-node xPyD is supported as long as the device has enough cores and memory for all prefill + decode instances.

Upstream vs. Plugin Responsibility#

Capability

Upstream vLLM

vllm-neuron

KV Connector framework (base interfaces)

✓

—

NixlConnector (RDMA transfer)

✓

—

LIBFABRIC/EFA backend

✓

—

Hybrid TP head splitting

✓

—

Bi-directional KV (D→P)

✓

✓ (config/enablement)

Streaming kv_transfer_params in SSE

✓ (≥ 0.21.0)

—

Component DP with DI

—

✓

Role-specific warmup skip

—

✓

Prefill/decode batch separation

—

✓ (NeuronScheduler)

Bucket padding for compiled NEFFs

—

✓ (NeuronScheduler)

Proxy server (toy)

✓ (examples)

✓ (examples)

Production orchestration

—

AIBrix, production-stack

Data Flow (Detailed)#

Prefill Server (data flow)#

Request arrives → NeuronScheduler
  → Pads to prefill bucket
  → Allocates KV blocks
  → NeuronModelRunner.execute_model()
    → _prepare_model_input() (input_ids, positions, slot_mapping)
    → with maybe_get_kv_connector_output():
        → model.forward()  [compiled NEFF]
        → connector.save_kv_layer() per attention layer
        → connector.wait_for_save()
    → sample one token (DISCARDED — only drives KV production; the first
      output token is produced by the decode server)
  → Return response (with kv_transfer_params for proxy)

Decode Server (data flow)#

Request arrives (from proxy) → NeuronScheduler
  → Request enters WAITING_FOR_REMOTE_KVS state
  → connector.start_load_kv() triggers NIXL RDMA READ
  → Polls until transfer complete
  → Request moves to RUNNING
  → NeuronModelRunner.execute_model()
    → _prepare_model_input() (uses transferred KV blocks)
    → model.forward() [compiled decode NEFF]
    → sample token
  → Stream tokens back through proxy

NIXL Transfer#

Decode worker                              Prefill worker
────────────                               ─────────────
1. register_kv_caches()                    1. register_kv_caches()
   → NIXL memory region registration         → NIXL memory region registration
   → Create block descriptors                → Create block descriptors

2. Handshake (one-time per P-D pair)
   ← Exchange: agent handles, block lens, base addrs

3. Transfer (per request)
   → Map block_ids to descriptor_ids
   → Issue RDMA READ (NIXL)                  → Device memory read by D
   → Poll for completion
   → Notify P to free blocks                 → Free KV blocks

Performance Characteristics#

Why DI Helps on Neuron#

  1. Prefill and decode use different compiled NEFFs — they can never share a batch anyway (unlike GPU which can mix). DI makes this explicit and eliminates scheduling contention.

  2. Dedicated decode instances maintain steady ITL — no prefill compilation or execution interrupts decode.

  3. Asymmetric parallelism exploits Neuron’s core partitioning — prefill can use fewer cores with higher utilization, decode can use more cores with DP for throughput.

Overheads#

Overhead

Source

Mitigation

TTFT increase

KV transfer latency (RDMA)

Write mode overlaps transfer with prefill

First-request latency

NIXL handshake (one-time)

Cached per P-D pair

Memory

Decode holds KV for all in-flight requests

Tune --max-num-seqs

Compilation time

Each role compiles only its buckets

VLLM_NEURON_SKIP_PREFILL_WARMUP / VLLM_NEURON_SKIP_DECODE_WARMUP

Future Directions#

  • Production orchestration — AIBrix and production-stack integration for autoscaling, health checks, and SLO-aware routing.

This document is relevant for: Trn2, Trn3