diff --git a/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py b/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py index 2a345ee..bc46105 100644 --- a/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py +++ b/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py @@ -30,6 +30,8 @@ from flash_mla import kv_cache as _kv # Rows of (q token x topk) gathered per chunk, to bound the fp32 score/kv intermediates. # Static per call, so each chunk is a static-shaped graph. _CHUNK_ELEMS = 1 << 22 +# Keys per online-softmax block (decode/sparse/head64/config.h, and the official reference kernel). +_KV_BLOCK = 64 @dataclasses.dataclass @@ -88,6 +90,10 @@ def _sparse_attn_rows( ``gather_kv(lo, hi)`` returns ``(kv[lo:hi] as bf16 [m, topk, d_qk], invalid_mask [m, topk])``. + Numerics follow the CUDA kernels (csrc/cuda_kernels/sm100/prefill/sparse/fwd/head128/phase1.cuh:238-253,324) + rather than tests/ref.py: online softmax over blocks of ``_KV_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. + Returns ``(out_fp32 [n, h_q, d_v], max_logits [n, h_q], lse [n, h_q])`` where ``lse`` is the attention log-sum-exp *before* the attention sink is merged, with ``+inf`` for rows that attend to nothing (tests/ref.py:52-58). @@ -106,30 +112,31 @@ def _sparse_attn_rows( kv = torch.where(kv != kv, torch.zeros_like(kv), kv) logits = torch.bmm(q[lo:hi].float(), kv.transpose(1, 2)) # [m, h_q, topk] - scale = sm_scale if rms_norm_scale_factor is not None: - logits = logits * (rms_norm_scale_factor[lo:hi].unsqueeze(-1) * scale) + logits = logits * (rms_norm_scale_factor[lo:hi].unsqueeze(-1) * sm_scale) else: - logits = logits * scale - neg_inf = torch.full_like(logits, float("-inf")) - logits = torch.where(invalid.unsqueeze(1), neg_inf, logits) - - lse = torch.logsumexp(logits, dim=-1) # [m, h_q] + logits = logits * sm_scale + logits = torch.where(invalid.unsqueeze(1), torch.full_like(logits, float("-inf")), logits) max_logits[lo:hi] = logits.max(dim=-1).values - lse_for_o = lse - if attn_sink is not None: - lse_for_o = torch.logsumexp( - torch.stack([lse, attn_sink.float().view(1, h_q).expand_as(lse)], dim=0), dim=0 - ) - # So that the corresponding O becomes 0 (tests/ref.py:48). - lse_for_o = torch.where(lse_for_o == float("-inf"), torch.full_like(lse_for_o, float("inf")), lse_for_o) - probs = torch.exp(logits - lse_for_o.unsqueeze(-1)) - chunk_out = torch.bmm(probs, kv[..., :d_v]) + 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, _KV_BLOCK): + s_blk = logits[..., t0 : t0 + _KV_BLOCK] + m_new = torch.maximum(m, s_blk.amax(-1)) + rescale = torch.exp(m - m_new) + p = torch.exp(s_blk - m_new.unsqueeze(-1)) + l = l * rescale + p.sum(-1) + acc = acc * rescale.unsqueeze(-1) + torch.bmm(p.to(torch.bfloat16).float(), kv[:, t0 : t0 + _KV_BLOCK, :d_v]) + m = m_new - lonely = (lse == float("-inf")).unsqueeze(-1) - out[lo:hi] = torch.where(lonely, torch.zeros_like(chunk_out), chunk_out) - lse_out[lo:hi] = torch.where(lse == float("-inf"), torch.full_like(lse, float("inf")), lse) + # +inf sink -> zero output, -inf sink -> no-op; rows that attend to nothing -> 0 (tests/ref.py:48). + denom = l + torch.exp(attn_sink.float().view(1, h_q) - m) if attn_sink is not None else l + 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) + lse_out[lo:hi] = torch.where(lonely, torch.full_like(l, float("inf")), m + torch.log(l)) return out, max_logits, lse_out diff --git a/examples/neuron/README.md b/examples/neuron/README.md index 2d93f79..7238907 100644 --- a/examples/neuron/README.md +++ b/examples/neuron/README.md @@ -30,8 +30,20 @@ under a different CPU thread count (FP8 activation quantization amplifies sub-ul | 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) | +| short (11) | 9/9 | 0.054 (0.032) | 0.9985 (0.9995) | +| medium (172, wraps the 128 SWA window) | 8/9 * | 0.041 (0.039) | 0.9992 (0.9993) | +| long (715, candidate-block filter active) | 9/9 | 0.059 (0.060) | 0.9983 (0.9982) | -Per-step / per-layer numbers: `golden/deepinfer_cpu_vs_golden.json`. +\* medium step 6 has a top-1/top-2 margin of 0.0096; the reference flips it against itself too. + +`probe_layer0.py` compares ~35 named ops per layer between the reference and DeepInfer. Layer 0 is +bit-exact once three rounding points are aligned: + +- attention P rounded to bf16 before P@V (online softmax, 64-key blocks) — the CUDA FlashMLA kernels and + the reference do this; the Neuron flash_mla shim now does too (was fp32, a port bug); +- SwiGLU output quantized to FP8 straight from fp32 (the reference rounds to bf16 first) and +- the MoE combine rounding the routed sum to bf16 before adding the shared expert (the reference adds in + fp32) — both are DeepInfer's own GPU/Ascend behaviour, kept as is. + +The remaining logits difference comes from those two, amplified by FP8 activation quantization layer by +layer. Per-step / per-layer numbers before the attention fix: `golden/deepinfer_cpu_vs_golden.json`. diff --git a/examples/neuron/probe_layer0.py b/examples/neuron/probe_layer0.py new file mode 100644 index 0000000..3a8a493 --- /dev/null +++ b/examples/neuron/probe_layer0.py @@ -0,0 +1,448 @@ +"""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]))) + + 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() diff --git a/examples/neuron/run_deepinfer_tiny.py b/examples/neuron/run_deepinfer_tiny.py index a3cbe88..1e3ce86 100644 --- a/examples/neuron/run_deepinfer_tiny.py +++ b/examples/neuron/run_deepinfer_tiny.py @@ -16,8 +16,8 @@ logits [steps+1, vocab] fp32, optional hidden), ready for compare_golden.py. 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. +NOTE-3 (fixed KV pool instead of the host-RAM-sized profile) -- the other bring-up bugs are fixed in-repo; the list +of workarounds that actually triggered is printed at the end. Set RUN_TINY_DEBUG=1 to print every sampler call. """ from __future__ import annotations @@ -45,6 +45,7 @@ def parse_args(): 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") + ap.add_argument("--attn-p-bf16", action="store_true", help="experiment: patch the flash_mla shim with the proposed bf16-P fix") return ap.parse_args() @@ -70,81 +71,10 @@ def setup_env(device: str) -> None: WORKAROUNDS: list[str] = [] -def install_workarounds(num_kv_blocks: int) -> None: +def install_workarounds(num_kv_blocks: int, attn_p_bf16: bool = False) -> 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 + # BUG-1 (engram gather_sf on gloo), BUG-2 (get_model is_cuda), BUG-4 (paged MQA logits page ids) and NOTE-5 + # (bind_cpu) were fixed in-repo on 2026-10-20; their workarounds are gone. Only NOTE-3 remains. # 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 @@ -161,6 +91,14 @@ def install_workarounds(num_kv_blocks: int) -> None: WorkerBase.profile_num_available_blocks = profile_num_available_blocks + # EXPERIMENT (proposed fix, see probe_layer0.py): FlashMLA-style P rounding to bf16 in the Neuron flash_mla shim + if attn_p_bf16: + import flash_mla.flash_mla_interface as fmi + from probe_layer0 import sparse_attn_rows_p_bf16 + + fmi._sparse_attn_rows = sparse_attn_rows_p_bf16 + WORKAROUNDS.append("EXPERIMENT: flash_mla _sparse_attn_rows with bf16 P (proposed fix)") + class Recorder: """Collects the lm-head logits of every sampler call (one request in flight -> one row per call).""" @@ -269,7 +207,7 @@ def main(): 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) + install_workarounds(args.num_kv_blocks, args.attn_p_bf16) rec = Recorder() install_logits_hook(rec)