| |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(REPO)) |
| sys.path.insert(0, str(REPO / "scripts")) |
|
|
| from model import AttentionMixer |
|
|
|
|
| class _Params: |
| def __init__(self, seqlen_offset=0): |
| self.seqlen_offset = seqlen_offset |
| self.key_value_memory_dict = {} |
|
|
|
|
| def _mixer(d=128): |
| torch.manual_seed(0) |
| return AttentionMixer(d_model=d, head_dim=64, device="cpu", dtype=torch.float32, |
| layer_idx=3).eval() |
|
|
|
|
| def test_uncached_forward_is_unchanged_and_prefill_matches_it(): |
| m = _mixer() |
| x = torch.randn(1, 11, 128) |
| with torch.no_grad(): |
| ref = m(x) |
| pre = m(x, inference_params=_Params(0)) |
| assert torch.allclose(ref, pre, atol=1e-6) |
|
|
|
|
| def test_stepwise_decode_equals_full_forward(): |
| m = _mixer() |
| x = torch.randn(1, 13, 128) |
| with torch.no_grad(): |
| ref = m(x) |
| p = _Params(0) |
| outs = [m(x[:, :5], inference_params=p)] |
| p.seqlen_offset = 5 |
| for i in range(5, 13): |
| outs.append(m(x[:, i:i + 1], inference_params=p)) |
| p.seqlen_offset = i + 1 |
| step = torch.cat(outs, dim=1) |
| assert torch.allclose(ref, step, atol=1e-5), (ref - step).abs().max() |
| assert p.key_value_memory_dict[3][0].shape[2] == 13 |
|
|
|
|
| def test_multi_position_step_onto_cache_is_causal_and_correct(): |
| m = _mixer() |
| x = torch.randn(1, 10, 128) |
| with torch.no_grad(): |
| ref = m(x) |
| p = _Params(0) |
| a = m(x[:, :4], inference_params=p); p.seqlen_offset = 4 |
| b = m(x[:, 4:], inference_params=p) |
| assert torch.allclose(ref, torch.cat([a, b], dim=1), atol=1e-5) |
|
|
|
|
| def test_layer_idx_is_required_for_the_cache(): |
| m = AttentionMixer(d_model=128, head_dim=64, device="cpu", dtype=torch.float32) |
| with pytest.raises(RuntimeError): |
| m(torch.randn(1, 3, 128), inference_params=_Params(0)) |
|
|
|
|
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA (mamba_ssm kernels)") |
| class TestDepthGenerateIdentityGPU: |
| @pytest.fixture(scope="class") |
| def tdc(self): |
| import test_depth_checkpoint as t |
| return t |
|
|
| def _model(self, hybrid, tdc): |
| from model import MambaCoTModel |
| torch.manual_seed(0) |
| m = MambaCoTModel(model_name="state-spaces/mamba2-130m", device="cuda", |
| dtype=torch.float32, mtp_num_heads=0) |
| if hybrid: |
| m.enable_hybrid_attention(top_k=2) |
| m.enable_depth_module(num_levels=tdc.NUM_LEVELS, codebook_size=64, |
| d_depth=64, depth_layers=1, depth_arch="gru") |
| return m.eval() |
|
|
| def _prefix(self, m): |
| torch.manual_seed(1) |
| text = torch.randint(5, 1000, (1, 9)) |
| return torch.cat([text, torch.tensor([[m.token_registry.audio_start_id]])], dim=1).cuda() |
|
|
| def _run(self, m, tdc, cached, sg=None, n=24): |
| torch.manual_seed(2) |
| with torch.no_grad(): |
| return tdc.depth_generate( |
| m, self._prefix(m), m.token_registry, max_frames=n, min_frames=n, |
| allow_stop=False, feedback="all", cross_frame_decode=False, |
| device="cuda", cached=cached, silence_guard=sg) |
|
|
| @pytest.mark.parametrize("hybrid", [False, True], ids=["pure", "hybrid"]) |
| def test_cached_equals_recompute_greedy(self, tdc, hybrid): |
| m = self._model(hybrid, tdc) |
| a = self._run(m, tdc, cached=False) |
| b = self._run(m, tdc, cached=True) |
| assert a is not None and b is not None |
| assert a == b, f"first divergence at frame {next(i for i in range(len(a[0])) if any(a[k][i] != b[k][i] for k in range(len(a))))}" |
|
|
| def test_rewind_reprefill_matches_recompute(self, tdc): |
| m = self._model(False, tdc) |
| base = self._run(m, tdc, cached=False) |
| pause = sorted(set(base[0][:6])) |
| mk = lambda: tdc.SilenceGuard(codes=pause, budget=2, stop_after=99, |
| retries=1, retry_temp=0.0) |
| g0, g1 = mk(), mk() |
| a = self._run(m, tdc, cached=False, sg=g0) |
| b = self._run(m, tdc, cached=True, sg=g1) |
| if g0.stats().get("rewinds", 0) == 0: |
| pytest.skip("guard did not fire on this init; rewind path not exercised") |
| assert g0.stats() == g1.stats() |
| assert a == b |
|
|
| def test_cached_plan_injection_matches_recompute_teacher_forced(self, tdc): |
| m = self._model(False, tdc) |
| m.enable_plan_injection() |
| m.eval() |
| torch.manual_seed(3) |
| torch.nn.init.normal_(m.plan_injector.proj.weight, std=0.05) |
| torch.nn.init.normal_(m.plan_injector.proj.bias, std=0.05) |
| reg = m.token_registry |
| prefix = self._prefix(m) |
| P, n, NL = prefix.shape[1], 24, tdc.NUM_LEVELS |
| sizes = m.plan_injector.bin_sizes |
| g = torch.Generator().manual_seed(4) |
| bins = torch.stack([torch.randint(0, s, (3,), generator=g) for s in sizes], |
| dim=-1).unsqueeze(0).cuda() |
| sched = tdc.build_decode_plan_schedule([6, 9, 9], prefix_len=P, max_frames=n, |
| mode="pack", device="cuda") |
| backbone = m.backbone.backbone |
| emb = backbone.embedding |
| gen = torch.Generator(device="cuda").manual_seed(7) |
| frames = [torch.randint(0, 64, (NL,), generator=gen, device="cuda") for _ in range(n)] |
|
|
| def ids_codes(t): |
| ids = torch.cat([prefix, torch.full((1, t), reg.audio_start_id, |
| dtype=torch.long, device="cuda")], 1) |
| codes = torch.full((1, P + t, NL), -1, dtype=torch.long, device="cuda") |
| for i in range(t): |
| codes[0, P + i] = frames[i] |
| return ids, codes |
|
|
| def stage(ids, codes, lo, plan): |
| emb.set_codes(codes[:, lo:]) |
| if plan: |
| emb.set_plan(m.plan_injector(bins, sched[:, :ids.shape[1]])[:, lo:]) |
|
|
| def run(plan): |
| re, ca = [], [] |
| with torch.no_grad(): |
| for t in range(n): |
| ids, codes = ids_codes(t) |
| stage(ids, codes, 0, plan) |
| re.append(backbone(ids)[:, -1, :].float().clone()) |
| ip = tdc._fresh_inference_cache(P + n + 8); n_fed = 0 |
| for t in range(n): |
| ids, codes = ids_codes(t) |
| lo = n_fed |
| stage(ids, codes, lo, plan) |
| ip.seqlen_offset = n_fed |
| h = backbone(ids[:, lo:], inference_params=ip) |
| n_fed = ids.shape[1] |
| ca.append(h[:, -1, :].float().clone()) |
| lm = sum(int(m.backbone.lm_head(a).argmax()) == int(m.backbone.lm_head(b).argmax()) |
| for a, b in zip(re, ca)) |
| dp = sum(list(m.depth_module.generate_frame(a)) == list(m.depth_module.generate_frame(b)) |
| for a, b in zip(re, ca)) |
| rel = [float((a - b).norm() / (a.norm() + 1e-9)) for a, b in zip(re, ca)] |
| return re, rel, lm, dp |
|
|
| with_plan, rel, lm, dp = run(True) |
| assert max(rel) < 5e-3, f"plan path: hidden rel {max(rel):.2e} (per frame {rel})" |
| assert lm == n and dp == n, f"plan path: lm-head {lm}/{n}, depth {dp}/{n}" |
| no_plan, rel0, _, _ = run(False) |
| moved = max(float((a - b).norm() / (a.norm() + 1e-9)) for a, b in zip(with_plan, no_plan)) |
| assert moved > 1e-2, f"the randomised injector moved the hidden state by only {moved:.2e}" |
| assert max(rel0) < 5e-3 |
|
|
