diff --git a/mamba_ssm/modules/mamba3.py b/mamba_ssm/modules/mamba3.py index 34c119c2c..6ef771a41 100644 --- a/mamba_ssm/modules/mamba3.py +++ b/mamba_ssm/modules/mamba3.py @@ -207,6 +207,11 @@ def forward(self, u, seq_idx=None, cu_seqlens=None, inference_params=None): # Apply Mamba-3 kernel if self.is_mimo: + input_states = ( + (angle_dt_state, ssm_state, k_state, v_state) + if ssm_state is not None + else None + ) y = mamba3_mimo_combined( Q=C, K=B, @@ -226,6 +231,7 @@ def forward(self, u, seq_idx=None, cu_seqlens=None, inference_params=None): rotary_dim_divisor=self.rotary_dim_divisor, dtype=x.dtype, return_state=ssm_state is not None, + Input_States=input_states, cu_seqlens=cu_seqlens, fuse_pregate_headwise_rms_norm=self.fuse_pregate_headwise_norm, outproj_norm_weight=self.norm.weight if self.fuse_pregate_headwise_norm else None, @@ -246,6 +252,11 @@ def forward(self, u, seq_idx=None, cu_seqlens=None, inference_params=None): y = torch.einsum("blrhp,hrp->blhp", y, self.mimo_o) y = rearrange(y, "b l h p -> b l (h p)") else: + input_states = ( + (angle_dt_state, ssm_state, k_state.squeeze(1), v_state) + if ssm_state is not None + else None + ) y = mamba3_siso_combined( Q=C.squeeze(2), K=B.squeeze(2), @@ -259,7 +270,7 @@ def forward(self, u, seq_idx=None, cu_seqlens=None, inference_params=None): D=self.D, Z=z if not self.is_outproj_norm else None, chunk_size=self.chunk_size, - Input_States=None, + Input_States=input_states, return_final_states=ssm_state is not None, cu_seqlens=cu_seqlens, ) diff --git a/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py b/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py index aabbda1f4..c94ffaa67 100644 --- a/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py +++ b/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py @@ -48,6 +48,10 @@ def forward( Angles: Tensor, D: Tensor, Z: Tensor, + Input_Angle_State: Optional[Tensor], + Input_SSM_State: Optional[Tensor], + Input_K_State: Optional[Tensor], + Input_V_State: Optional[Tensor], chunk_size: int, rotary_dim_divisor: int, dtype: torch.dtype, @@ -62,10 +66,26 @@ def forward( ctx.dtype = dtype ctx.fuse_pregate_headwise_rms_norm = fuse_pregate_headwise_rms_norm ctx.outproj_norm_eps = outproj_norm_eps - (Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z) = tuple( + ctx.has_input_state = Input_SSM_State is not None + all_states_present = ( + Input_Angle_State is not None + and Input_SSM_State is not None + and Input_K_State is not None + and Input_V_State is not None + ) + all_states_absent = ( + Input_Angle_State is None + and Input_SSM_State is None + and Input_K_State is None + and Input_V_State is None + ) + assert all_states_present or all_states_absent, "Input states must be provided together or all be None." + (Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z, + Input_Angle_State, Input_SSM_State, Input_K_State, Input_V_State) = tuple( t.contiguous() if t is not None else None for t in ( Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z, + Input_Angle_State, Input_SSM_State, Input_K_State, Input_V_State, ) ) # Kernels require cu_seqlens as int32; torch.cumsum promotes int32→int64 on CUDA @@ -75,6 +95,7 @@ def forward( # Compute cumulative angles (varlen-aware) Angles_Cumsum, angle_output_state = angle_dt_fwd( Angles, DT, + init_state=Input_Angle_State, chunk_size=chunk_size, return_output_state=True, cu_seqlens=cu_seqlens, @@ -89,6 +110,7 @@ def forward( Z, D, MIMO_Z, Angles_Cumsum, DA_CS, DA_CS_REV, DT, Trap, Segsum, cu_seqlens=cu_seqlens, + initial_states=(Input_SSM_State, Input_K_State, Input_V_State) if all_states_present else None, return_state=return_state, chunk_size=chunk_size, rotary_dim_divisor=rotary_dim_divisor, dtype=dtype, @@ -103,6 +125,7 @@ def forward( Q, K, V, Q_bias, K_bias, MIMO_V, MIMO_Out, Z, D, MIMO_Z, Angles_Cumsum, DA_CS, DA_CS_REV, DT, Trap, Segsum, + initial_states=(Input_SSM_State, Input_K_State, Input_V_State) if all_states_present else None, return_state=return_state, chunk_size=chunk_size, rotary_dim_divisor=rotary_dim_divisor, dtype=dtype, @@ -124,7 +147,14 @@ def forward( else: Final_SSM_State = Final_SSM_State.permute(0, 1, 3, 2).contiguous().detach() Final_K = Final_K.contiguous().detach() - Final_V = V[:, -1, :, :].contiguous().detach() + if cu_seqlens is None: + Final_V = torch.empty((V.shape[0], V.shape[2], V.shape[3]), device=V.device, dtype=V.dtype) + Final_V.copy_(V[:, -1, :, :]) + else: + final_v_indices = cu_seqlens[1:].to(torch.long) - 1 + Final_V = torch.empty((final_v_indices.shape[0], V.shape[2], V.shape[3]), device=V.device, dtype=V.dtype) + Final_V.copy_(V[0, final_v_indices, :, :]) + Final_V = Final_V.detach() ctx.mark_non_differentiable(Final_Angle, Final_SSM_State, Final_K, Final_V) return Out, Final_Angle, Final_SSM_State, Final_K, Final_V @@ -137,6 +167,11 @@ def backward(ctx, dout, *args) -> tuple: "Backward called but forward ran without gradient tracking. " "Ensure inputs require grad or run under torch.enable_grad()." ) + if ctx.has_input_state: + raise NotImplementedError( + "Mamba-3 MIMO backward with Input_States is not implemented; " + "use input states only for inference/prefill." + ) dout = dout.contiguous() (Q, K, V, ADT, DT, Trap, Q_bias, K_bias, Angles, Angles_Cumsum, @@ -236,6 +271,7 @@ def backward(ctx, dout, *args) -> tuple: dAngles, dD, dZ, + None, None, None, None, None, None, None, None, None, None, None, ) @@ -264,6 +300,7 @@ def mamba3_mimo( dtype: torch.dtype, return_state: bool = False, cu_seqlens: Optional[Tensor] = None, + Input_States: Optional[Tuple[Tensor, Tensor, Tensor, Tensor]] = None, fuse_pregate_headwise_rms_norm: bool = False, outproj_norm_weight: Optional[Tensor] = None, outproj_norm_eps: float = 1e-5, @@ -289,6 +326,10 @@ def mamba3_mimo( rotary_dim_divisor: Divisor for rotary embedding dimensions (default: 4, meaning angles have 1/4 of headdim_qk) dtype: Data type for lower-precision computation (e.g., torch.bfloat16) return_state: Whether to return final state for autoregressive decoding (default: False) + Input_States: Optional tuple of initial states (angle, ssm, k, v). Dense shapes are + (batch, nheads, angle_dim), (batch, nheads, headdim_v, headdim_qk), + (batch, mimo_rank, nheads, headdim_qk), and (batch, nheads, headdim_v). + Varlen mode uses num_sequences as the leading dimension. cu_seqlens: Optional tensor of cumulative sequence lengths for variable-length sequences. If provided, should be a tensor of shape (num_seq + 1,) where cu_seqlens[i] is the cumulative sequence length up to sequence i. This is used for efficient processing of @@ -324,6 +365,25 @@ def mamba3_mimo( assert nheads % nheads_qk == 0, f"nheads ({nheads}) must be divisible by nheads_qk ({nheads_qk})" assert headdim_qk % 2 == 0, f"headdim_qk ({headdim_qk}) must be even for rotary embeddings" assert rotary_dim_divisor in [2, 4], f"currently only supports rotary embedding on entire or half of headdim_qk" + num_sequences = cu_seqlens.shape[0] - 1 if cu_seqlens is not None else batch + if Input_States is None: + Input_Angle_State, Input_SSM_State, Input_K_State, Input_V_State = None, None, None, None + else: + assert len(Input_States) == 4, "Input_States must be a tuple of (angle_state, ssm_state, k_state, v_state)" + Input_Angle_State, Input_SSM_State, Input_K_State, Input_V_State = Input_States + angle_dim = Angles.shape[-1] + assert Input_Angle_State.shape == (num_sequences, nheads, angle_dim), ( + f"Input angle state shape mismatch: expected {(num_sequences, nheads, angle_dim)}, got {Input_Angle_State.shape}" + ) + assert Input_SSM_State.shape == (num_sequences, nheads, headdim_v, headdim_qk), ( + f"Input SSM state shape mismatch: expected {(num_sequences, nheads, headdim_v, headdim_qk)}, got {Input_SSM_State.shape}" + ) + assert Input_K_State.shape == (num_sequences, mimo_rank, nheads, headdim_qk), ( + f"Input K state shape mismatch: expected {(num_sequences, mimo_rank, nheads, headdim_qk)}, got {Input_K_State.shape}" + ) + assert Input_V_State.shape == (num_sequences, nheads, headdim_v), ( + f"Input V state shape mismatch: expected {(num_sequences, nheads, headdim_v)}, got {Input_V_State.shape}" + ) # NOTE: the following (headdim_qk, headdim_v) values currently can result in compilation errors: (16, 32), (256, 128) if headdim_qk not in [16, 32, 64, 128, 256]: print(f"WARNING: The value headdim_qk={headdim_qk} has not been tested. " +\ @@ -354,6 +414,10 @@ def mamba3_mimo( Angles, D, Z, + Input_Angle_State, + Input_SSM_State, + Input_K_State, + Input_V_State, chunk_size, rotary_dim_divisor, dtype, @@ -361,4 +425,4 @@ def mamba3_mimo( cu_seqlens, fuse_pregate_headwise_rms_norm, outproj_norm_eps, - ) \ No newline at end of file + ) diff --git a/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py b/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py index e1a536767..ef205813e 100755 --- a/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py +++ b/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py @@ -48,6 +48,7 @@ def mamba_mimo_fwd( reduceO, fuse_pregate_headwise_rms_norm=False, return_final_state=False, + has_initial_state=False, chunk_size: int = 16, rotary_dim_divisor = 4, dtype: str = 'bfloat16', @@ -61,6 +62,9 @@ def mamba_mimo_fwd( nchunks = tilelang.cdiv(S, chunk_size) tail_len = S % chunk_size fused_chunk_size = chunk_size * R + Init_SSM_State_shape = (B, H, P, N) + Init_K_State_shape = (B, R, H, N) + Init_V_State_shape = (B, H, P) if reduceO: O_shape = (B, S, H, P) @@ -88,6 +92,9 @@ def mamba_mimo_fwd_kernel( DT: T.Tensor([B, H, S], T.float32), # type: ignore TRAP: T.Tensor([B, H, S], dtype), # type: ignore SEGSUM: T.Tensor([B, H, nchunks, chunk_size, chunk_size], T.float32), # type: ignore + INIT_SSM_STATE: T.Tensor(Init_SSM_State_shape, T.float32), # type: ignore + INIT_K_STATE: T.Tensor(Init_K_State_shape, dtype), # type: ignore + INIT_V_STATE: T.Tensor(Init_V_State_shape, dtype), # type: ignore FINAL_STATE: T.Tensor([B, H, N, P], T.float32), # type: ignore FINAL_K: T.Tensor([B, R, H, N], dtype) # type: ignore @@ -189,6 +196,22 @@ def mamba_mimo_fwd_kernel( T.copy(Q_BIAS[i_h, :, :], q_bias_frag) T.copy(K_BIAS[i_h, :, :], k_bias_frag) + if has_initial_state: + boundary_scale = T.alloc_var(T.float32) + boundary_trap = T.alloc_var(dtype) + T.copy(DT[i_b, i_h, 0], boundary_scale) + T.copy(TRAP[i_b, i_h, 0], boundary_trap) + boundary_scale *= T.sigmoid(-boundary_trap) + for n, p in T.Parallel(N, P): + states_frag[n, p] = INIT_SSM_STATE[i_b, i_h, p, n] + for r in T.serial(R): + states_frag[n, p] += ( + INIT_K_STATE[i_b, r, i_h, n] + * INIT_V_STATE[i_b, i_h, p] + * MIMO_V[i_h, r, p] + * boundary_scale + ) + # --- Chunk Loop --- for i in T.Pipelined(0, nchunks, num_stages=num_stages): chunk_start = i * chunk_size @@ -477,6 +500,7 @@ def mamba_mimo_forward(q, k, v, segsum, chunk_size, rotary_dim_divisor, dtype, return_state=False, + initial_states=None, fuse_pregate_headwise_rms_norm=False, outproj_norm_weight=None, outproj_norm_eps=1e-5, @@ -504,6 +528,7 @@ def mamba_mimo_forward(q, k, v, reduceO, fuse_pregate_headwise_rms_norm, return_final_state=return_state, + has_initial_state=initial_states is not None, chunk_size=chunk_size, rotary_dim_divisor=rotary_dim_divisor, dtype=tl_dtype, @@ -540,6 +565,28 @@ def mamba_mimo_forward(q, k, v, z_arg = z D_arg = D mimo_z_arg = mimo_z + if initial_states is None: + init_ssm_state_arg, init_k_state_arg, init_v_state_arg = None, None, None + else: + init_ssm_state_arg, init_k_state_arg, init_v_state_arg = initial_states + if init_ssm_state_arg.shape != (B, H, P, N): + raise ValueError( + f"Expected initial SSM state shape {(B, H, P, N)}, " + f"got {tuple(init_ssm_state_arg.shape)}" + ) + if init_k_state_arg.shape != (B, R, H, N): + raise ValueError( + f"Expected initial K state shape {(B, R, H, N)}, " + f"got {tuple(init_k_state_arg.shape)}" + ) + if init_v_state_arg.shape != (B, H, P): + raise ValueError( + f"Expected initial V state shape {(B, H, P)}, " + f"got {tuple(init_v_state_arg.shape)}" + ) + init_ssm_state_arg = init_ssm_state_arg.contiguous() + init_k_state_arg = init_k_state_arg.contiguous() + init_v_state_arg = init_v_state_arg.contiguous() h = torch.empty((B, H, N, P), device='cuda', dtype=torch.float32) if return_state else None k_final = torch.empty((B, R, H, N), device='cuda', dtype=dtype) if return_state else None @@ -557,6 +604,9 @@ def mamba_mimo_forward(q, k, v, dt, trap, segsum, + init_ssm_state_arg, + init_k_state_arg, + init_v_state_arg, h, k_final ) diff --git a/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd_varlen.py b/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd_varlen.py index fc286b8a4..92975cfd7 100644 --- a/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd_varlen.py +++ b/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd_varlen.py @@ -69,6 +69,7 @@ def mamba_mimo_fwd( fuse_pregate_headwise_rms_norm=False, isVarlen: bool = True, return_final_state=False, + has_initial_state=False, chunk_size: int = 16, rotary_dim_divisor = 4, dtype: str = 'bfloat16', @@ -116,10 +117,16 @@ def mamba_mimo_fwd( max_nchunks = (S//chunk_size) + NS Final_State_shape = (NS, H, N, P) Final_K_shape = (NS, R, H, N) + Init_SSM_State_shape = (NS, H, P, N) + Init_K_State_shape = (NS, R, H, N) + Init_V_State_shape = (NS, H, P) else: max_nchunks = tilelang.cdiv(S, chunk_size) Final_State_shape = (B, H, N, P) Final_K_shape = (B, R, H, N) + Init_SSM_State_shape = (B, H, P, N) + Init_K_State_shape = (B, R, H, N) + Init_V_State_shape = (B, H, P) fused_chunk_size = chunk_size * R if reduceO: @@ -149,6 +156,9 @@ def mamba_mimo_fwd_kernel( TRAP: T.Tensor([B, H, S], dtype), # type: ignore SEGSUM: T.Tensor([B, H, max_nchunks, chunk_size, chunk_size], T.float32), # type: ignore CU_SEQLENS: T.Tensor([NS+1], dtype=T.int32), # type: ignore + INIT_SSM_STATE: T.Tensor(Init_SSM_State_shape, T.float32), # type: ignore + INIT_K_STATE: T.Tensor(Init_K_State_shape, dtype), # type: ignore + INIT_V_STATE: T.Tensor(Init_V_State_shape, dtype), # type: ignore FINAL_STATE: T.Tensor(Final_State_shape, T.float32), # type: ignore FINAL_K: T.Tensor(Final_K_shape, dtype) # type: ignore @@ -285,6 +295,33 @@ def mamba_mimo_fwd_kernel( if tail_len > 0: full_nchunks += 1 + if has_initial_state: + boundary_scale = T.alloc_var(T.float32) + boundary_trap = T.alloc_var(dtype) + T.copy(DT[i_b, i_h, start_seq_ind], boundary_scale) + T.copy(TRAP[i_b, i_h, start_seq_ind], boundary_trap) + boundary_scale *= T.sigmoid(-boundary_trap) + if isVarlen: + for n, p in T.Parallel(N, P): + states_frag[n, p] = INIT_SSM_STATE[i_ns, i_h, p, n] + for r in T.serial(R): + states_frag[n, p] += ( + INIT_K_STATE[i_ns, r, i_h, n] + * INIT_V_STATE[i_ns, i_h, p] + * MIMO_V[i_h, r, p] + * boundary_scale + ) + else: + for n, p in T.Parallel(N, P): + states_frag[n, p] = INIT_SSM_STATE[i_b, i_h, p, n] + for r in T.serial(R): + states_frag[n, p] += ( + INIT_K_STATE[i_b, r, i_h, n] + * INIT_V_STATE[i_b, i_h, p] + * MIMO_V[i_h, r, p] + * boundary_scale + ) + # --- Chunk Loop --- for i in T.Pipelined(0, full_nchunks, num_stages=num_stages): chunk_start = start_seq_ind + i * chunk_size @@ -596,6 +633,7 @@ def mamba_mimo_forward_varlen(q, k, v, chunk_size, rotary_dim_divisor, dtype, cu_seqlens=None, return_state=False, + initial_states=None, fuse_pregate_headwise_rms_norm=False, outproj_norm_weight=None, outproj_norm_eps=1e-5, @@ -682,6 +720,7 @@ def mamba_mimo_forward_varlen(q, k, v, fuse_pregate_headwise_rms_norm, isVarlen=cu_seqlens is not None, return_final_state=return_state, + has_initial_state=initial_states is not None, chunk_size=chunk_size, rotary_dim_divisor=rotary_dim_divisor, dtype=tl_dtype, @@ -718,6 +757,29 @@ def mamba_mimo_forward_varlen(q, k, v, z_arg = z D_arg = D mimo_z_arg = mimo_z + if initial_states is None: + init_ssm_state_arg, init_k_state_arg, init_v_state_arg = None, None, None + else: + init_ssm_state_arg, init_k_state_arg, init_v_state_arg = initial_states + expected_sequences = NS if cu_seqlens is not None else B + if init_ssm_state_arg.shape != (expected_sequences, H, P, N): + raise ValueError( + f"Expected initial SSM state shape {(expected_sequences, H, P, N)}, " + f"got {tuple(init_ssm_state_arg.shape)}" + ) + if init_k_state_arg.shape != (expected_sequences, R, H, N): + raise ValueError( + f"Expected initial K state shape {(expected_sequences, R, H, N)}, " + f"got {tuple(init_k_state_arg.shape)}" + ) + if init_v_state_arg.shape != (expected_sequences, H, P): + raise ValueError( + f"Expected initial V state shape {(expected_sequences, H, P)}, " + f"got {tuple(init_v_state_arg.shape)}" + ) + init_ssm_state_arg = init_ssm_state_arg.contiguous() + init_k_state_arg = init_k_state_arg.contiguous() + init_v_state_arg = init_v_state_arg.contiguous() if cu_seqlens is not None: h = torch.empty((NS, H, N, P), device='cuda', dtype=torch.float32) if return_state else None @@ -740,6 +802,9 @@ def mamba_mimo_forward_varlen(q, k, v, trap, segsum, cu_seqlens, + init_ssm_state_arg, + init_k_state_arg, + init_v_state_arg, h, k_final ) diff --git a/tests/modules/test_mamba3_inference.py b/tests/modules/test_mamba3_inference.py new file mode 100644 index 000000000..1b7d6976c --- /dev/null +++ b/tests/modules/test_mamba3_inference.py @@ -0,0 +1,58 @@ +""" +Mamba-3 module-level inference tests. + +Copyright (c) 2026, Dao AI Lab, Goombalab. +""" + +import pytest +import torch + +from mamba_ssm.modules.mamba3 import Mamba3 +from mamba_ssm.utils.generation import InferenceParams + + +def _require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + +def _make_model(is_mimo: bool, **kwargs) -> Mamba3: + defaults = dict( + d_model=128, + d_state=64, + expand=2, + headdim=64, + ngroups=1, + is_mimo=is_mimo, + mimo_rank=2, + chunk_size=16, + layer_idx=0, + device="cuda", + dtype=torch.bfloat16, + ) + defaults.update(kwargs) + return Mamba3(**defaults).eval() + + +@pytest.mark.parametrize("is_mimo", [False, True]) +def test_mamba3_inference_params_prefill_uses_cached_states(is_mimo): + """Prefill should consume cached states when seqlen_offset stays at 0.""" + _require_cuda() + torch.manual_seed(2) + + model = _make_model(is_mimo=is_mimo) + d_model = model.d_model + split = 32 + seqlen = 48 + u = torch.randn(1, seqlen, d_model, device="cuda", dtype=torch.bfloat16) + inference_params = InferenceParams(max_seqlen=seqlen, max_batch_size=1) + + with torch.no_grad(): + full = model(u) + first = model(u[:, :split], inference_params=inference_params) + second = model(u[:, split:], inference_params=inference_params) + + combined = torch.cat([first, second], dim=1) + assert torch.allclose(combined, full, atol=8e-2, rtol=8e-2), ( + f"max diff = {(combined - full).abs().max():.4f}" + ) diff --git a/tests/ops/tilelang/test_mamba3_mimo.py b/tests/ops/tilelang/test_mamba3_mimo.py index e3cf08794..2b9b39473 100644 --- a/tests/ops/tilelang/test_mamba3_mimo.py +++ b/tests/ops/tilelang/test_mamba3_mimo.py @@ -1637,6 +1637,157 @@ def test_mamba_mimo_smoke_forward_backward(mods: SimpleNamespace) -> None: assert torch.isfinite(grad).all(), f"Non-finite gradient detected for {name}" +def test_mamba3_mimo_initial_states_split_prefill_matches_full(mods: SimpleNamespace) -> None: + chunk_size = 16 + split = 32 + inputs = make_smoke_inputs( + batch=1, + seqlen=64, + mimo_rank=4, + nheads_qk=1, + nheads=8, + headdim_qk=128, + headdim_v=64, + chunk_size=chunk_size, + rotary_dim_divisor=FIXED_ROTARY_DIM_DIVISOR, + device="cuda", + dtype=FIXED_DTYPE, + seed=123, + ) + + def copied(t: Tensor) -> Tensor: + out = torch.empty(t.shape, device=t.device, dtype=t.dtype) + out.copy_(t) + return out + + def sliced(start: int, end: int) -> dict: + out = dict(inputs) + for name in ("Q", "K", "V", "Angles", "Z"): + out[name] = copied(inputs[name][:, start:end]) + for name in ("ADT", "DT", "Trap"): + out[name] = copied(inputs[name][:, :, start:end]) + return out + + with torch.no_grad(): + full_out, full_angle, full_ssm, full_k, full_v = mods.top.mamba3_mimo( + **inputs, return_state=True + ) + first_out, first_angle, first_ssm, first_k, first_v = mods.top.mamba3_mimo( + **sliced(0, split), return_state=True + ) + second_out, final_angle, final_ssm, final_k, final_v = mods.top.mamba3_mimo( + **sliced(split, 64), + return_state=True, + Input_States=(first_angle, first_ssm, first_k, first_v), + ) + + split_out = torch.cat([first_out, second_out], dim=1) + cfg = f"B=1, S=64, split={split}, H=8, P=64, N=128, R=4, C={chunk_size}" + assert_stable_rel(split_out, full_out, label="initial_state_split_out", cfg=cfg) + assert_stable_rel(final_angle, full_angle, label="initial_state_split_angle", cfg=cfg) + assert_stable_rel(final_ssm, full_ssm, label="initial_state_split_ssm", cfg=cfg) + assert_stable_rel(final_k, full_k, label="initial_state_split_k", cfg=cfg) + assert_stable_rel(final_v, full_v, label="initial_state_split_v", cfg=cfg) + + +def test_mamba3_mimo_varlen_initial_states_split_prefill_matches_full(mods: SimpleNamespace) -> None: + chunk_size = 16 + seqlens = [31, 42, 47] + splits = [16, 18, 21] + suffix_lens = [seqlen - split for seqlen, split in zip(seqlens, splits)] + full_starts = [0] + for seqlen in seqlens[:-1]: + full_starts.append(full_starts[-1] + seqlen) + s_total = sum(seqlens) + inputs = make_smoke_inputs( + batch=1, + seqlen=s_total, + mimo_rank=4, + nheads_qk=1, + nheads=8, + headdim_qk=128, + headdim_v=64, + chunk_size=chunk_size, + rotary_dim_divisor=FIXED_ROTARY_DIM_DIVISOR, + device="cuda", + dtype=FIXED_DTYPE, + seed=321, + ) + + def cu_seqlens(lengths: list[int]) -> Tensor: + return torch.tensor( + [0] + list(torch.cumsum(torch.tensor(lengths, dtype=torch.int32), dim=0).tolist()), + device="cuda", + dtype=torch.int32, + ) + + def packed_spans(starts: list[int], ends: list[int]) -> dict: + out = dict(inputs) + for name in ("Q", "K", "V", "Angles", "Z"): + out[name] = torch.cat( + [inputs[name][:, start:end] for start, end in zip(starts, ends)], + dim=1, + ) + for name in ("ADT", "DT", "Trap"): + out[name] = torch.cat( + [inputs[name][:, :, start:end] for start, end in zip(starts, ends)], + dim=2, + ) + return out + + full_cu = cu_seqlens(seqlens) + prefix_cu = cu_seqlens(splits) + suffix_cu = cu_seqlens(suffix_lens) + prefix_starts = full_starts + prefix_ends = [start + split for start, split in zip(full_starts, splits)] + suffix_starts = prefix_ends + suffix_ends = [start + seqlen for start, seqlen in zip(full_starts, seqlens)] + prefix_inputs = packed_spans(prefix_starts, prefix_ends) + suffix_inputs = packed_spans(suffix_starts, suffix_ends) + + with torch.no_grad(): + full_out, full_angle, full_ssm, full_k, full_v = mods.top.mamba3_mimo( + **inputs, + return_state=True, + cu_seqlens=full_cu, + ) + prefix_out, prefix_angle, prefix_ssm, prefix_k, prefix_v = mods.top.mamba3_mimo( + **prefix_inputs, + return_state=True, + cu_seqlens=prefix_cu, + ) + suffix_out, final_angle, final_ssm, final_k, final_v = mods.top.mamba3_mimo( + **suffix_inputs, + return_state=True, + cu_seqlens=suffix_cu, + Input_States=(prefix_angle, prefix_ssm, prefix_k, prefix_v), + ) + + combined_parts = [] + prefix_offset = 0 + suffix_offset = 0 + for split, suffix_len in zip(splits, suffix_lens): + combined_parts.append( + torch.cat( + [ + prefix_out[:, prefix_offset:prefix_offset + split], + suffix_out[:, suffix_offset:suffix_offset + suffix_len], + ], + dim=1, + ) + ) + prefix_offset += split + suffix_offset += suffix_len + split_out = torch.cat(combined_parts, dim=1) + + cfg = f"seqlens={seqlens}, splits={splits}, H=8, P=64, N=128, R=4, C={chunk_size}" + assert_stable_rel(split_out, full_out, label="varlen_initial_state_split_out", cfg=cfg) + assert_stable_rel(final_angle, full_angle, label="varlen_initial_state_split_angle", cfg=cfg) + assert_stable_rel(final_ssm, full_ssm, label="varlen_initial_state_split_ssm", cfg=cfg) + assert_stable_rel(final_k, full_k, label="varlen_initial_state_split_k", cfg=cfg) + assert_stable_rel(final_v, full_v, label="varlen_initial_state_split_v", cfg=cfg) + + def test_mamba_mimo_smoke_forward_backward_varlen(mods: SimpleNamespace) -> None: """Smoke test for the varlen forward+backward path through ``mamba3_mimo``.