This document is relevant for: Trn2, Trn3

DCP (Decode Context Parallelism)#

Overview#

DCP shards the KV cache sequence dimension across ranks, reducing per-rank memory for long contexts. Two modes:

  1. DCP Prefill (apply_prefill_dcp=True and DCP size greater than one): Replicates attention weights across DCP sub-groups and shards the input sequence during prefill. Each DCP rank processes different token chunks.

  2. DCP Decode (DCP size greater than one): Shards the KV cache within the DCP replica set during decode. It gathers Q across DCP peers, computes attention against each local KV shard, and combines the partial results with LSE correction.

DCP Prefill#

Configuration#

vllm serve <model> \
    --tensor-parallel-size 8 \
    --decode-context-parallel-size 4 \
    --cp-kv-cache-interleave-size 16 \
    --additional-config '{"neuron_config": {"apply_prefill_dcp": true}}'

Constraints#

  • TP % DCP == 0

  • apply_prefill_dcp requires DI (kv_role=kv_producer) and DCP > 1

  • cp_kv_cache_interleave_size must equal block_size for DI compatibility

Process Groups (TP=8, DCP=4)#

Full TP group: [0, 1, 2, 3, 4, 5, 6, 7]  (size 8)

dcp_tp_group (size TP/DCP = 2):
  [0, 1], [2, 3], [4, 5], [6, 7]
  → SP gather/scatter, QKV/O weight sharding

cp_kv_group (size DCP = 4):
  [0, 2, 4, 6], [1, 3, 5, 7]
  → KV AllGather (same heads, different tokens)

cp_rank = tp_rank // (TP/DCP):
  Rank 0,1 → cp_rank 0 (tokens in blocks 0,4,8,...)
  Rank 2,3 → cp_rank 1 (tokens in blocks 1,5,9,...)
  Rank 4,5 → cp_rank 2 (tokens in blocks 2,6,10,...)
  Rank 6,7 → cp_rank 3 (tokens in blocks 3,7,11,...)

Compute Pattern#

1. SP All-Gather (world_group): S/TP → S tokens
2. Interleave slice to owned positions: S → S/DCP tokens
3. QKV Projection (weights sharded across dcp_tp_group)
4. KV Cache Write (only local owned positions via slot_mapping)
5. Q AllGather (cp_kv_group): Q [Nh_q, S/DCP] → [Nh_q, S]
6. Unshuffle Q to global position order
7. Segmented attention: full_Q × local_KV (prior cache + current) with LSE correction
8. ReduceScatter (cp_kv_group): output [Nh_q, S] → [Nh_q, S/DCP]
9. O Projection + Reduce-Scatter (dcp_tp_group): S/DCP → S/TP

Weight Loading#

Component

Sharding

Rank

QKV, O

dcp_tp_group.world_size (TP/DCP)

dcp_tp_group.rank_in_group

MLP, Embedding, LM Head

Full TP

Full TP rank

DCP Decode#

Decode Configuration#

vllm serve <model> \
    --tensor-parallel-size 16 \
    --decode-context-parallel-size 2 \
    --cp-kv-cache-interleave-size 16

