Text Generation
Transformers
RWKV
PyTorch
English
maba_sparse
maba
maba-v2
maba-v2-architecture
architecture
recurrent
dgda
decoupled-gated-delta-attention
gated-deltanet
linear-attention
linear-recurrence
sparse-attention
maba-sa
mla
multi-head-latent-attention
deepseek
qwen
minicpm
mamba
mamba-2
transformer
causal-lm
llm
nlp
long-context
1m-context
sub-quadratic
state-space-model
ssm
triton
flash-attention
on-device-ai
efficient-llm
nope
dg-indexer
centroid-indexing
hca
3-stream
swiglu
rmsnorm
speculative-decoding
mtp
multi-token-prediction
needle-in-a-haystack
scaling
100m
1b
3b
7b
30b
| import math | |
| from typing import Any, Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from maba_sparse.kernels.dispatcher import dispatch_stream_superposition | |
| from maba_sparse.layers.indexer import DGIndexer | |
| class MabaSparseAttention(nn.Module): | |
| def __init__( | |
| self, | |
| config: Optional[Any] = None, | |
| dim: int = 640, | |
| n_heads: int = 10, | |
| d_head: int = 64, | |
| d_c: int = 128, | |
| window_size: int = 128, | |
| block_size: int = 64, | |
| top_k: int = 32, | |
| hca_pool_size: int = 64, | |
| dist_lambda: float = 0.5, | |
| d_idx: int = 64, | |
| ablation_mode: str = "full", | |
| ) -> None: | |
| super().__init__() | |
| if config is not None: | |
| self.dim = getattr(config, "dim", getattr(config, "d_model", dim)) | |
| self.n_heads = getattr(config, "n_heads", getattr(config, "num_heads", n_heads)) | |
| self.d_head = getattr(config, "d_head", d_head) | |
| self.d_c = getattr(config, "d_c", d_c) | |
| self.window_size = getattr(config, "window_size", window_size) | |
| self.block_size = getattr(config, "block_size", block_size) | |
| self.top_k = getattr(config, "top_k", top_k) | |
| self.hca_pool_size = getattr(config, "hca_pool_size", hca_pool_size) | |
| self.dist_lambda = getattr(config, "dist_lambda", dist_lambda) | |
| self.n_sink_tokens = getattr(config, "n_sink_tokens", 4) | |
| d_idx = getattr(config, "d_idx", d_idx) | |
| self.ablation_mode = ablation_mode if ablation_mode is not None else getattr(config, "ablation_mode", "full") | |
| else: | |
| self.dim = dim | |
| self.n_heads = n_heads | |
| self.d_head = d_head | |
| self.d_c = d_c | |
| self.window_size = window_size | |
| self.block_size = block_size | |
| self.top_k = top_k | |
| self.hca_pool_size = hca_pool_size | |
| self.dist_lambda = dist_lambda | |
| self.n_sink_tokens = 4 | |
| self.ablation_mode = ablation_mode | |
| self.scale = 1.0 / math.sqrt(self.d_head) | |
| self.q_proj = nn.Linear(self.dim, self.n_heads * self.d_head, bias=False) | |
| self.kv_down_proj = nn.Linear(self.dim, self.d_c, bias=False) | |
| self.k_up_proj = nn.Linear(self.d_c, self.n_heads * self.d_head, bias=False) | |
| self.v_up_proj = nn.Linear(self.d_c, self.n_heads * self.d_head, bias=False) | |
| self.o_proj = nn.Linear(self.n_heads * self.d_head, self.dim, bias=False) | |
| self.stream_gate = nn.Linear(self.dim, 3, bias=True) | |
| self.indexer = DGIndexer( | |
| dim=self.dim, | |
| d_idx=d_idx, | |
| block_size=self.block_size, | |
| top_k=self.top_k, | |
| dist_lambda=self.dist_lambda, | |
| ) | |
| self._cached_k: Optional[torch.Tensor] = None | |
| self._cached_v: Optional[torch.Tensor] = None | |
| def reset_cache(self) -> None: | |
| self.indexer.reset_cache() | |
| self._cached_k = None | |
| self._cached_v = None | |
| def get_gate_weights(self, x: torch.Tensor) -> torch.Tensor: | |
| return F.softmax(self.stream_gate(x), dim=-1) | |
| def _compute_local_attention( | |
| self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, L_full: int | |
| ) -> torch.Tensor: | |
| b, h, lq, _ = q.shape | |
| lkv = k.shape[2] | |
| if lq == 1 and lkv > (self.window_size + self.n_sink_tokens): | |
| sink_idx = torch.arange(min(self.n_sink_tokens, lkv), device=q.device) | |
| win_start = max(min(self.n_sink_tokens, lkv), lkv - self.window_size) | |
| win_idx = torch.arange(win_start, lkv, device=q.device) | |
| sel_idx = torch.cat([sink_idx, win_idx]) | |
| k_sel = k[:, :, sel_idx, :] | |
| v_sel = v[:, :, sel_idx, :] | |
| a = torch.matmul(q, k_sel.transpose(-1, -2)) * self.scale | |
| p = F.softmax(a, dim=-1) | |
| return torch.matmul(p, v_sel) | |
| a = torch.matmul(q, k.transpose(-1, -2)) * self.scale | |
| qo = lkv - lq | |
| i = (qo + torch.arange(lq, device=q.device)).unsqueeze(1) | |
| j = torch.arange(lkv, device=q.device).unsqueeze(0) | |
| c = j <= i | |
| w = (i - j) < self.window_size | |
| s = j < min(self.n_sink_tokens, lkv) | |
| m = c & (w | s) | |
| a = a.masked_fill(~m.unsqueeze(0).unsqueeze(0), float("-inf")) | |
| p = torch.nan_to_num(F.softmax(a, dim=-1), nan=0.0) | |
| return torch.matmul(p, v) | |
| def _compute_sparse_attention( | |
| self, | |
| q: torch.Tensor, | |
| k: torch.Tensor, | |
| v: torch.Tensor, | |
| top_indices: torch.Tensor, | |
| L_full: int, | |
| ) -> torch.Tensor: | |
| b, h, lq, dh = q.shape | |
| lkv = k.shape[2] | |
| nb = (lkv + self.block_size - 1) // self.block_size | |
| if nb == 0 or lq == 0: | |
| return torch.zeros(b, h, lq, dh, device=q.device, dtype=q.dtype) | |
| ak = top_indices.shape[-1] | |
| pad_len = nb * self.block_size - lkv | |
| kp = F.pad(k, (0, 0, 0, pad_len)) if pad_len > 0 else k | |
| vp = F.pad(v, (0, 0, 0, pad_len)) if pad_len > 0 else v | |
| kb = kp.view(b, h, nb, self.block_size, dh) | |
| vb = vp.view(b, h, nb, self.block_size, dh) | |
| qo = lkv - lq | |
| if lq == 1: | |
| bi = top_indices.squeeze(1) | |
| gather_idx = bi.view(b, 1, ak, 1, 1).expand(b, h, ak, self.block_size, dh) | |
| k_sel = torch.gather(kb, 2, gather_idx).reshape(b, h, ak * self.block_size, dh) | |
| v_sel = torch.gather(vb, 2, gather_idx).reshape(b, h, ak * self.block_size, dh) | |
| a = torch.matmul(q, k_sel.transpose(-1, -2)) * self.scale | |
| token_offsets = ( | |
| bi.unsqueeze(-1) * self.block_size | |
| + torch.arange(self.block_size, device=q.device).view(1, 1, self.block_size) | |
| ).reshape(b, 1, ak * self.block_size) | |
| valid = (token_offsets < lkv) & (token_offsets <= (lkv - 1)) | |
| a = a.masked_fill(~valid.unsqueeze(1), float("-inf")) | |
| p = torch.nan_to_num(F.softmax(a, dim=-1), nan=0.0) | |
| return torch.matmul(p, v_sel) | |
| os_chunks = [] | |
| chunk_sz = 128 | |
| for q_start in range(0, lq, chunk_sz): | |
| q_end = min(q_start + chunk_sz, lq) | |
| c_len = q_end - q_start | |
| qc = q[:, :, q_start:q_end, :] | |
| bi = top_indices[:, q_start:q_end, :] | |
| gather_idx = bi.view(b, 1, c_len, ak, 1, 1).expand(b, h, c_len, ak, self.block_size, dh) | |
| kb_exp = kb.view(b, h, 1, nb, self.block_size, dh).expand(b, h, c_len, nb, self.block_size, dh) | |
| vb_exp = vb.view(b, h, 1, nb, self.block_size, dh).expand(b, h, c_len, nb, self.block_size, dh) | |
| k_sel = torch.gather(kb_exp, 3, gather_idx).reshape(b, h, c_len, ak * self.block_size, dh) | |
| v_sel = torch.gather(vb_exp, 3, gather_idx).reshape(b, h, c_len, ak * self.block_size, dh) | |
| ac = (qc.unsqueeze(-2) @ k_sel.transpose(-1, -2)).squeeze(-2) * self.scale | |
| pos = (qo + torch.arange(q_start, q_end, device=q.device)).view(1, c_len, 1) | |
| token_offsets = ( | |
| bi.unsqueeze(-1) * self.block_size | |
| + torch.arange(self.block_size, device=q.device).view(1, 1, 1, self.block_size) | |
| ).reshape(b, c_len, ak * self.block_size) | |
| valid = (token_offsets < lkv) & (token_offsets <= pos) | |
| ac = ac.masked_fill(~valid.unsqueeze(1), float("-inf")) | |
| pc = torch.nan_to_num(F.softmax(ac, dim=-1), nan=0.0) | |
| oc = (pc.unsqueeze(-2) @ v_sel).squeeze(-2) | |
| os_chunks.append(oc) | |
| return torch.cat(os_chunks, dim=2) | |
| def _compute_hca_attention( | |
| self, q: torch.Tensor, c_kv: torch.Tensor, L_full: int | |
| ) -> torch.Tensor: | |
| b, h, lq, dh = q.shape | |
| lkv = c_kv.shape[1] | |
| r = self.hca_pool_size | |
| if lkv == 0 or lq == 0: | |
| return torch.zeros(b, h, lq, dh, device=q.device, dtype=q.dtype) | |
| qo = lkv - lq | |
| pos = qo + torch.arange(lq, device=q.device) | |
| num_complete = lkv // r | |
| early_mask = pos < (r - 1) | |
| out = torch.empty(b, h, lq, dh, device=q.device, dtype=q.dtype) | |
| if early_mask.any(): | |
| early_indices = torch.where(early_mask)[0] | |
| early_pos = pos[early_mask] | |
| max_early_pos = early_pos.max().item() | |
| c_prefix = c_kv[:, : max_early_pos + 1] | |
| c_cum = torch.cumsum(c_prefix, dim=1) | |
| denom = torch.arange(1, max_early_pos + 2, device=c_kv.device).view(1, -1, 1) | |
| c_mean = c_cum / denom | |
| c_pool_early = c_mean[:, early_pos, :] | |
| vh_early = self.v_up_proj(c_pool_early).view(b, len(early_pos), h, dh).transpose(1, 2) | |
| out[:, :, early_indices, :] = vh_early | |
| if (~early_mask).any() and num_complete > 0: | |
| late_indices = torch.where(~early_mask)[0] | |
| late_pos = pos[~early_mask] | |
| q_late = q[:, :, late_indices, :] | |
| c_complete = c_kv[:, : num_complete * r].view(b, num_complete, r, self.d_c).mean(dim=2) | |
| kh = self.k_up_proj(c_complete).view(b, num_complete, h, dh).transpose(1, 2) | |
| vh = self.v_up_proj(c_complete).view(b, num_complete, h, dh).transpose(1, 2) | |
| a = torch.matmul(q_late, kh.transpose(-1, -2)) * self.scale | |
| qi = late_pos.unsqueeze(1) | |
| bs = torch.arange(num_complete, device=q.device) + 1 | |
| ei = (bs * r - 1).unsqueeze(0) | |
| c = ei <= qi | |
| a = a.masked_fill(~c.unsqueeze(0).unsqueeze(0), float("-inf")) | |
| p = torch.nan_to_num(F.softmax(a, dim=-1), nan=0.0) | |
| out[:, :, late_indices, :] = torch.matmul(p, vh) | |
| return out | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| past_c_kv: Optional[torch.Tensor] = None, | |
| past_centroids: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| b, l, d = x.shape | |
| if past_c_kv is None or ( | |
| self._cached_k is not None and self._cached_k.shape[2] != past_c_kv.shape[1] | |
| ) or ( | |
| self.indexer._cached_k_idx is not None and self.indexer._cached_k_idx.shape[1] != past_c_kv.shape[1] | |
| ): | |
| self.reset_cache() | |
| c_kv = self.kv_down_proj(x) | |
| c_full = torch.cat([past_c_kv, c_kv], dim=1) if past_c_kv is not None else c_kv | |
| lf = c_full.shape[1] | |
| q = self.q_proj(x).view(b, l, self.n_heads, self.d_head).transpose(1, 2) | |
| if ( | |
| not self.training | |
| and past_c_kv is not None | |
| and self._cached_k is not None | |
| and self._cached_k.shape[0] == b | |
| and self._cached_k.shape[2] == past_c_kv.shape[1] | |
| ): | |
| k_new = self.k_up_proj(c_kv).view(b, l, self.n_heads, self.d_head).transpose(1, 2) | |
| v_new = self.v_up_proj(c_kv).view(b, l, self.n_heads, self.d_head).transpose(1, 2) | |
| k = torch.cat([self._cached_k, k_new], dim=2) | |
| v = torch.cat([self._cached_v, v_new], dim=2) | |
| else: | |
| k = self.k_up_proj(c_full).view(b, lf, self.n_heads, self.d_head).transpose(1, 2) | |
| v = self.v_up_proj(c_full).view(b, lf, self.n_heads, self.d_head).transpose(1, 2) | |
| if not self.training: | |
| self._cached_k = k | |
| self._cached_v = v | |
| else: | |
| self._cached_k = None | |
| self._cached_v = None | |
| ol = self._compute_local_attention(q, k, v, lf) | |
| idx, _ = self.indexer( | |
| x, return_scores=False, past_centroids=past_centroids | |
| ) | |
| os = self._compute_sparse_attention(q, k, v, idx, lf) | |
| if self.ablation_mode == "no_hca": | |
| oh = torch.zeros_like(ol) | |
| else: | |
| oh = self._compute_hca_attention(q, c_full, lf) | |
| gate_logits = self.stream_gate(x) | |
| o = dispatch_stream_superposition(ol, os, oh, gate_logits=gate_logits) | |
| o = o.transpose(1, 2).contiguous().view(b, l, self.n_heads * self.d_head).to(q.dtype) | |
| return self.o_proj(o), c_full | |
| MABASALayer = MabaSparseAttention | |