| """CPU integration checks against real tiny Transformers Qwen3.5 layers.""" |
| import importlib |
| import sys |
| from pathlib import Path |
| import torch |
| from torch import nn |
| from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig |
| from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel |
|
|
| root = Path(__file__).parent |
| repo = root / sys.argv[1] if len(sys.argv) > 1 else root.parent |
| sys.path.insert(0, str(repo)) |
| from dqwen.core import DiffusionLM, patch_bidirectional |
| from dqwen.kvcache import BlockCache, _segment, _patch, enable_kv_cache, disable_kv_cache, stream_kv |
| from dqwen import hub |
|
|
| torch.manual_seed(123) |
| torch.set_num_threads(1) |
| cfg = Qwen3_5TextConfig(vocab_size=32, hidden_size=16, intermediate_size=32, |
| num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, head_dim=8, |
| linear_num_key_heads=1, linear_num_value_heads=2, linear_key_head_dim=4, |
| linear_value_head_dim=4, linear_conv_kernel_dim=4, |
| layer_types=['linear_attention', 'full_attention'], |
| rope_parameters={'rope_type': 'default', 'rope_theta': 10000.0, |
| 'mrope_section': [1, 1, 2]}) |
| cfg._attn_implementation = 'sdpa' |
| lm = nn.Module() |
| lm.text = Qwen3_5TextModel(cfg).eval() |
| lm.lm_head = nn.Linear(16, 32, bias=False) |
| patch_bidirectional(lm) |
| for m in lm.modules(): |
| if hasattr(m, 'rev_gate'): |
| m.rev_gate.data.fill_(0.25) |
| dlm = DiffusionLM(lm, 31).eval() |
| dlm.eos_id = 30 |
| dlm.bart = None |
| ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) |
|
|
| with torch.no_grad(): |
| baseline = dlm.hidden(ids) |
| _patch(dlm) |
| |
| for lo, hi in [(1, 4), (3, 6), (3, 8)]: |
| cache = BlockCache(lo, hi, ids.shape[1]) |
| pos = torch.arange(ids.shape[1]) |
| record = dlm.text(inputs_embeds=dlm.embed(ids), |
| attention_mask={'full_attention': None, 'linear_attention': None}, |
| position_ids=pos[None], past_key_values=cache, use_cache=True).last_hidden_state |
| torch.testing.assert_close(record, baseline, atol=2e-5, rtol=2e-4) |
| saved = {i: tuple(t.clone() if t is not None else None for t in s) |
| for i, s in cache.f_state.items()} |
| cache.mode = 'step' |
| pos = torch.arange(lo - 1, hi) |
| for _ in range(2): |
| active = dlm.text(inputs_embeds=dlm.embed(ids[:, lo-1:hi]), |
| attention_mask={'full_attention': None, 'linear_attention': None}, |
| position_ids=pos[None], past_key_values=cache, use_cache=True).last_hidden_state |
| torch.testing.assert_close(active, record[:, lo-1:hi], atol=2e-5, rtol=2e-4) |
| for i, state in saved.items(): |
| for a, b in zip(state, cache.f_state[i]): |
| if a is not None: |
| torch.testing.assert_close(a, b, atol=0, rtol=0) |
|
|
| class Tok: |
| def apply_chat_template(self, *args, **kwargs): return 'prompt' |
| def __call__(self, *args, **kwargs): return {'input_ids': [1, 2, 3]} |
| def decode(self, ids, **kwargs): return str(ids) |
|
|
| tok = Tok() |
| enable_kv_cache(dlm) |
| results = list(hub.stream(dlm, tok, 'test', 4, 1.0)) |
| assert results and results[-1][2] <= 4 |
| |
| assert isinstance(hub.generate(dlm, tok, 'test', gen_len=4), str) |
| if not hub._kv_legacy_list: |
| assert list(hub.stream(dlm, tok, 'test', 4, 1.0, 2)) |
| |
| refreshed = list(stream_kv(dlm, tok, 'test', 4, 1.0, block_len=2, refresh=1)) |
| assert refreshed |
| |
| reference = torch.tensor([[1, 2, 3, 31, 31, 31, 31]]) |
| expected = [] |
| with torch.no_grad(): |
| for lo in (3, 5): |
| hi = lo + 2 |
| while (reference[0, lo:hi] == dlm.mask_id).any(): |
| h = dlm.hidden(reference) |
| pos = torch.arange(lo, hi)[reference[0, lo:hi] == dlm.mask_id] |
| logits = dlm.head(h[0, pos-1]).float() |
| logits[:, dlm.mask_id] = -torch.inf |
| conf, cand = logits.softmax(-1).max(-1) |
| select = conf >= 1.0 |
| if not select.any(): select[conf.argmax()] = True |
| reference[0, pos[select]] = cand[select] |
| expected.append(reference[:, 3:].clone()) |
| if (reference[0, 3:hi] == dlm.eos_id).any(): |
| reference[0, hi:] = dlm.eos_id |
| expected.append(reference[:, 3:].clone()) |
| break |
| assert len(expected) == len(refreshed) |
| for want, (got, _, _) in zip(expected, refreshed): |
| torch.testing.assert_close(want, got) |
| assert list(stream_kv(dlm, tok, 'test', 4, 1.0, max_steps=1))[-1][2] == 1 |
| disable_kv_cache(dlm) |
| assert not any(hasattr(m, '_dq_kv_old') for m in lm.modules()) |
| torch.testing.assert_close(dlm.hidden(ids), baseline, atol=2e-5, rtol=2e-4) |
| enable_kv_cache(dlm, refresh=1) |
| assert list(hub.stream(dlm, tok, 'test', 2, 1.0)) |
| disable_kv_cache() |
| print('PASS:', repo.name, 'record/active equivalence, empty segments, immutable state, stream/generate, restore') |
|
|
