| import torch |
|
|
| from grimboost.config import DrafterConfig |
| from grimboost.model.transformer import GrimBoostLM |
|
|
|
|
| def tiny_config() -> DrafterConfig: |
| return DrafterConfig( |
| vocab_size=256, |
| hidden_size=64, |
| num_hidden_layers=2, |
| num_attention_heads=4, |
| num_key_value_heads=2, |
| head_dim=16, |
| intermediate_size=128, |
| max_position_embeddings=64, |
| rope_theta=10_000.0, |
| ) |
|
|
|
|
| def make_model() -> GrimBoostLM: |
| torch.manual_seed(1) |
| return GrimBoostLM(tiny_config()).eval() |
|
|
|
|
| def test_incremental_equals_full(): |
| model = make_model() |
| ids = torch.randint(0, 256, (1, 12)) |
| with torch.no_grad(): |
| full = model(ids) |
| cache = model.new_cache(max_len=64, dtype=torch.float32) |
| _ = model(ids[:, :-1], past_key_values=cache, use_cache=True) |
| step = model(ids[:, -1:], past_key_values=cache, use_cache=True) |
| assert cache.length == 12 |
| assert torch.allclose(step[:, -1], full[:, -1], atol=1e-5) |
|
|
|
|
| def test_multi_token_block_after_cache(): |
| model = make_model() |
| ids = torch.randint(0, 256, (1, 14)) |
| with torch.no_grad(): |
| full = model(ids) |
| cache = model.new_cache(max_len=64, dtype=torch.float32) |
| _ = model(ids[:, :8], past_key_values=cache, use_cache=True) |
| block = model(ids[:, 8:], past_key_values=cache, use_cache=True) |
| assert cache.length == 14 |
| assert torch.allclose(block, full[:, 8:], atol=1e-5) |
|
|
|
|
| def test_newest_suffix_rollback_is_exact(): |
| model = make_model() |
| ids = torch.randint(0, 256, (1, 14)) |
| alt = torch.randint(0, 256, (1, 3)) |
| with torch.no_grad(): |
| cache = model.new_cache(max_len=64, dtype=torch.float32) |
| _ = model(ids[:, :11], past_key_values=cache, use_cache=True) |
| _ = model(ids[:, 11:], past_key_values=cache, use_cache=True) |
| assert cache.length == 14 |
| cache.rollback(3) |
| assert cache.length == 11 |
| block = model(alt, past_key_values=cache, use_cache=True) |
| reference = model(torch.cat([ids[:, :11], alt], dim=1)) |
| assert torch.allclose(block[:, -1], reference[:, -1], atol=1e-5) |
|
|
|
|
| def test_rebuild_with_positions_from_zero(): |
| model = make_model() |
| ids = torch.randint(0, 256, (1, 10)) |
| with torch.no_grad(): |
| full = model(ids) |
| cache = model.new_cache(max_len=64, dtype=torch.float32) |
| rebuilt = model(ids, past_key_values=cache, use_cache=True) |
| assert cache.length == 10 |
| assert torch.allclose(rebuilt, full, atol=1e-5) |
|
|
|
|
| def test_left_crop_is_not_equivalent(): |
| """Documents why the engine never left-crops a full-attention cache.""" |
| model = make_model() |
| ids = torch.randint(0, 256, (1, 12)) |
| with torch.no_grad(): |
| full_suffix = model(ids[:, 5:], position_ids=torch.arange(5, 12).unsqueeze(0)) |
| cache = model.new_cache(max_len=64, dtype=torch.float32) |
| _ = model(ids[:, :-1], past_key_values=cache, use_cache=True) |
| for layer in range(model.config.num_hidden_layers): |
| cache.k[layer] = cache.k[layer][:, :, 5:, :] |
| cache.v[layer] = cache.v[layer][:, :, 5:, :] |
| cache.length -= 5 |
| cropped = model(ids[:, -1:], past_key_values=cache, use_cache=True, |
| position_ids=torch.tensor([[11]])) |
| assert not torch.allclose(cropped[:, -1], full_suffix[:, -1], atol=1e-4) |
|
|