Decode Constraints#

  • tp > num_kv_heads (KV must be replicated across DCP group)

  • dcp <= tp // num_kv_heads

  • (num_q_heads // num_kv_heads) % dcp == 0

  • tp % dcp == 0

  • DCP cannot currently be combined with attention DP

Process Groups (TP=16, DCP=2)#

In this example, vLLM creates consecutive DCP rank pairs: [0,1], [2,3], ..., [14,15]. Each pair shares the same replicated KV head and splits the sequence.

dcp_rank = tp_rank % dcp_size
  Rank 0 → dcp_rank 0 (even blocks)
  Rank 1 → dcp_rank 1 (odd blocks)

Decode Compute Pattern#

1. QKV Projection (weights sharded by full TP)
2. AllGather Q across DCP group: q_heads_local → q_heads_local * dcp_size
3. Local attention: gathered Q against local KV shard (S/DCP tokens)
4. Extract partial LSE (logsumexp of local scores)
5. AllGather LSE across DCP group
6. Compute correction: weight_i = exp(local_lse - global_lse)
7. Apply correction to local attention output
8. ReduceScatter across DCP group: combine corrected outputs, scatter heads back

local_filled Computation#

The attention mask uses local_filled_slots to determine how many prior tokens are in this rank’s KV cache shard:

vblock = block_size * dcp_size
stride = interleave_size * dcp_size
rank_start = dcp_rank * interleave_size

local_filled = floor(N / vblock) * block_size
remaining = N - floor(N / vblock) * vblock
local_filled += floor(remaining / stride) * interleave_size
leftover = remaining - floor(remaining / stride) * stride
local_filled += clamp(clamp(leftover - rank_start, min=0), max=interleave_size)

Active Token Handling and dcp_active_mask#

During decode, all DCP ranks project the same input hidden states and produce the same K/V for the active token. The cache write uses slot_mapping:

  • Owning rank (slot_mapping >= 0): writes to the correct cache slot.

  • Non-owning ranks (slot_mapping = -1): writes to a garbage slot (last block).

After AllGather Q, each rank computes attention over its local KV shard. Without masking, every rank would include the active token’s K/V in its local attention — the LSE correction would then count the active token dcp_size times (once per rank).

dcp_active_mask prevents this. The attention mask has two parts:

  • prior_mask: slots [0..local_filled) — prior tokens this rank owns

  • active_slots: the current decode token’s slot, gated by dcp_active_mask

Only the owning rank sets dcp_active_mask = 1; others set it to 0. This ensures the active token contributes exactly once to the combined output.

pos_in_vblock = N - floor(N / vblock) * vblock
owner = floor(pos_in_vblock / interleave_size) % dcp_size
dcp_active_mask = (owner == dcp_rank).float()

cp_rank: Prefill vs Decode#

Both modes use the same --decode-context-parallel-size flag and the same slot mapping function, but the group topology differs, producing different cp_rank formulas:

DCP Prefill: cp_rank = tp_rank // (TP / DCP)#

DCP sub-divides TP into DCP token groups of TP/DCP consecutive ranks each. Ranks within the same group share the same tokens but have different heads.

TP=8, DCP=4:
  Group 0: ranks [0, 1]  → cp_rank 0 (token group 0)
  Group 1: ranks [2, 3]  → cp_rank 1 (token group 1)
  Group 2: ranks [4, 5]  → cp_rank 2 (token group 2)
  Group 3: ranks [6, 7]  → cp_rank 3 (token group 3)

Semantic: "which token chunk do I own?"
All ranks with the same cp_rank store identical token positions (different heads).

DCP Decode: cp_rank = tp_rank % DCP#

DCP forms TP/DCP independent replica sets of DCP consecutive ranks each. Ranks within the same set share the same (replicated) KV heads but own different tokens.

TP=16, DCP=2:
  Replica set 0: ranks [0, 1]   → cp_rank 0, 1
  Replica set 1: ranks [2, 3]   → cp_rank 0, 1
  ...
  Replica set 7: ranks [14, 15] → cp_rank 0, 1

Semantic: "which interleave position do I own within my replica set?"
cp_rank 0 owns even blocks, cp_rank 1 owns odd blocks.

Why the Difference#

In prefill, DCP replicates weights and shards sequence — the “group” is the set of ranks that process different heads for the same tokens. The group index (cp_rank) identifies the token chunk.

In decode, DCP replicates KV heads and shards sequence — the “group” is the replica set where each rank owns a different sequence slice. The position within the group (cp_rank) identifies which interleave slice.

Both formulas produce the correct cp_rank for _compute_slot_mapping_cpu: the slot mapping assigns positions where (pos_in_vblock // interleave_size) % cp_world_size == cp_rank to the local rank, and -1 to all others.

Slot Mapping and KV Cache#

Both DCP prefill and decode use CP-aware slot mapping:

_compute_slot_mapping_cpu(
    block_table, slot_mapping, positions, req_indices,
    block_size,
    cp_world_size=dcp_size,
    cp_rank=cp_rank,
    cp_kv_cache_interleave_size=interleave_size,
)

With interleave_size=block_size=16 and DCP=4:

  • Block 0 (positions 0-15) → cp_rank 0

  • Block 1 (positions 16-31) → cp_rank 1

  • Block 2 (positions 32-47) → cp_rank 2

  • Block 3 (positions 48-63) → cp_rank 3

  • Block 4 (positions 64-79) → cp_rank 0 (wraps)

Non-local positions get slot_mapping = -1.

Disaggregated Inference (NIXL Transfer)#

NeuronNixlConnector#

The NIXL connector is NeuronNixlConnector, a standalone connector specified in --kv-transfer-config:

"kv_connector": "NeuronNixlConnector"

The platform auto-injects kv_connector_module_path — users only set the name. Required when using DCP with DI (validated at startup).

Internally uses NeuronNixlConnectorWorker (subclass of NixlConnectorWorker) with DCP-specific overrides:

  • register_kv_caches: Stores local DCP rank/size; on prefill, re-encodes metadata with NeuronNixlAgentMetadata (includes dcp_size field); patches tp_ratio for match case

  • _validate_remote_agent_handshake: Skips strict validation for DCP prefill engines

  • add_remote_agent: Routes to split (head offset), merge (smaller descriptors), or match (passthrough)

  • _nixl_handshake: Connects to ALL remote ranks when TP differs

  • _read_blocks_for_req: Routes to DCP-specific block reading logic

NeuronNixlAgentMetadata#

Extends NixlAgentMetadata with a dcp_size: int = 0 field. The prefill advertises its DCP degree in the handshake metadata. The decode reads it during detection to compute correct remote rank mappings.

Transfer Topologies#

DCP Prefill → Standard Decode#

Each decode rank reads from P_DCP prefill ranks (all cp_ranks with matching head):

P:TP=4/DCP=2 → D:TP=2:
  Decode rank 0 (heads 0-3) ← Prefill ranks [0, 2] (cp_ranks 0,1, heads 0-3)
  Decode rank 1 (heads 4-7) ← Prefill ranks [1, 3] (cp_ranks 0,1, heads 4-7)

Local blocks are interleaved across remote ranks by token group.

DCP Prefill → DCP Decode (Same DCP Degree)#

1:1 cp_rank mapping. Each decode dcp_rank reads from exactly one prefill cp_rank:

P:TP=32/DCP=2 → D:TP=16/DCP=2:
  Decode rank 0 (dcp_rank=0) ← Prefill rank 0 (cp_rank=0)
  Decode rank 1 (dcp_rank=1) ← Prefill rank 16 (cp_rank=1)

DCP Prefill → DCP Decode (Different DCP)#

When P_DCP > D_DCP, each decode dcp_rank reads from P_DCP/D_DCP prefill cp_ranks:

P:TP=32/DCP=4 → D:TP=16/DCP=2:
  Decode rank 0 (dcp_rank=0) ← Prefill ranks [0, 16] (cp_ranks 0,2)
  Decode rank 1 (dcp_rank=1) ← Prefill ranks [8, 24] (cp_ranks 1,3)

Filter formula: decode dcp_rank r reads from prefill cp_ranks {r, r+D_DCP, r+2*D_DCP, ...}.

Standard Prefill → DCP Decode#

When prefill has no DCP (standard TP, same TP size), the decode filters remote blocks to its interleaved subset:

P:TP=16 → D:TP=16/DCP=2:
  Decode rank 0 (dcp_rank=0): keeps even-indexed remote blocks
  Decode rank 1 (dcp_rank=1): keeps odd-indexed remote blocks

Head Splitting / Merging#

Topology

Condition

Action

Matching

remote_block_len == local_block_len

No offset, standard interleave

Splitting

remote_block_len > local_block_len

Head offset on remote descriptors

Merging

remote_block_len < local_block_len

Custom remote descriptors, split local handles

Remote Rank Computation#

# num_target_ranks = tp_ratio * head_ratio (accounts for all topologies)
tp_ratio = P_TP // D_TP
head_ratio = max(1, remote_block_len // local_block_len)
num_target_ranks = tp_ratio * head_ratio

if inverse_head_ratio > 1:
    # Merge: multiple head positions × cp_ranks
    p_dcp = num_target_ranks // inverse_head_ratio
    tp_pair_size = P_TP // p_dcp
    remote_ranks = [ep + cp * tp_pair_size
                    for hp in range(inverse_head_ratio)
                    for cp in range(p_dcp)
                    for ep in [decode_tp_rank * inverse_head_ratio + hp]]
else:
    # Split/Match: one head position × all cp_ranks
    stride = P_TP // num_target_ranks
    effective_pos = decode_tp_rank * (P_TP // num_target_ranks) // D_TP
    remote_ranks = [effective_pos + i * stride for i in range(num_target_ranks)]

# Filter for decode DCP alignment
if D_DCP > 1:
    remote_ranks = [r for i, r in enumerate(remote_ranks) if i % D_DCP == dcp_rank]

Notification Handling#

When a decode DCP rank has no blocks to transfer (short sequence), it still sends a notification to the prefill rank via _read_blocks(local_block_ids=[]). If no transfers are issued at all, an empty entry is added to _recving_transfers so get_finished() can complete the request.

Model Runner (neuron_model_runner.py)#

Slot Mapping for DCP Prefill#

The model receives S/DCP owned tokens after the interleave slice in forward. The slot_mapping is extracted from the full S-sized mapping by selecting owned positions using the same interleave pattern:

slot_mapping = slot_mapping.cpu().view(S // (W * I), W, I)[:, R, :].contiguous().reshape(-1).to(device)

This runs in _build_attention_metadata (runtime) and _build_warmup_attention_metadata (warmup) to ensure consistent shapes. The output is always S/DCP entries.

Segmented Attention CP (segmented_attention_cp)#

Purpose#

Handles DCP prefill attention for both first-chunk and multi-chunk (chunked prefill with prior). Each rank has:

  • Full Q (AllGathered, unshuffled to global order)

  • Local KV: prior from cache (S_prior/DCP tokens) + current from projection (S_chunk/DCP tokens)

API#

segmented_attention_cp(
    q,              # [Nh_q, S_total, Dh] — AllGathered Q
    k_local,        # [Nh_kv, S_local, Dh] — current chunk's local KV
    v_local,
    k_cache,        # paged prior cache
    v_cache,
    block_tables,   # [1, max_blocks_per_seq]
    prior_tokens,   # [[local_prior_count]] (0 on first chunk)
    block_size,
    cp_rank, cp_world_size, cp_kv_cache_interleave_size,
    cp_group,       # for AllGather LSE + ReduceScatter
    scale, tp_q, tp_out,
)

Algorithm#

  1. Read prior KV from cache (static padded shape, validity masked by slot index)

  2. Concatenate prior + current local KV

  3. Build causal mask using actual global positions:

    • Prior slot s → global pos (s // I) * (W * I) + R * I + (s % I)

    • Current token j → global pos prior_global + (j // I) * (W * I) + R * I + (j % I)

    • Q position i → global pos prior_global + i

  4. Compute partial attention (Q × local_KV) + extract LSE

  5. AllGather LSE across cp_group → global LSE

  6. Weight local output by exp(local_lse - global_lse)

  7. ReduceScatter weighted output (dim=1) → each rank gets its S/DCP output slice

When prior_tokens = 0, all prior slots are masked out and contribute nothing. No dynamic branches — the mask handles it statically.

This document is relevant for: Trn2, Trn3