mamba2_surgery

Generic packed-tensor surgery for Mamba2-style mixers.

Classes

Mamba2TensorLayout

Logical parameter names and geometry for the common Mamba2 layout.

Functions

mamba2_projected_prefix_mask

Mask for [gate, x, B, C, dt] emitted by a Mamba2 in projection.

slice_mamba2_state_dict

Return a shallow state-dict copy with Mamba logical prefixes sliced.

sort_mamba2_state_dict

Sort Mamba heads, per-head dimensions, and grouped state channels.
class Mamba2TensorLayout

Bases: object

Logical parameter names and geometry for the common Mamba2 layout.

__init__(in_proj_key, out_proj_key, conv_weight_key, conv_bias_key, norm_key, a_log_key, d_key, dt_bias_key, num_heads, head_dim, num_groups, state_dim, in_proj_bias_key=None)
Parameters:

  • in_proj_key (str)

  • out_proj_key (str | None)

  • conv_weight_key (str | None)

  • conv_bias_key (str | None)

  • norm_key (str | None)

  • a_log_key (str | None)

  • d_key (str | None)

  • dt_bias_key (str | None)

  • num_heads (int)

  • head_dim (int)

  • num_groups (int)

  • state_dim (int)

  • in_proj_bias_key (str | None)

Return type:

None

a_log_key: str | None
property bias_key: str
conv_bias_key: str | None
property conv_size: int
conv_weight_key: str | None
d_key: str | None
dt_bias_key: str | None
property grouped_state_size: int
head_dim: int
in_proj_bias_key: str | None = None
in_proj_key: str
property in_proj_size: int
property inner_size: int
norm_key: str | None
num_groups: int
num_heads: int
out_proj_key: str | None
state_dim: int
validate(state_dict, *, require_in_proj=True)
Parameters:

  • state_dict (Mapping[str, Tensor])

  • require_in_proj (bool)

Return type:

None

mamba2_projected_prefix_mask(layout, *, target_heads=None, target_head_dim=None, target_state_dim=None, device=None)

Mask for [gate, x, B, C, dt] emitted by a Mamba2 in projection.

Parameters:

  • layout (Mamba2TensorLayout)

  • target_heads (int | None)

  • target_head_dim (int | None)

  • target_state_dim (int | None)

Return type:

Tensor

slice_mamba2_state_dict(state_dict, layout, *, target_heads=None, target_head_dim=None, target_state_dim=None)

Return a shallow state-dict copy with Mamba logical prefixes sliced.

Parameters:

  • state_dict (Mapping[str, Tensor])

  • layout (Mamba2TensorLayout)

  • target_heads (int | None)

  • target_head_dim (int | None)

  • target_state_dim (int | None)

Return type:

dict[str, Tensor]

sort_mamba2_state_dict(state_dict, layout, *, head_scores=None, head_dim_scores=None, state_scores=None)

Sort Mamba heads, per-head dimensions, and grouped state channels.

Parameters:

  • state_dict (Mapping[str, Tensor])

  • layout (Mamba2TensorLayout)

  • head_scores (Tensor | None)

  • head_dim_scores (Tensor | None)

  • state_scores (Tensor | None)

Return type:

tuple[dict[str, Tensor], dict[str, Tensor]]