Is this page helpful?

`core.inference.disaggregation.mamba_reshard`#

Heterogeneous TP/PP reshard of Mamba conv/ssm state between prefill and decode shard layouts (the Mamba analog of the attention KV reshard).

Module Contents#

Classes#

`MambaStateDims`	The model’s (global, unsharded) Mamba structural dims.
`MambaShardLayout`	One rank’s Mamba-state ownership: which global layers + TP rank, plus the model’s structural dims (:class:`MambaStateDims`). Per-rank locals follow by dividing by `tp_size`.
`MambaReshardTransfer`	One sub-block move for the reshard.

Functions#

plan_mamba_reshard

Plan the conv/ssm 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.mamba_reshard._CONV_BANDS#: (‘x’, ‘B’, ‘C’)

class core.inference.disaggregation.mamba_reshard.MambaStateDims#

The model’s (global, unsharded) Mamba structural dims.

These belong to the MambaMixer / model config – carried as one unit (rather than loose constants spread across the layout) so there’s a single source and they can’t drift apart. The producer should read them straight from the model config (e.g. ngroups = config.mamba_num_groups) rather than reverse-deriving 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.mamba_reshard.MambaShardLayout#

One rank’s Mamba-state ownership: which global layers + TP rank, plus the model’s structural dims (:class:MambaStateDims). Per-rank locals 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.mamba_reshard.MambaStateDims#: None

__post_init__() → None#

property nheads: int#: Global (unsharded) number of Mamba heads.

property headdim: int#: Dimension of each Mamba head.

property d_state: int#: SSM state size per head.

property ngroups: int#: Global (unsharded) number of B/C groups.

property d_conv: int#: Convolution kernel width.

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

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

property nheads_local: int#: Number of Mamba heads held by this TP rank.

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

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

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

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

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

(global_total, local_size, conv_local_offset) for a band.

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

class core.inference.disaggregation.mamba_reshard.MambaReshardTransfer#

One sub-block move for the reshard.

band is "x"/"B"/"C" (conv channel axis) or "ssm" (head axis). src_layer/dst_layer are local layer indices on each side; *_lo/*_hi are the local channel/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 ssm.

core.inference.disaggregation.mamba_reshard.plan_mamba_reshard( src_layouts: List[core.inference.disaggregation.mamba_reshard.MambaShardLayout], dst_layouts: List[core.inference.disaggregation.mamba_reshard.MambaShardLayout], ) → List[core.inference.disaggregation.mamba_reshard.MambaReshardTransfer]#: Plan the conv/ssm 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.

core.inference.disaggregation.mamba_reshard#