1"""Phase 7: cached incremental decoding == full-prefix recomputation (tolerance-based), bypass bitwise."""2import torch34from prisme.cache import RequestBatchCache5from prisme.generation import step_logits, generate6from prisme.flops import action_cost_proxy789def test_incremental_equals_full_prefix_forced(model_nonzero, cfg, tokens):10 m = model_nonzero11 g = torch.Generator().manual_seed(3)12 acts = torch.randint(0, cfg.n_actions + 1, (3, 12, 4), generator=g)13 full = m(tokens, acts)["logits"]14 cache = RequestBatchCache(layers=m.backbone.new_cache())15 outs = []16 # prefill 5 tokens then decode one at a time17 lg, _, _ = step_logits(m, tokens[:, :5], cache, 4, forced_actions=acts[:, :5])18 outs.append(lg)19 for t in range(5, 12):20 lg, _, _ = step_logits(m, tokens[:, t:t + 1], cache, 4, forced_actions=acts[:, t:t + 1])21 outs.append(lg)22 inc = torch.cat(outs, 1)23 assert torch.allclose(inc, full, atol=1e-4, rtol=1e-4)24 assert cache.length == 12 and cache.bank.length == 12252627def test_incremental_equals_full_prefix_policy(model_nonzero, cfg, tokens):28 m = model_nonzero29 with torch.no_grad():30 m.controller.heads.bias[0] = 0.5 # make the controller edit sometimes31 m.controller.heads.weight.normal_(std=0.5)32 cost = torch.tensor(action_cost_proxy(cfg, 12))33 full = m.forward_policy(tokens, 3, cost)34 cache = RequestBatchCache(layers=m.backbone.new_cache())35 outs, traces = [], []36 for t in range(12):37 lg, tr, _ = step_logits(m, tokens[:, t:t + 1], cache, 3, cost=cost)38 outs.append(lg); traces.append(tr)39 inc = torch.cat(outs, 1)40 assert torch.allclose(inc, full["logits"], atol=1e-4, rtol=1e-4)41 assert torch.equal(torch.cat(traces, 1), full["actions"])42 assert int(full["n_edits"].sum()) > 0434445def test_budget_zero_bypass_bitwise_and_no_workspace(model_nonzero, cfg, tokens):46 m = model_nonzero47 ref = m.backbone.head(m.backbone(tokens))48 cache = RequestBatchCache(layers=m.backbone.new_cache())49 outs = []50 for t in range(12):51 lg, tr, ne = step_logits(m, tokens[:, t:t + 1], cache, 0)52 assert tr is None and int(ne.sum()) == 053 outs.append(lg)54 inc = torch.cat(outs, 1)55 assert torch.allclose(inc, ref, atol=1e-4)56 # the explicit bypass branch of the full forward is bitwise identical to the backbone57 assert torch.equal(m(tokens, None)["logits"], ref)585960def test_cache_isolation_between_requests(model_nonzero, cfg, tokens):61 m = model_nonzero62 c1 = RequestBatchCache(layers=m.backbone.new_cache())63 c2 = RequestBatchCache(layers=m.backbone.new_cache())64 step_logits(m, tokens[:1, :6], c1, 2, forced_actions=torch.zeros(1, 6, 2, dtype=torch.long))65 lg_a, _, _ = step_logits(m, tokens[1:2, :6], c2, 2, forced_actions=torch.zeros(1, 6, 2, dtype=torch.long))66 c3 = RequestBatchCache(layers=m.backbone.new_cache())67 lg_b, _, _ = step_logits(m, tokens[1:2, :6], c3, 2, forced_actions=torch.zeros(1, 6, 2, dtype=torch.long))68 assert torch.equal(lg_a, lg_b)69 assert c1.length == 6 and c2.length == 670 c1.reset(); assert c1.length == 0 and c1.bank.length == 0 and c2.length == 6717273def test_generate_runs_and_caps_budget(model_nonzero, cfg, tokens):74 m = model_nonzero75 cost = torch.tensor(action_cost_proxy(cfg, 12))76 out = generate(m, tokens[:, :4], 6, max_budget=8, cost=cost)77 assert out["tokens"].shape == (3, 10)78 assert out["n_edits"] is None or int(out["n_edits"].max()) <= cfg.r_max79 out0 = generate(m, tokens[:, :4], 6, max_budget=0)80 assert out0["tokens"].shape == (3, 10)81