| """Model/runtime contract tests on a tiny random Qwen3.5 (CPU, exact torch reference |
| kernels) plus one GPU test on the real base model. |
| |
| * the functional backbone equals the HF forward, and a fork equals the concatenated read; |
| * field independence is bitwise: a field's head outputs are identical alone, reordered, |
| next to other fields, and in split vs merged schemas; |
| * the span pointer only addresses context tokens; |
| * only x-generate fields reach the LM head; |
| * responses validate against the schema; constrained decoding respects length/pattern. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import random |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| transformers = pytest.importorskip("transformers") |
|
|
| from deg import compile_schema, validate_value |
| from deg.infer import Constraint, field_query, fill |
| from deg.model import ( |
| base_snapshot, |
| load_model, |
| load_tokenizer, |
| token_span, |
| trim_span, |
| ) |
| from deg.request import FillOptions |
| from deg.response import FillError |
|
|
| try: |
| SNAPSHOT = base_snapshot() |
| except FileNotFoundError: |
| SNAPSHOT = None |
|
|
| needs_tokenizer = pytest.mark.skipif(SNAPSHOT is None, reason="Qwen3.5-0.8B not cached") |
|
|
| CONTEXT = ("Ticket from dana@example.com: order AB-1234 was charged twice on 2026-03-05, " |
| "total $1,250.50. Please refund today, I am furious. Items: 2 chairs, 1 desk.") |
| SCHEMA = { |
| "type": "object", |
| "properties": { |
| "department": {"enum": ["billing", "technical", "sales"], |
| "description": "Which team should handle this?"}, |
| "urgent": {"type": "boolean"}, |
| "anger": {"type": "integer", "minimum": 0, "maximum": 4}, |
| "email": {"type": ["string", "null"], "format": "email"}, |
| "order_id": {"type": ["string", "null"]}, |
| "total": {"type": ["number", "null"]}, |
| "charged_on": {"type": ["string", "null"], "format": "date"}, |
| "items": {"type": "array", "maxItems": 4, |
| "items": {"type": "object", |
| "properties": {"name": {"type": ["string", "null"]}, |
| "qty": {"type": "integer", "minimum": 1, |
| "maximum": 5}}, |
| "required": ["name", "qty"]}}, |
| "reply": {"type": "string", "x-generate": True, "maxLength": 40, "minLength": 3, |
| "pattern": "^[^{}]*$"}, |
| }, |
| "required": ["department", "urgent", "anger", "email", "order_id", "total", "charged_on", |
| "items", "reply"], |
| } |
|
|
|
|
| def tiny_lm(seed: int = 0): |
| from transformers import Qwen3_5ForCausalLM, Qwen3_5TextConfig |
|
|
| torch.manual_seed(seed) |
| cfg = Qwen3_5TextConfig( |
| vocab_size=248320, hidden_size=64, intermediate_size=128, num_hidden_layers=4, |
| layer_types=["linear_attention", "linear_attention", "linear_attention", |
| "full_attention"], |
| num_attention_heads=2, num_key_value_heads=1, head_dim=32, |
| linear_num_key_heads=2, linear_num_value_heads=2, linear_key_head_dim=16, |
| linear_value_head_dim=16, linear_conv_kernel_dim=4, tie_word_embeddings=True, |
| rope_parameters={"rope_type": "default", "rope_theta": 10000.0, |
| "partial_rotary_factor": 0.25, "mrope_section": [2, 1, 1], |
| "mrope_interleaved": True}, |
| ) |
| lm = Qwen3_5ForCausalLM(cfg).float().eval() |
| with torch.no_grad(): |
| for name, p in lm.named_parameters(): |
| if "norm" in name or "A_log" in name or "dt_bias" in name: |
| p.normal_(0.0, 0.3) |
| return lm |
|
|
|
|
| @pytest.fixture(scope="module") |
| def model(): |
| if SNAPSHOT is None: |
| pytest.skip("Qwen3.5-0.8B tokenizer not cached") |
| m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(), |
| tokenizer=load_tokenizer(SNAPSHOT)) |
| torch.manual_seed(1) |
| with torch.no_grad(): |
| for p in m.heads.parameters(): |
| p.normal_(0.0, 0.5) |
| return m |
|
|
|
|
| |
|
|
|
|
| def test_trim_span_strips_token_punctuation(): |
| text = 'Buy "Desk (White)." now, [Marcus Thorne] and Inc.,' |
| a = text.index('"') |
| assert text.__getitem__(slice(*trim_span(text, a, text.index(" now")))) == "Desk (White)." |
| a = text.index("[") |
| assert text.__getitem__(slice(*trim_span(text, a, text.index(" and")))) == "Marcus Thorne" |
| a = text.index("Inc") |
| assert text.__getitem__(slice(*trim_span(text, a, len(text)))) == "Inc." |
|
|
|
|
| def _cut(text: str, sub: str, extra: int = 0) -> str: |
| a = text.index(sub) |
| return text[slice(*trim_span(text, a, a + len(sub) + extra))] |
|
|
|
|
| def test_trim_span_keeps_value_stars_and_strips_wrappers(): |
| |
| t = "on card ****4821. It" |
| assert _cut(t, "****4821") == "****4821" |
| |
| t = "Status: **bold** today" |
| assert _cut(t, "**bold**") == "bold" |
| assert _cut(t, "**bold") == "bold" |
| assert _cut(t, "bold**") == "bold" |
| assert _cut("x *it*, y", "*it*,") == "it" |
| |
| assert _cut('vendor "ACME Corp" signed', '"ACME Corp"') == "ACME Corp" |
| |
| assert _cut("Grade: AB*, final", "AB*,") == "AB*" |
|
|
|
|
| def _pointer_logits(offsets, i, j): |
| T = len(offsets) |
| start = torch.full((T + 1,), -20.0) |
| end = torch.full((T + 1,), -20.0) |
| start[i] = end[j] = 10.0 |
| return start, end |
|
|
|
|
| def test_extract_result_tries_raw_span_when_trimmed_fails_constraints(): |
| from deg.infer import extract_result |
| from deg.model import EncodedContext |
|
|
| def run(ctx_text, prop, value_sub): |
| fld = compile_schema({"type": "object", "properties": {"v": prop}, |
| "required": ["v"]}).fields[0] |
| words = [(m.start(), m.end()) for m in __import__("re").finditer(r"\S+", ctx_text)] |
| a = ctx_text.index(value_sub) |
| i, j = token_span(tuple(words), a, a + len(value_sub)) |
| ctx = EncodedContext(ids=tuple(range(len(words))), ctx_start=0, offsets=tuple(words)) |
| start, end = _pointer_logits(words, i, j) |
| return extract_result(fld, ctx_text, ctx, start, end, torch.tensor(1.0), |
| FillOptions()).value |
|
|
| card = {"type": "string", "pattern": "^\\*{4}[0-9]{4}$", "minLength": 8} |
| assert run("Paid on card ****4821, it cleared.", card, "****4821,") == "****4821" |
| |
| quoted = {"type": "string", "pattern": '^"[A-Z]+"$'} |
| assert run('Use code "ABC" now', quoted, '"ABC"') == '"ABC"' |
| |
| assert run('vendor "ACME Corp" signed', {"type": "string"}, '"ACME Corp"') == "ACME Corp" |
| assert run("Grade: AB*, final", {"type": "string", "pattern": "^[A-Z]+\\*$"}, |
| "AB*,") == "AB*" |
|
|
|
|
| def test_constraint_prefix_and_final(): |
| fld = compile_schema({"type": "object", "properties": { |
| "code": {"type": "string", "x-generate": True, "pattern": "^[A-Z]{2}-\\d+$", |
| "maxLength": 8, "minLength": 4}}}).fields[0] |
| rule = Constraint(fld) |
| assert rule.prefix_ok("A") and rule.prefix_ok("AB-") and rule.prefix_ok("AB-12") |
| assert not rule.prefix_ok("a") and not rule.prefix_ok("AB-123456") |
| assert not rule.final_ok("AB-") and rule.final_ok("AB-12") |
|
|
|
|
| @needs_tokenizer |
| def test_token_span_covers_char_span(): |
| tok = load_tokenizer(SNAPSHOT) |
| enc = tok(CONTEXT, add_special_tokens=False, return_offsets_mapping=True) |
| offsets = tuple(tuple(x) for x in enc["offset_mapping"]) |
| a = CONTEXT.index("AB-1234") |
| i, j = token_span(offsets, a, a + 7) |
| assert offsets[i][0] <= a and offsets[j][1] >= a + 7 |
| got = trim_span(CONTEXT, offsets[i][0], offsets[j][1]) |
| assert CONTEXT[got[0]:got[1]] == "AB-1234" |
|
|
|
|
| |
|
|
|
|
| def test_backbone_matches_hf_and_fork_equals_concatenation(model, monkeypatch): |
| from transformers.models.qwen3_5 import modeling_qwen3_5 as hf |
|
|
| from deg.model import delta_rule_reference |
|
|
| def chunk_ref(q, k, v, g, beta, initial_state=None, output_final_state=False, **_): |
| return delta_rule_reference(q, k, v, g, beta, initial_state) |
|
|
| monkeypatch.setattr(hf, "torch_chunk_gated_delta_rule", chunk_ref) |
| rend = model.renderer |
| a = rend.enc(CONTEXT) |
| b = rend.enc("FIELD /order_id\nKIND: copy\nANSWER:") |
| with torch.no_grad(): |
| ref = model.backbone.text(torch.tensor([a + b])).last_hidden_state[0] |
| h1, st = model.backbone(torch.tensor([a]), torch.tensor([len(a)])) |
| h2, _ = model.backbone(torch.tensor([b]), torch.tensor([len(b)]), st) |
| |
| pad = torch.full((2, len(a) + 5), rend.pad_id) |
| pad[0, :len(a)] = torch.tensor(a) |
| pad[1, :3] = torch.tensor(a[:3]) |
| _, st2 = model.backbone(pad, torch.tensor([len(a), 3])) |
| fork = torch.full((1, len(b) + 9), rend.pad_id) |
| fork[0, :len(b)] = torch.tensor(b) |
| h3, _ = model.backbone(fork, torch.tensor([len(b)]), st2.select(torch.tensor([0]))) |
| mine = torch.cat([h1[0], h2[0]]) |
| torch.testing.assert_close(mine, ref, atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(h3[0, :len(b)], h2[0], atol=2e-4, rtol=2e-4) |
|
|
|
|
| |
|
|
|
|
| def _raw_outputs(model, schema): |
| """pointer -> raw head tensors (logits or start/end) for every parallel leaf.""" |
| from deg.infer import LeafInstance, walk |
|
|
| plan = compile_schema(schema) |
| ectx = model.renderer.context(CONTEXT) |
| with torch.no_grad(): |
| hidden, state = model.trunk([ectx]) |
| leaves = [x for x in walk(plan.root, {}) if isinstance(x, LeafInstance) |
| and x.field.phase == "parallel"] |
| queries = [model.renderer.query(field_query(x.field, x.ptr, list(x.items))) |
| for x in leaves] |
| outs = model.answer_batched(state, hidden, ectx, queries) |
| return {x.ptr: (o.logits if o.logits is not None else torch.cat([o.start, o.end])) |
| for x, o in zip(leaves, outs)} |
|
|
|
|
| def _sub(names, schema=SCHEMA): |
| return {"type": "object", |
| "properties": {n: schema["properties"][n] for n in names}, |
| "required": [n for n in names if n in schema["required"]]} |
|
|
|
|
| def test_field_outputs_are_bitwise_independent(model): |
| flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")] |
| merged = _raw_outputs(model, _sub(flat)) |
| assert len(merged) == len(flat) |
| reordered = _raw_outputs(model, _sub(list(reversed(flat)))) |
| |
| padded = {n: {"type": ["string", "null"], "description": f"filler {n}"} |
| for n in ("f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "f9")} |
| big_schema = {"type": "object", "properties": {**padded, **_sub(flat)["properties"]}, |
| "required": _sub(flat)["required"]} |
| big = _raw_outputs(model, big_schema) |
| assert len(big) > 8 |
| for name in flat: |
| alone = _raw_outputs(model, _sub([name])) |
| ptr = f"/{name}" |
| for other in (merged, reordered, big): |
| assert torch.equal(alone[ptr], other[ptr]), name |
|
|
|
|
| def test_fill_meta_is_identical_in_split_and_merged_schemas(model): |
| names = ["department", "urgent", "anger", "email", "order_id", "total", "charged_on"] |
| merged = fill(model, CONTEXT, _sub(names)) |
| for half in (names[:3], names[3:]): |
| part = fill(model, CONTEXT, _sub(half)) |
| for n in half: |
| assert part["meta"][f"/{n}"] == merged["meta"][f"/{n}"] |
| assert part["value"].get(n, "<omitted>") == merged["value"].get(n, "<omitted>") |
|
|
|
|
| |
|
|
|
|
| def test_span_pointer_addresses_context_tokens_only(model): |
| ectx = model.renderer.context(CONTEXT) |
| fld = compile_schema(_sub(["order_id"])).fields[0] |
| q = model.renderer.query(field_query(fld, "/order_id", [])) |
| with torch.no_grad(): |
| hidden, state = model.trunk([ectx]) |
| (out,) = model.answer_batched(state, hidden, ectx, [q]) |
| |
| assert out.start.shape[0] == len(ectx.offsets) + 1 == out.end.shape[0] |
| assert len(ectx.offsets) < len(ectx.ids) |
| for _ in range(5): |
| resp = fill(model, CONTEXT, _sub(["order_id", "email", "total", "charged_on"])) |
| for ptr, meta in resp["meta"].items(): |
| if meta.get("span") is not None: |
| a, b = meta["span"] |
| assert 0 <= a < b <= len(CONTEXT) |
|
|
|
|
| |
|
|
|
|
| def test_only_x_generate_fields_use_the_lm_head(model, monkeypatch): |
| calls = [] |
| real = model.backbone.lm_logits |
|
|
| def spy(hidden): |
| calls.append(hidden.shape) |
| return real(hidden) |
|
|
| monkeypatch.setattr(model.backbone, "lm_logits", spy) |
| no_gen = {**SCHEMA, "properties": {k: v for k, v in SCHEMA["properties"].items() |
| if k != "reply"}, |
| "required": [r for r in SCHEMA["required"] if r != "reply"]} |
| resp = fill(model, CONTEXT, no_gen) |
| assert calls == [] and resp["usage"]["generated_tokens"] == 0 |
| resp = fill(model, CONTEXT, SCHEMA) |
| gen = [p for p, m in resp["meta"].items() if m.get("mode") == "generate"] |
| assert gen == ["/reply"] |
| assert len(calls) == resp["meta"]["/reply"]["generated_tokens"] > 0 |
| assert resp["usage"]["generated_tokens"] == len(calls) |
|
|
|
|
| def test_response_validates_and_generation_respects_constraints(model): |
| for seed in range(3): |
| opts = FillOptions(temperature=0.8, seed=seed) if seed else FillOptions() |
| resp = fill(model, CONTEXT, SCHEMA, opts) |
| validate_value(compile_schema(SCHEMA), resp["value"]) |
| reply = resp["value"]["reply"] |
| assert 3 <= len(reply) <= 40 and "{" not in reply and "}" not in reply |
| assert resp["model"] == "deg-qwen3.5-0.8b-0.1" |
| assert set(resp["value"]) >= {"department", "urgent", "anger", "items", "reply"} |
|
|
|
|
| def test_sampled_generation_is_reproducible_with_seed(model): |
| opts = FillOptions(temperature=1.0, seed=7) |
| a = fill(model, CONTEXT, SCHEMA, opts) |
| b = fill(model, CONTEXT, SCHEMA, opts) |
| assert a == b |
|
|
|
|
| def test_required_field_without_any_parsable_span_is_a_fill_error(model): |
| schema = {"type": "object", "properties": {"when": {"type": "string", |
| "format": "date-time"}}, |
| "required": ["when"]} |
| with pytest.raises(FillError): |
| fill(model, "no timestamps in here at all", schema) |
|
|
|
|
| def _with_note(schema=SCHEMA, budget=None): |
| note = {"type": "string", "x-generate": True, "maxLength": 30, |
| "description": "internal note for the agent"} |
| if budget is not None: |
| note["x-max-tokens"] = budget |
| return {**schema, "properties": {**schema["properties"], "note": note}, |
| "required": [*schema["required"], "note"]} |
|
|
|
|
| def test_generate_fields_are_independent_of_each_other(model): |
| """Generate fields condition on decided/extracted values only and decode as |
| batch-invariant rows: /reply is the same alone and next to another generate field.""" |
| for opts in (FillOptions(), FillOptions(temperature=0.9, seed=3)): |
| alone = fill(model, CONTEXT, SCHEMA, opts) |
| both = fill(model, CONTEXT, _with_note(), opts) |
| assert both["value"]["reply"] == alone["value"]["reply"] |
| assert both["meta"]["/reply"] == alone["meta"]["/reply"] |
| assert "note" in both["value"] |
|
|
|
|
| def test_per_field_generate_budget(model): |
| resp = fill(model, CONTEXT, _with_note(budget=3), FillOptions(max_new_tokens=64)) |
| assert resp["meta"]["/note"]["generated_tokens"] <= 3 |
| assert resp["meta"]["/reply"]["generated_tokens"] > 3 or "reply" in resp["value"] |
|
|
|
|
| def test_fill_many_matches_fill_alone(model): |
| from deg.infer import fill_many |
|
|
| other = "Dana again: the desk arrived broken. Order AB-9 total $80." |
| jobs = [(CONTEXT, SCHEMA, None), (other, _sub(["department", "order_id", "total"]), None), |
| (CONTEXT, {"type": "object", "properties": {"x": {"type": "object"}}, |
| "required": ["x"], "bogus": 1}, None), |
| (other, _with_note(), FillOptions(temperature=0.7, seed=5))] |
| many = fill_many(model, jobs) |
| assert len(many) == len(jobs) |
| for job, got in zip(jobs, many): |
| try: |
| want = fill(model, *job) |
| except Exception as exc: |
| assert type(got) is type(exc) |
| continue |
| assert got == want |
|
|
|
|
| def test_score_level_labels_become_option_lines(): |
| from deg.infer import outcome_texts |
|
|
| schema = {"type": "object", "properties": {"anger": { |
| "type": "integer", "minimum": 0, "maximum": 2, |
| "x-levels": {"0": "calm", "2": "furious"}}}, "required": ["anger"]} |
| fld = compile_schema(schema).fields[0] |
| assert fld.level_labels == ("calm", None, "furious") |
| assert outcome_texts(fld) == ("0: calm", "1", "2: furious") |
|
|
|
|
| |
|
|
|
|
| class _LargestTensor: |
| """Records the largest tensor (numel) any torch op produces while active.""" |
|
|
| def __enter__(self): |
| from torch.utils._python_dispatch import TorchDispatchMode |
|
|
| rec = self |
|
|
| class Mode(TorchDispatchMode): |
| def __torch_dispatch__(self, func, types, args=(), kwargs=None): |
| out = func(*args, **(kwargs or {})) |
| for t in (out if isinstance(out, (tuple, list)) else (out,)): |
| if isinstance(t, torch.Tensor): |
| rec.largest = max(rec.largest, t.numel()) |
| return out |
|
|
| self.largest = 0 |
| self.mode = Mode() |
| self.mode.__enter__() |
| return self |
|
|
| def __exit__(self, *exc): |
| self.mode.__exit__(*exc) |
|
|
|
|
| def test_long_prefill_streams_in_chunks_without_quadratic_tensors(model, monkeypatch): |
| """No O(L^2) tensor anywhere in a long prefill (no dense L x L causal mask, no |
| [B, L, P+L] mask), and the chunked prefill equals the one-shot read.""" |
| import deg.model as dm |
|
|
| rend = model.renderer |
| ids = rend.context(" ".join([CONTEXT] * 50)).ids |
| L = len(ids) |
| assert L > 2000 |
| x, n = torch.tensor([ids]), torch.tensor([L]) |
| with torch.no_grad(): |
| h_one, st_one = model.backbone(x, n) |
| monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 256) |
| monkeypatch.setattr(dm, "PREFILL_CHUNK", 256) |
| with _LargestTensor() as big: |
| h, st = model.backbone(x, n) |
| assert big.largest < L * L // 4, big.largest |
| torch.testing.assert_close(h, h_one, atol=2e-4, rtol=2e-4) |
| for (r, t), (r1, t1) in zip(st.lin, st_one.lin): |
| torch.testing.assert_close(r, r1, atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(t, t1, atol=2e-4, rtol=2e-4) |
| for (k, v), (k1, v1) in zip(st.kv, st_one.kv): |
| torch.testing.assert_close(k, k1, atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(v, v1, atol=2e-4, rtol=2e-4) |
| assert torch.equal(st.kv_mask, st_one.kv_mask) and torch.equal(st.lengths, st_one.lengths) |
|
|
|
|
| def test_long_prefill_padded_batch_and_prefix_state(model, monkeypatch): |
| """Chunked prefill over right-padded rows and on top of an existing prefix state |
| equals the one-shot read for every row's real tokens.""" |
| import deg.model as dm |
|
|
| rend = model.renderer |
| a = rend.enc(" ".join([CONTEXT] * 12)) |
| b = rend.enc(" ".join([CONTEXT] * 5)) |
| pre = rend.enc("Note: ") |
| x = torch.full((2, len(a)), rend.pad_id) |
| x[0], x[1, :len(b)] = torch.tensor(a), torch.tensor(b) |
| n = torch.tensor([len(a), len(b)]) |
| with torch.no_grad(): |
| _, st0 = model.backbone(torch.tensor([pre, pre]), torch.tensor([len(pre)] * 2)) |
| h_one, s_one = model.backbone(x, n, st0) |
| monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 64) |
| monkeypatch.setattr(dm, "PREFILL_CHUNK", 64) |
| h, s = model.backbone(x, n, st0) |
| for row, m in enumerate((len(a), len(b))): |
| torch.testing.assert_close(h[row, :m], h_one[row, :m], atol=2e-4, rtol=2e-4) |
| for (r, t), (r1, t1) in zip(s.lin, s_one.lin): |
| torch.testing.assert_close(r, r1, atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(t, t1, atol=2e-4, rtol=2e-4) |
| assert torch.equal(s.kv_mask, s_one.kv_mask) |
| for (k, v), (k1, v1) in zip(s.kv, s_one.kv): |
| for row in range(2): |
| real = s.kv_mask[row] |
| torch.testing.assert_close(k[row][:, real], k1[row][:, real], atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(v[row][:, real], v1[row][:, real], atol=2e-4, rtol=2e-4) |
|
|
|
|
| def test_fork_batch_split_on_oom_gives_the_same_outputs(model, monkeypatch): |
| """A fork batch that runs out of memory is split in half and retried; the outputs are |
| the same bits as the unsplit batch.""" |
| from deg.model import Backbone |
|
|
| flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")] |
| want = _raw_outputs(model, _sub(flat)) |
| orig = Backbone.fork_forward |
| calls = [] |
|
|
| def tight(self, ids, lengths, fs, keep=False, decode=False, sel=None): |
| calls.append(ids.shape[0]) |
| if ids.shape[0] > 2: |
| raise torch.OutOfMemoryError("synthetic") |
| return orig(self, ids, lengths, fs, keep, decode, sel) |
|
|
| monkeypatch.setattr(Backbone, "fork_forward", tight) |
| got = _raw_outputs(model, _sub(flat)) |
| assert max(calls) > 2 and any(c <= 2 for c in calls) |
| assert got.keys() == want.keys() |
| for k in want: |
| assert torch.equal(got[k], want[k]), k |
|
|
|
|
| def test_blocked_span_logits_match_one_shot(model, monkeypatch): |
| """Span logits computed in answer-row micro-batches and context-row blocks (the path |
| for long contexts / big schemas) equal the one-shot path, and big sets wait in host |
| memory, not on the device.""" |
| import deg.model as dm |
|
|
| flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")] |
| want = _raw_outputs(model, _sub(flat)) |
| monkeypatch.setattr(dm, "HEAD_ROWS", 7) |
| monkeypatch.setattr(dm, "HEAD_LOGIT_BYTES", 8 * 3 * 64) |
| got = _raw_outputs(model, _sub(flat)) |
| assert got.keys() == want.keys() |
| for k in want: |
| torch.testing.assert_close(got[k], want[k], atol=1e-5, rtol=1e-5) |
|
|
|
|
| def test_fork_hidden_keeps_only_head_positions(model): |
| """Between the fork pass and the heads only the answer/option rows of each fork are |
| held (not every fork token), and outputs are unchanged.""" |
| from deg.model import _Picked |
|
|
| h = torch.arange(12.0).reshape(6, 2) |
| p = _Picked(h, (5, 1, 3)) |
| assert p.rows.shape == (3, 2) |
| for i in (1, 3, 5): |
| assert torch.equal(p[i], h[i]) |
|
|
|
|
| def test_consume_extends_state_without_holding_two_copies(model): |
| """forward(consume=True) gives the same outputs as a normal extension and releases the |
| given state's K/V (so a long context is never held twice); same for the chunked |
| path.""" |
| import deg.model as dm |
|
|
| rend = model.renderer |
| ctx = torch.tensor([rend.enc(" ".join([CONTEXT] * 3))]) |
| more = torch.tensor([rend.enc("KNOWN VALUES:\n/a = 1\n")]) |
| with torch.no_grad(): |
| _, st = model.backbone(ctx, torch.tensor([ctx.shape[1]])) |
| h_ref, s_ref = model.backbone(more, torch.tensor([more.shape[1]]), st) |
| assert all(kv is not None for kv in st.kv) |
| h, s = model.backbone(more, torch.tensor([more.shape[1]]), st, consume=True) |
| assert all(kv is None for kv in st.kv) |
| assert torch.equal(h, h_ref) |
| for (k, v), (k1, v1) in zip(s.kv, s_ref.kv): |
| assert torch.equal(k, k1) and torch.equal(v, v1) |
| _, st = model.backbone(ctx, torch.tensor([ctx.shape[1]])) |
| old = dm.DENSE_PREFILL_MAX, dm.PREFILL_CHUNK |
| dm.DENSE_PREFILL_MAX = dm.PREFILL_CHUNK = 4 |
| try: |
| h2, s2 = model.backbone(more, torch.tensor([more.shape[1]]), st, consume=True) |
| finally: |
| dm.DENSE_PREFILL_MAX, dm.PREFILL_CHUNK = old |
| assert all(kv is None for kv in st.kv) |
| torch.testing.assert_close(h2, h_ref, atol=2e-4, rtol=2e-4) |
|
|
|
|
| def test_span_candidates_band_is_linear_and_unchanged(): |
| """Span candidates over a long context never build the T x T grid; on a short one |
| they equal the dense-grid reference.""" |
| from deg.infer import MAX_SPAN_TOKENS, SPAN_CANDIDATES, span_candidates |
|
|
| def dense(start, end, temperature): |
| ls = torch.log_softmax(start.float() / temperature, -1) |
| le = torch.log_softmax(end.float() / temperature, -1) |
| T = ls.shape[0] - 1 |
| grid = ls[:T, None] + le[None, :T] |
| idx = torch.arange(T) |
| ok = (idx[None] >= idx[:, None]) & (idx[None] - idx[:, None] < MAX_SPAN_TOKENS) |
| grid = grid.masked_fill(~ok, float("-inf")) |
| vals, flat = grid.flatten().topk(min(SPAN_CANDIDATES, int(ok.sum()))) |
| return [(float(v), int(f) // T, int(f) % T) for v, f in zip(vals, flat)] |
|
|
| g = torch.Generator().manual_seed(0) |
| temp = torch.tensor(1.3) |
| for T in (1, 5, 127, 128, 129, 400): |
| s, e = torch.randn(T + 1, generator=g), torch.randn(T + 1, generator=g) |
| got, null = span_candidates(s, e, temp) |
| assert got == dense(s, e, temp), T |
| T = 20000 |
| s, e = torch.randn(T + 1, generator=g), torch.randn(T + 1, generator=g) |
| with _LargestTensor() as big: |
| got, _ = span_candidates(s, e, temp) |
| assert big.largest <= T * MAX_SPAN_TOKENS + MAX_SPAN_TOKENS |
| assert all(0 <= j - i < MAX_SPAN_TOKENS for _, i, j in got) |
|
|
|
|
| def test_token_span_bisect_matches_linear_scan(): |
| offsets = ((0, 3), (3, 4), (4, 4), (4, 9), (9, 12), (12, 12), (13, 20)) |
| for start in range(0, 20): |
| for end in range(start + 1, 21): |
| try: |
| i = next(k for k, (_, b) in enumerate(offsets) if b > start) |
| j = max(k for k, (a, _) in enumerate(offsets) if a < end) |
| except (StopIteration, ValueError): |
| continue |
| assert token_span(offsets, start, end) == (i, j), (start, end) |
|
|
|
|
| |
|
|
|
|
| def test_training_example_and_losses(model): |
| from deg.train import deg_example, forward_losses, jev_example |
|
|
| a = CONTEXT.index("AB-1234") |
| rec = {"uid": "t1", "schema": _sub(["department", "order_id", "items", "reply"]), |
| "context": CONTEXT, |
| "value": {"department": "billing", "order_id": "AB-1234", |
| "items": [{"name": "chairs", "qty": 2}], "reply": "Refund issued."}, |
| "spans": {"/order_id": [a, a + 7], |
| "/items/0/name": [CONTEXT.index("chairs"), |
| CONTEXT.index("chairs") + 6]}, |
| "span_alternatives": {}, "span_probs": {}, |
| "decide_probs": {"/department": [0.8, 0.15, 0.05]}} |
| model.meta["array_mode"] = "index" |
| try: |
| ex_index = deg_example(model, rec, random.Random(0), shuffle_p=0.0) |
| finally: |
| model.meta["array_mode"] = "anchor" |
| assert sorted(it.kind for it in ex_index.items) == [ |
| "choice", "count", "extract", "extract", "generate", "score"] |
| ex = deg_example(model, rec, random.Random(0), shuffle_p=0.0) |
| kinds = sorted(it.kind for it in ex.items) |
| assert kinds == ["anchor", "choice", "extract", "extract", "generate", "score"] |
| oid = next(it for it in ex.items if it.ptr == "/order_id") |
| lo = ex.ctx.offsets[oid.starts[0]][0] |
| hi = ex.ctx.offsets[oid.ends[0]][1] |
| assert CONTEXT[slice(*trim_span(CONTEXT, lo, hi))] == "AB-1234" |
| jev = jev_example(model, {"uid": "j", "state": {"a": 1}, "questions": [ |
| {"type": "noul", "instructions": "Is a set?", "criteria": {"true": "yes"}, |
| "label": {"probs": [0.9, 0.1]}}, |
| {"type": "score", "instructions": "How big?", "criteria": ["small", "big"], |
| "label": {"hard": 1}}]}, random.Random(0)) |
| assert jev is not None and [it.kind for it in jev.items] == ["noul", "score"] |
| model.backbone.train() |
| try: |
| losses, stats = forward_losses(model, [ex, ex_index, jev], 4096, 16) |
| finally: |
| model.backbone.eval() |
| assert {"decide", "extract", "count", "anchor", "lm"} <= set(losses) |
| total = sum(losses[k] for k in ("decide", "extract", "count", "anchor", "lm")) |
| assert torch.isfinite(total) |
| total.backward() |
| assert model.heads.span.q_start.weight.grad is not None |
| assert model.backbone.text.layers[0].linear_attn.in_proj_qkv.weight.grad is not None |
| for p in [*model.heads.parameters(), *model.backbone.parameters()]: |
| p.grad = None |
|
|
|
|
| |
|
|
|
|
| @pytest.mark.skipif(SNAPSHOT is None or not torch.cuda.is_available(), |
| reason="needs CUDA and the Qwen3.5-0.8B snapshot") |
| def test_gpu_bitwise_independence_real_model(): |
| m = load_model(None, device="cuda") |
| flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")] |
| merged = _raw_outputs(m, _sub(flat)) |
| reordered = _raw_outputs(m, _sub(list(reversed(flat)))) |
| for name in flat: |
| alone = _raw_outputs(m, _sub([name])) |
| assert torch.equal(alone[f"/{name}"], merged[f"/{name}"]), name |
| assert torch.equal(alone[f"/{name}"], reordered[f"/{name}"]), name |
| big = {"type": "object", "properties": { |
| **{f"f{i}": {"type": ["string", "null"], "description": f"filler {i}"} |
| for i in range(12)}, **_sub(flat)["properties"]}, |
| "required": _sub(flat)["required"]} |
| big_out = _raw_outputs(m, big) |
| for name in flat: |
| assert torch.equal(_raw_outputs(m, _sub([name]))[f"/{name}"], big_out[f"/{name}"]) |
| resp = fill(m, CONTEXT, SCHEMA) |
| validate_value(compile_schema(SCHEMA), resp["value"]) |
| |
| both = fill(m, CONTEXT, _with_note()) |
| assert both["value"]["reply"] == resp["value"]["reply"] |
| assert both["meta"]["/reply"] == resp["meta"]["/reply"] |
| |
| from deg.infer import fill_many |
|
|
| other = "Dana again: the desk arrived broken. Order AB-9 total $80." |
| jobs = [(CONTEXT, SCHEMA, None), (other, _with_note(), None), |
| (other, _sub(["department", "items"]), None)] |
| for job, got in zip(jobs, fill_many(m, jobs)): |
| assert got == fill(m, *job) |
|
|
|
|
| @pytest.mark.skipif(SNAPSHOT is None or not torch.cuda.is_available(), |
| reason="needs CUDA and the Qwen3.5-0.8B snapshot") |
| def test_gpu_long_prefill_chunked_matches_and_is_linear_memory(monkeypatch): |
| """Real model on CUDA: the chunked prefill (flash-style fork kernel over the K/V |
| buffer) matches the one-shot read within bf16 noise for padded rows, and a 32k-token |
| prefill's peak memory is linear (K/V + hidden + one chunk's work, no L x L mask).""" |
| import deg.model as dm |
|
|
| m = load_model(None, device="cuda") |
| rend = m.renderer |
| a = rend.context(" ".join([CONTEXT] * 70)).ids |
| b = rend.context(" ".join([CONTEXT] * 30)).ids |
| x = torch.full((2, len(a)), rend.pad_id, device="cuda") |
| x[0], x[1, :len(b)] = torch.tensor(a), torch.tensor(b) |
| n = torch.tensor([len(a), len(b)], device="cuda") |
| with torch.no_grad(): |
| h_one, s_one = m.backbone(x, n) |
| monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 1024) |
| monkeypatch.setattr(dm, "PREFILL_CHUNK", 1024) |
| h, s = m.backbone(x, n) |
| for row, L in enumerate((len(a), len(b))): |
| |
| |
| cos = torch.nn.functional.cosine_similarity(h[row, :L].float(), |
| h_one[row, :L].float(), -1) |
| assert float(cos.min()) > 0.99 and float(cos.mean()) > 0.9995, row |
| assert torch.equal(s.kv_mask, s_one.kv_mask) |
| monkeypatch.undo() |
| ctx = rend.context(" ".join([CONTEXT] * 750)) |
| L = len(ctx.ids) |
| assert L > 32000 |
| torch.cuda.synchronize() |
| torch.cuda.reset_peak_memory_stats() |
| base = torch.cuda.memory_allocated() |
| with torch.no_grad(): |
| hid, st = m.trunk([ctx]) |
| torch.cuda.synchronize() |
| peak = torch.cuda.max_memory_allocated() - base |
| kept = sum(k.numel() * k.element_size() * 2 for k, _ in st.kv) + hid.numel() * 2 |
| work = peak - kept |
| assert work < 640 * 2**20, (peak / 2**20, kept / 2**20) |
| assert L * L > 640 * 2**20 |
| |
| from deg.kernels import bi_linear |
|
|
| ks = m.heads.span.k_start |
| rows = hid[0, :20000].float() |
| whole = bi_linear(rows, ks.weight, ks.bias) |
| blocks = torch.cat([bi_linear(rows[i:i + 8192], ks.weight, ks.bias) |
| for i in range(0, rows.shape[0], 8192)]) |
| assert torch.equal(whole, blocks) |
|
|
|
|
| @pytest.mark.skipif(SNAPSHOT is None or not torch.cuda.is_available(), |
| reason="needs CUDA and the Qwen3.5-0.8B snapshot") |
| def test_gpu_peak_memory_does_not_scale_with_fields_times_context(monkeypatch): |
| """A long context (~24k tokens) with a big schema: the fill's peak device memory above |
| the context's own K/V + hidden states is bounded by one fork batch + one head |
| micro-batch, independent of how many fields there are (fields x context logits wait |
| in host memory; fork hidden states keep only head rows), and blocked span logits are |
| bitwise the one-shot ones on CUDA.""" |
| import deg.model as dm |
|
|
| m = load_model(None, device="cuda") |
| monkeypatch.setattr(dm, "FORK_TOKENS", 2048) |
| ctx = " ".join([CONTEXT] * 550) |
| L = len(m.renderer.context(ctx).ids) |
| assert L > 20000 |
|
|
| def schema(n): |
| return {"type": "object", "properties": { |
| f"f{i}": {"type": ["string", "null"], "description": f"value {i} of the ticket"} |
| for i in range(n)}} |
|
|
| def peak(n): |
| fill(m, ctx, schema(n)) |
| torch.cuda.synchronize() |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
| base = torch.cuda.memory_allocated() |
| fill(m, ctx, schema(n)) |
| torch.cuda.synchronize() |
| return torch.cuda.max_memory_allocated() - base |
|
|
| small, big = peak(20), peak(400) |
| old_logits = 400 * (L + 1) * 8 |
| assert big - small < 96 * 2**20, (small / 2**20, big / 2**20) |
| assert old_logits > 64 * 2**20 |
| |
| ectx = m.renderer.context(ctx) |
| with torch.no_grad(): |
| hid, st = m.trunk([ectx]) |
| ch = hid[0, ectx.ctx_start: ectx.ctx_start + len(ectx.offsets)] |
| a = torch.randn(5, hid.shape[-1], device="cuda") |
| blocked = m._span_logits_blocked(a, ch) |
| monkeypatch.setattr(dm, "HEAD_ROWS", 10**9) |
| monkeypatch.setattr(dm, "HEAD_LOGIT_BYTES", 2**40) |
| whole = m._span_logits_blocked(a, ch) |
| for (s1, e1), (s2, e2) in zip(blocked, whole): |
| assert torch.equal(s1, s2) and torch.equal(e1, e2) |
|
|
