Spaces:
Runtime error
Runtime error
| """Triton-3.2-compatible Mamba3 SISO wrapper. | |
| The upstream Mamba3 SISO Triton kernel in the release graft uses | |
| ``tl.make_tensor_descriptor``. HF A10 Jobs require Triton 3.2.0 for working | |
| driver discovery with torch 2.6, and Triton 3.2.0 does not expose that API. | |
| This module keeps the public ``mamba3_siso_combined`` entrypoint intact and | |
| uses regular Torch ops for the SISO path. It is intentionally conservative: | |
| no silent CPU fallback, no fake outputs, and no dependency on newer Triton APIs. | |
| """ | |
| from __future__ import annotations | |
| from typing import Optional, Tuple | |
| import torch | |
| from torch import Tensor | |
| import torch.nn.functional as F | |
| def _expand_qk_heads(x: Tensor, nheads: int) -> Tensor: | |
| """Expand grouped Q/K heads to value-head count.""" | |
| nheads_qk = x.shape[2] | |
| if nheads_qk == nheads: | |
| return x | |
| if nheads % nheads_qk != 0: | |
| raise ValueError(f"nheads={nheads} must be divisible by nheads_qk={nheads_qk}") | |
| repeat = nheads // nheads_qk | |
| return x.repeat_interleave(repeat, dim=2) | |
| def _rotate_pairwise(x: Tensor, angles: Tensor) -> Tensor: | |
| """Apply adjacent-pair RoPE to the first 2 * angles.shape[-1] channels.""" | |
| rotary_pairs = angles.shape[-1] | |
| rotary_dim = rotary_pairs * 2 | |
| if rotary_dim == 0: | |
| return x | |
| if x.shape[-1] < rotary_dim: | |
| raise ValueError(f"rotary_dim={rotary_dim} exceeds q/k dim={x.shape[-1]}") | |
| x_rot = x.clone() | |
| x0 = x[..., :rotary_dim:2] | |
| x1 = x[..., 1:rotary_dim:2] | |
| cos = torch.cos(angles.to(torch.float32)).to(dtype=x.dtype) | |
| sin = torch.sin(angles.to(torch.float32)).to(dtype=x.dtype) | |
| x_rot[..., :rotary_dim:2] = x0 * cos - x1 * sin | |
| x_rot[..., 1:rotary_dim:2] = x0 * sin + x1 * cos | |
| return x_rot | |
| def _torch_siso_scan( | |
| Q: Tensor, | |
| K: Tensor, | |
| V: Tensor, | |
| ADT: Tensor, | |
| DT: Tensor, | |
| Trap: Tensor, | |
| Q_bias: Tensor, | |
| K_bias: Tensor, | |
| Angles: Tensor, | |
| D: Optional[Tensor], | |
| Z: Optional[Tensor], | |
| Input_States: Optional[Tuple[Tensor, Tensor, Tensor]], | |
| return_final_states: bool, | |
| ) -> Tensor | Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: | |
| if not (Q.is_cuda and K.is_cuda and V.is_cuda): | |
| raise RuntimeError("Mamba3 Torch SISO fallback requires CUDA tensors") | |
| batch, seqlen, _, qk_dim = Q.shape | |
| _, _, nheads, headdim_v = V.shape | |
| q = _expand_qk_heads(Q, nheads) | |
| k = _expand_qk_heads(K, nheads) | |
| q = q + Q_bias.to(dtype=q.dtype, device=q.device).view(1, 1, nheads, qk_dim) | |
| k = k + K_bias.to(dtype=k.dtype, device=k.device).view(1, 1, nheads, qk_dim) | |
| # Upstream angle_dt computes cumulative angle rates weighted by dt. | |
| dt_seq = DT.transpose(1, 2).to(torch.float32) # (B, L, H) | |
| angle_cumsum = torch.cumsum(Angles.to(torch.float32) * dt_seq.unsqueeze(-1), dim=1) | |
| q = _rotate_pairwise(q, angle_cumsum) | |
| k_unscaled = _rotate_pairwise(k, angle_cumsum) | |
| trap = torch.sigmoid(Trap.to(torch.float32)) | |
| dt = DT.to(torch.float32) | |
| shifted_dt = torch.cat([dt[:, :, 1:], torch.zeros_like(dt[:, :, :1])], dim=2) | |
| shifted_trap = torch.cat([trap[:, :, 1:], torch.zeros_like(trap[:, :, :1])], dim=2) | |
| gamma = dt * trap | |
| scale = shifted_dt * (1.0 - shifted_trap) + gamma | |
| k_scaled = k_unscaled * scale.transpose(1, 2).unsqueeze(-1).to(k_unscaled.dtype) | |
| if Input_States is not None: | |
| state, k_state, v_state = Input_States | |
| state_acc = state.to(torch.float32).clone() | |
| first_dt = dt[:, :, 0].unsqueeze(-1).unsqueeze(-1) | |
| first_trap = trap[:, :, 0].unsqueeze(-1).unsqueeze(-1) | |
| state_acc = state_acc + v_state.to(torch.float32).unsqueeze(-1) * k_state.to(torch.float32).unsqueeze(-2) * first_dt * (1.0 - first_trap) | |
| else: | |
| state_acc = torch.zeros(batch, nheads, headdim_v, qk_dim, device=V.device, dtype=torch.float32) | |
| d = torch.zeros(nheads, device=V.device, dtype=torch.float32) if D is None else D.to(torch.float32) | |
| outputs: list[Tensor] = [] | |
| for t in range(seqlen): | |
| q_t = q[:, t].to(torch.float32) # (B, H, N) | |
| k_t = k_scaled[:, t].to(torch.float32) | |
| k_raw_t = k_unscaled[:, t].to(torch.float32) | |
| v_t = V[:, t].to(torch.float32) # (B, H, P) | |
| decayed = state_acc * torch.exp(ADT[:, :, t].to(torch.float32)).unsqueeze(-1).unsqueeze(-1) | |
| state_out = torch.einsum("bhn,bhpn->bhp", q_t, decayed) | |
| qk_dot = (q_t * k_raw_t).sum(dim=-1) * gamma[:, :, t] | |
| y_t = state_out + (d.view(1, nheads) + qk_dot).unsqueeze(-1) * v_t | |
| if Z is not None: | |
| y_t = y_t * F.silu(Z[:, t].to(torch.float32)) | |
| outputs.append(y_t.to(V.dtype)) | |
| state_acc = decayed + torch.einsum("bhp,bhn->bhpn", v_t, k_t) | |
| out = torch.stack(outputs, dim=1) | |
| if not return_final_states: | |
| return out | |
| final_angle = angle_cumsum[:, -1] | |
| final_k = k_unscaled[:, -1] | |
| final_v = V[:, -1] | |
| return out, final_angle, state_acc, final_k, final_v | |
| def mamba3_siso_combined( | |
| Q: Tensor, | |
| K: Tensor, | |
| V: Tensor, | |
| ADT: Tensor, | |
| DT: Tensor, | |
| Trap: Tensor, | |
| Q_bias: Tensor, | |
| K_bias: Tensor, | |
| Angles: Tensor, | |
| D: Optional[Tensor] = None, | |
| Z: Optional[Tensor] = None, | |
| Input_States: Optional[Tuple[Tensor, Tensor, Tensor, Tensor]] = None, | |
| chunk_size: int = 64, | |
| return_final_states: bool = False, | |
| cu_seqlens: Optional[Tensor] = None, | |
| ) -> Tensor | Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: | |
| del chunk_size | |
| if cu_seqlens is not None: | |
| raise NotImplementedError("Torch Mamba3 SISO fallback does not support varlen cu_seqlens") | |
| angle_state, ssm_state, k_state, v_state = ( | |
| Input_States if Input_States is not None else (None, None, None, None) | |
| ) | |
| if any(x is not None for x in (angle_state, ssm_state, k_state, v_state)): | |
| if not all(x is not None for x in (angle_state, ssm_state, k_state, v_state)): | |
| raise ValueError("Mamba3 input states must be all present or all absent") | |
| torch_states = (ssm_state, k_state, v_state) | |
| else: | |
| torch_states = None | |
| return _torch_siso_scan( | |
| Q=Q, | |
| K=K, | |
| V=V, | |
| ADT=ADT, | |
| DT=DT, | |
| Trap=Trap, | |
| Q_bias=Q_bias, | |
| K_bias=K_bias, | |
| Angles=Angles, | |
| D=D, | |
| Z=Z, | |
| Input_States=torch_states, | |
| return_final_states=return_final_states, | |
| ) | |