core.inference.disaggregation.kv_reshard

TP/PP/EP/ETP KV-shard layouts and the range-intersection reshard planner.

Module Contents

Classes

KVShardLayout	A worker’s KV-cache ownership within the global model.
KVReshardTransfer	One sub-block exchange between a (src, dst) rank pair.

Functions

plan_kv_reshard

Full reshard plan: every sub-block that must move src -> dst.

API

class core.inference.disaggregation.kv_reshard.KVShardLayout

A worker’s KV-cache ownership within the global model.

num_layers / num_heads are the global attention layer count and KV-head count (for GQA, the number of KV heads). global_rank is the worker’s torch rank (used as the transport peer id).

When tp_size exceeds num_heads, each KV head is replicated on tp_size // num_heads consecutive TP ranks and every rank holds one head, matching how Attention partitions GQA heads.

num_layers: int

None

num_heads: int

None

tp_size: int

None

tp_rank: int

None

pp_size: int

None

pp_rank: int

None

global_rank: int

None

ep_size: int

1

ep_rank: int

0

etp_size: int

1

etp_rank: int

0

layer_start: Optional[int]

None

num_local_layers: Optional[int]

None

__post_init__() → None

kv_shard_key() → Tuple[int, int]

This rank’s attention position: (tp_rank, pp_rank). Ranks sharing a key are EP/ETP replicas holding identical KV.

kv_replicas() → int

Number of TP ranks holding each of this rank’s KV heads.

kv_replica_index() → int

This rank’s index among the TP ranks replicating its KV heads.

layer_range() → Tuple[int, int]

Global attention-layer range [lo, hi) owned by this rank.

head_range() → Tuple[int, int]

Global KV-head range [lo, hi) owned by this rank.

local_num_layers() → int

Number of attention layers held locally by this rank.

local_num_heads() → int

Number of KV heads held locally by this rank.

class core.inference.disaggregation.kv_reshard.KVReshardTransfer

One sub-block exchange between a (src, dst) rank pair.

Global coords identify the intersection; the local-slice helpers convert to each side’s buffer offsets. There is at most one transfer per (src, dst) pair (each owns a contiguous rectangle, so the intersection is a single rectangle).

src_rank: int

None

dst_rank: int

None

global_layer_lo: int

None

global_layer_hi: int

None

global_head_lo: int

None

global_head_hi: int

None

src_layer_slice(

src: core.inference.disaggregation.kv_reshard.KVShardLayout,

) → slice

Local layer slice on the source side for this transfer.

src_head_slice(

src: core.inference.disaggregation.kv_reshard.KVShardLayout,

) → slice

Local KV-head slice on the source side for this transfer.

dst_layer_slice(

dst: core.inference.disaggregation.kv_reshard.KVShardLayout,

) → slice

Local layer slice on the destination side for this transfer.

dst_head_slice(

dst: core.inference.disaggregation.kv_reshard.KVShardLayout,

) → slice

Local KV-head slice on the destination side for this transfer.

core.inference.disaggregation.kv_reshard.plan_kv_reshard(

srcs: List[core.inference.disaggregation.kv_reshard.KVShardLayout],

dsts: List[core.inference.disaggregation.kv_reshard.KVShardLayout],

) → List[core.inference.disaggregation.kv_reshard.KVReshardTransfer]

Full reshard plan: every sub-block that must move src -> dst.

Both sides compute the same plan from the same layouts and filter to their own rank (transfers_for_src / transfers_for_dst).

KV is replicated across the EP and ETP dimensions, so each attention position (tp_rank, pp_rank) may be held by several source ranks. We source each position from exactly one of them – the smallest global_rank – which avoids duplicate sends and is independent of how EP/ETP map onto ranks.

When tp_size > num_heads several TP ranks also hold the same KV heads. Each destination reads from the source replica whose index matches its own (modulo the source replica count), which spreads reads over the replicas and maps a homogeneous layout onto itself. The choice depends only on the (src, dst) pair, so a side that plans against a single peer layout computes the same transfers.