| """Native Mamba-2 blocks and bounded absolute-position suffix-cache support. |
| |
| Training uses upstream Mamba2.forward equations with fixed native kernel tiles |
| and ordered native backward reductions. Decoding uses upstream step. |
| Continuation chunks use the upstream SSD scan with an initial state and the |
| upstream multi-token causal-convolution update over retained inputs. |
| This path requires real CUDA comparison against full forward before admission. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def validate_config(cfg): |
| if (cfg.variant != 'taildrop' or cfg.history_policy != 'suffix' |
| or cfg.recent_block or cfg.minimum_history != 32): |
| raise ValueError('Mamba outer stacks require unchanged unprotected S-min32') |
| if cfg.dense_ends < 2 or cfg.dense_ends % 2 or cfg.layers <= 2 * cfg.dense_ends: |
| raise ValueError('Mamba stacks replace dense ends two-for-one around a nonempty reader') |
| if cfg.hidden * cfg.mamba_expand % cfg.mamba_headdim: |
| raise ValueError('Mamba inner width is not divisible by head width') |
| if min(cfg.mamba_d_state, cfg.mamba_expand, cfg.mamba_d_conv) < 1: |
| raise ValueError('Mamba dimensions must be positive') |
| if cfg.mamba_state_dtype != 'float32': |
| raise ValueError('This trajectory requires FP32 recurrent inference state') |
| if cfg.mamba_headdim > 64 or cfg.mamba_d_state > 128: |
| raise ValueError('Ordered native backward supports at most head64/state128') |
|
|
|
|
| def validate_cache_call(model, tokens, lengths, cap, cache): |
| """Reject incompatible calls before any outer recurrent state is modified.""" |
| if model.training or lengths is not None: |
| raise ValueError('Persistent inference requires eval mode and a fixed suffix cap') |
| if type(cap) is not int or cap < 1: |
| raise ValueError('Mamba/suffix inference requires a positive fixed cap') |
| if tokens.ndim != 2 or tokens.shape[1] < 1: |
| raise ValueError('Cached input must contain at least one token') |
| if len(cache.states) != model.cfg.layers or len(set(cache.seen)) != 1: |
| raise ValueError('Cache layer count or token cursors are inconsistent') |
| end = cache.seen[0] + tokens.shape[1] |
| if cache.max_length is not None and end > cache.max_length: |
| raise ValueError('Decode would exceed the declared cache capacity') |
| identity = ('mamba2-suffix-v1', cap, tokens.shape[0], str(tokens.device)) |
| if cache.seen[0] and getattr(cache, 'mamba_policy', None) != identity: |
| raise ValueError('A populated cache cannot change its retained-history policy or batch') |
| cache.mamba_policy = identity |
|
|
|
|
| class BoundedRotary(nn.Module): |
| """Identical absolute RoPE frequencies; cache only the current position slice. |
| |
| The same instance is shared by all 20 middle layers. Single-token decoding |
| caches one position, independent of total stream length. Long uncached calls |
| may allocate temporary tables but never retain more than max_cached rows. |
| """ |
| def __init__(self, dim, base=10000.0, max_cached=2048, native=True): |
| super().__init__() |
| self.dim, self.base, self.max_cached, self.native = dim, base, max_cached, native |
| self.interleaved = False |
| self.register_buffer('inv_freq', self._frequencies('cpu'), persistent=False) |
| self._key = self._cos_cached = self._sin_cached = None |
|
|
| def _frequencies(self, device): |
| return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)) |
|
|
| def _apply(self, fn): |
| super()._apply(fn) |
| |
| self.inv_freq = self._frequencies(self.inv_freq.device) |
| self._key = self._cos_cached = self._sin_cached = None |
| return self |
|
|
| def forward(self, q, k, *, seqlen_offset=0, max_seqlen=None): |
| if not isinstance(seqlen_offset, int) or seqlen_offset < 0 or q.shape[1] != k.shape[1]: |
| raise ValueError('RoPE requires a common nonnegative absolute position slice') |
| length = q.shape[1] |
| key = (seqlen_offset, length, q.device, q.dtype) |
| if key == self._key: |
| cos, sin = self._cos_cached, self._sin_cached |
| else: |
| positions = torch.arange(seqlen_offset, seqlen_offset + length, device=q.device, dtype=torch.float32) |
| angles = torch.outer(positions, self.inv_freq) |
| cos, sin = angles.cos().to(q.dtype), angles.sin().to(q.dtype) |
| if length <= self.max_cached: |
| self._key, self._cos_cached, self._sin_cached = key, cos, sin |
| else: |
| self._key = self._cos_cached = self._sin_cached = None |
| if self.native: |
| from fla.modules.rotary.ops import rotary_embedding |
| return tuple(rotary_embedding(x=x, cos=cos, sin=sin, interleaved=False, |
| seqlen_offsets=0) for x in (q, k)) |
| |
| def rotate(x): |
| first, second = x.chunk(2, dim=-1) |
| c, s = cos[None, :, None], sin[None, :, None] |
| return torch.cat((first * c - second * s, first * s + second * c), dim=-1) |
| return rotate(q), rotate(k) |
|
|
|
|
| class MambaOuterBlock(nn.Module): |
| def __init__(self, cfg, index): |
| super().__init__() |
| from mamba_ssm.modules.mamba2 import Mamba2 |
| from pretraining_mamba2_outer.deterministic_backward import install |
| self.native_execution = install() |
| self.index, self.cfg = index, cfg |
| self.attn_norm = nn.RMSNorm(cfg.hidden, eps=1e-6) |
| key = f'{cfg.seed}:native-mamba2-outer:{index}'.encode() |
| seed = int.from_bytes(hashlib.sha256(key).digest()[:8], 'little') % (2**63 - 1) |
| |
| with torch.random.fork_rng(devices=[]): |
| torch.random.default_generator.manual_seed(seed) |
| self.attn = Mamba2(d_model=cfg.hidden, d_state=cfg.mamba_d_state, |
| d_conv=cfg.mamba_d_conv, expand=cfg.mamba_expand, |
| headdim=cfg.mamba_headdim, ngroups=1, layer_idx=index, |
| bias=False, conv_bias=True, rmsnorm=True, |
| D_has_hdim=False, norm_before_gate=False) |
| from pretraining_mamba2_outer.native_initialization import preserve_audited_initialization |
| self.scratch_rounding = preserve_audited_initialization(cfg, index, self.attn) |
|
|
| def _cached_chunk(self, u, conv_state, ssm_state, *, initial): |
| from causal_conv1d import causal_conv1d_fn, causal_conv1d_update |
| from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined |
| m = self.attn |
| b, t, _ = u.shape |
| projected = m.in_proj(u) |
| z, xbc, dt = torch.split(projected, [m.d_inner, m.d_inner + 2 * m.d_state, m.nheads], dim=-1) |
| conv_input = xbc.transpose(1, 2) |
| if initial: |
| |
| |
| |
| |
| |
| conv_state.copy_(F.pad(conv_input, (m.d_conv - t, 0))) |
| if (conv_input.stride(1) == 1 and conv_input.stride(2) > 1 |
| and (conv_input.shape[1] % 8 or conv_input.stride(0) % 8 |
| or conv_input.stride(2) % 8)): |
| conv_input = conv_input.contiguous() |
| mixed = causal_conv1d_fn(conv_input, m.conv1d.weight.squeeze(1), |
| m.conv1d.bias, activation=m.activation) |
| else: |
| mixed = causal_conv1d_update(conv_input, conv_state, |
| m.conv1d.weight.squeeze(1), m.conv1d.bias, |
| activation=m.activation) |
| x, B, C = torch.split(mixed.transpose(1, 2), [m.d_inner, m.d_state, m.d_state], dim=-1) |
| y, final = mamba_chunk_scan_combined( |
| x.reshape(b, t, m.nheads, m.headdim), dt, |
| -torch.exp(m.A_log.float()), B.unsqueeze(2), C.unsqueeze(2), |
| chunk_size=m.chunk_size, D=m.D, dt_bias=m.dt_bias, |
| dt_softplus=True, initial_states=None if initial else ssm_state, |
| return_final_states=True) |
| ssm_state.copy_(final) |
| return m.out_proj(m.norm(y.reshape(b, t, m.d_inner), z)) |
|
|
| def forward(self, x, *, lengths=None, block_mask=None, cap=None, cache=None): |
| u = self.attn_norm(x) |
| if cache is None: |
| return x + self.attn(u) |
| offset = cache.get_seq_length(self.index) |
| state = cache.states[self.index] |
| if state is None: |
| conv, ssm = self.attn.allocate_inference_cache(x.shape[0], 0) |
| ssm = ssm.float() |
| else: |
| conv, ssm = state['conv_state'], state['ssm_state'] |
| if x.shape[1] == 1: |
| output, _, _ = self.attn.step(u, conv, ssm) |
| else: |
| output = self._cached_chunk(u, conv, ssm, initial=(offset == 0)) |
| cache.update(layer_idx=self.index, offset=x.shape[1], conv_state=conv, ssm_state=ssm) |
| return x + output |
|
|
