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
DeepInfer/examples/neuron/probe_layer0.py
whn09 760e3cadc0 Neuron: test suite for new boxes, device-vs-host differential tool 
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
  on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
  tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
  tilelang_ops interface tests (basic-mode stubs reported as skips via
  nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
  process group, device primitive checks, end-to-end vs golden on CPU and
  on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
  the port depends on or works around (fp8 strided-view indexing, uint64,
  1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
  run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
  tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-21 09:46:29 +08:00

453 lines
21 KiB
Python
Raw Blame History

"""Op-by-op capture of the first decoder layers, in the official reference and in DeepInfer, for one prefill.
# reference (pure CPU, torch kernels from ref_kernels_torch), twice with different thread counts = floor
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I examples/neuron/probe_layer0.py ref --out ~/di/probe/ref16.pt --threads 16
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I examples/neuron/probe_layer0.py ref --out ~/di/probe/ref7.pt --threads 7
# DeepInfer engine (CPU mode), same prompt, prefill only
python examples/neuron/probe_layer0.py deepinfer --out ~/di/probe/di.pt
# table
python -I examples/neuron/probe_layer0.py compare ~/di/probe/ref16.pt ~/di/probe/di.pt --floor ~/di/probe/ref7.pt
Captured names are "L{layer}.{op}"; every tensor is [seqlen, ...] (batch dim dropped), fp32 on CPU. The prompt is
the `short` golden prompt (BOS + 10 tokens); the table reports token 0 (BOS), token 1 and all tokens.
"""
from __future__ import annotations
import argparse
import functools
import json
import os
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
LAYERS = (0, 1, 2) # layer 2 only to catch layer 1's FFN hc_post (fused into layer 2's pre in DeepInfer)
COMPARE_LAYERS = (0, 1)
class Cap:
def __init__(self):
self.t = {}
self.layer = None
self.count = {}
self.active = True
def put(self, name, x, layer=None):
import torch
layer = self.layer if layer is None else layer
if not self.active or layer is None or layer not in LAYERS:
return
key = f"L{layer}.{name}"
n = self.count.get(key, 0)
self.count[key] = n + 1
if n:
key = f"{key}#{n}"
if isinstance(x, (tuple, list)):
x = x[0]
x = x.detach().float().cpu().clone()
self.t[key] = x
# ------------------------------------------------------------------------------------------------ reference
def run_ref(args):
import torch
torch.set_num_threads(args.threads)
sys.path.insert(0, str(HERE))
import ref_kernels_torch
sys.modules["kernel"] = ref_kernels_torch
sys.path.insert(1, os.path.expanduser(args.ref_dir))
import engram as ref_engram
import model as M
import reference_golden as G
from safetensors.torch import load_file
from transformers import AutoTokenizer
ref_engram.build_compressed_token_map = functools.lru_cache(None)(ref_engram.build_compressed_token_map)
hf = json.load(open(os.path.join(os.path.expanduser(args.hf_ckpt), "config.json")))
ma = G.hf_to_model_args(hf)
ma.update(max_batch_size=1, max_seq_len=2048, temperature=0.0, vision_n_layers=0, dspark_block_size=0)
conv = os.path.expanduser(args.converted)
tok = AutoTokenizer.from_pretrained(conv)
sd = load_file(os.path.join(conv, "model0-mp1.safetensors"))
sd = {k: v for k, v in sd.items() if not (k.startswith(("vision.", "aligner.", "mtp.", "image_")) or k.endswith(".bias_vl"))}
torch.set_default_dtype(torch.bfloat16)
model = M.Transformer(M.ModelArgs(**ma), tok)
model.load_state_dict(sd, strict=False)
cap = Cap()
def put_hook(name):
return lambda mod, i, o: cap.put(name, o)
model.embed.register_forward_hook(lambda m, i, o: cap.put("embed", o[0], layer=0))
for li, blk in enumerate(model.layers):
blk.register_forward_pre_hook(lambda m, i, li=li: setattr(cap, "layer", li))
a = blk.attn
blk.attn_norm.register_forward_hook(put_hook("attn_norm"))
blk.ffn_norm.register_forward_hook(put_hook("ffn_norm"))
a.wq_a.register_forward_hook(put_hook("wq_a"))
a.q_norm.register_forward_hook(put_hook("q_norm"))
a.wq_b.register_forward_hook(put_hook("wq_b"))
a.wkv.register_forward_hook(put_hook("wkv"))
a.kv_norm.register_forward_hook(put_hook("kv_norm"))
a.wo_b.register_forward_pre_hook(lambda m, i: cap.put("wo_b_in", i[0]))
a.wo_b.register_forward_hook(put_hook("attn_out"))
g = blk.ffn.gate
g.register_forward_pre_hook(lambda m, i: cap.put("gate_logits", torch.nn.functional.linear(i[0].float(), m.weight.float())))
g.register_forward_hook(lambda m, i, o: (cap.put("topk_weights", o[0]), cap.put("topk_idx", o[1].float()), None)[-1])
se = blk.ffn.shared_experts
se.register_forward_hook(put_hook("shared_out"))
se.w1.register_forward_hook(put_hook("shared_w1"))
se.w3.register_forward_hook(put_hook("shared_w3"))
se.w2.register_forward_pre_hook(lambda m, i: cap.put("shared_act", i[0]))
blk.ffn.register_forward_hook(put_hook("moe_out"))
if blk.engram is not None: # runs before the block (Transformer.forward), so pass the layer explicitly
blk.engram.register_forward_hook(lambda m, i, o, li=li: cap.put("engram_out", o[0], layer=li))
orig_split = M.hc_split_sinkhorn
def hc_split_sinkhorn(mixes, *a, **k):
pre, post, comb = orig_split(mixes, *a, **k)
cap.put("mixes", mixes[0])
cap.put("pre_mix", pre[0])
cap.put("post_mix", post[0])
cap.put("comb_mix", comb[0])
return pre, post, comb
M.hc_split_sinkhorn = hc_split_sinkhorn
orig_pre, orig_post = M.Block.hc_pre, M.Block.hc_post
M.Block.hc_pre = lambda self, x, m: (lambda y: (cap.put("hc_pre", y[0]), y)[1])(orig_pre(self, x, m))
M.Block.hc_post = lambda self, x, r, p, c: (lambda y: (cap.put("hc_post", y[0]), y)[1])(orig_post(self, x, r, p, c))
orig_sa = M.sparse_attn
def sparse_attn(q, kv, sink, idx, scale):
o = orig_sa(q, kv, sink, idx, scale)
cap.put("q_rope", q[0])
cap.put("kv_rope_quant", kv[0, : q.size(1)]) # prefill: kv rows == tokens (window part)
cap.put("attn_core", o[0])
return o
M.sparse_attn = sparse_attn
# MoE.forward (model.py:891-906), unchanged numerics, additionally recording the fp32 routed sum and the
# DeepEP-combine-style result (routed partial rounded to bf16 before the shared expert is added)
def moe_forward(self, x, image_mask=None):
shape = x.size()
x = x.view(-1, self.dim)
weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten())
y = torch.zeros_like(x, dtype=torch.float32)
counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
for i in range(self.experts_start_idx, self.experts_end_idx):
if counts[i] == 0:
continue
idx, top = torch.where(indices == i)
y[idx] += self.experts[i](x[idx], weights[idx, top, None])
cap.put("routed_fp32", y)
shared = self.shared_experts(x)
cap.put("moe_out_combine_style", (y.to(torch.bfloat16).float() + shared.float()).to(torch.bfloat16))
y += shared
return y.type_as(x).view(shape)
M.MoE.forward = moe_forward
gold = torch.load(os.path.join(os.path.expanduser(args.golden_dir), "short.pt"), weights_only=False)
ids = gold["prompt_ids"]
with torch.inference_mode():
model(torch.tensor([ids]), 0)
torch.save({"side": "ref", "ids": ids, "t": cap.t}, os.path.expanduser(args.out))
print(f"[probe] ref: {len(cap.t)} tensors -> {args.out}")
# ------------------------------------------------------------------------------------------------ proposed fix
def sparse_attn_rows_p_bf16(q, gather_kv, topk, sm_scale, d_v, attn_sink, rms_norm_scale_factor, block=64):
"""Drop-in for neuron_ops/vendor/flash_mla/flash_mla_interface.py:_sparse_attn_rows with the numerics of the
CUDA FlashMLA sparse kernels (e.g. csrc/cuda_kernels/sm100/prefill/sparse/fwd/head128/phase1.cuh:238-253,324):
online softmax over blocks of `block` keys, P = exp(s - running_max) rounded to bf16 before the P@V GEMM,
running sum and the final 1 / (l + exp(sink - m)) in fp32. Same return contract as the original."""
import torch
n, h_q, d_qk = q.shape
out = torch.empty((n, h_q, d_v), dtype=torch.float32, device=q.device)
max_logits = torch.empty((n, h_q), dtype=torch.float32, device=q.device)
lse_out = torch.empty((n, h_q), dtype=torch.float32, device=q.device)
rows = max(1, min(n, (1 << 26) // max(1, topk * d_qk)))
for lo in range(0, n, rows):
hi = min(lo + rows, n)
kv, invalid = gather_kv(lo, hi)
kv = kv.float()
kv = torch.where(kv != kv, torch.zeros_like(kv), kv)
logits = torch.bmm(q[lo:hi].float(), kv.transpose(1, 2))
if rms_norm_scale_factor is not None:
logits = logits * (rms_norm_scale_factor[lo:hi].unsqueeze(-1) * sm_scale)
else:
logits = logits * sm_scale
logits = torch.where(invalid.unsqueeze(1), torch.full_like(logits, float("-inf")), logits)
m = torch.full((hi - lo, h_q), -1e30, dtype=torch.float32, device=q.device)
l = torch.zeros_like(m)
acc = torch.zeros((hi - lo, h_q, d_v), dtype=torch.float32, device=q.device)
for t0 in range(0, topk, block):
sblk = logits[..., t0 : t0 + block]
new = torch.maximum(m, sblk.amax(-1))
sc = torch.exp(m - new)
p = torch.exp(sblk - new.unsqueeze(-1))
l = l * sc + p.sum(-1)
acc = acc * sc.unsqueeze(-1) + torch.bmm(p.to(torch.bfloat16).float(), kv[:, t0 : t0 + block, :d_v])
m = new
denom = l + (torch.exp(attn_sink.float().view(1, h_q) - m) if attn_sink is not None else 0)
lonely = l == 0
o = acc / torch.where(denom == 0, torch.ones_like(denom), denom).unsqueeze(-1)
out[lo:hi] = torch.where(lonely.unsqueeze(-1), torch.zeros_like(o), o)
max_logits[lo:hi] = logits.max(dim=-1).values
lse_out[lo:hi] = torch.where(lonely, torch.full_like(l, float("inf")), m + torch.log(l))
return out, max_logits, lse_out
# ------------------------------------------------------------------------------------------------ DeepInfer
def run_deepinfer(args):
sys.path.insert(0, str(HERE))
import run_deepinfer_tiny as R
R.setup_env(args.device)
import torch
torch.set_num_threads(args.threads)
import deepinfer # noqa: F401
from deepinfer import PLATFORM
assert str(PLATFORM.device) == "cpu" or args.device == "neuron"
cap = Cap()
cap.active = False
rec = R.Recorder()
R.install_logits_hook(rec) # also tags placeholder steps
import deepinfer.model_executor.models.v4_modules.attn as attn_mod
import deepinfer.model_executor.models.v4_modules.hc as hc_mod
import deepinfer.model_executor.models.v4_modules.mlp as mlp_mod
def real_step():
return rec.step_kind is not None and rec.step_kind[0] and not rec.step_kind[1]
def wrap(mod, fname, fn_out):
orig = getattr(mod, fname)
def w(*a, **k):
out = orig(*a, **k)
if real_step():
fn_out(out, a, k)
return out
setattr(mod, fname, w)
wrap(hc_mod, "mhc_pre_apply_mix", lambda o, a, k: cap.put("hc_pre", o.reshape(-1, o.shape[-1])))
wrap(hc_mod, "mhc_pre_norm_fn", lambda o, a, k: cap.put("mixes", o.reshape(-1, o.shape[-1])))
def split_out(o, a, k):
pre, post, comb = o
cap.put("pre_mix", pre.reshape(-1, pre.shape[-2]))
cap.put("post_mix", post.reshape(-1, post.shape[-2]))
wrap(hc_mod, "mhc_pre_split_mixes", split_out)
wrap(hc_mod, "sinkhorn_normalize", lambda o, a, k: cap.put("comb_mix", o.reshape(-1, *o.shape[-2:])))
wrap(hc_mod, "mhc_post", lambda o, a, k: cap.put("hc_post_raw", o.reshape(-1, *o.shape[-2:])))
wrap(hc_mod, "expand_to_mhc", lambda o, a, k: None)
def fused_out(o, a, k):
if o is None:
return
x = a[0]
if x.dim() == 3 and x.shape[1] > 1: # q [s, heads, head_dim]
cap.put("q_rope", o)
else:
cap.put("kv_rope_quant", o.reshape(o.shape[0], -1))
wrap(attn_mod, "fused_norm_rope_hadamard_cast_store_kvcache", fused_out)
wrap(attn_mod, "prefill_attn_varlen_sparse", lambda o, a, k: cap.put("attn_core", o[0]))
import tile_kernels.quant.swiglu_forward_kernel as swk
orig_swiglu = mlp_mod.swiglu_forward_and_per_token_cast
def swiglu(x, *a, **k):
if real_step() and k.get("topk_weights") is None and (len(a) < 4 or a[3] is None):
inter = x.shape[-1] // 2
cap.put("shared_w1", x[:, :inter])
cap.put("shared_w3", x[:, inter:])
cap.put("shared_act", swk.swiglu_compute(x, None, None, None, k.get("clamp_value"), None))
return orig_swiglu(x, *a, **k)
mlp_mod.swiglu_forward_and_per_token_cast = swiglu
if args.round_act_bf16:
# experiment (not a proposed fix): round the SwiGLU output to bf16 before the FP8 cast, like the reference's
# Expert.forward (`self.w2(x.to(dtype))`), to confirm that this rounding point explains the FFN deviation
orig_compute = swk.swiglu_compute
swk.swiglu_compute = lambda *a, **k: orig_compute(*a, **k).to(torch.bfloat16).float()
print("[probe] SwiGLU output rounded to bf16 before the FP8 cast (experiment)")
if args.fix_attn:
import flash_mla.flash_mla_interface as fmi
fmi._sparse_attn_rows = sparse_attn_rows_p_bf16
print("[probe] using P-bf16 sparse attention (proposed fix)")
wrap(mlp_mod, "moe_topk_gate", lambda o, a, k: (cap.put("topk_idx", o[0].float()), cap.put("topk_weights", o[1])))
# Fixed KV pool: the engine's profile run (a max_num_batched_tokens dummy prefill) is slow on the host and
# exhausts HBM on the NeuronCore with the pure-torch kernels.
R.install_workarounds(32)
from deepinfer.entrypoints.llm import LLM
from deepinfer.sampling_params import SamplingParams
llm = LLM(model=os.path.expanduser(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=2048,
max_num_batched_tokens=1024, max_num_seqs=1, gpu_memory_utilization=0.5)
model = llm.llm_engine.driver_worker.model_runner.model
q_lora = None
def hook(name, fn=None):
def h(m, i, o):
if real_step():
cap.put(name, fn(o) if fn else o)
return h
def bf16_of(o):
# norms return ((fp8, sf), bf16) or (fp8, bf16) when return_bf16=True
return o[1] if isinstance(o, tuple) else o
model.model.embed_tokens.register_forward_hook(lambda m, i, o: cap.put("embed", o, layer=0) if real_step() else None)
for li, layer in enumerate(model.model.layers):
layer.register_forward_pre_hook(lambda m, i, li=li: setattr(cap, "layer", li) if real_step() else None)
a = layer.self_attn
q_lora = a.q_lora_rank
layer.attn_layernorm.register_forward_hook(hook("attn_norm", bf16_of))
layer.ffn_layernorm.register_forward_hook(hook("ffn_norm", bf16_of))
a.q_a_kv_swa_proj.register_forward_hook(hook("q_a_kv"))
a.q_a_norm.register_forward_hook(hook("q_norm", bf16_of))
a.q_b_proj.register_forward_hook(hook("wq_b"))
a.out_proj.register_forward_pre_hook(lambda m, i: cap.put("wo_b_in", i[0]) if real_step() else None)
a.out_proj.register_forward_hook(hook("attn_out"))
mlp = layer.mlp
mlp.gate.register_forward_hook(hook("gate_logits"))
orig_shared = mlp.shared_forward
def shared_forward(*aa, _o=orig_shared, **kk):
out = _o(*aa, **kk)
if real_step() and out is not None:
cap.put("shared_out", out)
return out
mlp.shared_forward = shared_forward
mlp.register_forward_hook(hook("moe_out"))
if layer.engram is not None:
layer.engram.register_forward_hook(hook("engram_out"))
gold = __import__("torch").load(os.path.join(os.path.expanduser(args.golden_dir), "short.pt"), weights_only=False)
ids = gold["prompt_ids"]
cap.active = True
llm.generate(prompt_token_ids=[ids], sampling_params=SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True,
golden_token_ids=gold["tokens"][:1]), use_tqdm=False)
t = cap.t
# split the fused q_a|kv projection like the reference's separate wq_a / wkv
for li in LAYERS:
k = f"L{li}.q_a_kv"
if k in t:
t[f"L{li}.wq_a"], t[f"L{li}.wkv"] = t[k][:, :q_lora], t[k][:, q_lora:]
k = f"L{li}.wq_b"
if k in t:
H, hd = model.model.layers[li].self_attn.num_heads, model.model.layers[li].self_attn.qk_head_dim
t[k] = t[k][:, : H * hd]
# mhc_post order: inside layer L, call #0 is L's attention post; L's FFN post runs inside layer L+1 (fused
# post+pre), i.e. "L{L+1}.hc_post_raw" call #0. Re-key to the reference's per-block hc_post / hc_post#1.
raw = {k: v for k, v in t.items() if ".hc_post_raw" in k}
for k in raw:
t.pop(k)
for li in LAYERS:
attn_post = raw.get(f"L{li}.hc_post_raw" if li == 0 else f"L{li}.hc_post_raw#1")
ffn_post = raw.get(f"L{li + 1}.hc_post_raw")
if attn_post is not None:
t[f"L{li}.hc_post"] = attn_post
if ffn_post is not None:
t[f"L{li}.hc_post#1"] = ffn_post
torch_ = __import__("torch")
torch_.save({"side": "deepinfer", "ids": ids, "t": t}, os.path.expanduser(args.out))
print(f"[probe] deepinfer: {len(t)} tensors -> {args.out}; workarounds {R.WORKAROUNDS}")
llm.close()
# ------------------------------------------------------------------------------------------------ compare
ORDER = ["embed", "mixes", "pre_mix", "post_mix", "comb_mix", "hc_pre", "attn_norm", "wq_a", "q_norm", "wq_b", "q_rope",
"wkv", "kv_rope_quant", "attn_core", "wo_b_in", "attn_out", "hc_post", "mixes#1", "pre_mix#1", "post_mix#1",
"comb_mix#1", "hc_pre#1", "ffn_norm", "gate_logits", "topk_idx", "topk_weights", "shared_w1", "shared_w3", "shared_act", "shared_out", "routed_fp32", "moe_out", "moe_out_combine_style", "hc_post#1",
"engram_out"]
def run_compare(args):
import torch
A = torch.load(os.path.expanduser(args.a), weights_only=False)["t"]
B = torch.load(os.path.expanduser(args.b), weights_only=False)["t"]
F = torch.load(os.path.expanduser(args.floor), weights_only=False)["t"] if args.floor else None
def rel(x, y):
x, y = x.float(), y.float()
if x.shape != y.shape:
return f"shape {tuple(x.shape)}!={tuple(y.shape)}"
return ((x - y).norm() / x.norm().clamp_min(1e-30)).item()
def fmt(v):
return f"{v:9.2e}" if isinstance(v, float) else f"{v:>9}"
print(f"{'op':22s} {'shape':>18s} {'BOS':>9s} {'tok1':>9s} {'all':>9s} {'floor(all)':>10s}")
def sort_topk(T):
for k in list(T):
if k.endswith("topk_idx"):
idx = T[k]
order = idx.argsort(-1)
T[k] = idx.gather(-1, order)
w = k.replace("topk_idx", "topk_weights")
if w in T:
T[w] = T[w].gather(-1, order)
for T in (A, B, F):
if T is not None:
sort_topk(T)
for li in COMPARE_LAYERS:
names = [n for n in ORDER if f"L{li}.{n}" in A]
names += sorted(k.split(".", 1)[1] for k in A if k.startswith(f"L{li}.") and k.split(".", 1)[1] not in ORDER)
for n in names:
k = f"L{li}.{n}"
if k not in B:
print(f"{k:22s} {'(missing in candidate)':>18s}")
continue
a = A[k][0] if A[k].dim() >= 3 and A[k].shape[0] == 1 else A[k]
b = B[k].reshape(a.shape) if B[k].numel() == a.numel() else B[k]
f = F[k][0] if F is not None and k in F and F[k].dim() >= 3 and F[k].shape[0] == 1 else (F[k] if F is not None and k in F else None)
cols = [rel(a[:1], b[:1]) if a.shape == b.shape else rel(a, b), rel(a[1:2], b[1:2]) if a.shape == b.shape else "", rel(a, b)]
fl = rel(a, f) if f is not None else ""
print(f"{k:22s} {str(tuple(a.shape)):>18s} " + " ".join(fmt(c) for c in cols) + f" {fmt(fl):>10s}")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("side", choices=["ref", "deepinfer", "compare"])
ap.add_argument("a", nargs="?")
ap.add_argument("b", nargs="?")
ap.add_argument("--floor", default=None)
ap.add_argument("--out", default="~/di/probe/out.pt")
ap.add_argument("--threads", type=int, default=16)
ap.add_argument("--device", default="cpu")
ap.add_argument("--fix-attn", action="store_true", help="deepinfer side: use the proposed P-bf16 attention")
ap.add_argument("--round-act-bf16", action="store_true", help="deepinfer side: bf16-round SwiGLU output before FP8 cast")
ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt")
ap.add_argument("--converted", default="~/di/tiny_ckpt_ref_mp1")
ap.add_argument("--ref-dir", default="~/di/ref_dsv41/inference")
ap.add_argument("--golden-dir", default="~/di/golden")
args = ap.parse_args()
if args.side != "compare":
os.makedirs(os.path.dirname(os.path.expanduser(args.out)), exist_ok=True)
{"ref": run_ref, "deepinfer": run_deepinfer, "compare": run_compare}[args.side](args)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink