File size: 1,447 Bytes
00c9d45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import numpy as np
from mla.model import Config, Model
from mla.kvcache import KVCache, forward_cached
def _tiny():
return Config(vocab_size=64, d_model=32, n_layers=2, n_heads=4,
n_kv_heads=2, head_dim=8, swiglu_hidden=48, seq_len=16)
def test_cached_matches_full_forward():
cfg = _tiny()
model = Model(cfg)
rng = np.random.default_rng(0)
T = 10
ids = rng.integers(0, cfg.vocab_size, size=T)
full = model(np.array([ids], dtype=np.int64)).data[0]
cache = KVCache(cfg.n_layers)
step_logits = []
for t in range(T):
lg = forward_cached(model, np.array([[ids[t]]], dtype=np.int64), cache)
step_logits.append(np.asarray(lg)[0, -1])
inc = np.stack(step_logits)
assert cache.length() == T
assert np.allclose(full, inc, atol=1e-6, rtol=1e-6)
def test_cached_chunk_prefill_matches():
cfg = _tiny()
model = Model(cfg)
rng = np.random.default_rng(1)
T = 8
ids = rng.integers(0, cfg.vocab_size, size=T)
full = model(np.array([ids], dtype=np.int64)).data[0]
cache = KVCache(cfg.n_layers)
prefill = forward_cached(model, np.array([ids[:5]], dtype=np.int64), cache)
assert np.allclose(np.asarray(prefill)[0], full[:5], atol=1e-6, rtol=1e-6)
for t in range(5, T):
lg = forward_cached(model, np.array([[ids[t]]], dtype=np.int64), cache)
assert np.allclose(np.asarray(lg)[0, -1], full[t], atol=1e-6, rtol=1e-6)
|