diff --git a/deepinfer/model_executor/model_loader.py b/deepinfer/model_executor/model_loader.py index 21253e6..6b05403 100644 --- a/deepinfer/model_executor/model_loader.py +++ b/deepinfer/model_executor/model_loader.py @@ -12,7 +12,7 @@ from pathlib import Path import torch import torch.nn as nn -from deepinfer import AttrDict, envs +from deepinfer import PLATFORM, AttrDict, envs from deepinfer.config import ModelConfig from deepinfer.hugging_face import HFCoverage from deepinfer.logger import init_logger @@ -73,10 +73,12 @@ def get_model( model_config.model_cfg, skip_load, ) + # The model device is "cuda", or the mapped Neuron target (neuron:0, or the host in Neuron CPU mode). + device_type = PLATFORM.device.type if PLATFORM.is_neuron else "cuda" for n, p in model.named_parameters(): - assert p.is_cuda, f"{n} not on cuda" + assert p.device.type == device_type, f"{n} not on {device_type}" if hasattr(p, "sf"): - assert p.sf.is_cuda, f"{n}.sf not on cuda" + assert p.sf.device.type == device_type, f"{n}.sf not on {device_type}" return model.eval().requires_grad_(False) diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/attention.py b/deepinfer/neuron_ops/vendor/deep_gemm/attention.py index afc6c3e..7797fee 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/attention.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/attention.py @@ -250,7 +250,9 @@ def fp8_fp4_paged_mqa_logits( positions = torch.arange(c0, c1, device=device, dtype=torch.int64)[None, :] valid = (positions < lens[:, None]) & (positions < max_positions) safe = positions.clamp(max=max(0, max_positions - 1)).expand(num_q_tokens, -1) - page_ids = table.gather(1, safe // page_kv) + # Block-table slots past the context may hold stale/garbage page ids: read page 0 there + # (masked to -inf below) instead of dereferencing them. + page_ids = torch.where(valid, table.gather(1, safe // page_kv), 0) keys = _decode_paged_bytes(*_gather_paged_bytes(kv_cache, page_ids, safe % page_kv, is_fp4), is_fp4) values = _scores_per_query_keys(q_f32, keys, weights_f32) logits[:, c0:c1] = torch.where(valid, values, torch.full_like(values, _NEG_INF)).to(torch.bfloat16) @@ -438,7 +440,7 @@ def fp8_fp4_paged_sparse_mqa_logits( if metadata._kv_snapshot is None: safe = tokens.clamp(0, max(0, max_positions - 1)) - page_ids = table.gather(1, safe // page_kv) + page_ids = torch.where(valid, table.gather(1, safe // page_kv), 0) # see the dense paged kernel metadata._kv_snapshot = _gather_paged_bytes(kv_cache, page_ids, safe % page_kv, is_fp4) data_bytes, sf_bytes = metadata._kv_snapshot diff --git a/deepinfer/utils.py b/deepinfer/utils.py index 10810c7..491b1bf 100644 --- a/deepinfer/utils.py +++ b/deepinfer/utils.py @@ -141,6 +141,11 @@ def allow_cuda_sync() -> AbstractContextManager[None]: def bind_cpu(local_rank: int) -> None: + from deepinfer import PLATFORM + + if PLATFORM.is_neuron and envs.DEEPINFER_NEURON_DEVICE == "cpu": + # Neuron CPU mode computes on the host: a 1-core share (96 cores / 64 NeuronCores) would serialize it. + return phy_cpus = psutil.cpu_count(logical=False) ngpus = get_actual_num_gpus() num_cpus_per_gpu = phy_cpus // ngpus diff --git a/examples/neuron/README.md b/examples/neuron/README.md new file mode 100644 index 0000000..2d93f79 --- /dev/null +++ b/examples/neuron/README.md @@ -0,0 +1,37 @@ +# DeepInfer on AWS Trainium — end-to-end check with a tiny checkpoint + +Design and conventions of the port: [deepinfer/neuron_ops/NEURON_PORT.md](../../deepinfer/neuron_ops/NEURON_PORT.md). + +The ground truth is DeepSeek's official reference implementation (`inference/` in the +`deepseek-ai/DeepSeek-V4.1-Flash-1017` HF repo) running on CPU, with its TileLang kernels replaced by +the pure-torch versions in `ref_kernels_torch.py`. + +| Script | Purpose | +|---|---| +| `make_tiny_checkpoint.py` | Random tiny V4.1 checkpoint in the real HF layout (FP8 e4m3 + ue8m0 32x32 scales, FP4 e2m1 experts with 1x32 scales, engram tables) | +| `check_tiny_checkpoint.py` | DeepInfer loader coverage / shape / dtype checks on the meta device | +| `ref_kernels_torch.py`, `test_ref_kernels_torch.py` | Torch replacements for the reference `kernel.py` | +| `reference_golden.py` | Reference prefill + greedy decode; saves logits, tokens and per-layer hidden states | +| `run_deepinfer_tiny.py` | Same prompts through the DeepInfer engine (`LLM` offline API), teacher-forced | +| `compare_golden.py` | Candidate vs golden: argmax match, logits rel / cos / KL, per-layer hidden error | + +Tiny config: 10 layers + 3 nextn, hidden 1024, 8 heads x 512 (MLA, rope 64), 16 experts / top-4, +indexer 4 x 128 with top-64, engram on layers 1 and 4, real tokenizer (vocab 129280). + +```bash +export DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu # or =neuron for the NeuronCore +python examples/neuron/run_deepinfer_tiny.py --threads 32 --hidden --compare +``` + +## Results so far (CPU mode, basic mode: eager, no DSpark/MTP, EP=DP=CP=1) + +Teacher-forced, prefill + 8 decode steps per prompt; "floor" is the reference compared with itself +under a different CPU thread count (FP8 activation quantization amplifies sub-ulp differences). + +| prompt (tokens) | argmax match | max logits rel (floor) | min cos (floor) | +|---|---|---|---| +| short (11) | 9/9 | 0.055 (0.032) | 0.9985 (0.9995) | +| medium (172, wraps the 128 SWA window) | 8/9 | 0.049 (0.039) | 0.9988 (0.9993) | +| long (715, candidate-block filter active) | 8/9 | 0.056 (0.060) | 0.9985 (0.9982) | + +Per-step / per-layer numbers: `golden/deepinfer_cpu_vs_golden.json`. diff --git a/examples/neuron/check_tiny_checkpoint.py b/examples/neuron/check_tiny_checkpoint.py new file mode 100644 index 0000000..7c7a7be --- /dev/null +++ b/examples/neuron/check_tiny_checkpoint.py @@ -0,0 +1,142 @@ +"""Check that DeepInfer's HF loader accepts a (tiny) V4.1 checkpoint, on the host CPU only. + +Runs the same validation chain as DeepseekForCausalLM.load_model / model_loader.load_hf_checkpoint, but +without materializing or loading weights: + load_hf_model_config -> prepare_hf_coverage -> V41Builder on meta -> cover_tied_parameters + -> "every model parameter is covered" -> validate_hf_shapes (shapes + dtypes from the safetensors headers) + +Usage (on the box): + cd ~/di/DeepInfer && DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu \ + ~/di/venv/bin/python examples/neuron/check_tiny_checkpoint.py ~/di/tiny_ckpt +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron") +os.environ.setdefault("DEEPINFER_NEURON_DEVICE", "cpu") +os.environ.setdefault("MASTER_ADDR", "127.0.0.1") +os.environ.setdefault("MASTER_PORT", str(29000 + os.getpid() % 1000)) + + +_VENDOR = ("deep_gemm", "deep_ep", "deep_select", "flash_mla", "tile_kernels", "tilelang", "flash_attn") + + +def _stub_missing_vendor_packages() -> None: + """The pure-torch vendor packages are being written in parallel; building parameter containers on meta + only needs them to be importable. Stub any that are still missing (attribute access returns a dummy).""" + import importlib + import importlib.abc + import importlib.machinery + import types + + class _Dummy: + def __init__(self, name): + self._name = name + + def __getattr__(self, item): + return _Dummy(f"{self._name}.{item}") + + def __call__(self, *a, **k): + raise RuntimeError(f"stubbed vendor function {self._name} called") + + class _StubModule(types.ModuleType): + def __getattr__(self, item): + if item.startswith("__"): + raise AttributeError(item) + return _Dummy(f"{self.__name__}.{item}") + + class _Finder(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def find_spec(self, fullname, path, target=None): + if fullname.split(".")[0] in stubbed: + return importlib.machinery.ModuleSpec(fullname, self, is_package=True) + return None + + def create_module(self, spec): + m = _StubModule(spec.name) + m.__path__ = [] + return m + + def exec_module(self, module): + pass + + stubbed = set() + for name in _VENDOR: + try: + importlib.import_module(name) + except Exception as exc: # noqa: BLE001 + stubbed.add(name) + print(f"[check] vendor package {name!r} not importable ({type(exc).__name__}: {exc}); stubbing it") + if stubbed: + sys.meta_path.insert(0, _Finder()) + + +def main(hf_dir: str) -> None: + import torch + import torch.distributed as dist + + import deepinfer + from deepinfer import PLATFORM # noqa: F401 (installs transfer_to_neuron -> cpu) + + _stub_missing_vendor_packages() + + # transfer_to_neuron replaces torch.empty/zeros/... with Python wrappers. torch.utils._device caches the + # set of factory functions it redirects inside `with torch.device("meta")` on first use, i.e. after the + # patch, so it holds only the wrappers and the original functions they call are NOT redirected: the + # "meta" build then allocates real tensors (and prepare_meta_fp8 asserts). Work around it here by also + # registering the wrapped originals. (Bug to be fixed in deepinfer/neuron_ops/transfer_to_neuron.py.) + import torch.utils._device as _tdev + + _ctors = set(_tdev._device_constructors()) + _ctors |= {f.__wrapped__ for f in list(_ctors) if hasattr(f, "__wrapped__")} + _tdev._device_constructors = lambda: _ctors + + assert str(PLATFORM.device) == "cpu", f"refusing to run on {PLATFORM.device}; set DEEPINFER_NEURON_DEVICE=cpu" + if not dist.is_initialized(): + dist.init_process_group("gloo", rank=0, world_size=1) + from deepinfer.model_executor.parallel_utils.parallel_state import PG + + world = dist.group.WORLD + for attr in ("dp", "cp", "ep", "ep_in_dp", "enp", "world", "local", "cpu_dp_group", "cpu_cp_group", "cpu_ep_in_dp"): + setattr(PG, attr, world) + PG.rank = 0 + + from deepinfer.hugging_face import cover_tied_parameters, load_hf_model_config, model_params, prepare_hf_coverage, validate_hf_shapes + from deepinfer.model_executor.model_loader import build_model_on_meta + + hf = Path(hf_dir).expanduser() + model_cfg = load_hf_model_config(hf / "config.json") + gpt = model_cfg.gpt + # runtime fields normally filled by ModelConfig / finalize_model_config + gpt.setdefault("runtime_max_model_len", 4096) + gpt.setdefault("max_num_seqs", 4) + gpt.setdefault("max_num_batched_tokens", 4096) + gpt.moe.setdefault("n_redundant_experts", 0) + if gpt.get("dflash_cfg") and gpt.dflash_cfg.get("moe"): + gpt.dflash_cfg.moe.setdefault("n_redundant_experts", 0) + + coverage = prepare_hf_coverage(gpt, hf_dir=hf) + model = build_model_on_meta(model_cfg) + params = model_params(model) + coverage = cover_tied_parameters(coverage, set(params), n_layer=gpt.num_hidden_layers) + missing = sorted(set(params) - coverage.model_names()) + unused = sorted(set(coverage.model_names()) - set(params)) + print(f"model params: {len(params)}, covered names: {len(coverage.model_names())}") + print(f"missing (model params without HF source): {len(missing)} {missing[:20]}") + print(f"unused HF groups (no model param): {len(unused)} {unused[:20]}") + validate_hf_shapes(params, coverage, hf_dir=hf) + n_exp = {k: len(v) for k, v in coverage.expert_tensors.items()} + print(f"expert groups: {len(n_exp)} (experts per group: {sorted(set(n_exp.values()))}); engram tables: {coverage.engram_tables}") + for name in sorted(n for n in params if "engram_embedding.embedding.embedding" in n): + p = params[name] + print(f" {name}: runtime rows {tuple(p.shape)} (row_start {getattr(p, 'row_start', None)}), head sizes {getattr(p, 'head_vocab_sizes', None)}") + assert not missing, "DeepInfer would reject this checkpoint (uncovered parameters)" + print("OK: DeepInfer HF loader accepts", hf) + dist.destroy_process_group() + + +if __name__ == "__main__": + main(sys.argv[1] if len(sys.argv) > 1 else "~/di/tiny_ckpt") diff --git a/examples/neuron/compare_golden.py b/examples/neuron/compare_golden.py new file mode 100644 index 0000000..79206eb --- /dev/null +++ b/examples/neuron/compare_golden.py @@ -0,0 +1,106 @@ +"""Compare a candidate run against a reference golden (both in the reference_golden.py .pt format). + + python -I examples/neuron/compare_golden.py ~/di/golden/medium.pt .pt [--json out.json] + +The candidate must have been run TEACHER-FORCED on the golden's tokens (prompt_ids, then golden `tokens[:-1]` as +decode inputs); otherwise decode-step logits are compared across different contexts. Required keys in the +candidate: prompt_ids, logits [steps+1, vocab]; optional: fed_tokens, hidden (list per forward call of +[n_layers, seqlen, hc_mult, dim]), indexer_topk. + +Metrics per forward call (step 0 = prefill, which predicts the first token): + argmax match, rel = ||a-b|| / ||a|| of the fp32 logits, cosine, max |a-b|, top-5 overlap, KL(p_ref || p_cand), + and the golden top1-top2 margin (an argmax mismatch where the margin is below the noise floor is not a bug). +Per layer (if both have `hidden`): relative L2 error of the block output over all tokens, and of the worst token. + +The reference is not bit-stable even against itself: FP8 activation quantization turns sub-ulp differences +(e.g. a different CPU thread count) into ~1e-2 relative per-layer differences, and top-k routing / index +selection compound them. golden/noise_floor_*.json records the reference-vs-reference spread; judge a +candidate against that floor, not against zero. +""" + +from __future__ import annotations + +import argparse +import json + +import torch + + +def rel(a, b): + a, b = a.float(), b.float() + return ((a - b).norm() / a.norm().clamp_min(1e-30)).item() + + +def compare(ref: dict, cand: dict) -> dict: + assert list(ref["prompt_ids"]) == list(cand["prompt_ids"]), "different prompts" + if "fed_tokens" in cand: + exp = list(ref["tokens"][: len(cand["fed_tokens"])]) + if list(cand["fed_tokens"]) != exp: + print("WARNING: candidate was not teacher-forced on the golden tokens; decode steps are not comparable") + la, lb = ref["logits"].float(), cand["logits"].float() + n = min(la.size(0), lb.size(0)) + steps = [] + for i in range(n): + a, b = la[i], lb[i] + top2 = a.topk(2).values + pa, lpb = a.log_softmax(-1), b.log_softmax(-1) + steps.append( + dict( + step=i, + argmax_ref=int(a.argmax()), + argmax_cand=int(b.argmax()), + match=bool(a.argmax() == b.argmax()), + margin_ref=round((top2[0] - top2[1]).item(), 5), + rel=rel(a, b), + cos=torch.nn.functional.cosine_similarity(a, b, dim=0).item(), + maxabs=(a - b).abs().max().item(), + top5_overlap=len(set(a.topk(5).indices.tolist()) & set(b.topk(5).indices.tolist())), + kl=(pa.exp() * (pa - lpb)).sum().item(), + ) + ) + layers = [] + if ref.get("hidden") and cand.get("hidden"): + for i in range(min(len(ref["hidden"]), len(cand["hidden"]))): + ha, hb = ref["hidden"][i].float(), cand["hidden"][i].float() + if ha.shape != hb.shape: + layers.append(dict(step=i, error=f"shape {tuple(ha.shape)} vs {tuple(hb.shape)}")) + continue + dims = tuple(range(2, ha.dim())) + per_tok = (ha - hb).pow(2).sum(dims).sqrt() / ha.pow(2).sum(dims).sqrt().clamp_min(1e-30) # [L, s] + layers.append( + dict( + step=i, + rel=[rel(ha[l], hb[l]) for l in range(ha.size(0))], + worst_token_rel=per_tok.max(-1).values.tolist(), + ) + ) + return dict(steps=steps, layers=layers) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("ref") + ap.add_argument("cand") + ap.add_argument("--json", default=None) + args = ap.parse_args() + ref = torch.load(args.ref, weights_only=False) + cand = torch.load(args.cand, weights_only=False) + res = compare(ref, cand) + print(f"{'step':>4} {'match':>5} {'margin':>7} {'rel':>9} {'cos':>9} {'maxabs':>8} {'top5':>4} {'KL':>9}") + for s in res["steps"]: + print( + f"{s['step']:>4} {str(s['match']):>5} {s['margin_ref']:>7.3f} {s['rel']:>9.2e} {s['cos']:>9.6f} " + f"{s['maxabs']:>8.3f} {s['top5_overlap']:>4} {s['kl']:>9.2e}" + ) + for L in res["layers"]: + if "error" in L: + print(f"hidden step {L['step']}: {L['error']}") + else: + print(f"hidden step {L['step']}: layer rel " + " ".join(f"{v:.1e}" for v in L["rel"])) + if args.json: + with open(args.json, "w") as fh: + json.dump(res, fh, indent=1) + + +if __name__ == "__main__": + main() diff --git a/examples/neuron/golden/deepinfer_cpu_vs_golden.json b/examples/neuron/golden/deepinfer_cpu_vs_golden.json new file mode 100644 index 0000000..eefcb65 --- /dev/null +++ b/examples/neuron/golden/deepinfer_cpu_vs_golden.json @@ -0,0 +1,316 @@ +{ + "what": "DeepInfer LLM engine (DEEPINFER_NEURON_DEVICE=cpu, EP=DP=CP=1, eager, nextn=0) teacher-forced on the reference golden, vs golden; floor = reference vs itself", + "date": "2026-10-20", + "prompts": { + "short": { + "argmax_match": 9, + "steps": 9, + "mismatches": [], + "max_logits_rel": 0.05466, + "floor_max_logits_rel": 0.031786, + "min_logits_cos": 0.998507, + "floor_min_logits_cos": 0.999495, + "max_kl": 0.001507, + "floor_max_kl": 0.000499, + "max_hidden_rel_per_layer": [ + 0.00545, + 0.01526, + 0.02523, + 0.03109, + 0.03846, + 0.04597, + 0.0491, + 0.05428, + 0.05774, + 0.06081 + ], + "floor_max_hidden_rel_per_layer": [ + 6e-05, + 0.00047, + 0.004855, + 0.009456, + 0.012744, + 0.018298, + 0.02, + 0.021909, + 0.030419, + 0.031869 + ], + "per_step": [ + { + "step": 0, + "match": true, + "rel": 0.04924, + "cos": 0.998789, + "kl": 0.0012128 + }, + { + "step": 1, + "match": true, + "rel": 0.04901, + "cos": 0.998799, + "kl": 0.0011874 + }, + { + "step": 2, + "match": true, + "rel": 0.05466, + "cos": 0.998507, + "kl": 0.0015074 + }, + { + "step": 3, + "match": true, + "rel": 0.04741, + "cos": 0.998877, + "kl": 0.001123 + }, + { + "step": 4, + "match": true, + "rel": 0.04551, + "cos": 0.998965, + "kl": 0.0010243 + }, + { + "step": 5, + "match": true, + "rel": 0.0474, + "cos": 0.998878, + "kl": 0.0011334 + }, + { + "step": 6, + "match": true, + "rel": 0.05391, + "cos": 0.998548, + "kl": 0.0014566 + }, + { + "step": 7, + "match": true, + "rel": 0.04517, + "cos": 0.99898, + "kl": 0.0010055 + }, + { + "step": 8, + "match": true, + "rel": 0.04621, + "cos": 0.998934, + "kl": 0.001059 + } + ] + }, + "medium": { + "argmax_match": 8, + "steps": 9, + "mismatches": [ + { + "step": 6, + "golden_margin": 0.00963 + } + ], + "max_logits_rel": 0.04872, + "floor_max_logits_rel": 0.038841, + "min_logits_cos": 0.998814, + "floor_min_logits_cos": 0.999248, + "max_kl": 0.001187, + "floor_max_kl": 0.000751, + "max_hidden_rel_per_layer": [ + 0.02296, + 0.02309, + 0.02499, + 0.02653, + 0.02806, + 0.03014, + 0.03367, + 0.03765, + 0.04512, + 0.04879 + ], + "floor_max_hidden_rel_per_layer": [ + 3e-05, + 0.000146, + 0.00754, + 0.010748, + 0.015295, + 0.020652, + 0.030513, + 0.032413, + 0.036607, + 0.0383 + ], + "per_step": [ + { + "step": 0, + "match": true, + "rel": 0.03771, + "cos": 0.99929, + "kl": 0.0007106 + }, + { + "step": 1, + "match": true, + "rel": 0.04001, + "cos": 0.9992, + "kl": 0.0008061 + }, + { + "step": 2, + "match": true, + "rel": 0.04275, + "cos": 0.999087, + "kl": 0.0009217 + }, + { + "step": 3, + "match": true, + "rel": 0.0362, + "cos": 0.999346, + "kl": 0.0006492 + }, + { + "step": 4, + "match": true, + "rel": 0.04872, + "cos": 0.998814, + "kl": 0.0011875 + }, + { + "step": 5, + "match": true, + "rel": 0.03131, + "cos": 0.999511, + "kl": 0.0004906 + }, + { + "step": 6, + "match": false, + "rel": 0.03696, + "cos": 0.999318, + "kl": 0.0006812 + }, + { + "step": 7, + "match": true, + "rel": 0.04059, + "cos": 0.999177, + "kl": 0.0008209 + }, + { + "step": 8, + "match": true, + "rel": 0.03795, + "cos": 0.999281, + "kl": 0.0007192 + } + ] + }, + "long": { + "argmax_match": 8, + "steps": 9, + "mismatches": [ + { + "step": 7, + "golden_margin": 0.1591 + } + ], + "max_logits_rel": 0.05571, + "floor_max_logits_rel": 0.059876, + "min_logits_cos": 0.998449, + "floor_min_logits_cos": 0.998208, + "max_kl": 0.001541, + "floor_max_kl": 0.001755, + "max_hidden_rel_per_layer": [ + 0.01959, + 0.02303, + 0.02454, + 0.02814, + 0.03133, + 0.03751, + 0.04054, + 0.04661, + 0.05095, + 0.05564 + ], + "floor_max_hidden_rel_per_layer": [ + 0.000172, + 0.000188, + 0.000175, + 0.00383, + 0.022313, + 0.046369, + 0.050459, + 0.054834, + 0.05773, + 0.059775 + ], + "per_step": [ + { + "step": 0, + "match": true, + "rel": 0.05293, + "cos": 0.9986, + "kl": 0.001395 + }, + { + "step": 1, + "match": true, + "rel": 0.05107, + "cos": 0.998697, + "kl": 0.0013076 + }, + { + "step": 2, + "match": true, + "rel": 0.05101, + "cos": 0.9987, + "kl": 0.0013051 + }, + { + "step": 3, + "match": true, + "rel": 0.05571, + "cos": 0.998449, + "kl": 0.0015413 + }, + { + "step": 4, + "match": true, + "rel": 0.04936, + "cos": 0.998783, + "kl": 0.0012217 + }, + { + "step": 5, + "match": true, + "rel": 0.05548, + "cos": 0.998462, + "kl": 0.0015173 + }, + { + "step": 6, + "match": true, + "rel": 0.05265, + "cos": 0.998616, + "kl": 0.0013786 + }, + { + "step": 7, + "match": false, + "rel": 0.05182, + "cos": 0.998658, + "kl": 0.0013534 + }, + { + "step": 8, + "match": true, + "rel": 0.05312, + "cos": 0.99859, + "kl": 0.0013851 + } + ] + } + } +} \ No newline at end of file diff --git a/examples/neuron/golden/long.json b/examples/neuron/golden/long.json new file mode 100644 index 0000000..a76c53a --- /dev/null +++ b/examples/neuron/golden/long.json @@ -0,0 +1,900 @@ +{ + "name": "long", + "prompt_text": "In the early morning the harbor was quiet. Fishing boats rocked gently against the wooden piers, their ropes creaking as the tide came in. An old man sat on an overturned crate, mending a net with slow, careful hands. He had worked these waters for more than fifty years, and he could read the weather in the color of the sky and the shape of the clouds. Today the sky was pale gray, and a thin mist hung over the water. Somewhere a gull cried, and the sound echoed off the stone walls of the warehouses. A young girl ran down the pier, her boots clattering on the boards. She stopped beside the old man and asked him whether the boats would go out today. He looked at the horizon for a long moment before answering. The sea, he said, decides that for us. We only listen. The history of computing is a story of abstraction. Early machines were programmed by rewiring panels and flipping switches, and every program was tied to the physical layout of the hardware. Assembly languages gave names to instructions, and compilers later translated whole expressions into sequences of those instructions. Operating systems hid the details of disks and terminals behind files and streams, while virtual memory let each process pretend it owned the entire machine. Networks added another layer. Packets travel through routers that know nothing about the applications sending them, and protocols stacked on top of one another turn unreliable links into dependable conversations. A web page that appears in a fraction of a second may involve dozens of servers, caches and databases spread across several continents, each one following rules that were agreed upon decades ago by people who could not have imagined how the system would grow. Mathematics offers a similar lesson. A proof about prime numbers does not care whether the numbers are written in decimal or binary, on paper or in silicon. The same theorem holds in every notation, and that independence is exactly what makes it useful. Engineers borrow this idea when they define interfaces: as long as both sides respect the contract, either side can be rewritten, optimized or replaced without the other noticing. Cooking, perhaps surprisingly, works the same way. A recipe names its ingredients and steps but leaves the choice of pan, stove and kitchen to the cook. Bread rises because yeast consumes sugar and releases gas, whether the dough is kneaded by hand in a village bakery or by a machine in a factory. Understanding the underlying chemistry lets a baker adapt when the flour is different, the air is humid, or the oven runs hot. Music, too, depends on structure. A melody can be transposed into another key, played on a different instrument, or slowed down to half its tempo, and listeners will still recognize it. Composers rely on patterns of tension and release, repetition and variation, to guide the ear through a piece. The notes on the page are only instructions; the performance brings them to life in a particular room on a particular evening, never exactly the same twice. Gardens change with every season. In spring the first green shoots push through the cold soil, and by summer the beds are crowded with tomatoes, beans and sunflowers that lean toward the light. Autumn brings the harvest and the slow work of clearing dead stems, while winter rests under frost and snow. A patient gardener learns that each year is different and that the soil remembers what was done to it. At the end of the day the harbor grew quiet again. The boats returned one by one, low in the water with their catch, and the old man folded his net and walked slowly home along the shore.", + "prompt_ids": [ + 0, + 1124, + 270, + 3808, + 7469, + 270, + 58523, + 515, + 12426, + 16, + 75989, + 36474, + 107943, + 25680, + 2765, + 270, + 26166, + 121446, + 14, + 786, + 78139, + 1914, + 1690, + 412, + 270, + 41949, + 4127, + 295, + 16, + 1417, + 3072, + 836, + 3825, + 377, + 411, + 106554, + 71040, + 14, + 292, + 3261, + 260, + 3150, + 418, + 6470, + 14, + 16372, + 6099, + 16, + 1166, + 936, + 7911, + 1305, + 16985, + 362, + 850, + 1099, + 12099, + 1737, + 14, + 305, + 468, + 1494, + 1733, + 270, + 9670, + 295, + 270, + 3605, + 294, + 270, + 12709, + 305, + 270, + 7633, + 294, + 270, + 26316, + 16, + 13253, + 270, + 12709, + 515, + 24754, + 22799, + 14, + 305, + 260, + 12220, + 9470, + 15179, + 1060, + 270, + 2237, + 16, + 5183, + 4779, + 260, + 108711, + 32080, + 14, + 305, + 270, + 5706, + 61415, + 1375, + 270, + 13754, + 13654, + 294, + 270, + 103677, + 16, + 334, + 4047, + 6482, + 11693, + 1910, + 270, + 21561, + 14, + 969, + 39182, + 667, + 46354, + 377, + 270, + 26730, + 16, + 2566, + 12852, + 24570, + 270, + 3072, + 836, + 305, + 4869, + 1440, + 3920, + 270, + 36474, + 1114, + 807, + 798, + 4316, + 16, + 1166, + 6273, + 509, + 270, + 31383, + 362, + 260, + 1606, + 4712, + 1936, + 37298, + 16, + 455, + 8343, + 14, + 468, + 1359, + 14, + 37581, + 396, + 362, + 550, + 16, + 1350, + 1353, + 12440, + 16, + 455, + 3980, + 294, + 17117, + 344, + 260, + 4093, + 294, + 57015, + 16, + 15225, + 17341, + 881, + 55346, + 513, + 84331, + 6694, + 25940, + 305, + 94587, + 38252, + 14, + 305, + 1750, + 2305, + 515, + 21234, + 304, + 270, + 5131, + 22343, + 294, + 270, + 17491, + 16, + 20298, + 10555, + 6718, + 6362, + 304, + 12846, + 14, + 305, + 675, + 54934, + 3859, + 25914, + 4445, + 18224, + 1055, + 18323, + 294, + 1948, + 12846, + 16, + 40030, + 3977, + 25342, + 270, + 6744, + 294, + 59050, + 305, + 43478, + 5573, + 9405, + 305, + 28010, + 14, + 1901, + 10814, + 6464, + 2513, + 1660, + 1699, + 45077, + 436, + 18998, + 270, + 5221, + 8126, + 16, + 31956, + 5144, + 2658, + 7104, + 16, + 21250, + 1697, + 5901, + 1407, + 84413, + 396, + 1153, + 5760, + 943, + 270, + 6532, + 19849, + 1066, + 14, + 305, + 24749, + 57129, + 377, + 2656, + 294, + 834, + 2658, + 2859, + 78122, + 10348, + 1055, + 103360, + 22213, + 16, + 334, + 6545, + 4366, + 396, + 9238, + 295, + 260, + 9078, + 294, + 260, + 1957, + 1142, + 15546, + 39085, + 294, + 29260, + 14, + 274, + 5813, + 305, + 25496, + 7670, + 3881, + 3586, + 58748, + 14, + 1660, + 834, + 2502, + 7027, + 396, + 881, + 12141, + 4465, + 12706, + 5084, + 513, + 1482, + 995, + 1494, + 554, + 611, + 34814, + 1192, + 270, + 1487, + 1114, + 2334, + 16, + 20835, + 7070, + 260, + 4032, + 12316, + 16, + 334, + 12745, + 943, + 4309, + 3737, + 1918, + 554, + 2312, + 3920, + 270, + 3737, + 477, + 4866, + 295, + 7767, + 469, + 11680, + 14, + 377, + 4468, + 469, + 295, + 31077, + 16, + 455, + 1975, + 23916, + 12927, + 295, + 1750, + 16584, + 14, + 305, + 396, + 19088, + 344, + 9045, + 1205, + 4491, + 436, + 6923, + 16, + 47363, + 17242, + 566, + 5090, + 1082, + 851, + 11348, + 32042, + 28, + 412, + 1606, + 412, + 1952, + 10210, + 3678, + 270, + 6085, + 14, + 4338, + 3317, + 588, + 366, + 92769, + 14, + 37874, + 469, + 15484, + 2503, + 270, + 915, + 69258, + 16, + 71921, + 14, + 8808, + 36177, + 14, + 2984, + 270, + 1975, + 1722, + 16, + 334, + 17144, + 6362, + 1009, + 15833, + 305, + 6531, + 790, + 9120, + 270, + 6749, + 294, + 5569, + 14, + 54599, + 305, + 15350, + 304, + 270, + 6674, + 16, + 57216, + 33234, + 1754, + 34757, + 84592, + 11288, + 305, + 30985, + 5606, + 14, + 3920, + 270, + 30476, + 344, + 14343, + 16474, + 513, + 1893, + 295, + 260, + 13671, + 95042, + 469, + 513, + 260, + 8126, + 295, + 260, + 21727, + 16, + 15565, + 270, + 13716, + 20555, + 28162, + 260, + 111555, + 6708, + 1082, + 270, + 16668, + 344, + 1688, + 14, + 270, + 3525, + 344, + 50750, + 14, + 469, + 270, + 23366, + 12122, + 6025, + 16, + 14388, + 14, + 2727, + 14, + 11776, + 377, + 4456, + 16, + 334, + 54736, + 588, + 366, + 1331, + 20094, + 1055, + 2658, + 3077, + 14, + 8178, + 377, + 260, + 1688, + 11471, + 14, + 469, + 53544, + 1910, + 304, + 4747, + 1009, + 21901, + 14, + 305, + 42971, + 759, + 2413, + 14122, + 436, + 16, + 125854, + 387, + 18276, + 377, + 8804, + 294, + 17599, + 305, + 7917, + 14, + 42039, + 305, + 13882, + 14, + 304, + 8003, + 270, + 12765, + 1407, + 260, + 8521, + 16, + 455, + 9575, + 377, + 270, + 4366, + 477, + 1353, + 12846, + 29, + 270, + 4197, + 14257, + 1066, + 304, + 1988, + 295, + 260, + 3191, + 4772, + 377, + 260, + 3191, + 13044, + 14, + 3214, + 9045, + 270, + 1975, + 13796, + 16, + 51710, + 2757, + 418, + 1750, + 5029, + 16, + 660, + 9623, + 270, + 1257, + 6726, + 57158, + 9702, + 1407, + 270, + 8147, + 7842, + 14, + 305, + 513, + 8012, + 270, + 31352, + 477, + 41161, + 418, + 36123, + 14, + 26327, + 305, + 6029, + 77423, + 396, + 20543, + 7212, + 270, + 2900, + 16, + 67901, + 14257, + 270, + 18481, + 305, + 270, + 6470, + 1116, + 294, + 40445, + 7299, + 33106, + 14, + 1901, + 11369, + 52169, + 1202, + 38571, + 305, + 12494, + 16, + 334, + 5828, + 101816, + 47734, + 396, + 1660, + 1141, + 344, + 1688, + 305, + 396, + 270, + 7842, + 59676, + 1205, + 515, + 3989, + 304, + 436, + 16, + 3256, + 270, + 1522, + 294, + 270, + 2173, + 270, + 58523, + 15021, + 12426, + 1820, + 16, + 455, + 36474, + 10274, + 834, + 513, + 834, + 14, + 3157, + 295, + 270, + 2237, + 418, + 786, + 8764, + 14, + 305, + 270, + 3072, + 836, + 44056, + 793, + 3150, + 305, + 13577, + 13257, + 2680, + 3758, + 270, + 27918, + 16 + ], + "n_prompt_tokens": 715, + "tokens": [ + 110226, + 112553, + 15227, + 85761, + 86959, + 85315, + 27946, + 44597, + 106053 + ], + "decoded": "\u8f66\u7a97yat\u8282\u70b9 zostologiques\u90ae\u7f16 leaki\u010d Kond", + "top1_top2_margin": [ + 0.16406, + 0.23452, + 0.2282, + 0.55923, + 0.37373, + 0.11843, + 0.18562, + 0.1591, + 0.42782 + ], + "top5_ids": [ + [ + 110226, + 85792, + 98669, + 80461, + 9362 + ], + [ + 112553, + 17629, + 36136, + 101971, + 106529 + ], + [ + 15227, + 67648, + 129083, + 92665, + 52707 + ], + [ + 85761, + 10766, + 65263, + 16068, + 62747 + ], + [ + 86959, + 17925, + 62702, + 71188, + 49377 + ], + [ + 85315, + 66522, + 85353, + 46669, + 107743 + ], + [ + 27946, + 79720, + 9625, + 116206, + 71526 + ], + [ + 44597, + 56846, + 35818, + 42526, + 67572 + ], + [ + 106053, + 35940, + 20015, + 121068, + 24048 + ] + ], + "top5_logits": [ + [ + 4.38729, + 4.22324, + 4.11084, + 3.98992, + 3.97544 + ], + [ + 4.69323, + 4.45871, + 4.29074, + 4.22572, + 4.13341 + ], + [ + 4.31865, + 4.09045, + 3.96603, + 3.91204, + 3.90222 + ], + [ + 4.78987, + 4.23064, + 4.12185, + 3.95752, + 3.88381 + ], + [ + 4.70271, + 4.32898, + 4.00672, + 3.94672, + 3.91344 + ], + [ + 4.15846, + 4.04003, + 3.90236, + 3.87659, + 3.867 + ], + [ + 4.71181, + 4.52619, + 4.30112, + 4.25565, + 4.02916 + ], + [ + 4.44295, + 4.28385, + 4.22194, + 4.10899, + 3.94827 + ], + [ + 5.09891, + 4.67109, + 4.37312, + 3.96956, + 3.94636 + ] + ], + "logits_l2": [ + 359.677, + 360.695, + 359.8685, + 359.4601, + 359.465, + 358.5112, + 359.0301, + 359.3313, + 358.0074 + ], + "logits_mean": [ + 0.003591, + 0.001506, + -0.000184, + -0.003338, + 0.000688, + 0.00179, + 0.001405, + 0.002071, + 0.002439 + ], + "full_golden": "/home/ubuntu/di/golden/long.pt", + "full_golden_sha256": "4e00fd6f9bc63306a7aa7ade8125c00cd7f9bfcd9f4ea298b0dc7284e0983088", + "seconds": 4.0 +} \ No newline at end of file diff --git a/examples/neuron/golden/medium.json b/examples/neuron/golden/medium.json new file mode 100644 index 0000000..9aea047 --- /dev/null +++ b/examples/neuron/golden/medium.json @@ -0,0 +1,357 @@ +{ + "name": "medium", + "prompt_text": "In the early morning the harbor was quiet. Fishing boats rocked gently against the wooden piers, their ropes creaking as the tide came in. An old man sat on an overturned crate, mending a net with slow, careful hands. He had worked these waters for more than fifty years, and he could read the weather in the color of the sky and the shape of the clouds. Today the sky was pale gray, and a thin mist hung over the water. Somewhere a gull cried, and the sound echoed off the stone walls of the warehouses. A young girl ran down the pier, her boots clattering on the boards. She stopped beside the old man and asked him whether the boats would go out today. He looked at the horizon for a long moment before answering. The sea, he said, decides that for us. We only listen.", + "prompt_ids": [ + 0, + 1124, + 270, + 3808, + 7469, + 270, + 58523, + 515, + 12426, + 16, + 75989, + 36474, + 107943, + 25680, + 2765, + 270, + 26166, + 121446, + 14, + 786, + 78139, + 1914, + 1690, + 412, + 270, + 41949, + 4127, + 295, + 16, + 1417, + 3072, + 836, + 3825, + 377, + 411, + 106554, + 71040, + 14, + 292, + 3261, + 260, + 3150, + 418, + 6470, + 14, + 16372, + 6099, + 16, + 1166, + 936, + 7911, + 1305, + 16985, + 362, + 850, + 1099, + 12099, + 1737, + 14, + 305, + 468, + 1494, + 1733, + 270, + 9670, + 295, + 270, + 3605, + 294, + 270, + 12709, + 305, + 270, + 7633, + 294, + 270, + 26316, + 16, + 13253, + 270, + 12709, + 515, + 24754, + 22799, + 14, + 305, + 260, + 12220, + 9470, + 15179, + 1060, + 270, + 2237, + 16, + 5183, + 4779, + 260, + 108711, + 32080, + 14, + 305, + 270, + 5706, + 61415, + 1375, + 270, + 13754, + 13654, + 294, + 270, + 103677, + 16, + 334, + 4047, + 6482, + 11693, + 1910, + 270, + 21561, + 14, + 969, + 39182, + 667, + 46354, + 377, + 270, + 26730, + 16, + 2566, + 12852, + 24570, + 270, + 3072, + 836, + 305, + 4869, + 1440, + 3920, + 270, + 36474, + 1114, + 807, + 798, + 4316, + 16, + 1166, + 6273, + 509, + 270, + 31383, + 362, + 260, + 1606, + 4712, + 1936, + 37298, + 16, + 455, + 8343, + 14, + 468, + 1359, + 14, + 37581, + 396, + 362, + 550, + 16, + 1350, + 1353, + 12440, + 16 + ], + "n_prompt_tokens": 172, + "tokens": [ + 110226, + 112553, + 15227, + 85761, + 86959, + 11563, + 10257, + 126318, + 127265 + ], + "decoded": "\u8f66\u7a97yat\u8282\u70b9 zostologiques\u0e18\ufffd\ufffd Nasdaq\u5408\u7406\u5316", + "top1_top2_margin": [ + 0.09576, + 0.176, + 0.47041, + 0.55217, + 0.16517, + 0.02413, + 0.00963, + 0.06668, + 0.38849 + ], + "top5_ids": [ + [ + 110226, + 85792, + 80461, + 98669, + 57424 + ], + [ + 112553, + 36136, + 17629, + 64921, + 30344 + ], + [ + 15227, + 124367, + 22709, + 551, + 11173 + ], + [ + 85761, + 39686, + 27013, + 118390, + 127694 + ], + [ + 86959, + 17925, + 49377, + 114808, + 28304 + ], + [ + 11563, + 119471, + 93205, + 85315, + 85353 + ], + [ + 10257, + 103677, + 81986, + 54394, + 64408 + ], + [ + 126318, + 88760, + 31927, + 89605, + 47916 + ], + [ + 127265, + 106011, + 78074, + 6497, + 81795 + ] + ], + "top5_logits": [ + [ + 4.49845, + 4.40269, + 4.19878, + 3.98309, + 3.9216 + ], + [ + 4.44214, + 4.26613, + 4.10688, + 4.03814, + 3.96695 + ], + [ + 4.63011, + 4.15971, + 3.94564, + 3.83741, + 3.83039 + ], + [ + 4.67657, + 4.1244, + 4.04267, + 4.00452, + 3.98093 + ], + [ + 4.51708, + 4.35191, + 4.09455, + 4.063, + 3.97634 + ], + [ + 3.93672, + 3.91259, + 3.88577, + 3.85805, + 3.83134 + ], + [ + 4.16555, + 4.15592, + 4.15041, + 4.10569, + 3.88085 + ], + [ + 4.37949, + 4.31282, + 4.25638, + 3.9865, + 3.96823 + ], + [ + 4.95869, + 4.5702, + 4.0482, + 4.04001, + 3.99707 + ] + ], + "logits_l2": [ + 359.7198, + 360.4961, + 360.233, + 359.8513, + 359.4761, + 358.8099, + 359.3246, + 360.7809, + 358.6741 + ], + "logits_mean": [ + 0.004896, + 0.002201, + -0.000475, + -0.002406, + 0.0013, + 0.001889, + -0.005621, + 0.001583, + 0.00249 + ], + "full_golden": "/home/ubuntu/di/golden/medium.pt", + "full_golden_sha256": "73df855690a8423e4117d1178ce167b1720dbdd42e9e37e54cef307be81746bf", + "seconds": 3.3 +} \ No newline at end of file diff --git a/examples/neuron/golden/noise_floor.json b/examples/neuron/golden/noise_floor.json new file mode 100644 index 0000000..5a22cac --- /dev/null +++ b/examples/neuron/golden/noise_floor.json @@ -0,0 +1,73 @@ +{ + "what": "reference vs itself: torch CPU threads 16 (golden) vs 7 and 32, teacher-forced on the golden tokens; use as the tolerance floor when comparing DeepInfer against the golden (compare_golden.py)", + "prompts": { + "short": { + "max_logits_rel": 0.031786, + "min_logits_cos": 0.999495, + "max_kl": 0.000499, + "min_top5_overlap": 4, + "argmax_mismatches": [], + "max_hidden_rel_per_layer": [ + 6e-05, + 0.00047, + 0.004855, + 0.009456, + 0.012744, + 0.018298, + 0.02, + 0.021909, + 0.030419, + 0.031869 + ] + }, + "medium": { + "max_logits_rel": 0.038841, + "min_logits_cos": 0.999248, + "max_kl": 0.000751, + "min_top5_overlap": 4, + "argmax_mismatches": [ + { + "threads": 7, + "step": 6, + "golden_margin": 0.00963 + }, + { + "threads": 32, + "step": 6, + "golden_margin": 0.00963 + } + ], + "max_hidden_rel_per_layer": [ + 3e-05, + 0.000146, + 0.00754, + 0.010748, + 0.015295, + 0.020652, + 0.030513, + 0.032413, + 0.036607, + 0.0383 + ] + }, + "long": { + "max_logits_rel": 0.059876, + "min_logits_cos": 0.998208, + "max_kl": 0.001755, + "min_top5_overlap": 4, + "argmax_mismatches": [], + "max_hidden_rel_per_layer": [ + 0.000172, + 0.000188, + 0.000175, + 0.00383, + 0.022313, + 0.046369, + 0.050459, + 0.054834, + 0.05773, + 0.059775 + ] + } + } +} \ No newline at end of file diff --git a/examples/neuron/golden/short.json b/examples/neuron/golden/short.json new file mode 100644 index 0000000..8b5bb5c --- /dev/null +++ b/examples/neuron/golden/short.json @@ -0,0 +1,196 @@ +{ + "name": "short", + "prompt_text": "The quick brown fox jumps over the lazy dog.", + "prompt_ids": [ + 0, + 671, + 4787, + 13769, + 46012, + 54994, + 1060, + 270, + 41638, + 6397, + 16 + ], + "n_prompt_tokens": 11, + "tokens": [ + 85792, + 78870, + 110111, + 9689, + 101353, + 107157, + 84220, + 19174, + 82767 + ], + "decoded": "\u6301\u3061))\n\n\ncombined Pract@media acuity\u7ea4\u7ef4\u7d20 Sin\u89c4\u8303\u7684", + "top1_top2_margin": [ + 0.10489, + 0.18469, + 0.09769, + 0.10715, + 0.06664, + 0.23189, + 0.33355, + 0.29129, + 0.39137 + ], + "top5_ids": [ + [ + 85792, + 109217, + 110226, + 91841, + 87149 + ], + [ + 78870, + 62985, + 4375, + 37823, + 71835 + ], + [ + 110111, + 62076, + 14882, + 127459, + 10901 + ], + [ + 9689, + 75659, + 30795, + 6070, + 49485 + ], + [ + 101353, + 46574, + 36355, + 20609, + 21109 + ], + [ + 107157, + 5900, + 31951, + 75048, + 62107 + ], + [ + 84220, + 16903, + 26762, + 104561, + 106094 + ], + [ + 19174, + 94517, + 53537, + 46354, + 43228 + ], + [ + 82767, + 70793, + 124801, + 69604, + 40493 + ] + ], + "top5_logits": [ + [ + 4.30327, + 4.19838, + 4.06756, + 4.05694, + 4.0096 + ], + [ + 4.72153, + 4.53684, + 4.06865, + 4.00015, + 3.96316 + ], + [ + 4.20809, + 4.1104, + 4.08111, + 4.05279, + 4.05153 + ], + [ + 4.17989, + 4.07274, + 3.99208, + 3.97799, + 3.95346 + ], + [ + 4.39166, + 4.32503, + 3.96568, + 3.95226, + 3.93489 + ], + [ + 4.5328, + 4.30091, + 4.14662, + 4.01601, + 3.98233 + ], + [ + 4.48246, + 4.14891, + 4.05826, + 4.0559, + 4.02218 + ], + [ + 4.82395, + 4.53265, + 4.26447, + 4.00352, + 3.94851 + ], + [ + 4.66934, + 4.27798, + 3.91097, + 3.8789, + 3.85056 + ] + ], + "logits_l2": [ + 359.5489, + 359.0275, + 360.0101, + 359.5485, + 358.1219, + 360.3374, + 359.3568, + 358.4936, + 358.7865 + ], + "logits_mean": [ + 0.002626, + 0.004761, + -0.003585, + 0.002637, + -0.00203, + -0.000569, + 0.000388, + -0.001356, + 0.003071 + ], + "full_golden": "/home/ubuntu/di/golden/short.pt", + "full_golden_sha256": "59f13484afa07472500d07958212c54203e7bd99b89de3fb936e10af93deb9a6", + "seconds": 3.6 +} \ No newline at end of file diff --git a/examples/neuron/make_tiny_checkpoint.py b/examples/neuron/make_tiny_checkpoint.py new file mode 100644 index 0000000..1b20c54 --- /dev/null +++ b/examples/neuron/make_tiny_checkpoint.py @@ -0,0 +1,465 @@ +"""Generate a tiny, random, deterministic DeepSeek-V4.1-Flash HF checkpoint. + +Same tensor names, dtypes and layout conventions as the released checkpoint +(deepseek-ai/DeepSeek-V4.1-Flash-1017); only the sizes shrink. The result is meant to be loadable by +both DeepInfer's HF loader (deepinfer/hugging_face/*) and the official reference +(inference/convert.py + inference/model.py). + +Conventions (taken from the official inference/ code): + * FP8 weights: F8_E4M3 [out, in] + F8_E8M0 scale [ceil(out/32), ceil(in/32)] (32x32 blocks). + scale = 2**ceil(log2(amax/448)), q = e4m3(clamp(w/scale, +-448)); dequant = q * scale. + * Routed experts: FP4 e2m1, two codes per byte stored as I8 [out, in/2]; element 2j is the LOW nibble, + element 2j+1 the HIGH nibble (convert.py cast_e2m1fn_to_e4m3fn: low = x & 0xF -> first, high -> second); + code table = FP4_TABLE in convert.py (bit 3 = sign, 0..7 = 0,.5,1,1.5,2,3,4,6). + F8_E8M0 scale [out, in/32] (1x32 blocks along K), scale = 2**ceil(log2(amax/6)), RNE to the e2m1 grid. + * Engram table: F8_E4M3 [rows, head_dim] + F8_E8M0 [rows, head_dim/32] (1x32 blocks), rows = + sum of the per-(ngram, head) prime bucket sizes (EngramLayout in DeepInfer and in inference/engram.py). + * hc_* and attn_sink and gate biases are F32; norms, compressor, indexer wk/weights_proj, gate weight, + embed/head, vision, markov/confidence heads are BF16 (as in the real index). + +Random weights: W ~ N(0, 1/in_features) in fp32 -> bf16 (-> quantized), norms = 1, biases/sinks small; the +projections that write into the residual stream get an extra gain (attn wo_b: --attn-out-scale 0.5; shared/routed +expert w2 and engram wkv: --ffn-out-scale 1/16) so that the reference is numerically stable enough to be a golden. +Every tensor uses its own generator seeded from (seed, crc32(name)), so the content of one tensor does not +depend on which other tensors exist. + +Usage: + python make_tiny_checkpoint.py --out ~/di/tiny_ckpt --tokenizer-dir \ + --real-config +""" + +from __future__ import annotations + +import argparse +import copy +import json +import math +import os +import shutil +import zlib +from collections import OrderedDict + +import torch +from safetensors.torch import save_file + +# ---------------------------------------------------------------------------------------------- +# Tiny config. Keys not listed keep the value from the real HF config. +# ---------------------------------------------------------------------------------------------- +TINY_TEXT_OVERRIDES = dict( + hidden_size=1024, + moe_intermediate_size=256, + num_hidden_layers=10, + num_attention_heads=8, + head_dim=512, + qk_rope_head_dim=64, + q_lora_rank=256, + o_lora_rank=256, + o_groups=2, + n_routed_experts=16, + num_experts_per_tok=4, + # 10 backbone layers + 3 nextn layers. Layout mirrors the real one: 2 SWA-only layers, a ratio-2 + # segment, a ratio-1 segment starting at the candidate source, 3 SWA-only MTP layers. + compress_ratios=[0, 0, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0], + kv_source_layer_ids=[2, 4, 6], + index_source_layer_ids=[2, 4, 6, 8], + candidate_source_layer_id=6, + # DeepInfer: candidate_tokens * sizeof(bf16 psum) % 1024 == 0 (index_attn.py:109) -> topk_blocks*8 % 512 == 0, + # so 64 is the smallest legal value with block size 8. + candidate_topk_blocks=64, + index_n_heads=4, + index_head_dim=128, + index_topk=64, + engram_layer_ids=[1, 4], + engram_vocab_size=4096, + engram_n_heads=2, + engram_head_dim=256, + num_nextn_predict_layers=3, + dspark_target_layer_ids=[7, 8, 9], + dspark_n_routed_experts=8, + dspark_num_experts_per_tok=2, +) +# DeepInfer's VisionTransformer hard-codes DEPTH=32 / dim 1024 / 16 heads / mlp 2816 +# (deepinfer/model_executor/models/vision_encoder/vit.py), always builds it, and the loader requires every +# model parameter to be covered (model_loader.load_hf_checkpoint). So the vision tower keeps its real size. +TINY_VISION_OVERRIDES = dict() + +FP8_MAX = 448.0 +FP4_MAX = 6.0 +BLOCK = 32 +FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) + + +# ---------------------------------------------------------------------------------------------- +# Quantization helpers (match inference/kernel.py fast_round_scale: 2**ceil(log2(amax * (1/max)))) +# ---------------------------------------------------------------------------------------------- +def pow2_ceil_scale(amax: torch.Tensor, max_inv: float) -> torch.Tensor: + """2**ceil(log2(amax * max_inv)) computed from the fp32 bit pattern, like kernel.py fast_log2_ceil.""" + v = (amax.float() * torch.tensor(max_inv, dtype=torch.float32)).contiguous() + bits = v.view(torch.int32) + exp = (bits >> 23) & 0xFF + man = bits & ((1 << 23) - 1) + e = exp - 127 + (man != 0).to(torch.int32) + return torch.ldexp(torch.ones_like(v), e) + + +def to_e8m0(scale: torch.Tensor) -> torch.Tensor: + """Power-of-two fp32 -> float8_e8m0fnu (exact).""" + bits = (scale.float().contiguous().view(torch.int32) >> 23) & 0xFF + return bits.to(torch.uint8).view(torch.float8_e8m0fnu) + + +def quant_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """[out, in] -> e4m3 weight + e8m0 scale [ceil(out/32), ceil(in/32)].""" + out_f, in_f = w.shape + bo, bi = math.ceil(out_f / BLOCK), math.ceil(in_f / BLOCK) + wp = torch.zeros(bo * BLOCK, bi * BLOCK, dtype=torch.float32) + wp[:out_f, :in_f] = w.float() + blocks = wp.view(bo, BLOCK, bi, BLOCK) + amax = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-4) + s = pow2_ceil_scale(amax, 1.0 / FP8_MAX) + q = (blocks / s[:, None, :, None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + q = q.view(bo * BLOCK, bi * BLOCK)[:out_f, :in_f].contiguous() + return q, to_e8m0(s) + + +def round_e2m1_codes(v: torch.Tensor) -> torch.Tensor: + """fp32 (already scaled, |v|<=6) -> 4-bit e2m1 codes, round-to-nearest-even.""" + a = v.abs() + code = torch.zeros_like(a, dtype=torch.int32) + # midpoints whose tie goes DOWN to an even code: 0.25(0|1) 1.25(2|3) 2.5(4|5) 5.0(6|7) + for b in (0.25, 1.25, 2.5, 5.0): + code += (a > b).to(torch.int32) + # midpoints whose tie goes UP to an even code: 0.75(1|2) 1.75(3|4) 3.5(5|6) + for b in (0.75, 1.75, 3.5): + code += (a >= b).to(torch.int32) + sign = (v < 0) | ((v == 0) & torch.signbit(v)) + return code | (sign.to(torch.int32) << 3) + + +def quant_fp4_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """[out, in] -> packed int8 [out, in/2] (low nibble = even k) + e8m0 scale [out, in/32].""" + out_f, in_f = w.shape + assert in_f % BLOCK == 0 + blocks = w.float().view(out_f, in_f // BLOCK, BLOCK) + amax = blocks.abs().amax(-1).clamp_min(FP4_MAX * 2.0**-126) + s = pow2_ceil_scale(amax, 1.0 / FP4_MAX) + codes = round_e2m1_codes((blocks / s[..., None]).clamp(-FP4_MAX, FP4_MAX)).view(out_f, in_f) + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).to(torch.uint8).view(torch.int8).contiguous() + return packed, to_e8m0(s) + + +def dequant_fp4_row(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + x = packed.view(torch.uint8) + lo, hi = (x & 0xF).long(), (x >> 4).long() + vals = torch.stack([FP4_TABLE[lo], FP4_TABLE[hi]], dim=-1).flatten(-2) + return vals * scale.float().repeat_interleave(BLOCK, dim=-1) + + +def quant_fp8_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Engram table: [rows, d] -> e4m3 + e8m0 [rows, d/32] (1x32 blocks).""" + rows, d = w.shape + blocks = w.float().view(rows, d // BLOCK, BLOCK) + amax = blocks.abs().amax(-1).clamp_min(1e-4) + s = pow2_ceil_scale(amax, 1.0 / FP8_MAX) + q = (blocks / s[..., None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn).view(rows, d) + return q.contiguous(), to_e8m0(s) + + +# ---------------------------------------------------------------------------------------------- +# Engram layout (identical in DeepInfer engram_metadata.EngramLayout.from_cfg and inference/engram.py) +# ---------------------------------------------------------------------------------------------- +def engram_primes(layer_ids, max_ngram_size, n_heads, vocab_size): + from sympy import isprime, nextprime + + # reference (inference/engram.py find_next_prime) + seen, ref = set(), [] + for _ in layer_ids: + per_ngram = [] + for _ in range(max_ngram_size - 1): + sizes, current = [], vocab_size - 1 + for _ in range(n_heads): + c = current + 1 + while not isprime(c) or c in seen: + c += 1 + current = c + seen.add(c) + sizes.append(c) + per_ngram.append(sizes) + ref.append(per_ngram) + # DeepInfer (engram_metadata.EngramLayout.from_cfg) + seen, di = set(), [] + for _ in layer_ids: + per_ngram = [] + for _ in range(max_ngram_size - 1): + start, sizes = vocab_size - 1, [] + for _ in range(n_heads): + p = int(nextprime(start)) + while p in seen: + p = int(nextprime(p)) + seen.add(p) + sizes.append(p) + start = p + per_ngram.append(sizes) + di.append(per_ngram) + assert ref == di, (ref, di) + return ref + + +# ---------------------------------------------------------------------------------------------- +# Tensor factory +# ---------------------------------------------------------------------------------------------- +class Factory: + def __init__(self, seed: int, norm_noise: float): + self.seed = seed + self.norm_noise = norm_noise + self.tensors: "OrderedDict[str, torch.Tensor]" = OrderedDict() + + def gen(self, name: str) -> torch.Generator: + g = torch.Generator() + g.manual_seed((self.seed * 1_000_003 + zlib.crc32(name.encode())) % (2**63)) + return g + + def randn(self, name, shape, std): + return torch.randn(shape, generator=self.gen(name), dtype=torch.float32) * std + + def add(self, name, t): + assert name not in self.tensors, name + self.tensors[name] = t.contiguous() + + def bf16(self, name, out_f, in_f, std=None): + std = (1.0 / in_f) ** 0.5 if std is None else std + self.add(name, self.randn(name, (out_f, in_f), std).bfloat16()) + + def norm(self, name, dim): + w = torch.ones(dim) + if self.norm_noise: + w = w + self.randn(name, (dim,), self.norm_noise) + self.add(name, w.bfloat16()) + + def bias(self, name, dim, std=0.02, dtype=torch.bfloat16): + self.add(name, self.randn(name, (dim,), std).to(dtype)) + + def fp8(self, prefix, out_f, in_f, gain=1.0): + w = self.randn(prefix + ".weight", (out_f, in_f), gain * (1.0 / in_f) ** 0.5).bfloat16() + q, s = quant_fp8_block(w) + self.add(prefix + ".weight", q) + self.add(prefix + ".scale", s) + + def fp4(self, prefix, out_f, in_f, gain=1.0): + w = self.randn(prefix + ".weight", (out_f, in_f), gain * (1.0 / in_f) ** 0.5).bfloat16() + q, s = quant_fp4_row(w) + self.add(prefix + ".weight", q) + self.add(prefix + ".scale", s) + + def f32(self, name, t): + self.add(name, t.float()) + + +def build_tensors(cfg: dict, seed: int, norm_noise: float, attn_out_scale: float, ffn_out_scale: float) -> "OrderedDict[str, torch.Tensor]": + t = cfg["text_config"] + v = cfg["vision_config"] + f = Factory(seed, norm_noise) + dim, H, hd = t["hidden_size"], t["num_attention_heads"], t["head_dim"] + qr, olr, og = t["q_lora_rank"], t["o_lora_rank"], t["o_groups"] + inter, hc = t["moe_intermediate_size"], t["hc_mult"] + n_layer, nextn = t["num_hidden_layers"], t["num_nextn_predict_layers"] + ratios = t["compress_ratios"] + kv_src, idx_src = set(t["kv_source_layer_ids"]), set(t["index_source_layer_ids"]) + ih, ihd = t["index_n_heads"], t["index_head_dim"] + vocab = t["vocab_size"] + mix_hc = (2 + hc) * hc + + # ---- vision tower + aligner (BF16), real shapes ---- + vd, vinter, ps = v["hidden_size"], v["intermediate_size"], v["patch_size"] + f.bf16("vision.patch_embed.proj.weight", vd, 3 * ps * ps) + f.bias("vision.patch_embed.proj.bias", vd) + for i in range(v["num_hidden_layers"]): + p = f"vision.blocks.{i}." + f.norm(p + "norm1.weight", vd) + f.bf16(p + "attn.wqkv.weight", 3 * vd, vd) + f.bias(p + "attn.wqkv.bias", 3 * vd) + f.bf16(p + "attn.wo.weight", vd, vd) + f.bias(p + "attn.wo.bias", vd) + f.norm(p + "norm2.weight", vd) + f.bf16(p + "mlp.w1.weight", 2 * vinter, vd) + f.bf16(p + "mlp.w2.weight", vd, vinter) + f.norm("vision.norm.weight", vd) + ds = v["downsample_ratio"] + f.bf16("aligner.w1.weight", dim, vd * ds * ds) + f.bias("aligner.w1.bias", dim) + f.bf16("aligner.w2.weight", dim, dim) + f.bias("aligner.w2.bias", dim) + + # ---- global ---- + f.bf16("embed.weight", vocab, dim, std=1.0) + for n in ("image_start", "image_end", "image_newline"): + f.add(n, f.randn(n, (dim,), 1.0).bfloat16()) + f.norm("norm.weight", dim) + f.bf16("head.weight", vocab, dim) + + def hc_params(p): + for which in ("attn", "ffn"): + f.f32(p + f"hc_{which}_fn", f.randn(p + f"hc_{which}_fn", (mix_hc, hc * dim), (1.0 / (hc * dim)) ** 0.5)) + f.f32(p + f"hc_{which}_scale", torch.full((3,), 0.3) + f.randn(p + f"hc_{which}_scale", (3,), 0.05)) + base = f.randn(p + f"hc_{which}_base", (mix_hc,), 0.5) + base[2 * hc :] += 2.0 * torch.eye(hc).flatten() # comb biased towards identity + f.f32(p + f"hc_{which}_base", base) + + def attn_params(p, layer_id): + f.f32(p + "attn.attn_sink", f.randn(p + "attn.attn_sink", (H,), 0.1)) + f.fp8(p + "attn.wq_a", qr, dim) + f.norm(p + "attn.q_norm.weight", qr) + f.fp8(p + "attn.wq_b", H * hd, qr) + f.fp8(p + "attn.wkv", hd, dim) + f.norm(p + "attn.kv_norm.weight", hd) + f.fp8(p + "attn.wo_a", og * olr, H * hd // og) + f.fp8(p + "attn.wo_b", dim, og * olr, gain=attn_out_scale) + backbone = layer_id is not None + if backbone and layer_id in kv_src: + f.bf16(p + "attn.compressor.wkv.weight", hd, dim) + if ratios[layer_id] > 1: + f.bf16(p + "attn.compressor.wgate.weight", hd, dim) + f.norm(p + "attn.compressor.norm.weight", hd) + if backbone and layer_id in idx_src: + f.fp8(p + "attn.indexer.wq_b", ih * ihd, qr) + f.bf16(p + "attn.indexer.weights_proj.weight", ih, dim) + if layer_id in kv_src: # only index-K owners carry wk / k_norm + f.bf16(p + "attn.indexer.wk.weight", ihd, hd) + f.norm(p + "attn.indexer.k_norm.weight", ihd) + + def ffn_params(p, n_exp): + f.norm(p + "attn_norm.weight", dim) + f.norm(p + "ffn_norm.weight", dim) + f.bf16(p + "ffn.gate.weight", n_exp, dim) + f.bias(p + "ffn.gate.bias", n_exp, std=0.02, dtype=torch.float32) + f.bias(p + "ffn.gate.bias_vl", n_exp, std=0.02, dtype=torch.float32) + f.fp8(p + "ffn.shared_experts.w1", inter, dim) + f.fp8(p + "ffn.shared_experts.w3", inter, dim) + f.fp8(p + "ffn.shared_experts.w2", dim, inter, gain=ffn_out_scale) + for e in range(n_exp): + q = p + f"ffn.experts.{e}." + f.fp4(q + "w1", inter, dim) + f.fp4(q + "w3", inter, dim) + f.fp4(q + "w2", dim, inter, gain=ffn_out_scale) + + primes = engram_primes(t["engram_layer_ids"], t["engram_max_ngram_size"], t["engram_n_heads"], t["engram_vocab_size"]) + n_hash_cols = (t["engram_max_ngram_size"] - 1) * t["engram_n_heads"] + ehd = t["engram_head_dim"] + + for L in range(n_layer): + p = f"layers.{L}." + hc_params(p) + attn_params(p, L) + ffn_params(p, t["n_routed_experts"]) + if L in t["engram_layer_ids"]: + li = t["engram_layer_ids"].index(L) + rows = sum(sum(x) for x in primes[li]) + assert rows == t["engram_num_embeddings"][li] + w = f.randn(p + "engram.embed.weight", (rows, ehd), 1.0).bfloat16() + q, s = quant_fp8_row(w) + f.add(p + "engram.embed.weight", q) + f.add(p + "engram.embed.scale", s) + f.add(p + "engram.q_weight", (1.0 + f.randn(p + "engram.q_weight", (hc, dim), 0.1)).bfloat16()) + f.add(p + "engram.k_weight", (1.0 + f.randn(p + "engram.k_weight", (hc, dim), 0.1)).bfloat16()) + # key half is RMS-normalised inside the gate, so the gain only shrinks the value written to the stream + f.fp8(p + "engram.wkv", dim * (hc + 1), n_hash_cols * ehd, gain=ffn_out_scale) + + for i in range(nextn): + p = f"mtp.{i}." + hc_params(p) + attn_params(p, None) + ffn_params(p, t["dspark_n_routed_experts"]) + if i == 0: + f.fp8(p + "main_proj", dim, dim * len(t["dspark_target_layer_ids"])) + f.norm(p + "main_norm.weight", dim) + if i == nextn - 1: + r = t["dspark_markov_rank"] + f.norm(p + "norm.weight", dim) + f.bf16(p + "markov_head.embed.weight", vocab, r, std=1.0) + f.bf16(p + "markov_head.head.weight", vocab, r) + f.bf16(p + "confidence_head.proj.weight", 1, dim + r) + return f.tensors + + +def make_config(real: dict) -> dict: + cfg = copy.deepcopy(real) + t = cfg["text_config"] + for k, val in TINY_TEXT_OVERRIDES.items(): + assert k in t, k + t[k] = val + for k, val in TINY_VISION_OVERRIDES.items(): + cfg["vision_config"][k] = val + assert len(t["compress_ratios"]) == t["num_hidden_layers"] + t["num_nextn_predict_layers"] + primes = engram_primes(t["engram_layer_ids"], t["engram_max_ngram_size"], t["engram_n_heads"], t["engram_vocab_size"]) + t["engram_num_embeddings"] = [sum(sum(x) for x in per_layer) for per_layer in primes] + return cfg + + +def shard_of(name: str) -> int: + if name.startswith(("vision.", "aligner.")): + return 0 + if name.startswith(("embed.", "head.", "norm.", "image_")): + return 1 + if name.startswith("layers."): + return 2 if int(name.split(".")[1]) < 5 else 3 + return 4 # mtp.* + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--out", required=True) + ap.add_argument("--tokenizer-dir", required=True) + ap.add_argument("--real-config", required=True, help="the released HF config.json") + ap.add_argument("--seed", type=int, default=20261020) + ap.add_argument("--norm-noise", type=float, default=0.0, help="std of N(1, s) norm weights (0 = all ones)") + # Gains on the projections that write into the residual stream. With plain N(0, 1/in) everywhere (1.0, 1.0) + # the reference is chaotic: FP8 activation quantization + top-k routing turn sub-ulp differences (e.g. another + # CPU thread count) into 20-50% logit differences. Small residual updates mimic trained models; the MoE/engram + # branch (discrete routing) is damped most, attention less so that the logits still depend on the context. + # Measured (threads 16 vs 7, teacher forced): logits rel err <= 0.04 while short/medium/long contexts ending in + # the same token differ by 0.24-0.49. + ap.add_argument("--attn-out-scale", type=float, default=0.5, help="gain of attn.wo_b") + ap.add_argument("--ffn-out-scale", type=float, default=1.0 / 16, help="gain of shared/routed expert w2 and engram wkv") + args = ap.parse_args() + torch.set_num_threads(min(16, os.cpu_count() or 1)) + + real = json.load(open(args.real_config)) + cfg = make_config(real) + os.makedirs(args.out, exist_ok=True) + tensors = build_tensors(cfg, args.seed, args.norm_noise, args.attn_out_scale, args.ffn_out_scale) + + # self-check of the FP4 packing against convert.py's decoding convention + for name, ten in tensors.items(): + if name.endswith("experts.0.w1.weight") and name.startswith("layers.0."): + deq = dequant_fp4_row(ten, tensors[name.replace("weight", "scale")]) + ref = torch.randn((ten.shape[0], ten.shape[1] * 2), generator=Factory(args.seed, 0).gen(name)) * (1.0 / (ten.shape[1] * 2)) ** 0.5 + rel = (deq - ref.bfloat16().float()).norm() / ref.norm() + print(f"fp4 self-check {name}: rel err {rel:.4f}") + assert rel < 0.2 + + n_shards = 5 + shards = [OrderedDict() for _ in range(n_shards)] + for name, ten in tensors.items(): + shards[shard_of(name)][name] = ten + weight_map, total = {}, 0 + for i, sd in enumerate(shards): + fname = f"model-{i + 1:05d}-of-{n_shards:05d}.safetensors" + save_file(sd, os.path.join(args.out, fname), metadata={"format": "pt"}) + for name, ten in sd.items(): + weight_map[name] = fname + total += ten.numel() * ten.element_size() + with open(os.path.join(args.out, "model.safetensors.index.json"), "w") as fh: + json.dump({"metadata": {"total_size": total}, "weight_map": dict(sorted(weight_map.items()))}, fh, indent=2) + with open(os.path.join(args.out, "config.json"), "w") as fh: + json.dump(cfg, fh, indent=2) + with open(os.path.join(args.out, "tiny_checkpoint.json"), "w") as fh: + json.dump({"generator": "examples/neuron/make_tiny_checkpoint.py", "seed": args.seed, "norm_noise": args.norm_noise, + "attn_out_scale": args.attn_out_scale, "ffn_out_scale": args.ffn_out_scale, "n_tensors": len(tensors), "total_bytes": total}, fh, indent=2) + for fn in ("tokenizer.json", "tokenizer_config.json"): + shutil.copyfile(os.path.join(args.tokenizer_dir, fn), os.path.join(args.out, fn)) + print(f"wrote {len(tensors)} tensors, {total / 2**30:.2f} GiB to {args.out}") + print("engram_num_embeddings", cfg["text_config"]["engram_num_embeddings"]) + + +if __name__ == "__main__": + main() diff --git a/examples/neuron/ref_kernels_torch.py b/examples/neuron/ref_kernels_torch.py new file mode 100644 index 0000000..0297b5a --- /dev/null +++ b/examples/neuron/ref_kernels_torch.py @@ -0,0 +1,246 @@ +"""Pure-torch drop-in for the official DeepSeek-V4.1 reference `inference/kernel.py` (TileLang GPU kernels). + +Mirrors the public API of kernel.py (act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn, +hc_split_sinkhorn) with the same numerics, so the unmodified reference `model.py` can run on CPU: + + import ref_kernels_torch, sys + sys.modules["kernel"] = ref_kernels_torch # before `import model` + +Numerics reproduced from kernel.py (line numbers refer to the downloaded kernel.py): + * fast_round_scale (22-37): scale = 2**ceil(log2(amax * fp32(1/max))) taken from the fp32 bit pattern. + * act_quant (41-124): per (row, block) amax clamped to >= 1e-4; e4m3 RNE cast after clamp(+-448); + inplace=True writes back dtype(fp32(e4m3(x/s)) * s) (a quant->dequant round trip). + * fp4_act_quant (128-204): e8m0 scales: amax >= 6*2**-126, pow2-ceil(amax/6); e4m3 scales: amax >= 6*2**-9, + s = fp32(e4m3(amax / 6)); values RNE to the e2m1 grid after clamp(+-6); inplace => dtype(fp32(e2m1)*s). + * fp8_gemm / fp4_gemm (208-307, 478-591): per 32-wide K block, fp32 partial GEMM, then + acc += partial * scale_a[m, k] * scale_b[n(/32), k] in K-block order; output in torch default dtype. + FP4 weights: float4_e2m1fn_x2 [N, K/2], element 2j = low nibble, 2j+1 = high nibble. + * sparse_attn (311-403): FlashAttention-style online softmax over blocks of 64 gathered indices + (-1 = masked, gathered as zeros with score -inf), running max initialised to -1e30, the probabilities + are cast to bf16 before the P@V GEMM (acc_s_cast), sink added as exp(sink - max) to the denominator. + * hc_split_sinkhorn (407-474): sigmoid pre (+eps), 2*sigmoid post, softmax(+eps) then column + normalisation, then (iters-1) x (row, column) normalisation, each with +eps in the denominator. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +FP8_MAX = 448.0 +FP4_MAX = 6.0 +FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) +_GEMM_BLOCK_K = 32 # kernel.py: block_K = group_size (32 for this model) / fp4 sub_K = 32 +_ATTN_BLOCK = 64 # kernel.py sparse_attn_kernel: block = 64 + + +def _fast_round_scale(amax: torch.Tensor, max_inv: float) -> torch.Tensor: + """kernel.py fast_round_scale: 2**ceil(log2(amax * max_inv)), via the IEEE-754 bits.""" + v = (amax.float() * torch.tensor(max_inv, dtype=torch.float32)).contiguous() + bits = v.view(torch.int32) + exp = (bits >> 23) & 0xFF + man = bits & ((1 << 23) - 1) + e = exp - 127 + (man != 0).to(torch.int32) + # fast_pow2: (e + 127) << 23 reinterpreted as fp32 + return ((e + 127) << 23).to(torch.int32).view(torch.float32) + + +def _e2m1_round(v: torch.Tensor) -> torch.Tensor: + """fp32 in [-6, 6] -> nearest e2m1 value (round half to even), returned as fp32.""" + a = v.abs() + code = torch.zeros_like(a, dtype=torch.int64) + for b in (0.25, 1.25, 2.5, 5.0): # ties go down (to the even code) + code += a > b + for b in (0.75, 1.75, 3.5): # ties go up (to the even code) + code += a >= b + mag = FP4_TABLE.to(v.device)[code] + return torch.where(v < 0, -mag, mag) + + +def _e2m1_codes(v: torch.Tensor) -> torch.Tensor: + a = v.abs() + code = torch.zeros_like(a, dtype=torch.int32) + for b in (0.25, 1.25, 2.5, 5.0): + code += (a > b).to(torch.int32) + for b in (0.75, 1.75, 3.5): + code += (a >= b).to(torch.int32) + return code | ((v < 0).to(torch.int32) << 3) + + +def _unpack_fp4(b: torch.Tensor) -> torch.Tensor: + """float4_e2m1fn_x2 / int8 [..., K/2] -> fp32 [..., K]; low nibble first.""" + x = b.contiguous().view(torch.uint8) + table = FP4_TABLE.to(x.device) + lo, hi = (x & 0x0F).long(), (x >> 4).long() + return torch.stack([table[lo], table[hi]], dim=-1).flatten(-2) + + +def _to_scale_dtype(s: torch.Tensor, scale_dtype: torch.dtype) -> torch.Tensor: + if scale_dtype == torch.float8_e8m0fnu: + # s is a power of two here: its exponent byte is the e8m0 encoding + return ((s.contiguous().view(torch.int32) >> 23) & 0xFF).to(torch.uint8).view(torch.float8_e8m0fnu) + return s.to(scale_dtype) + + +def act_quant( + x: torch.Tensor, + block_size: int = 128, + scale_fmt: Optional[str] = None, + scale_dtype: torch.dtype = torch.float32, + inplace: bool = False, +): + """Block-wise FP8 quantization; inplace=True does the fused quant+dequant back to x.dtype.""" + N = x.size(-1) + assert N % block_size == 0 + z = x.contiguous() + xf = z.float().unflatten(-1, (N // block_size, block_size)) + amax = xf.abs().amax(-1).clamp_min(1e-4) + if scale_fmt is not None: + s = _fast_round_scale(amax, 1.0 / FP8_MAX) + else: + s = amax * torch.tensor(1.0 / FP8_MAX, dtype=torch.float32) + q = (xf / s.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + if inplace: + y = (q.float() * s.unsqueeze(-1)).to(x.dtype).flatten(-2) + x.copy_(y) + return x + return q.flatten(-2), _to_scale_dtype(s, scale_dtype) + + +def fp4_act_quant( + x: torch.Tensor, + block_size: int = 32, + inplace: bool = False, + scale_dtype: torch.dtype = torch.float8_e8m0fnu, +): + """FP4 with E8M0 scales (indexer) or E4M3 scales (compressed KV); inplace writes the dequantized values.""" + assert scale_dtype in (torch.float8_e8m0fnu, torch.float8_e4m3fn) + N = x.size(-1) + assert N % block_size == 0 + z = x.contiguous() + xf = z.float().unflatten(-1, (N // block_size, block_size)) + amax = xf.abs().amax(-1) + if scale_dtype == torch.float8_e4m3fn: + amax = amax.clamp_min(FP4_MAX * 2.0**-9) + s = (amax / FP4_MAX).to(torch.float8_e4m3fn).float() + else: + amax = amax.clamp_min(FP4_MAX * 2.0**-126) + s = _fast_round_scale(amax, 1.0 / FP4_MAX) + v = (xf / s.unsqueeze(-1)).clamp(-FP4_MAX, FP4_MAX) + if inplace: + y = (_e2m1_round(v) * s.unsqueeze(-1)).to(x.dtype).flatten(-2) + x.copy_(y) + return x + codes = _e2m1_codes(v).flatten(-2) + packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).to(torch.uint8).view(torch.float4_e2m1fn_x2) + s_out = _to_scale_dtype(s, scale_dtype) if scale_dtype == torch.float8_e8m0fnu else s.to(torch.float8_e4m3fn) + return packed, s_out + + +def _blocked_gemm(A: torch.Tensor, sa: torch.Tensor, B: torch.Tensor, sb_rows: torch.Tensor, act_group: int) -> torch.Tensor: + """A [M,K] fp32 (exact fp8 values), sa [M, K/act_group]; B [N,K] fp32 (exact), sb_rows [N, K/32]. + acc += (A_k @ B_k^T) * sa[:, k // n_sub] * sb[:, k] for each 32-wide K block, in order.""" + M, K = A.shape + N = B.shape[0] + n_sub = act_group // _GEMM_BLOCK_K + acc = torch.zeros(M, N, dtype=torch.float32, device=A.device) + for k in range(K // _GEMM_BLOCK_K): + sl = slice(k * _GEMM_BLOCK_K, (k + 1) * _GEMM_BLOCK_K) + part = A[:, sl] @ B[:, sl].T + acc += part * sa[:, k // n_sub].unsqueeze(1) * sb_rows[:, k].unsqueeze(0) + return acc + + +def fp8_gemm( + a: torch.Tensor, + a_s: torch.Tensor, + b: torch.Tensor, + b_s: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, + block_size: int = 128, +) -> torch.Tensor: + """C[M,N] = A[M,K] @ B[N,K]^T with per-block FP8 scaling (weight scales per (N block, K block)).""" + assert a.is_contiguous() and b.is_contiguous() + assert a_s.is_contiguous() and b_s.is_contiguous() + assert block_size == _GEMM_BLOCK_K, "this model only uses 32x32 blocks" + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + assert K % block_size == 0 + assert a_s.numel() == M * (K // block_size) + assert b_s.shape == ((N + block_size - 1) // block_size, K // block_size) + sb_rows = b_s.float().repeat_interleave(block_size, dim=0)[:N] + acc = _blocked_gemm(a.view(M, K).float(), a_s.view(M, -1).float(), b.float(), sb_rows, block_size) + return acc.to(torch.get_default_dtype()).view(*a.size()[:-1], N) + + +def fp4_gemm( + a: torch.Tensor, + a_s: torch.Tensor, + b: torch.Tensor, + b_s: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, + act_block_size: int = 128, +) -> torch.Tensor: + """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T; B is [N, K/2] float4_e2m1fn_x2 with per-32 E8M0 scales.""" + assert a.is_contiguous() and b.is_contiguous() + assert a_s.is_contiguous() and b_s.is_contiguous() + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + assert act_block_size in (32, 128) + assert K % act_block_size == 0 + assert a_s.numel() == M * (K // act_block_size) + assert b_s.shape == (N, K // 32) + acc = _blocked_gemm(a.view(M, K).float(), a_s.view(M, -1).float(), _unpack_fp4(b), b_s.float(), act_block_size) + return acc.to(torch.get_default_dtype()).view(*a.size()[:-1], N) + + +def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor: + """q [b,m,h,d] bf16, kv [b,n,d] bf16, attn_sink [h] fp32, topk_idxs [b,m,topk] int32 (-1 = none).""" + b, m, h, d = q.shape + topk = topk_idxs.size(-1) + qf = q.float() + sink = attn_sink.float() + scores_max = torch.full((b, m, h), -1e30, dtype=torch.float32, device=q.device) + sum_exp = torch.zeros((b, m, h), dtype=torch.float32, device=q.device) + acc_o = torch.zeros((b, m, h, d), dtype=torch.float32, device=q.device) + batch_ix = torch.arange(b, device=q.device).view(b, 1, 1) + n_blocks = (topk + _ATTN_BLOCK - 1) // _ATTN_BLOCK + for t in range(n_blocks): + idx = topk_idxs[..., t * _ATTN_BLOCK : (t + 1) * _ATTN_BLOCK].long() + if idx.size(-1) < _ATTN_BLOCK: # positions past topk read as -1 + idx = torch.cat([idx, idx.new_full((b, m, _ATTN_BLOCK - idx.size(-1)), -1)], dim=-1) + valid = idx != -1 + kv_blk = kv[batch_ix, idx.clamp_min(0)] # [b, m, 64, d] + kv_blk = torch.where(valid.unsqueeze(-1), kv_blk, torch.zeros((), dtype=kv.dtype)).float() + s = torch.einsum("bmhd,bmkd->bmhk", qf, kv_blk) + s = torch.where(valid.unsqueeze(2), s, torch.tensor(-torch.inf)) + s = s * softmax_scale + prev = scores_max + scores_max = torch.maximum(prev, s.amax(-1)) + scale = torch.exp(prev - scores_max) + p = torch.exp(s - scores_max.unsqueeze(-1)) + sum_exp = sum_exp * scale + p.sum(-1) + acc_o = acc_o * scale.unsqueeze(-1) + torch.einsum("bmhk,bmkd->bmhd", p.to(torch.bfloat16).float(), kv_blk) + sum_exp = sum_exp + torch.exp(sink.view(1, 1, h) - scores_max) + return (acc_o / sum_exp.unsqueeze(-1)).to(q.dtype) + + +def hc_split_sinkhorn( + mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6 +): + hc = hc_mult + m = mixes.float() + sc, base = hc_scale.float(), hc_base.float() + pre = torch.sigmoid(m[..., :hc] * sc[0] + base[:hc]) + eps + post = 2 * torch.sigmoid(m[..., hc : 2 * hc] * sc[1] + base[hc : 2 * hc]) + comb = m[..., 2 * hc :].unflatten(-1, (hc, hc)) * sc[2] + base[2 * hc :].view(hc, hc) + comb = torch.exp(comb - comb.amax(-1, keepdim=True)) + comb = comb / comb.sum(-1, keepdim=True) + eps + comb = comb / (comb.sum(-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(-1, keepdim=True) + eps) + comb = comb / (comb.sum(-2, keepdim=True) + eps) + return pre, post, comb diff --git a/examples/neuron/reference_golden.py b/examples/neuron/reference_golden.py new file mode 100644 index 0000000..bab419d --- /dev/null +++ b/examples/neuron/reference_golden.py @@ -0,0 +1,341 @@ +"""Ground truth for the Neuron port: run the OFFICIAL DeepSeek-V4.1 reference model on the host CPU. + +The downloaded reference (`inference/model.py`, `engram.py`, `convert.py`, ...) is used unmodified; only its +TileLang `kernel` module is swapped for the pure-torch `ref_kernels_torch` (same numerics, see that file). +DSpark/MTP (`dspark_block_size=0`) and vision (`vision_n_layers=0`) are disabled: DeepInfer runs text-only with +DEEPINFER_USE_NEXTN_DECODE=0, and neither affects the backbone logits. + +Steps: + 1. convert the HF checkpoint with the official convert.py (MP=1, --expert-dtype fp4) unless already done; + 2. build reference ModelArgs from the HF config.json (same mapping as the official inference/config.json); + 3. per prompt (batch 1, fresh model): prefill the whole prompt, then N greedy decode steps (temperature 0, + no EOS stop), saving logits (fp32), greedy tokens, per-layer block outputs, indexer top-k and engram hashes. + +Usage (on the box, host CPU only; TORCH_DEVICE_BACKEND_AUTOLOAD=0 keeps torch_neuronx/NRT out of the process): + TORCH_DEVICE_BACKEND_AUTOLOAD=0 ~/di/venv/bin/python -I examples/neuron/reference_golden.py \ + --hf-ckpt ~/di/tiny_ckpt --ref-dir ~/di/ref_dsv41/inference --converted ~/di/tiny_ckpt_ref_mp1 \ + --out ~/di/golden --summary-dir examples/neuron/golden +""" + +from __future__ import annotations + +import argparse +import dataclasses +import functools +import hashlib +import json +import os +import subprocess +import sys +import time + +HERE = os.path.dirname(os.path.abspath(__file__)) + +PROMPTS = { + # short: no window wrap, no compressed top-k selection + "short": "The quick brown fox jumps over the lazy dog.", + # medium (~200 tokens): wraps the 128-token sliding window and makes both indexer top-k levels select + "medium": ( + "In the early morning the harbor was quiet. Fishing boats rocked gently against the wooden piers, " + "their ropes creaking as the tide came in. An old man sat on an overturned crate, mending a net with " + "slow, careful hands. He had worked these waters for more than fifty years, and he could read the " + "weather in the color of the sky and the shape of the clouds. Today the sky was pale gray, and a thin " + "mist hung over the water. Somewhere a gull cried, and the sound echoed off the stone walls of the " + "warehouses. A young girl ran down the pier, her boots clattering on the boards. She stopped beside " + "the old man and asked him whether the boats would go out today. He looked at the horizon for a long " + "moment before answering. The sea, he said, decides that for us. We only listen." + ), +} +# long (~700 tokens): more than candidate_topk_blocks * candidate_block_size = 512 compressed positions in the +# ratio-1 layers, so the candidate block pre-filter is active as well +PROMPTS["long"] = " ".join( + [ + PROMPTS["medium"], + "The history of computing is a story of abstraction. Early machines were programmed by rewiring " + "panels and flipping switches, and every program was tied to the physical layout of the hardware. " + "Assembly languages gave names to instructions, and compilers later translated whole expressions " + "into sequences of those instructions. Operating systems hid the details of disks and terminals " + "behind files and streams, while virtual memory let each process pretend it owned the entire machine.", + "Networks added another layer. Packets travel through routers that know nothing about the " + "applications sending them, and protocols stacked on top of one another turn unreliable links into " + "dependable conversations. A web page that appears in a fraction of a second may involve dozens of " + "servers, caches and databases spread across several continents, each one following rules that were " + "agreed upon decades ago by people who could not have imagined how the system would grow.", + "Mathematics offers a similar lesson. A proof about prime numbers does not care whether the numbers " + "are written in decimal or binary, on paper or in silicon. The same theorem holds in every notation, " + "and that independence is exactly what makes it useful. Engineers borrow this idea when they define " + "interfaces: as long as both sides respect the contract, either side can be rewritten, optimized or " + "replaced without the other noticing.", + "Cooking, perhaps surprisingly, works the same way. A recipe names its ingredients and steps but " + "leaves the choice of pan, stove and kitchen to the cook. Bread rises because yeast consumes sugar " + "and releases gas, whether the dough is kneaded by hand in a village bakery or by a machine in a " + "factory. Understanding the underlying chemistry lets a baker adapt when the flour is different, the " + "air is humid, or the oven runs hot.", + "Music, too, depends on structure. A melody can be transposed into another key, played on a different " + "instrument, or slowed down to half its tempo, and listeners will still recognize it. Composers rely on " + "patterns of tension and release, repetition and variation, to guide the ear through a piece. The notes " + "on the page are only instructions; the performance brings them to life in a particular room on a " + "particular evening, never exactly the same twice.", + "Gardens change with every season. In spring the first green shoots push through the cold soil, and by " + "summer the beds are crowded with tomatoes, beans and sunflowers that lean toward the light. Autumn " + "brings the harvest and the slow work of clearing dead stems, while winter rests under frost and snow. " + "A patient gardener learns that each year is different and that the soil remembers what was done to it.", + "At the end of the day the harbor grew quiet again. The boats returned one by one, low in the water " + "with their catch, and the old man folded his net and walked slowly home along the shore.", + ] +) + + +def hf_to_model_args(hf: dict) -> dict: + """HF config.json -> reference ModelArgs kwargs (field mapping as in the official inference/config.json).""" + t = hf["text_config"] + v = hf["vision_config"] + rs = t["rope_scaling"] + return dict( + vocab_size=t["vocab_size"], + dim=t["hidden_size"], + moe_inter_dim=t["moe_intermediate_size"], + n_layers=t["num_hidden_layers"], + n_mtp_layers=t["num_nextn_predict_layers"], + dspark_block_size=t["dspark_block_size"], + dspark_noise_token_id=t["dspark_noise_token_id"], + dspark_target_layer_ids=tuple(t["dspark_target_layer_ids"]), + dspark_markov_rank=t["dspark_markov_rank"], + dspark_n_routed_experts=t["dspark_n_routed_experts"], + dspark_n_activated_experts=t["dspark_num_experts_per_tok"], + n_heads=t["num_attention_heads"], + n_routed_experts=t["n_routed_experts"], + n_shared_experts=t["n_shared_experts"], + n_activated_experts=t["num_experts_per_tok"], + score_func=t["scoring_func"], + norm_topk_prob=t["norm_topk_prob"], + route_scale=t["routed_scaling_factor"], + swiglu_limit=t["swiglu_limit"], + q_lora_rank=t["q_lora_rank"], + head_dim=t["head_dim"], + rope_head_dim=t["qk_rope_head_dim"], + norm_eps=t["rms_norm_eps"], + o_groups=t["o_groups"], + o_lora_rank=t["o_lora_rank"], + window_size=t["sliding_window"], + kv_source_layers=tuple(t["kv_source_layer_ids"]), + index_source_layers=tuple(t["index_source_layer_ids"]), + original_seq_len=rs["original_max_position_embeddings"], + rope_theta=t["rope_theta"], + rope_factor=rs["factor"], + beta_fast=rs["beta_fast"], + beta_slow=rs["beta_slow"], + index_n_heads=t["index_n_heads"], + index_head_dim=t["index_head_dim"], + index_topk=t["index_topk"], + candidate_source_layer=t["candidate_source_layer_id"], + candidate_topk_blocks=t["candidate_topk_blocks"], + candidate_block_size=t["candidate_block_size"], + hc_mult=t["hc_mult"], + hc_sinkhorn_iters=t["hc_sinkhorn_iters"], + hc_eps=t["hc_eps"], + engram_layer_ids=tuple(t["engram_layer_ids"]), + engram_vocab_size=t["engram_vocab_size"], + engram_num_embeddings=tuple(t["engram_num_embeddings"]), + engram_max_ngram_size=t["engram_max_ngram_size"], + engram_pad_id=t["engram_pad_token_id"], + engram_compressed_vocab_size=t["engram_compressed_vocab_size"], + engram_n_heads=t["engram_n_heads"], + engram_head_dim=t["engram_head_dim"], + dtype="fp8" if hf["quantization_config"]["quant_method"] == "fp8" else "bf16", + expert_dtype=hf["quantization_config"]["expert_dtype"], + compress_rope_theta=t["compress_rope_theta"], + compress_ratios=tuple(t["compress_ratios"]), + vision_n_layers=v["num_hidden_layers"], + vision_dim=v["hidden_size"], + vision_n_heads=v["num_attention_heads"], + vision_inter_dim=v["intermediate_size"], + vision_patch_size=v["patch_size"], + vision_rope_theta=v["rope_theta"], + vision_downsample_ratio=v["downsample_ratio"], + vision_max_n_token=v["max_image_tokens"], + vision_min_pixels=v["min_pixels"], + vision_max_wh_ratio=v["max_wh_ratio"], + image_token_id=hf["image_token_id"], + ) + + +def convert(ref_dir: str, hf_ckpt: str, out: str) -> str: + path = os.path.join(out, "model0-mp1.safetensors") + if os.path.exists(path): + print(f"[golden] using existing converted checkpoint {path}") + return path + cmd = [sys.executable, "-I", os.path.join(ref_dir, "convert.py"), "--hf-ckpt-path", hf_ckpt, "--save-path", out, + "--model-parallel", "1", "--expert-dtype", "fp4"] + print("[golden] $", " ".join(cmd)) + subprocess.run(cmd, check=True, cwd="/tmp") + return path + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt") + ap.add_argument("--ref-dir", default="~/di/ref_dsv41/inference", help="copy of the official inference/ dir") + ap.add_argument("--converted", default="~/di/tiny_ckpt_ref_mp1") + ap.add_argument("--out", default="~/di/golden") + ap.add_argument("--summary-dir", default=os.path.join(HERE, "golden")) + ap.add_argument("--steps", type=int, default=8, help="greedy decode steps after the prefill") + ap.add_argument("--prompts", default=",".join(PROMPTS)) + ap.add_argument("--threads", type=int, default=16) + ap.add_argument("--max-seq-len", type=int, default=2048) + ap.add_argument( + "--force-tokens-dir", + default=None, + help="teacher forcing: feed the `tokens` of /.pt as decode inputs instead of this run's argmax " + "(used to measure the reference's own numerical noise floor against an existing golden)", + ) + args = ap.parse_args() + for k in ("hf_ckpt", "ref_dir", "converted", "out", "summary_dir"): + setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k)))) + if args.force_tokens_dir: + args.force_tokens_dir = os.path.abspath(os.path.expanduser(args.force_tokens_dir)) + os.makedirs(args.out, exist_ok=True) + os.makedirs(args.summary_dir, exist_ok=True) + + import torch + + torch.set_num_threads(args.threads) + assert "torch_neuronx" not in sys.modules, "run with TORCH_DEVICE_BACKEND_AUTOLOAD=0 (host CPU only)" + + # ---- reference code: real model.py / engram.py, torch kernels instead of TileLang ---- + sys.path.insert(0, HERE) + import ref_kernels_torch + + sys.modules["kernel"] = ref_kernels_torch + sys.path.insert(1, args.ref_dir) + import engram as ref_engram + import model as ref_model + from safetensors.torch import load_file + from transformers import AutoTokenizer + + assert ref_model.act_quant is ref_kernels_torch.act_quant + assert os.path.dirname(os.path.abspath(ref_model.__file__)) == args.ref_dir + # the compressed token map walks the whole vocab; build it once + ref_engram.build_compressed_token_map = functools.lru_cache(maxsize=None)(ref_engram.build_compressed_token_map) + + ckpt_file = convert(args.ref_dir, args.hf_ckpt, args.converted) + hf_cfg = json.load(open(os.path.join(args.hf_ckpt, "config.json"))) + margs = hf_to_model_args(hf_cfg) + margs.update( + max_batch_size=1, + max_seq_len=args.max_seq_len, + temperature=0.0, + vision_n_layers=0, # text only + dspark_block_size=0, # no DSpark / MTP modules + ) + model_args = ref_model.ModelArgs(**margs) + print("[golden]", model_args) + tokenizer = AutoTokenizer.from_pretrained(args.converted) + bos = hf_cfg["bos_token_id"] + + state = load_file(ckpt_file) + dropped = sorted(k for k in state if k.startswith(("vision.", "aligner.", "mtp.", "image_")) or k.endswith(".bias_vl")) + state = {k: v for k, v in state.items() if k not in set(dropped)} + print(f"[golden] {len(state)} tensors for the backbone; dropped {len(dropped)} (vision/aligner/mtp/image_*/bias_vl)") + + torch.set_default_dtype(torch.bfloat16) + for name in args.prompts.split(","): + text = PROMPTS[name] + ids = [bos] + tokenizer.encode(text, add_special_tokens=False) + t0 = time.time() + torch.manual_seed(0) + with torch.device("cpu"): + model = ref_model.Transformer(model_args, tokenizer) + missing, unexpected = model.load_state_dict(state, strict=False) + assert not missing and not unexpected, (missing, unexpected) + model.eval() + + rec = {"hidden": [], "indexer": [], "engram": []} + cur = {} + + def block_hook(i): + def fn(mod, inp, out): + cur.setdefault("hidden", {})[i] = out[0].detach().clone() + return fn + + def idx_hook(i): + def fn(mod, inp, out): + cur.setdefault("indexer", {})[i] = out.detach().clone() + return fn + + for i, layer in enumerate(model.layers): + layer.register_forward_hook(block_hook(i)) + if layer.attn.indexer is not None: + layer.attn.indexer.register_forward_hook(idx_hook(i)) + if model.engram_hash is not None: + model.engram_hash.register_forward_hook(lambda m, i, o: cur.__setitem__("engram", o.detach().clone())) + + forced = None + if args.force_tokens_dir: + ref = torch.load(os.path.join(args.force_tokens_dir, f"{name}.pt"), weights_only=False) + assert ref["prompt_ids"] == ids, "teacher-forcing golden was built from a different prompt" + forced = ref["tokens"] + tokens, fed, logits_all = [], [], [] + x = torch.tensor([ids], dtype=torch.long) + pos = 0 + for step in range(args.steps + 1): + cur.clear() + out_ids, logits, _ = model(x, pos) + pos += x.size(1) + nxt = int(logits[0].argmax()) + assert nxt == int(out_ids[0]) + tokens.append(nxt) + logits_all.append(logits[0].float().clone()) + rec["hidden"].append(torch.stack([cur["hidden"][i][0] for i in range(len(model.layers))]).bfloat16()) + rec["indexer"].append({k: v[0].int() for k, v in cur.get("indexer", {}).items()}) + rec["engram"].append(cur["engram"][0] if "engram" in cur else None) + feed = forced[step] if forced is not None else nxt + fed.append(feed) + x = torch.tensor([[feed]], dtype=torch.long) + logits_all = torch.stack(logits_all) + dt = time.time() - t0 + + top2 = logits_all.topk(2, dim=-1) + margins = (top2.values[:, 0] - top2.values[:, 1]).tolist() + payload = { + "name": name, + "prompt_text": text, + "prompt_ids": ids, + "tokens": tokens, # argmax of each step: tokens[0] from the prefill, tokens[1:] from the decode steps + "fed_tokens": fed[:-1], # decode step i+1 was fed fed_tokens[i] (== tokens[i] unless teacher-forced) + "teacher_forced_from": args.force_tokens_dir, + "logits": logits_all, # [steps+1, vocab] fp32; row i predicts tokens[i] + "start_pos": [0] + [len(ids) + i for i in range(args.steps)], + "hidden": rec["hidden"], # per forward call: [n_layers, seqlen, hc_mult, dim] bf16 block outputs + "indexer_topk": rec["indexer"], # per forward call: {layer: [seqlen, topk] int32 (-1 = none, +window offset)} + "engram_hashes": rec["engram"], # per forward call: [seqlen, n_engram_layers, n_hash_cols] int64 + "model_args": dataclasses.asdict(model_args), + "hf_config": hf_cfg, + } + pt = os.path.join(args.out, f"{name}.pt") + torch.save(payload, pt) + top5 = logits_all.topk(5, dim=-1) + summary = { + "name": name, + "prompt_text": text, + "prompt_ids": ids, + "n_prompt_tokens": len(ids), + "tokens": tokens, + "decoded": tokenizer.decode(tokens), + "top1_top2_margin": [round(m, 5) for m in margins], + "top5_ids": top5.indices.tolist(), + "top5_logits": [[round(v, 5) for v in row] for row in top5.values.tolist()], + "logits_l2": [round(v, 4) for v in logits_all.norm(dim=-1).tolist()], + "logits_mean": [round(v, 6) for v in logits_all.mean(dim=-1).tolist()], + "full_golden": pt, + "full_golden_sha256": hashlib.sha256(open(pt, "rb").read()).hexdigest(), + "seconds": round(dt, 1), + } + with open(os.path.join(args.summary_dir, f"{name}.json"), "w") as fh: + json.dump(summary, fh, indent=1) + print(f"[golden] {name}: {len(ids)} prompt tokens -> {tokens} ({tokenizer.decode(tokens)!r}); " + f"min top1-top2 margin {min(margins):.4f}; {dt:.1f}s; saved {pt}") + + +if __name__ == "__main__": + main() diff --git a/examples/neuron/run_deepinfer_tiny.py b/examples/neuron/run_deepinfer_tiny.py new file mode 100644 index 0000000..a3cbe88 --- /dev/null +++ b/examples/neuron/run_deepinfer_tiny.py @@ -0,0 +1,351 @@ +"""Run the tiny V4.1 checkpoint through the real DeepInfer engine, teacher-forced on the reference golden. + +Engine path: deepinfer.entrypoints.llm.LLM (scheduler, KV cache, model runner, sampler), EP=DP=CP=1, Neuron basic +mode defaults (eager, DEEPINFER_USE_NEXTN_DECODE=0). Teacher forcing uses the engine's own mechanism, +SamplingParams(golden_token_ids=...), so decode step i is fed the golden token i-1 exactly like the reference +(reference_golden.py --force-tokens-dir). Logits are captured by wrapping the sampler's `get_logits` (the fp32 +[num_reqs, vocab] lm-head output before any sampling transform); per-layer block outputs by forward hooks on the +decoder layers (best effort, see --hidden). + +Output: /.pt in the reference_golden.py format (prompt_ids, tokens = argmax per step, fed_tokens, +logits [steps+1, vocab] fp32, optional hidden), ready for compare_golden.py. + + cd ~/di/DeepInfer && ~/di/venv/bin/python examples/neuron/run_deepinfer_tiny.py \ + --hf-ckpt ~/di/tiny_ckpt --golden-dir ~/di/golden --out ~/di/deepinfer_out [--device cpu|neuron] \ + [--hidden] [--compare] [--threads 32] + +Placeholder engine steps (is_placeholder=True, no real sequence) also run the model and the sampler; they are +filtered out of the recorded logits/hidden. `install_workarounds` patches, at runtime and only in this process, +the bugs found in other modules while bringing this up (BUG-1/2/4, NOTE-3/5 in the code); the list of the ones +that actually triggered is printed at the end. Set RUN_TINY_DEBUG=1 to print every sampler call. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time +import traceback +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def parse_args(): + ap = argparse.ArgumentParser() + ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt") + ap.add_argument("--golden-dir", default="~/di/golden") + ap.add_argument("--out", default="~/di/deepinfer_out") + ap.add_argument("--prompts", default="short,medium,long") + ap.add_argument("--device", choices=["cpu", "neuron"], default="cpu", help="sets DEEPINFER_NEURON_DEVICE") + ap.add_argument("--max-model-len", type=int, default=2048) + ap.add_argument("--max-num-batched-tokens", type=int, default=1024) + ap.add_argument("--threads", type=int, default=16) + ap.add_argument("--num-kv-blocks", type=int, default=32, help="fixed KV block pool (128 tokens each); 0 = engine profiling") + ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)") + ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards") + return ap.parse_args() + + +def setup_env(device: str) -> None: + # Must happen before deepinfer is imported (see examples/_config.py). + os.environ["DEEPINFER_ACCELERATOR"] = "neuron" + os.environ["DEEPINFER_NEURON_DEVICE"] = device + os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0" + os.environ.setdefault("DEEPINFER_ENFORCE_EAGER", "1") + os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0") + os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") # one engine rank; mapped to the target device + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + os.environ.setdefault("DEEPINFER_USE_MULTIMODAL_POOL", "0") # text only; avoids an 8 GB pinned host pool + if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + here = str(Path(__file__).resolve().parent) + if here not in sys.path: + sys.path.insert(0, here) + + +WORKAROUNDS: list[str] = [] + + +def install_workarounds(num_kv_blocks: int) -> None: + """Work around bugs in other modules without editing them; each one is reported (see the module docstring).""" + import torch + import torch.distributed as dist + + # BUG-1: PlatformNeuron(cpu) uses gloo and fp8_sf_type=torch.int16; gloo collectives reject int16 + # ("RuntimeError: Invalid scalar type") in v4_modules/engram.py NgramEmbedding.gather_sf. Byte-reinterpret + # dtypes gloo lacks as uint8 for the gather-type collectives (lossless; concatenation is along dim 0). + orig_agit = dist.all_gather_into_tensor + unsupported = {torch.int16, torch.uint16, torch.float8_e4m3fn, torch.float8_e8m0fnu} + + def all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False): + if input_tensor.dtype in unsupported and dist.get_backend(group) == "gloo": + if "BUG-1" not in WORKAROUNDS: + WORKAROUNDS.append("BUG-1") + return orig_agit(output_tensor.view(torch.uint8), input_tensor.contiguous().view(torch.uint8), group=group, async_op=async_op) + return orig_agit(output_tensor, input_tensor, group=group, async_op=async_op) + + dist.all_gather_into_tensor = all_gather_into_tensor + + # BUG-2: model_executor/model_loader.py get_model asserts `p.is_cuda` for every parameter; on Neuron the + # parameters live on the transfer_to_neuron target (neuron:0 or cpu). Same function, device-agnostic check. + import deepinfer.model_executor.model_loader as ml + import deepinfer.worker.model_runner as mr + from deepinfer import PLATFORM + + def get_model(model_config, skip_load=False): + model_class = ml._get_model_architecture(model_config) + with ml._set_default_torch_dtype(model_config.dtype), torch.device("cuda"): + model = model_class.load_model(model_config.model, model_config.model_cfg, skip_load) + target = PLATFORM.device.type + for n, p in model.named_parameters(): + if p.device.type != target: + raise AssertionError(f"{n} on {p.device}, expected {target}") + if hasattr(p, "sf") and p.sf.device.type != target: + raise AssertionError(f"{n}.sf on {p.sf.device}, expected {target}") + if not any("BUG-2" == w for w in WORKAROUNDS): + WORKAROUNDS.append("BUG-2") + return model.eval().requires_grad_(False) + + ml.get_model = get_model + mr.get_model = get_model + + # NOTE-5: WorkerBase pins the worker to phys_cpus // get_actual_num_gpus() CPUs; on Trn2 that is 96 // 64 = 1 + # CPU, so CPU-mode runs are single-threaded (~10x slower). Keep the process affinity instead. + import deepinfer.worker.base as wb + + def bind_cpu(local_rank): + WORKAROUNDS.append("NOTE-5: skipped bind_cpu (1 CPU per rank on Trn2)") + + wb.bind_cpu = bind_cpu + + # BUG-4: vendor deep_gemm fp8_fp4_paged_mqa_logits gathers keys for every column < max_context_len and only + # masks afterwards, so block-table slots past a request's last page (unused entries, not valid page ids) are + # dereferenced -> IndexError. Sanitize out-of-range table entries (their columns are masked to -inf anyway) and + # flag it loudly if an out-of-range entry is ever needed by a valid position. + import deep_gemm + + orig_paged = deep_gemm.fp8_fp4_paged_mqa_logits + + def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, *a, **k): + num_pages, page_kv = kv_cache.shape[0], kv_cache.shape[1] + bad = (block_table < 0) | (block_table >= num_pages) + if bool(bad.any()): + lens = context_lens.reshape(-1).to(torch.int64) + needed_pages = (lens + page_kv - 1) // page_kv + col = torch.arange(block_table.shape[1], device=block_table.device)[None, :] + if bool((bad & (col < needed_pages[:, None])).any()): + WORKAROUNDS.append(f"BUG-4b: out-of-range page id inside the valid context: table={block_table.tolist()} lens={lens.tolist()}") + if "BUG-4" not in WORKAROUNDS: + WORKAROUNDS.append("BUG-4") + block_table = torch.where(bad, torch.zeros_like(block_table), block_table) + return orig_paged(q, kv_cache, weights, context_lens, block_table, *a, **k) + + deep_gemm.fp8_fp4_paged_mqa_logits = fp8_fp4_paged_mqa_logits + + # NOTE-3 (not a bug, a CPU-mode hazard): WorkerBase.profile_num_available_blocks runs a dummy prefill of + # max_num_batched_tokens tokens (903 s on the host for 2048 tokens) and then sizes the KV cache from + # mem_get_info(), which on cpu is host RAM -> 0.5 * 2 TB = 559 GiB of KV blocks. Use a fixed small pool instead. + if num_kv_blocks > 0: + import torch.distributed as _dist # noqa: F401 + from deepinfer.worker.base import WorkerBase + + def profile_num_available_blocks(self): + self.model_runner.set_block_size(self.cache_config.block_size) + self.cache_config.num_gpu_blocks = num_kv_blocks + WORKAROUNDS.append(f"NOTE-3: fixed {num_kv_blocks} KV blocks, no profile run") + return num_kv_blocks + + WorkerBase.profile_num_available_blocks = profile_num_available_blocks + + +class Recorder: + """Collects the lm-head logits of every sampler call (one request in flight -> one row per call).""" + + def __init__(self): + self.step_kind = None # (is_prompt, is_placeholder) of the engine step being executed + self.n_placeholder = 0 + self.logits = [] + self.hidden_calls = [] + self._cur_hidden = {} + + def reset(self): + self.n_placeholder = 0 + self.logits.clear() + self.hidden_calls.clear() + self._cur_hidden = {} + + +def install_logits_hook(rec: Recorder) -> None: + import deepinfer.model_executor.layers.sampler as sampler_mod + + orig = sampler_mod.get_logits + + def get_logits(hidden_states, embedding, vocab_size): + logits = orig(hidden_states, embedding, vocab_size) + if rec.step_kind is not None and rec.step_kind[1]: + rec.n_placeholder += 1 # placeholder step (no real sequence): not part of the request + return logits + rec.logits.append(logits.detach().float().cpu().clone()) + if os.environ.get("RUN_TINY_DEBUG"): + import traceback as _tb + + st = " <- ".join(f"{f.name}:{f.lineno}" for f in _tb.extract_stack()[-6:-1][::-1]) + print(f"[run] sampler call {len(rec.logits)}: logits {tuple(logits.shape)} argmax {int(logits[0].argmax())} " + f"hs {tuple(hidden_states.shape)} at {time.strftime('%H:%M:%S')} via {st}", flush=True) + return logits + + sampler_mod.get_logits = get_logits + + from deepinfer.worker.worker import Worker + + orig_step = Worker.execute_model_step + + def execute_model_step(self, is_prompt, is_placeholder, *a, **k): + rec.step_kind = (bool(is_prompt), bool(is_placeholder)) + try: + return orig_step(self, is_prompt, is_placeholder, *a, **k) + finally: + rec.step_kind = None + + Worker.execute_model_step = execute_model_step + + +def install_hidden_hooks(rec: Recorder, llm) -> str: + """Hook the decoder layers of the driver worker's model. Returns a status string.""" + try: + worker = llm.llm_engine.driver_worker + runner = getattr(worker, "model_runner", None) + model = getattr(runner, "model", None) if runner is not None else getattr(worker, "model", None) + layers = model.model.layers + except Exception as exc: # noqa: BLE001 + return f"hidden hooks not installed: {type(exc).__name__}: {exc}" + + from deepinfer.model_executor.models.v4_modules.hc import HCLayerMixin + + n = len(layers) + + def hook(i): + def fn(mod, inp, out): + if rec.step_kind is None or rec.step_kind[1]: + return None # placeholder / non-request forward + h, residual = out + h = h.move() if hasattr(h, "move") else h + try: + # 2D output = before hc_post; the reference block output is the expanded hc stream + x = HCLayerMixin.hc_post(h, residual) if h.dim() == 2 else h + except Exception: # noqa: BLE001 + x = h + rec._cur_hidden[i] = x.detach().float().cpu().clone() + if i == n - 1: + rec.hidden_calls.append(dict(rec._cur_hidden)) + rec._cur_hidden = {} + return None + + return fn + + for i, layer in enumerate(layers): + layer.register_forward_hook(hook(i)) + return f"hidden hooks on {n} layers of {type(model).__name__}" + + +def main(): + args = parse_args() + for k in ("hf_ckpt", "golden_dir", "out"): + setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k)))) + setup_env(args.device) + os.makedirs(args.out, exist_ok=True) + + import torch + + torch.set_num_threads(args.threads) + import deepinfer # noqa: F401 + from deepinfer import PLATFORM + + print(f"[run] platform device = {PLATFORM.device}", flush=True) + if args.device == "cpu": + assert str(PLATFORM.device) == "cpu", "refusing to run: target device is not the host CPU" + + install_workarounds(args.num_kv_blocks) + rec = Recorder() + install_logits_hook(rec) + + from deepinfer.entrypoints.llm import LLM + from deepinfer.sampling_params import SamplingParams + + t0 = time.time() + llm = LLM( + model=args.hf_ckpt, + trust_remote_code=True, + expert_parallel_size=1, + context_parallel_size=1, + data_parallel_size=1, + engine_rank=0, + chunked_prefill=True, + decode_sequence=False, + max_model_len=args.max_model_len, + max_num_batched_tokens=args.max_num_batched_tokens, + max_num_seqs=1, + gpu_memory_utilization=0.5, + ) + print(f"[run] engine up in {time.time() - t0:.1f}s", flush=True) + status = install_hidden_hooks(rec, llm) if args.hidden else "hidden not requested" + print(f"[run] {status}", flush=True) + + for name in args.prompts.split(","): + gold = torch.load(os.path.join(args.golden_dir, f"{name}.pt"), weights_only=False) + ids, gtok = list(gold["prompt_ids"]), list(gold["tokens"]) + steps = len(gtok) + rec.reset() + sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True, golden_token_ids=gtok) + t1 = time.time() + outputs = llm.generate(prompt_token_ids=[ids], sampling_params=sp, use_tqdm=False) + dt = time.time() - t1 + out_tokens = [o.token_id for o in outputs[0]] + rows = [r for r in rec.logits] + assert all(r.size(0) == 1 for r in rows), [tuple(r.shape) for r in rows] + logits = torch.cat(rows)[-steps:] if len(rows) >= steps else torch.cat(rows) + print(f"[run] {name}: {len(rows)} request sampler calls, {rec.n_placeholder} placeholder steps", flush=True) + if len(rows) != steps: + print(f"[run] WARNING {name}: {len(rows)} sampler calls for {steps} steps (chunked prefill?)", flush=True) + payload = { + "name": name, + "prompt_ids": ids, + "tokens": logits.argmax(-1).tolist(), + "engine_output_tokens": out_tokens, + "fed_tokens": gtok[: steps - 1], + "logits": logits, + "source": "deepinfer LLM engine, teacher-forced via SamplingParams.golden_token_ids", + "device": args.device, + } + if rec.hidden_calls: + hid = [] + for call in rec.hidden_calls[-steps:]: + hid.append(torch.stack([call[i].reshape(call[i].shape[0], -1, gold["hidden"][0].shape[-1]) for i in sorted(call)]).bfloat16()) + payload["hidden"] = hid + path = os.path.join(args.out, f"{name}.pt") + torch.save(payload, path) + match = sum(a == b for a, b in zip(payload["tokens"], gtok)) + print(f"[run] {name}: {len(ids)} prompt tokens, engine output {out_tokens} (golden-forced: {out_tokens == gtok}), " + f"argmax==golden {match}/{steps}, {dt:.1f}s -> {path}", flush=True) + if args.compare: + import compare_golden + + res = compare_golden.compare(gold, payload) + for s in res["steps"]: + print(f" step {s['step']}: match {s['match']} margin {s['margin_ref']:.3f} rel {s['rel']:.2e} cos {s['cos']:.6f} kl {s['kl']:.2e}") + for L in res["layers"][:1]: + print(" prefill hidden rel per layer:", L.get("error") or " ".join(f"{v:.1e}" for v in L["rel"])) + print(f"[run] workarounds triggered: {WORKAROUNDS}", flush=True) + llm.close() + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + sys.exit(1) diff --git a/examples/neuron/test_ref_kernels_torch.py b/examples/neuron/test_ref_kernels_torch.py new file mode 100644 index 0000000..7e3c856 --- /dev/null +++ b/examples/neuron/test_ref_kernels_torch.py @@ -0,0 +1,128 @@ +"""Self-consistency tests for ref_kernels_torch (the pure-torch stand-in for the reference kernel.py). + + TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I examples/neuron/test_ref_kernels_torch.py +""" + +import os +import sys + +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import ref_kernels_torch as K # noqa: E402 + + +def rel(a, b): + return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-30)).item() + + +def test_act_quant(): + x = torch.randn(7, 3, 256).bfloat16() * 3 + q, s = K.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu) + assert q.dtype == torch.float8_e4m3fn and s.dtype == torch.float8_e8m0fnu and s.shape == (7, 3, 8) + sf = s.float() + assert torch.all(torch.log2(sf) == torch.round(torch.log2(sf))) + deq = (q.float().unflatten(-1, (8, 32)) * sf.unsqueeze(-1)).flatten(-2) + y = x.clone() + K.act_quant(y, 32, "ue8m0", torch.float8_e8m0fnu, True) + assert torch.equal(y, deq.bfloat16()) + z = y.clone() + K.act_quant(z, 32, "ue8m0", torch.float8_e8m0fnu, True) + assert torch.equal(z, y), "quant->dequant must be idempotent" + assert rel(y, x) < 0.05 + # every block amax maps below 448 after scaling, and the scale is the smallest power of two that allows it + amax = x.float().unflatten(-1, (8, 32)).abs().amax(-1) + assert torch.all(amax / sf <= 448) and torch.all(amax / (sf / 2) > 448) + + +def test_fp8_gemm(): + torch.set_default_dtype(torch.bfloat16) + x = torch.randn(5, 4, 256).bfloat16() + w = torch.randn(96, 256) / 16 + wq, ws = [], [] + # weight 32x32 block quant, like the checkpoint + wb = w.view(3, 32, 8, 32) + s = K._fast_round_scale(wb.abs().amax((1, 3)), 1 / 448) + wq = (wb / s[:, None, :, None]).clamp(-448, 448).to(torch.float8_e4m3fn).view(96, 256) + ws = K._to_scale_dtype(s, torch.float8_e8m0fnu) + a, a_s = K.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu) + c = K.fp8_gemm(a, a_s, wq, ws, torch.float8_e8m0fnu, 32) + wdeq = (wq.float().view(3, 32, 8, 32) * s[:, None, :, None]).view(96, 256) + adeq = (a.float().unflatten(-1, (8, 32)) * a_s.float().unsqueeze(-1)).flatten(-2) + ref = adeq @ wdeq.T + assert c.dtype == torch.bfloat16 and c.shape == (5, 4, 96) + assert rel(c, ref) < 4e-3 + assert rel(c, x.float() @ w.float().T) < 0.08 + torch.set_default_dtype(torch.float32) + + +def test_fp4_gemm(): + torch.set_default_dtype(torch.bfloat16) + x = torch.randn(9, 128).bfloat16() + w = torch.randn(64, 128) / 11 + s = K._fast_round_scale(w.view(64, 4, 32).abs().amax(-1), 1 / 6) + codes = K._e2m1_codes((w.view(64, 4, 32) / s[..., None]).clamp(-6, 6)).view(64, 128) + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).to(torch.uint8).view(torch.float4_e2m1fn_x2) + # decode exactly like convert.py cast_e2m1fn_to_e4m3fn (low nibble first, FP4_TABLE) + FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) + u = packed.view(torch.uint8) + wdeq = torch.stack([FP4_TABLE[(u & 0xF).long()], FP4_TABLE[((u >> 4) & 0xF).long()]], -1).flatten(1) + wdeq = wdeq * s.repeat_interleave(32, -1) + a, a_s = K.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu) + c = K.fp4_gemm(a, a_s, packed, K._to_scale_dtype(s, torch.float8_e8m0fnu), torch.float8_e8m0fnu, 32) + adeq = (a.float().view(9, 4, 32) * a_s.float().unsqueeze(-1)).flatten(-2) + assert rel(c, adeq @ wdeq.T) < 4e-3 + assert rel(c, x.float() @ w.float().T) < 0.2 + torch.set_default_dtype(torch.float32) + + +def test_fp4_act_quant(): + x = torch.randn(3, 5, 128).bfloat16() + for sd, bs in ((torch.float8_e8m0fnu, 32), (torch.float8_e4m3fn, 16)): + y = x.clone() + K.fp4_act_quant(y, bs, True, sd) + packed, s = K.fp4_act_quant(x, bs, False, sd) + deq = (K._unpack_fp4(packed).unflatten(-1, (-1, bs)) * s.float().unsqueeze(-1)).flatten(-2) + assert torch.equal(deq.bfloat16(), y), sd + assert rel(y, x) < 0.15 + # RNE ties on the e2m1 grid + v = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, -0.25, -5.0, 5.01]) + assert K._e2m1_round(v).tolist() == [0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0, 0.0, -4.0, 6.0] + + +def test_sparse_attn(): + b, m, h, d, n, topk = 2, 5, 8, 512, 150, 140 + q = torch.randn(b, m, h, d).bfloat16() + kv = torch.randn(b, n, d).bfloat16() + sink = torch.randn(h) * 0.3 + idx = torch.stack([torch.randperm(n)[:topk] for _ in range(b * m)]).view(b, m, topk).int() + idx[0, 0, 100:] = -1 + idx[1, 2, :] = -1 # a row with nothing to attend to -> all-zero output + scale = d**-0.5 + o = K.sparse_attn(q, kv, sink, idx, scale) + # dense fp32 reference with an explicit sink column + kvg = kv[torch.arange(b).view(b, 1, 1), idx.long().clamp_min(0)].float() + s = torch.einsum("bmhd,bmkd->bmhk", q.float(), kvg) * scale + s = s.masked_fill((idx == -1).unsqueeze(2), -torch.inf) + s = torch.cat([s, sink.view(1, 1, h, 1).expand(b, m, h, 1)], -1) + p = torch.softmax(s, -1)[..., :-1] + ref = torch.einsum("bmhk,bmkd->bmhd", p, kvg) + assert rel(o, ref) < 1e-2, rel(o, ref) + assert torch.all(o[1, 2] == 0) + + +def test_sinkhorn(): + mixes = torch.randn(2, 3, 24) + pre, post, comb = K.hc_split_sinkhorn(mixes, torch.tensor([0.3, 0.3, 0.3]), torch.randn(24), 4, 20, 1e-6) + assert pre.shape == (2, 3, 4) and post.shape == (2, 3, 4) and comb.shape == (2, 3, 4, 4) + assert torch.allclose(comb.sum(-2), torch.ones(2, 3, 4), atol=1e-5) + assert torch.allclose(comb.sum(-1), torch.ones(2, 3, 4), atol=1e-3) + assert torch.all((post > 0) & (post < 2)) and torch.all(pre > 0) + + +if __name__ == "__main__": + torch.manual_seed(0) + for name, fn in list(globals().items()): + if name.startswith("test_"): + fn() + print("ok", name)