Is this page helpful?

`core.inference.disaggregation.ssm_reshard`#

Heterogeneous TP/PP reshard of SSM boundary-snapshot state between prefill and decode shard layouts (the SSM analog of attention KV resharding).

A snapshot’s conv state packs three channel bands, [x | B | C], on one axis: x is head-sharded (d_inner) and B/C are group-sharded (ngroups * d_state). The recurrent state is head-sharded. plan_ssm_reshard emits one transfer per (src rank, dst rank, global layer, band) whose layer and channel ranges overlap; both sides compute the same plan from the same layout lists, so the send and receive orders match.

Module Contents#

Classes#

`SSMStateDims`	The model’s global (unsharded) SSM structural dimensions.
`SSMShardLayout`	One rank’s SSM-state ownership: which global layers and TP rank, plus the model’s structural dims. Per-rank local sizes follow by dividing by tp_size.
`SSMReshardTransfer`	One sub-block move of the snapshot reshard.

Functions#

plan_ssm_reshard

Plan the conv/recurrent sub-block moves from the prefill (src) layouts to the decode (dst) layouts: one transfer per (src rank, dst rank, global layer, band) where both the layer ranges and the channel ranges overlap.

Data#

_CONV_BANDS

API#

core.inference.disaggregation.ssm_reshard._CONV_BANDS#: (‘x’, ‘B’, ‘C’)

class core.inference.disaggregation.ssm_reshard.SSMStateDims#

The model’s global (unsharded) SSM structural dimensions.

Carried as one unit so there is a single source and the dims cannot drift apart. The producer should read them from the model config rather than deriving them from tensor shapes. TP shards nheads/ngroups; the rest are unsharded.

nheads: int#: None

headdim: int#: None

d_state: int#: None

ngroups: int#: None

d_conv: int#: None

class core.inference.disaggregation.ssm_reshard.SSMShardLayout#

One rank’s SSM-state ownership: which global layers and TP rank, plus the model’s structural dims. Per-rank local sizes follow by dividing by tp_size.

global_rank: int#: None

tp_size: int#: None

tp_rank: int#: None

layer_start: int#: None

num_layers: int#: None

dims: core.inference.disaggregation.ssm_reshard.SSMStateDims#: None

__post_init__() → None#

property d_inner: int#: Global inner dimension (nheads * headdim).

property nheads_local: int#: SSM heads held by this TP rank.

property d_inner_local: int#: Local inner dimension for this TP rank.

property ngroups_local: int#: B/C groups held by this TP rank.

property conv_dim_local: int#: Total local conv channel width (x + B + C bands).

shard_key() → Tuple[int, int]#: The SSM shard this rank holds: (tp_rank, layer_start). Ranks sharing a key hold identical state (e.g. EP/DP replicas).

layer_range() → Tuple[int, int]#: Global SSM-layer range [lo, hi) owned by this rank.

band(name: str) → Tuple[int, int, int]#

Return (global_total, local_size, local_offset) for a band.

local_offset is the band’s start on the local conv channel axis; for the “recurrent” (head) band it is the start on the local head axis (always 0, heads are the whole tensor).

class core.inference.disaggregation.ssm_reshard.SSMReshardTransfer#

One sub-block move of the snapshot reshard.

band is “x”/”B”/”C” (conv channel axis) or “recurrent” (head axis). src_layer/dst_layer are local layer indices on each side; _lo/_hi are the local channel or head slice bounds.

src_rank: int#: None

dst_rank: int#: None

band: str#: None

global_layer: int#: None

src_layer: int#: None

dst_layer: int#: None

src_lo: int#: None

src_hi: int#: None

dst_lo: int#: None

dst_hi: int#: None

property is_conv: bool#: True if this transfer targets the conv state; False for recurrent.

core.inference.disaggregation.ssm_reshard.plan_ssm_reshard( src_layouts: List[core.inference.disaggregation.ssm_reshard.SSMShardLayout], dst_layouts: List[core.inference.disaggregation.ssm_reshard.SSMShardLayout], ) → List[core.inference.disaggregation.ssm_reshard.SSMReshardTransfer]#

Ranks sharing (tp_rank, layer_start) hold identical SSM state (e.g. EP/DP replicas), so each shard is sourced from exactly one of them, the smallest global_rank. Deterministic given the layout lists, so both sides enumerate the same transfers in the same order.

core.inference.disaggregation.ssm_reshard#