coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 23:07:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron basic mode — part 3: end-to-end check vs the official reference

Browse Source 
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout,
  pure-torch replacements for the official reference kernels, golden
  generation, a teacher-forced DeepInfer runner (LLM offline API) and a
  comparison tool. CPU mode: 25/27 greedy tokens match, worst logits
  cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060.
- model_loader: check parameters against the platform device, not is_cuda.
- deep_gemm paged MQA logits: do not dereference block-table slots past the
  context.
- bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would
  serialize host compute).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-20 20:20:02 +08:00
parent f4144fc3fc
commit 8a5209ee65
16 changed files with 3672 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
deepinfer/model_executor/model_loader.py
View File

@ -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)

6
deepinfer/neuron_ops/vendor/deep_gemm/attention.py vendored
View File

@ -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

5
deepinfer/utils.py
View File

@ -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

37
examples/neuron/README.md Normal file
View File

@ -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`.

142
examples/neuron/check_tiny_checkpoint.py Normal file
View File

@ -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")

106
examples/neuron/compare_golden.py Normal file
View File

@ -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 <candidate>.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()

316
examples/neuron/golden/deepinfer_cpu_vs_golden.json Normal file
View File

@ -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
}
]
}
}
}

900
examples/neuron/golden/long.json Normal file
View File

@ -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
}

357
examples/neuron/golden/medium.json Normal file
View File

@ -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
}

73
examples/neuron/golden/noise_floor.json Normal file
View File

@ -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
]
}
}
}

196
examples/neuron/golden/short.json Normal file
View File

@ -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
}

465
examples/neuron/make_tiny_checkpoint.py Normal file
View File

@ -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 <dir with tokenizer.json> \
--real-config <real HF config.json>
"""
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()

246
examples/neuron/ref_kernels_torch.py Normal file
View File

@ -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

341
examples/neuron/reference_golden.py Normal file
View File

@ -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 <dir>/<prompt>.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()

351
examples/neuron/run_deepinfer_tiny.py Normal file
View File

@ -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: <out>/<prompt>.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)

128
examples/neuron/test_ref_kernels_torch.py Normal file
View File

@ -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)