pub struct Mamba3SingleSsdCacheConfig {
pub batch: usize,
pub state_rank: usize,
pub per_head_dim: usize,
pub nheads: usize,
pub mimo_rank: usize,
pub num_rope_angles: usize,
pub rotation: RotationKind,
pub num_quat_blocks: usize,
}Expand description
Configuration / factory for a single Mamba3SingleSsdCache.
Fields§
§batch: usizeBatch size.
state_rank: usizeState rank.
per_head_dim: usizeHead dimension per_head_dim.
nheads: usizeNumber of SSM heads.
mimo_rank: usizeMIMO rank. 1 = SISO.
num_rope_angles: usizeNumber of RoPE angle pairs
(see crate::mamba3::double_ssd::cache::Mamba3DoubleSsdCacheConfig::num_rope_angles).
rotation: RotationKindWhich positional rotation the block uses (see
crate::mamba3::double_ssd::cache::Mamba3DoubleSsdCacheConfig::rotation).
num_quat_blocks: usizeNumber of quaternion blocks (rope_dim / 4); only used for
RotationKind::Quaternion4D.
Implementations§
Source§impl Mamba3SingleSsdCacheConfig
impl Mamba3SingleSsdCacheConfig
Sourcepub fn new(batch: usize, nheads: usize, num_rope_angles: usize) -> Self
pub fn new(batch: usize, nheads: usize, num_rope_angles: usize) -> Self
Create a new instance of the config.
§Arguments
§Required Arguments
§batch
Batch size.
§nheads
Number of SSM heads.
§num_rope_angles
Number of RoPE angle pairs
(see crate::mamba3::double_ssd::cache::Mamba3DoubleSsdCacheConfig::num_rope_angles).
§Default Arguments
§state_rank
State rank.
- Defaults to
128
§per_head_dim
Head dimension per_head_dim.
- Defaults to
64
§mimo_rank
MIMO rank. 1 = SISO.
- Defaults to
1
§rotation
Which positional rotation the block uses (see
crate::mamba3::double_ssd::cache::Mamba3DoubleSsdCacheConfig::rotation).
- Defaults to
"crate::mamba3::rotation::RotationKind::Complex2D"
§num_quat_blocks
Number of quaternion blocks (rope_dim / 4); only used for
RotationKind::Quaternion4D.
- Defaults to
1
Source§impl Mamba3SingleSsdCacheConfig
impl Mamba3SingleSsdCacheConfig
Sourcepub fn with_state_rank(self, state_rank: usize) -> Self
pub fn with_state_rank(self, state_rank: usize) -> Self
Sourcepub fn with_per_head_dim(self, per_head_dim: usize) -> Self
pub fn with_per_head_dim(self, per_head_dim: usize) -> Self
Sourcepub fn with_mimo_rank(self, mimo_rank: usize) -> Self
pub fn with_mimo_rank(self, mimo_rank: usize) -> Self
Sourcepub fn with_rotation(self, rotation: RotationKind) -> Self
pub fn with_rotation(self, rotation: RotationKind) -> Self
Sets the value for the field rotation.
Which positional rotation the block uses (see
crate::mamba3::double_ssd::cache::Mamba3DoubleSsdCacheConfig::rotation).
- Defaults to
"crate::mamba3::rotation::RotationKind::Complex2D"
Sourcepub fn with_num_quat_blocks(self, num_quat_blocks: usize) -> Self
pub fn with_num_quat_blocks(self, num_quat_blocks: usize) -> Self
Sets the value for the field num_quat_blocks.
Number of quaternion blocks (rope_dim / 4); only used for
RotationKind::Quaternion4D.
- Defaults to
1
Source§impl Mamba3SingleSsdCacheConfig
impl Mamba3SingleSsdCacheConfig
Sourcepub fn new_from_block_config(batch: usize, block_config: Mamba3Config) -> Self
pub fn new_from_block_config(batch: usize, block_config: Mamba3Config) -> Self
Derive cache shapes from a Mamba-3 block configuration plus a batch size.
Sourcepub fn init(&self, device: &Device) -> Mamba3SingleSsdCache
pub fn init(&self, device: &Device) -> Mamba3SingleSsdCache
Allocate zero/identity-initialised cache tensors on device.