Spaces:
Sleeping
Sleeping
| """Independent full-prefix NumPy oracle for the native, incremental decoder.""" | |
| from concurrent.futures import ThreadPoolExecutor | |
| import gc | |
| import importlib | |
| import os | |
| import threading | |
| import time | |
| import numpy as np | |
| import pytest | |
| # A direct CMake build may expose the extension as _native before installation. | |
| # Normal package/CI builds exercise the public cism._native installation path. | |
| native = importlib.import_module(os.environ.get("CISM_NATIVE_TEST_MODULE", "cism._native")) | |
| def tiny_model(model_type="llama", tied=False, hidden=24, intermediate=37): | |
| rng = np.random.default_rng(312) | |
| config = dict( | |
| model_type=model_type, | |
| hidden_size=hidden, | |
| intermediate_size=intermediate, | |
| num_hidden_layers=2, | |
| num_attention_heads=4, | |
| num_key_value_heads=2, | |
| head_dim=8, | |
| vocab_size=43, | |
| max_position_embeddings=80, | |
| rms_norm_eps=1e-5, | |
| rope_theta=10000.0, | |
| tie_word_embeddings=tied, | |
| hidden_act="silu", | |
| ) | |
| weights = {} | |
| def matrix(name, shape): | |
| weights[name] = rng.normal(0, 0.14, shape).astype(np.float32) | |
| def norm(name, width): | |
| weights[name] = rng.uniform(0.75, 1.25, width).astype(np.float32) | |
| matrix("model.embed_tokens.weight", (43, hidden)) | |
| if not tied: | |
| matrix("lm_head.weight", (43, hidden)) | |
| norm("model.norm.weight", hidden) | |
| for i in range(2): | |
| prefix = f"model.layers.{i}." | |
| norm(prefix + "input_layernorm.weight", hidden) | |
| norm(prefix + "post_attention_layernorm.weight", hidden) | |
| for name, shape in ( | |
| ("q_proj", (32, hidden)), | |
| ("k_proj", (16, hidden)), | |
| ("v_proj", (16, hidden)), | |
| ("o_proj", (hidden, 32)), | |
| ): | |
| matrix(prefix + f"self_attn.{name}.weight", shape) | |
| if model_type == "qwen3": | |
| norm(prefix + "self_attn.q_norm.weight", 8) | |
| norm(prefix + "self_attn.k_norm.weight", 8) | |
| for name, shape in ( | |
| ("gate_proj", (intermediate, hidden)), | |
| ("up_proj", (intermediate, hidden)), | |
| ("down_proj", (hidden, intermediate)), | |
| ): | |
| matrix(prefix + f"mlp.{name}.weight", shape) | |
| return config, weights | |
| E2M1_HALF = np.array([0, 1, 2, 3, 4, 6, 8, 12], dtype=np.float32) | |
| def encode_e4m3_numpy(values): | |
| """Round-to-nearest-even positive floats to E4M3, mirroring the C++ encoder.""" | |
| values = np.asarray(values, dtype=np.float32) | |
| bits = np.zeros(values.shape, dtype=np.int64) | |
| normal = values >= np.float32(2.0 ** -6) | |
| with np.errstate(divide="ignore"): | |
| exponent = np.floor(np.log2(np.where(normal, values, np.float32(1)))).astype(np.int64) | |
| mantissa = values / (2.0 ** exponent) | |
| quantized_m = np.rint((mantissa - np.float32(1)) * 8).astype(np.int64) # half-even | |
| field = exponent + 7 + (quantized_m >= 8) | |
| quantized_m = np.where(quantized_m >= 8, 0, quantized_m) | |
| bits[normal] = field[normal] * 8 + quantized_m[normal] | |
| subnormal = np.rint(values * 512.0).astype(np.int64) # value / 2^-9 | |
| bits[~normal] = np.minimum(subnormal[~normal], 0x7E) | |
| bits = np.minimum(bits, 0x7E) | |
| decoded = np.where(bits >= 8, (1 + (bits & 7) / 8) * 2.0 ** ((bits >> 3) - 7), (bits & 15) * 2.0 ** -9) | |
| return decoded.astype(np.float32) | |
| def reconstructed_weights(weights, precision): | |
| result = {} | |
| for name, original in weights.items(): | |
| if precision == "fp32" or original.ndim == 1: | |
| result[name] = original.copy() | |
| continue | |
| if precision == "hybrid-fp4" and ".mlp." in name: | |
| rows, cols = original.shape | |
| pad = (-cols) % 16 | |
| blocks = np.pad(original, ((0, 0), (0, pad))).reshape(rows, -1, 16) | |
| amax = np.abs(blocks).max(axis=-1) | |
| # No 2^-6 floor here — matches the native packer, whose E4M3 | |
| # encoder handles small scales via subnormals. | |
| ideal = np.where(amax == 0, np.float32(0), amax / 6) | |
| scale = encode_e4m3_numpy(ideal) | |
| magnitude = np.abs(blocks) / scale[..., None] | |
| # Compare against true E2M1 values (half table times one half). | |
| index = np.argmin(np.abs(magnitude[..., None] - E2M1_HALF * np.float32(0.5)), axis=-1) | |
| nibbles = np.where(np.sign(blocks) < 0, index | 8, index).astype(np.int64) | |
| decoded = np.where(np.sign(blocks) < 0, -E2M1_HALF[index], E2M1_HALF[index]) \ | |
| * np.float32(0.5) * scale[..., None] | |
| result[name] = decoded.reshape(rows, -1)[:, :cols] | |
| continue | |
| int4 = precision == "hybrid-int4" and ".mlp." in name | |
| if int4: | |
| decoded = np.empty_like(original) | |
| for start in range(0, original.shape[1], 32): | |
| values = original[:, start : start + 32] | |
| m = values.shape[1] | |
| for r in range(values.shape[0]): | |
| row = values[r] | |
| # First-max wins (matches the native strict-greater scan). | |
| idx = int(np.argmax(np.abs(row))) | |
| extreme = np.float32(row[idx]) | |
| if extreme == 0: | |
| d = np.float32(1.0) | |
| else: | |
| d = np.float32(extreme / np.float32(-8.0)) | |
| for t in range(m): | |
| x = float(np.float32(row[t]) / d + np.float32(8.5)) | |
| xi = int(x) | |
| if xi < 0: | |
| xi = 0 | |
| elif xi > 15: | |
| xi = 15 | |
| decoded[r, start + t] = np.float32( | |
| (np.float32(xi) - np.float32(8.0)) * d) | |
| result[name] = decoded | |
| continue | |
| block = original.shape[1] | |
| maximum = np.float32(127) | |
| decoded = np.empty_like(original) | |
| for start in range(0, original.shape[1], block): | |
| values = original[:, start : start + block] | |
| largest = np.max(np.abs(values), axis=1, keepdims=True) | |
| scale = np.where(largest == 0, np.float32(1), largest / maximum) | |
| scale = np.maximum(scale, np.nextafter(np.float32(0), np.float32(1))) | |
| divided = values / scale | |
| fraction, integral = np.modf(np.abs(divided)) | |
| rounded = np.copysign(integral + (fraction >= np.float32(0.5)), divided) | |
| quantized = np.clip(rounded, -maximum, maximum) | |
| decoded[:, start : start + block] = quantized * scale | |
| result[name] = decoded | |
| return result | |
| def reference_logits(config, weights, tokens): | |
| """Vectorized dense attention over every prefix position, without a KV cache.""" | |
| heads = config["num_attention_heads"] | |
| kv_heads = config["num_key_value_heads"] | |
| dim = config["head_dim"] | |
| length = len(tokens) | |
| def rms(values, weight): | |
| variance = np.mean(values.astype(np.float64) ** 2, axis=-1, keepdims=True) | |
| scale = (1 / np.sqrt(variance + config["rms_norm_eps"])).astype(np.float32) | |
| return (values * scale) * weight | |
| powers = np.arange(0, dim, 2, dtype=np.float32) / np.float32(dim) | |
| frequencies = np.float32(1) / (np.float32(config["rope_theta"]) ** powers) | |
| angles = np.arange(length, dtype=np.float32)[:, None] * frequencies[None, :] | |
| cosine, sine = np.cos(angles)[:, None, :], np.sin(angles)[:, None, :] | |
| def rotate(values): | |
| first, second = np.split(values, 2, axis=-1) | |
| return np.concatenate((first * cosine - second * sine, second * cosine + first * sine), axis=-1) | |
| x = weights["model.embed_tokens.weight"][tokens].copy() | |
| for layer in range(config["num_hidden_layers"]): | |
| p = f"model.layers.{layer}." | |
| normalized = rms(x, weights[p + "input_layernorm.weight"]) | |
| q = (normalized @ weights[p + "self_attn.q_proj.weight"].T).reshape(length, heads, dim) | |
| k = (normalized @ weights[p + "self_attn.k_proj.weight"].T).reshape(length, kv_heads, dim) | |
| v = (normalized @ weights[p + "self_attn.v_proj.weight"].T).reshape(length, kv_heads, dim) | |
| if config["model_type"] == "qwen3": | |
| q = rms(q, weights[p + "self_attn.q_norm.weight"]) | |
| k = rms(k, weights[p + "self_attn.k_norm.weight"]) | |
| q = rotate(q) | |
| k = np.repeat(rotate(k), heads // kv_heads, axis=1) | |
| v = np.repeat(v, heads // kv_heads, axis=1) | |
| scores = np.einsum("thd,shd->hts", q, k) / np.sqrt(np.float32(dim)) | |
| scores = np.where(np.triu(np.ones((length, length), dtype=bool), 1)[None], -np.inf, scores) | |
| probabilities = np.exp(scores - np.max(scores, axis=-1, keepdims=True)) | |
| probabilities /= np.sum(probabilities, axis=-1, keepdims=True) | |
| attention = np.einsum("hts,shd->thd", probabilities, v).reshape(length, heads * dim) | |
| x += attention @ weights[p + "self_attn.o_proj.weight"].T | |
| normalized = rms(x, weights[p + "post_attention_layernorm.weight"]) | |
| gate = normalized @ weights[p + "mlp.gate_proj.weight"].T | |
| up = normalized @ weights[p + "mlp.up_proj.weight"].T | |
| silu = gate / (np.float32(1) + np.exp(-gate)) | |
| x += (silu * up) @ weights[p + "mlp.down_proj.weight"].T | |
| final = rms(x, weights["model.norm.weight"]) | |
| head = weights["model.embed_tokens.weight"] if config["tie_word_embeddings"] else weights["lm_head.weight"] | |
| return final[-1] @ head.T | |
| def test_logits_and_cached_generation(architecture, precision, tied): | |
| config, weights = tiny_model(architecture, tied) | |
| model = native.Model(config, weights, precision) | |
| reconstructed = reconstructed_weights(weights, precision) | |
| prompt = [2, 7, 1, 19, 4] | |
| for prefix in (prompt[:1], prompt[:3], prompt): | |
| actual = model.logits(prefix) | |
| assert actual.dtype == np.float32 | |
| np.testing.assert_allclose(actual, reference_logits(config, reconstructed, prefix), atol=8e-6, rtol=8e-5) | |
| session = model.create_session(prompt, max_new_tokens=9) | |
| assert session.finish_reason == "" | |
| assert session.generated_tokens == 0 | |
| assert session.next_tokens(0) == [] | |
| expected = [] | |
| for _ in range(9): | |
| logits = reference_logits(config, reconstructed, prompt + expected) | |
| expected.append(int(np.argmax(logits))) | |
| actual = session.next_tokens(2) + session.next_tokens(3) + session.next_tokens(100) | |
| assert actual == expected | |
| assert session.generated_tokens == 9 | |
| assert session.finish_reason == "length" | |
| assert session.next_tokens(2) == [] | |
| for index, token in enumerate(actual): | |
| assert token == int(np.argmax(model.logits(prompt + actual[:index]))) | |
| def test_nll_matches_reference(architecture, precision): | |
| config, weights = tiny_model(architecture) | |
| model = native.Model(config, weights, precision) | |
| reconstructed = reconstructed_weights(weights, precision) | |
| tokens = [2, 7, 1, 19, 4, 11] | |
| # Row i of the reference (one prefix forward per position) predicts token i+1. | |
| rows = np.stack([reference_logits(config, reconstructed, tokens[: i + 1]) | |
| for i in range(len(tokens) - 1)]) | |
| maximum = rows.max(axis=-1) | |
| log_total = np.log(np.exp(rows - maximum[:, None]).sum(axis=-1)) | |
| targets = np.array(tokens[1:]) | |
| expected = float(np.mean(maximum + log_total - rows[np.arange(len(tokens) - 1), targets])) | |
| actual = model.nll(tokens) | |
| assert isinstance(actual, float) | |
| assert actual == model.nll(tokens) # deterministic | |
| assert abs(actual - expected) < 1e-4 | |
| with pytest.raises(ValueError, match="at least 2 tokens"): | |
| model.nll(tokens[:1]) | |
| def expected_bytes(weights, precision): | |
| total = 0 | |
| for name, array in weights.items(): | |
| if array.ndim == 1 or precision == "fp32": | |
| total += array.nbytes | |
| elif precision == "hybrid-int4" and ".mlp." in name: | |
| rows, cols = array.shape | |
| blocks = (cols + 31) // 32 | |
| total += rows * blocks * 16 + rows * blocks * 4 | |
| elif precision == "hybrid-fp4" and ".mlp." in name: | |
| rows, cols = array.shape | |
| total += rows * ((cols + 1) // 2) + rows * ((cols + 15) // 16) | |
| else: | |
| total += array.size + array.shape[0] * 4 | |
| return total | |
| def test_fp4_element_lut_is_exact(): | |
| expected = np.array([0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12]) | |
| np.testing.assert_array_equal(np.asarray(native.fp4_element_lut()), expected) | |
| def test_int8_activations_stay_close_and_deterministic(precision): | |
| # The pmaddwd quantized-activation path must engage (outputs move), stay | |
| # within a few percent of fp32 activations, and be bit-deterministic. | |
| # hidden=24/intermediate=37 exercise every partial-block tail. | |
| config, weights = tiny_model() | |
| prompt = [1, 3, 4, 2, 5] | |
| model = native.Model(config, weights, precision, 1) | |
| baseline = np.asarray(model.logits(prompt)) | |
| model.set_act_precision("int8") | |
| quantized = np.asarray(model.logits(prompt)) | |
| again = np.asarray(model.logits(prompt)) | |
| assert not np.array_equal(baseline, quantized) | |
| np.testing.assert_allclose(quantized, baseline, rtol=0.03, atol=0.03) | |
| np.testing.assert_array_equal(quantized, again) | |
| # The constructor kwarg must land in the same place as the setter. | |
| via_ctor = native.Model(config, weights, precision, 1, "int8").logits(prompt) | |
| np.testing.assert_array_equal(quantized, np.asarray(via_ctor)) | |
| def test_threaded_gemm_matches_single_thread(precision, act): | |
| # Wide enough rows*cols to engage the pooled path at threads=2; row | |
| # results are thread-independent, so threading must be bitwise exact. | |
| config, weights = tiny_model(hidden=256, intermediate=512) | |
| prompt = [1, 3, 4, 2, 5] | |
| single = np.asarray(native.Model(config, weights, precision, 1, act).logits(prompt)) | |
| multi = np.asarray(native.Model(config, weights, precision, 2, act).logits(prompt)) | |
| np.testing.assert_array_equal(single, multi) | |
| def test_spec_stats_count_proposed_and_accepted(): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights, "int8", 1) | |
| session = model.create_session([1, 2, 1, 2, 1, 2], max_new_tokens=16) | |
| assert session.spec_stats == (0, 0) | |
| first = session.next_tokens(16, 2) | |
| proposed, accepted = session.spec_stats | |
| assert proposed > 0 | |
| assert 0 <= accepted <= proposed | |
| # Same model, fresh session: deterministic tally and tokens. | |
| again = model.create_session([1, 2, 1, 2, 1, 2], max_new_tokens=16) | |
| assert again.next_tokens(16, 2) == first | |
| assert again.spec_stats == (proposed, accepted) | |
| # Direct verification calls are tallied too, bounded by construction. | |
| direct = model.create_session([1, 2, 3], max_new_tokens=8) | |
| direct.verify([1, 2, 3]) | |
| proposed2, accepted2 = direct.spec_stats | |
| assert proposed2 == 3 | |
| assert 0 <= accepted2 <= 3 | |
| def test_engine_reports_spec_totals(): | |
| from cism.engine import Engine | |
| class StubTokenizer: | |
| bos_token_id = 0 | |
| eos_token_id = 1 | |
| chat_template = None | |
| def encode(self, prompt, add_special_tokens=True): | |
| return list(prompt) | |
| def decode(self, ids, **kwargs): | |
| return "" | |
| config, weights = tiny_model() | |
| engine = Engine(native.Model(config, weights, "int8", 1), StubTokenizer(), config, "test") | |
| assert engine.generate([1, 2, 3], max_new_tokens=4, temperature=0).spec_proposed == 0 | |
| chunks = list(engine.stream([1, 2, 1, 2, 1, 2], max_new_tokens=8, temperature=0, spec_k=2)) | |
| assert sum(c.spec_proposed for c in chunks) > 0 | |
| result = engine.generate([1, 2, 1, 2, 1, 2], max_new_tokens=8, temperature=0, spec_k=2) | |
| assert result.spec_proposed == sum(c.spec_proposed for c in chunks) | |
| assert result.spec_accepted == sum(c.spec_accepted for c in chunks) | |
| assert 0 <= result.spec_accepted <= result.spec_proposed | |
| def test_fp4_e4m3_scale_encoding_matches_torch(): | |
| torch = pytest.importorskip("torch") | |
| rng = np.random.default_rng(77) | |
| values = rng.uniform(2 ** -9, 400, 4096).astype(np.float32) | |
| encoded = np.array([native.fp4_encode_scale(float(value)) for value in values], dtype=np.uint8) | |
| # Native decode must agree with torch's view of the same bit pattern. | |
| torch_view = np.array([ | |
| float(torch.from_numpy(np.array([bits], dtype=np.uint8)).view(torch.float8_e4m3fn)) | |
| for bits in encoded], dtype=np.float32) | |
| native_view = np.array([native.fp4_decode_scale(int(bits)) for bits in encoded], dtype=np.float32) | |
| np.testing.assert_allclose(native_view, torch_view, rtol=0, atol=0) | |
| # And encoding a value must round-trip within E4M3 relative resolution. | |
| np.testing.assert_allclose(native_view, values, rtol=0.07) | |
| assert native.fp4_encode_scale(0.0) == 0 | |
| assert native.fp4_encode_scale(1e30) == 0x7E | |
| assert native.fp4_decode_scale(0x7E) == 448.0 | |
| def test_owned_weights_and_accounting(precision): | |
| config, weights = tiny_model("qwen3", tied=True) | |
| expected = expected_bytes(weights, precision) | |
| model = native.Model(config, weights, precision) | |
| before = model.logits([1, 2, 3]) | |
| assert model.info["weight_bytes"] == expected | |
| assert model.info["precision"] == precision | |
| assert model.info["activation_dtype"] == "float32" | |
| assert model.info["act_precision"] == "fp32" | |
| assert model.info["threads"] == 1 | |
| assert model.info["kernel"] in {"scalar", "avx2", "neon"} | |
| assert set(model.info["kernels"]) == {"fp32", "int8", "int4", "fp4", "variant"} | |
| assert all(name == model.info["kernel"] for key, name in model.info["kernels"].items() | |
| if key != "variant") | |
| assert model.info["kernels"]["variant"] in {"scalar", "avx2", "avx-vnni", "avx512-vnni", "neon"} | |
| model.set_act_precision("int8") | |
| assert model.info["activation_dtype"] == "int8-act" | |
| assert model.info["act_precision"] == "int8" | |
| model.set_act_precision("fp32") | |
| if expected_kernel := os.environ.get("CISM_NATIVE_EXPECT_KERNEL"): | |
| assert model.info["kernel"] == expected_kernel | |
| supplied_head = dict(weights, **{"lm_head.weight": weights["model.embed_tokens.weight"].copy()}) | |
| assert native.Model(config, supplied_head, precision).info["weight_bytes"] == expected | |
| for array in weights.values(): | |
| array.fill(0) | |
| config["hidden_size"] = 1 | |
| np.testing.assert_array_equal(before, model.logits([1, 2, 3])) | |
| session = model.create_session([1, 2], max_new_tokens=3) | |
| del model | |
| gc.collect() | |
| assert len(session.next_tokens(9)) == 3 | |
| def test_sampling_chunks_seed_and_filters(): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights) | |
| options = dict(max_new_tokens=20, temperature=0.9, top_p=0.8, top_k=11, seed=876) | |
| first = model.create_session([3, 1], **options).next_tokens(20) | |
| session = model.create_session([3, 1], **options) | |
| second = session.next_tokens(1) + session.next_tokens(7) + session.next_tokens(12) | |
| assert first == second | |
| other = model.create_session([3, 1], **dict(options, seed=877)).next_tokens(20) | |
| assert first != other | |
| greedy = model.create_session([3], max_new_tokens=10).next_tokens(10) | |
| assert model.create_session([3], max_new_tokens=10, temperature=1, top_k=1).next_tokens(10) == greedy | |
| assert model.create_session([3], max_new_tokens=10, temperature=1, top_p=1e-12).next_tokens(10) == greedy | |
| assert model.create_session([3], max_new_tokens=10, temperature=0, top_k=1, top_p=0.2).next_tokens(10) == greedy | |
| prefix = [3, 1] | |
| for token in first: | |
| logits = model.logits(prefix) | |
| order = np.argsort(-logits, kind="stable")[:11] | |
| probabilities = np.exp((logits[order].astype(np.float64) - logits[order[0]]) / 0.9) | |
| probabilities /= probabilities.sum() | |
| count = np.searchsorted(np.cumsum(probabilities), 0.8) + 1 | |
| assert token in order[:count] | |
| prefix.append(token) | |
| def test_speculative_greedy_matches_plain_decode(architecture, precision, spec_k): | |
| config, weights = tiny_model(architecture, tied=True) | |
| model = native.Model(config, weights, precision) | |
| # A repetitive prompt makes the n-gram drafter fire and multi-accept. | |
| prompt = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4] | |
| plain = model.create_session(prompt, max_new_tokens=24, eos_token_ids=[]).next_tokens(24) | |
| session = model.create_session(prompt, max_new_tokens=24, eos_token_ids=[]) | |
| chunked = session.next_tokens(3, spec_k=spec_k) + session.next_tokens(21, spec_k=spec_k) | |
| assert chunked == plain | |
| assert session.finish_reason == "length" | |
| def test_speculative_draft_hits_and_misses(): | |
| config, weights = tiny_model("llama", tied=True) | |
| model = native.Model(config, weights) | |
| prompt = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4] | |
| def py_draft(history, k): | |
| """Independent reimplementation of the v4 prompt-lookup drafter: | |
| most-recent strictly-inside occurrence per position, iteratively | |
| extended (fixes v1 truncation at history end); stop when the | |
| context never occurred.""" | |
| out = [] | |
| size = len(history) | |
| if k == 0 or size < 2: | |
| return out | |
| for _ in range(k): | |
| total = size + len(out) | |
| placed = False | |
| for pattern in range(min(3, total - 1), 0, -1): | |
| if pattern >= size: | |
| continue | |
| ctx = (history + out)[total - pattern:total] | |
| for s in range(size - pattern - 1, -1, -1): | |
| if history[s:s + pattern] == ctx: | |
| out.append(history[s + pattern]) | |
| placed = True | |
| break | |
| if placed: | |
| break | |
| if not placed: | |
| break | |
| return out | |
| session = model.create_session(prompt, max_new_tokens=8, eos_token_ids=[]) | |
| history = list(prompt) | |
| for _ in range(3): | |
| generated = session.next_tokens(1) | |
| history.extend(generated) | |
| assert session.draft(3) == py_draft(history, 3) | |
| assert session.draft(0) == [] | |
| empty = model.create_session([1], max_new_tokens=4, eos_token_ids=[]) | |
| assert empty.draft(4) == [] # History too short to find a pattern. | |
| def test_speculative_verify_accept_and_reject(): | |
| config, weights = tiny_model("llama", tied=True) | |
| model = native.Model(config, weights) | |
| prompt = [1, 2, 3] | |
| plain = model.create_session(prompt, max_new_tokens=16, eos_token_ids=[]).next_tokens(16) | |
| # Feeding the true continuation must fully accept (plus bonus) and land on | |
| # exactly the same greedy sequence as plain decoding. | |
| session = model.create_session(prompt, max_new_tokens=16, eos_token_ids=[]) | |
| out = [] | |
| while len(out) < 16: | |
| take = min(3, 16 - len(out)) | |
| out += session.verify(plain[len(out):len(out) + take]) | |
| assert out == plain | |
| # Wrong candidates must be rejected and corrected, still matching greedy. | |
| session = model.create_session(prompt, max_new_tokens=16, eos_token_ids=[]) | |
| out = [] | |
| while len(out) < 16: | |
| take = min(3, 16 - len(out)) | |
| wrong = [(token + 7) % config["vocab_size"] for token in plain[len(out):len(out) + take]] | |
| out += session.verify(wrong) | |
| assert out == plain | |
| def test_speculative_rejection_path_and_boundaries(): | |
| config, weights = tiny_model("qwen3", tied=True) | |
| model = native.Model(config, weights) | |
| # spec_k validation happens at the session call boundary. | |
| session = model.create_session([1, 2], max_new_tokens=10, eos_token_ids=[]) | |
| with pytest.raises(ValueError, match="spec_k"): | |
| session.next_tokens(4, spec_k=1) | |
| with pytest.raises(ValueError, match="spec_k"): | |
| session.next_tokens(4, spec_k=17) | |
| # EOS inside an accepted speculative block must stop cleanly. | |
| eos = int(np.argmax(model.logits([1, 2]))) | |
| reference = model.create_session([1, 2], max_new_tokens=12, eos_token_ids=[eos]) | |
| expected = reference.next_tokens(12) | |
| assert reference.finish_reason in ("stop", "length") | |
| speculative = model.create_session([1, 2], max_new_tokens=12, eos_token_ids=[eos]) | |
| actual = speculative.next_tokens(12, spec_k=3) | |
| assert actual == expected | |
| assert speculative.finish_reason == reference.finish_reason | |
| # Token budget boundary: speculation never overshoots max_new_tokens. | |
| exact = model.create_session([1, 2], max_new_tokens=2, eos_token_ids=[]) | |
| tokens = exact.next_tokens(2, spec_k=4) | |
| assert len(tokens) == 2 and exact.generated_tokens == 2 | |
| assert exact.finish_reason in ("stop", "length") | |
| def test_eos_lengths_and_cancellation(): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights) | |
| eos = int(np.argmax(model.logits([1, 2]))) | |
| session = model.create_session([1, 2], max_new_tokens=1, eos_token_ids=[eos]) | |
| assert session.next_tokens(20) == [eos] | |
| assert session.finish_reason == "stop" | |
| assert session.generated_tokens == 1 | |
| session.cancel() | |
| assert session.finish_reason == "stop" | |
| empty = model.create_session([1], max_new_tokens=0) | |
| assert empty.finish_reason == "length" | |
| assert empty.next_tokens(1) == [] | |
| cancelled = model.create_session([1], max_new_tokens=7) | |
| cancelled.cancel() | |
| assert cancelled.finish_reason == "cancelled" | |
| assert cancelled.next_tokens(7) == [] | |
| assert cancelled.generated_tokens == 0 | |
| # EOS in the prompt is ordinary context, not a generated stopping event. | |
| session = model.create_session([eos], max_new_tokens=2, eos_token_ids=[eos]) | |
| assert session.finish_reason == "" | |
| assert len(session.next_tokens(1)) == 1 | |
| config["max_position_embeddings"] = 3 | |
| small = native.Model(config, weights) | |
| assert len(small.create_session([1, 2], max_new_tokens=1).next_tokens(2)) == 1 | |
| assert small.logits([1, 2, 3]).shape == (43,) | |
| with pytest.raises(ValueError, match="exceeds"): | |
| small.create_session([1, 2], max_new_tokens=2) | |
| def test_invalid_prompt(prompt): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights) | |
| with pytest.raises((ValueError, TypeError, OverflowError)): | |
| model.create_session(prompt, max_new_tokens=1) | |
| with pytest.raises((ValueError, TypeError, OverflowError)): | |
| model.logits(prompt) | |
| def test_invalid_sampling(options): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights) | |
| with pytest.raises((ValueError, TypeError, OverflowError)): | |
| model.create_session([1], **dict({"max_new_tokens": 1}, **options)) | |
| with pytest.raises(ValueError): | |
| model.create_session([1], max_new_tokens=1).next_tokens(-1) | |
| def test_invalid_config(change): | |
| config, weights = tiny_model() | |
| config.update(change) | |
| with pytest.raises((ValueError, TypeError, OverflowError)): | |
| native.Model(config, weights) | |
| def test_invalid_weights(mutation): | |
| config, weights = tiny_model() | |
| name = "model.embed_tokens.weight" | |
| if mutation == "missing": | |
| del weights[name] | |
| elif mutation == "extra": | |
| weights["model.unknown.weight"] = np.zeros(1, np.float32) | |
| elif mutation == "bias": | |
| weights["model.layers.0.self_attn.q_proj.bias"] = np.zeros(32, np.float32) | |
| elif mutation == "shape": | |
| weights[name] = weights[name][:-1].copy() | |
| elif mutation == "rank": | |
| weights[name] = weights[name].ravel() | |
| elif mutation == "dtype": | |
| weights[name] = weights[name].astype(np.float64) | |
| elif mutation == "noncontiguous": | |
| weights[name] = np.asfortranarray(weights[name]) | |
| elif mutation == "nan": | |
| weights[name][0, 0] = np.nan | |
| elif mutation == "inf": | |
| weights[name][0, 0] = np.inf | |
| elif mutation == "list": | |
| weights[name] = weights[name].tolist() | |
| elif mutation == "endian": | |
| weights[name] = weights[name].astype(">f4") | |
| with pytest.raises(ValueError): | |
| native.Model(config, weights) | |
| def test_tied_and_untied_head_validation(): | |
| config, weights = tiny_model(tied=True) | |
| weights["lm_head.weight"] = weights["model.embed_tokens.weight"].copy() | |
| weights["lm_head.weight"][0, 0] += 1 | |
| with pytest.raises(ValueError, match="tied"): | |
| native.Model(config, weights) | |
| del weights["lm_head.weight"] | |
| config["tie_word_embeddings"] = False | |
| with pytest.raises(ValueError, match="lm_head"): | |
| native.Model(config, weights) | |
| with pytest.raises(ValueError, match="precision"): | |
| native.Model(config, weights, "int4") | |
| def test_zero_weights_and_tie_breaking(precision): | |
| config, weights = tiny_model() | |
| for value in weights.values(): | |
| value.fill(0) | |
| model = native.Model(config, weights, precision) | |
| np.testing.assert_array_equal(model.logits([1, 2]), np.zeros(43, np.float32)) | |
| assert model.create_session([1], max_new_tokens=3).next_tokens(3) == [0, 0, 0] | |
| assert model.create_session([1], max_new_tokens=3, temperature=1, top_k=1).next_tokens(3) == [0, 0, 0] | |
| def test_unaligned_contiguous_input(): | |
| config, weights = tiny_model() | |
| name = "model.embed_tokens.weight" | |
| original = weights[name] | |
| buffer = bytearray(original.nbytes + 1) | |
| unaligned = np.ndarray(original.shape, dtype=np.float32, buffer=buffer, offset=1) | |
| unaligned[:] = original | |
| weights[name] = unaligned | |
| np.testing.assert_allclose(native.Model(config, weights).logits([1, 3]), reference_logits(config, weights, [1, 3]), atol=8e-6) | |
| def test_concurrent_sessions_and_serialized_calls(): | |
| config, weights = tiny_model("qwen3") | |
| model = native.Model(config, weights, "hybrid-int4") | |
| expected = model.create_session([1, 2], max_new_tokens=24).next_tokens(24) | |
| with ThreadPoolExecutor(max_workers=4) as executor: | |
| futures = [executor.submit(lambda: model.create_session([1, 2], max_new_tokens=24).next_tokens(24)) for _ in range(4)] | |
| assert all(future.result(timeout=20) == expected for future in futures) | |
| shared = model.create_session([1, 2], max_new_tokens=24) | |
| futures = [executor.submit(shared.next_tokens, 6) for _ in range(4)] | |
| results = [future.result(timeout=20) for future in futures] | |
| assert sorted(map(tuple, results)) == sorted(tuple(expected[i : i + 6]) for i in range(0, 24, 6)) | |
| assert shared.generated_tokens == 24 | |
| assert shared.finish_reason == "length" | |
| def test_gil_release_and_atomic_cancel(): | |
| config, weights = tiny_model() | |
| config["max_position_embeddings"] = 100000 | |
| model = native.Model(config, weights) | |
| session = model.create_session([1, 2], max_new_tokens=99990) | |
| started = threading.Event() | |
| def decode(): | |
| started.set() | |
| return session.next_tokens(99990) | |
| with ThreadPoolExecutor(max_workers=1) as executor: | |
| future = executor.submit(decode) | |
| assert started.wait(timeout=5) | |
| session.cancel() | |
| # This property must not wait on the mutex while holding the GIL. | |
| reason = session.finish_reason | |
| output = future.result(timeout=20) | |
| assert reason == "cancelled" | |
| assert session.generated_tokens == len(output) | |
| assert len(output) < 99990 | |
| def test_cancel_during_lazy_prefill(): | |
| config, weights = tiny_model() | |
| config["max_position_embeddings"] = 100000 | |
| model = native.Model(config, weights) | |
| session = model.create_session([1, 2] * 25000, max_new_tokens=2) | |
| started = threading.Event() | |
| def decode(): | |
| started.set() | |
| return session.next_tokens(2) | |
| with ThreadPoolExecutor(max_workers=1) as executor: | |
| future = executor.submit(decode) | |
| assert started.wait(timeout=5) | |
| # Allow native prefill to run; a 50k-token quadratic prefill cannot finish | |
| # in this interval. The monitor must be able to acquire the released GIL. | |
| time.sleep(0.02) | |
| assert not future.done() | |
| session.cancel() | |
| assert future.result(timeout=20) == [] | |
| assert session.finish_reason == "cancelled" | |
| assert session.generated_tokens == 0 | |
| assert session.next_tokens(2) == [] | |
| def test_lazy_prefill_owns_prompt(): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights) | |
| prompt = [1, 2, 3] | |
| expected = int(np.argmax(model.logits(prompt))) | |
| session = model.create_session(prompt, max_new_tokens=1) | |
| prompt[:] = [-1] | |
| assert session.next_tokens(1) == [expected] | |
| def test_kv_limit_rejected_before_allocation(): | |
| config, weights = tiny_model() | |
| config["max_position_embeddings"] = 1048576 | |
| config["num_hidden_layers"] = 256 | |
| for layer in range(2, 256): | |
| for name, value in list(weights.items()): | |
| if name.startswith("model.layers.0."): | |
| weights[name.replace("model.layers.0.", f"model.layers.{layer}.")] = value | |
| model = native.Model(config, weights) | |
| with pytest.raises(ValueError, match="KV cache"): | |
| model.create_session([1], max_new_tokens=1048575) | |
| def test_multi_query_and_multi_head_attention(architecture, kv_heads): | |
| config, weights = tiny_model(architecture) | |
| config["num_key_value_heads"] = kv_heads | |
| rng = np.random.default_rng(913) | |
| for layer in range(config["num_hidden_layers"]): | |
| for projection in ("k_proj", "v_proj"): | |
| weights[f"model.layers.{layer}.self_attn.{projection}.weight"] = rng.normal( | |
| 0, 0.14, (kv_heads * config["head_dim"], config["hidden_size"]) | |
| ).astype(np.float32) | |
| model = native.Model(config, weights) | |
| prompt = [1, 4, 6, 7] | |
| np.testing.assert_allclose(model.logits(prompt), reference_logits(config, weights, prompt), atol=8e-6, rtol=8e-5) | |
| expected = [] | |
| for _ in range(5): | |
| expected.append(int(np.argmax(reference_logits(config, weights, prompt + expected)))) | |
| assert model.create_session(prompt, max_new_tokens=5).next_tokens(5) == expected | |
| def test_inference_failure_terminates_session(): | |
| config, weights = tiny_model() | |
| model = native.Model(config, weights) | |
| # A finite FP32 input can still overflow a dot product. Refuse non-finite output. | |
| weights["model.layers.0.self_attn.q_proj.weight"].fill(np.finfo(np.float32).max) | |
| weights["model.layers.0.self_attn.k_proj.weight"].fill(np.finfo(np.float32).max) | |
| invalid = native.Model(config, weights) | |
| with pytest.raises(RuntimeError, match="non-finite"): | |
| invalid.logits([1, 2]) | |
| # Finite weights pass construction; the first real inference is lazy. | |
| session = invalid.create_session([1, 2], max_new_tokens=2) | |
| assert session.finish_reason == "" | |
| assert session.next_tokens(0) == [] | |
| with pytest.raises(RuntimeError, match="non-finite"): | |
| session.next_tokens(1) | |
| assert session.finish_reason == "cancelled" | |
| assert session.generated_tokens == 0 | |
| assert session.next_tokens(1) == [] | |
| cancelled = invalid.create_session([1, 2], max_new_tokens=2) | |
| cancelled.cancel() | |
| assert cancelled.next_tokens(1) == [] | |
| assert cancelled.finish_reason == "cancelled" | |
| empty = invalid.create_session([1, 2], max_new_tokens=0) | |
| assert empty.next_tokens(1) == [] | |
| assert empty.finish_reason == "length" | |
| assert np.isfinite(model.logits([1, 2])).all() | |
| def test_quantized_vector_boundaries(architecture, precision, width): | |
| config, weights = tiny_model(architecture, hidden=width, intermediate=width + 2) | |
| reconstructed = reconstructed_weights(weights, precision) | |
| model = native.Model(config, weights, precision) | |
| prompt = [1, 5, 2] | |
| expected = [] | |
| for _ in range(4): | |
| prefix = prompt + expected | |
| target = reference_logits(config, reconstructed, prefix) | |
| np.testing.assert_allclose(model.logits(prefix), target, rtol=0, atol=1e-4) | |
| expected.append(int(np.argmax(target))) | |
| session = model.create_session(prompt, max_new_tokens=4) | |
| assert session.next_tokens(1) + session.next_tokens(3) == expected | |
| def test_int4_signed_nibbles_and_distinct_block_scales(): | |
| config, weights = tiny_model("qwen3", hidden=65, intermediate=97) | |
| for name, weight in weights.items(): | |
| if ".mlp." not in name: | |
| continue | |
| # Include every representable quantized value, distinct adjacent nibbles, | |
| # negative/positive row extrema, zero blocks, and multiple block scales. | |
| rows, cols = np.indices(weight.shape) | |
| weight[:] = (((cols + rows) % 15) - 7) * np.where((cols // 32) % 2, 0.007, 0.02) | |
| weight[0, :32] = 0 | |
| weight[1, 32:64] = 0 | |
| reconstructed = reconstructed_weights(weights, "hybrid-int4") | |
| model = native.Model(config, weights, "hybrid-int4") | |
| for prompt in ([1], [3, 9, 2], [2, 7, 5, 6, 1]): | |
| np.testing.assert_allclose(model.logits(prompt), reference_logits(config, reconstructed, prompt), rtol=0, atol=1e-4) | |
| def test_multithreaded_matches_single_thread_bitwise(threads): | |
| config, weights = tiny_model("qwen3", tied=True) | |
| single = native.Model(config, weights, "hybrid-int4", threads=1) | |
| multi = native.Model(config, weights, "hybrid-int4", threads=threads) | |
| assert multi.info["threads"] == threads | |
| # Row slices execute identical per-row dot kernels; results are exact. | |
| for prompt in ([1], [3, 9, 2], [2, 7, 5, 6, 1]): | |
| np.testing.assert_array_equal(multi.logits(prompt), single.logits(prompt)) | |
| session_single = single.create_session([1, 5, 2], max_new_tokens=8, temperature=0.9, | |
| top_p=0.9, top_k=20, seed=42, eos_token_ids=[]) | |
| session_multi = multi.create_session([1, 5, 2], max_new_tokens=8, temperature=0.9, | |
| top_p=0.9, top_k=20, seed=42, eos_token_ids=[]) | |
| np.testing.assert_array_equal(session_multi.next_tokens(8), session_single.next_tokens(8)) | |
| def test_invalid_thread_counts(): | |
| config, weights = tiny_model("llama", tied=True) | |
| for threads in (0, 65, -1): | |
| with pytest.raises(ValueError, match="threads"): | |
| native.Model(config, weights, "fp32", threads=threads) | |
| # Floats are strictly rejected at the binding boundary (TypeError), while | |
| # booleans are integer subclasses: True is a valid single-threaded model. | |
| with pytest.raises(TypeError): | |
| native.Model(config, weights, "fp32", threads=2.5) | |
| assert native.Model(config, weights, "fp32", threads=True).info["threads"] == 1 | |
| def test_page_locking_and_touch(precision): | |
| config, weights = tiny_model("qwen3", tied=True) | |
| model = native.Model(config, weights, precision) | |
| info_bytes = model.info["weight_bytes"] | |
| assert not model.info["pages_locked"] | |
| # Touch sums every stored weight byte, including quantized buffers, scales, | |
| # and the rotary table, which weight_bytes excludes (head_dim/2 floats). | |
| first = model.touch() | |
| rotary_bytes = (config["head_dim"] // 2) * 4 | |
| assert info_bytes < first <= info_bytes + rotary_bytes | |
| assert model.touch() == first | |
| assert model.scan() == first | |
| assert model.scan() == model.scan() | |
| locked = model.lock_pages() | |
| if locked: | |
| assert locked <= info_bytes + rotary_bytes | |
| assert model.info["pages_locked"] | |
| with pytest.raises(ValueError, match="already locked"): | |
| model.lock_pages() | |
| # Page sharing makes exact byte equality impossible; small neighboring | |
| # regions occupy the same OS page, so unlock can only report <= locked. | |
| unlocked = model.unlock_pages() | |
| assert 0 < unlocked <= locked | |
| assert not model.info["pages_locked"] | |
| with pytest.raises(ValueError, match="not locked"): | |
| model.unlock_pages() | |
| # Inference still works after lock/unlock cycles on all storage formats. | |
| model.create_session([1], max_new_tokens=2, temperature=0, top_p=1, top_k=0, | |
| seed=0, eos_token_ids=[]).next_tokens(2) | |