Skip to content

Add support for input states to Mamba-3 forward pass #976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aakashlahoti wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add support for input states to Mamba-3 forward pass #976

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion mamba_ssm/modules/mamba3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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),
Expand All @@ -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,
)
Expand Down
70 changes: 67 additions & 3 deletions mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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

Expand All @@ -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,
Expand Down Expand Up @@ -236,6 +271,7 @@ def backward(ctx, dout, *args) -> tuple:
dAngles,
dD,
dZ,
None, None, None, None,
None, None, None, None, None, None, None,
)

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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. " +\
Expand Down Expand Up @@ -354,11 +414,15 @@ def mamba3_mimo(
Angles,
D,
Z,
Input_Angle_State,
Input_SSM_State,
Input_K_State,
Input_V_State,
chunk_size,
rotary_dim_divisor,
dtype,
return_state,
cu_seqlens,
fuse_pregate_headwise_rms_norm,
outproj_norm_eps,
)
)
50 changes: 50 additions & 0 deletions mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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
)
Expand Down
Loading