"""Mamba-2 (SSD) from the authors' code, running where Triton cannot. Upstream's Mamba-2 block calls fused Triton kernels that need CUDA. Off GPU the block takes upstream's own unfused path and the chunked scan is computed by their reference implementation (``ssd_minimal.py``, "the same as Listing 1 from the paper"), reached through our adapter in :mod:`torch_dimensions.mixers.mamba2_compat`. The adapter is ours, so it is the thing that has to be *proved*: the tests below check it against an independent naive SSM recurrence written from the definition, in float64, with the awkward cases turned on (a length that is not a whole number of chunks, fewer groups than heads, dt bias, the D skip and the z gate). The recurrence and the chunked algorithm are mathematically equal, so anything above float64 epsilon is a bug in the plumbing — which is exactly how a hardcoded ``.float()`` on ``dt_bias`` was caught. """ from __future__ import annotations import pytest import torch import torch.nn.functional as F import torch_dimensions as td pytest.importorskip("einops", reason="the vendored Mamba-2 needs the [upstream] extra") from torch_dimensions.mixers.mamba2_compat import mamba_chunk_scan_combined # noqa: E402 from torch_dimensions.mixers.upstream import UpstreamMamba2Mixer # noqa: E402 def _naive_ssd(x, dt, A, B, C, D=None, z=None, dt_bias=None): """The recurrence the chunked SSD algorithm claims to equal, written from the definition: state <- state * exp(dt*A) + dt * x B^T, y = state C.""" b, length, h, p = x.shape g = B.shape[2] dtp = F.softplus(dt + dt_bias) if dt_bias is not None else F.softplus(dt) be = B.repeat_interleave(h // g, dim=2) ce = C.repeat_interleave(h // g, dim=2) state = torch.zeros(b, h, p, B.shape[-1], dtype=x.dtype) ys = [] for t in range(length): decay = torch.exp(dtp[:, t] * A) # (b, h) state = state * decay[..., None, None] + torch.einsum( "bh,bhp,bhn->bhpn", dtp[:, t], x[:, t], be[:, t] ) ys.append(torch.einsum("bhpn,bhn->bhp", state, ce[:, t])) y = torch.stack(ys, dim=1) if D is not None: y = y + x * D[:, None] if z is not None: y = y * F.silu(z) return y @pytest.mark.parametrize("length,chunk,ngroups", [(32, 8, 4), (30, 8, 2), (32, 32, 4), (17, 4, 1)]) def test_adapter_equals_the_definition(length, chunk, ngroups): torch.manual_seed(0) b, h, p, n = 2, 4, 8, 16 kw = dict(dtype=torch.float64) x = torch.randn(b, length, h, p, **kw) dt = torch.randn(b, length, h, **kw) a = -torch.rand(h, **kw).exp() bb = torch.randn(b, length, ngroups, n, **kw) cc = torch.randn(b, length, ngroups, n, **kw) d = torch.randn(h, **kw) z = torch.randn(b, length, h, p, **kw) dt_bias = torch.randn(h, **kw) got = mamba_chunk_scan_combined( x, dt, a, bb, cc, chunk, D=d, z=z, dt_bias=dt_bias, dt_softplus=True ) want = _naive_ssd(x, dt, a, bb, cc, D=d, z=z, dt_bias=dt_bias) rel = ((got - want).abs().max() / want.abs().max()).item() assert rel < 1e-13, rel def test_adapter_keeps_float64_precision(): """The regression that motivated the parametrization above: a hardcoded float32 cast anywhere in the plumbing shows up here as ~1e-8.""" torch.manual_seed(1) b, length, h, p, n = 1, 16, 2, 4, 8 kw = dict(dtype=torch.float64) got = mamba_chunk_scan_combined( torch.randn(b, length, h, p, **kw), torch.zeros(b, length, h, **kw), -torch.ones(h, **kw), torch.ones(b, length, h, n, **kw), torch.ones(b, length, h, n, **kw), 8, dt_bias=torch.full((h,), 0.5, **kw), dt_softplus=True, ) assert got.dtype == torch.float64 def test_variable_length_is_refused_not_wrong(): """Ragged batches are a Triton-kernel feature; the reference path says so instead of quietly computing something else.""" kw = dict(dtype=torch.float32) args = ( torch.randn(1, 8, 2, 4, **kw), torch.randn(1, 8, 2, **kw), -torch.ones(2, **kw), torch.randn(1, 8, 2, 8, **kw), torch.randn(1, 8, 2, 8, **kw), 4, ) with pytest.raises(NotImplementedError, match="variable-length"): mamba_chunk_scan_combined(*args, seq_idx=torch.zeros(1, 8, dtype=torch.long)) def test_the_block_runs_and_sweeps_a_lattice(): torch.manual_seed(0) lat = td.Lattice(shape=(3, 4), names=("y", "x")) model = td.Mamba2(64, 2, lattice=lat, mixer_kwargs={"headdim": 32, "d_state": 16}).eval() x = torch.randn(2, 3, 4, 64) with torch.no_grad(): y = model(x) assert y.shape == x.shape assert torch.isfinite(y).all() def test_version_two_and_the_named_class_are_the_same_model(): lat = td.Lattice(shape=(3, 4), names=("y", "x")) kw = dict(mixer_kwargs={"headdim": 32, "d_state": 16}) by_flag = td.Mamba(64, 2, lattice=lat, version=2, **kw) by_name = td.Mamba2(64, 2, lattice=lat, **kw) assert type(by_flag.nd.mixers[0]) is type(by_name.nd.mixers[0]) assert by_flag.config["version"] == by_name.config["version"] == 2 def test_save_load_round_trips_the_version(tmp_path): torch.manual_seed(0) lat = td.Lattice(shape=(3, 4), names=("y", "x")) m = td.Mamba2(64, 2, lattice=lat, mixer_kwargs={"headdim": 32, "d_state": 16}).eval() p = tmp_path / "m2.td" m.save(p) r = td.load(p).eval() x = torch.randn(2, 3, 4, 64) with torch.no_grad(): assert torch.equal(m(x), r(x)) assert r.config["version"] == 2 def test_unsupported_versions_are_refused_with_a_reason(): # Versions 1, 2 and 3 all exist; anything else is a typo, not a model. for bad in (0, 4): with pytest.raises(ValueError, match="must be 1, 2 or 3"): td.Mamba(64, 1, version=bad) with pytest.raises(ValueError, match="no portable build of Mamba-2"): td.Mamba(64, 1, version=2, portable=True) @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="no MPS device") def test_mamba2_on_mps_matches_cpu(): torch.manual_seed(0) cpu = UpstreamMamba2Mixer(64, d_state=16, headdim=32).eval() mps = UpstreamMamba2Mixer(64, d_state=16, headdim=32).to("mps") mps.load_state_dict({k: v.to("mps") for k, v in cpu.state_dict().items()}) mps.eval() x = torch.randn(2, 24, 64) with torch.no_grad(): diff = (cpu(x) - mps(x.to("mps")).cpu()).abs().max().item() assert diff < 1e-5, diff xg = torch.randn(2, 24, 64, device="mps", requires_grad=True) mps.train() mps(xg).pow(2).mean().backward() assert torch.isfinite(xg.grad).all() # --- which implementation the block reaches for ------------------------------- def test_the_fused_path_is_chosen_per_tensor_not_per_machine(): """`use_mem_eff_path` used to default to `torch.cuda.is_available()`, which is a property of the *box*. On a CUDA machine that made every CPU-resident Mamba-2 ask for a kernel it could not reach, and the block refused on its first forward — so a CPU sanity check, a CPU test, or the CPU half of a device comparison all raised `NotImplementedError` on a machine that had a GPU, and only on such a machine. Found by running the suite on an RTX 5090. The predicate is now the input tensor's, evaluated in `forward`, matching how the rest of the library dispatches. """ mixer = UpstreamMamba2Mixer(64, d_state=16, headdim=32).eval() x = torch.randn(2, 16, 64) with torch.no_grad(): out = mixer(x) assert out.shape == x.shape assert torch.isfinite(out).all() # A CPU tensor can never take the fused path, whatever the box has. assert mixer.block.use_mem_eff_path is False def test_an_explicit_choice_is_left_alone(): """Only the *default* is deferred. A caller who names the flag is making a decision about upstream's own switch and must keep it.""" chosen = UpstreamMamba2Mixer(64, d_state=16, headdim=32, use_mem_eff_path=False) assert chosen._mem_eff_is_ours is False assert chosen.block.use_mem_eff_path is False default = UpstreamMamba2Mixer(64, d_state=16, headdim=32) assert default._mem_eff_is_ours is True def test_float64_takes_the_reference_path(): """The fused kernel has no float64 instantiation, and float64 is the control that separates reassociation from a different computation in the agreement benchmark — so it must not silently fail to run.""" mixer = UpstreamMamba2Mixer(64, d_state=16, headdim=32).double().eval() with torch.no_grad(): out = mixer(torch.randn(2, 16, 64, dtype=torch.float64)) assert out.dtype == torch.float64 assert torch.isfinite(out).all()