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).

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]

The attention shard this rank holds: (tp_rank, pp_rank). Ranks sharing a key hold identical KV (EP/ETP replicas of it).

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 shard (tp_rank, pp_rank) may be held by several source ranks. We source each shard from exactly one of them – the smallest global_rank – which avoids duplicate sends and is independent of how EP/ETP map onto ranks.