From 1dce2c49950478000433879d540f02dc82a0e580 Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 26 Oct 2026 03:16:54 +0800 Subject: [PATCH] Neuron bench: vllm-neuron#71's sonnet prompts, prompt/output dumps; DSpark draft ground truth - bench_neuron.py --dataset vllm-sonnet: vLLM's SonnetDataset as #71's bench builds it for the completions backend ("Pick as many lines..." + random poem lines, chat-formatted; --chat_tokenizer for a template the checkpoint lacks). --dataset sonnet: a contiguous text window. --dump: prompt/output token ids. - reference_dspark.py: the official inference/model.py forward_spec on the host CPU (tiny checkpoint, teacher-forced on reference_golden.py tokens): each decode step's 5 drafts and their logits. - run_deepinfer_tiny.py RUN_TINY_DRAFTS=1 prints every draft block. DeepInfer's drafts match the reference on host eager, on device and with attention TP: every divergence is a near tie (top-1/top-2 gap 0.001-0.055). With #71's prompt format the acceptance is #71's: 2.8 tokens a step at 512 (first position 65-79%); a raw text window gave 1.4 (continuations not in the context), random token ids 3.85 (repetitive outputs). Co-Authored-By: Claude Opus 5.5 --- examples/neuron/bench_neuron.py | 75 +++++++- examples/neuron/reference_dspark.py | 254 ++++++++++++++++++++++++++ examples/neuron/run_deepinfer_tiny.py | 14 ++ 3 files changed, 340 insertions(+), 3 deletions(-) create mode 100644 examples/neuron/reference_dspark.py diff --git a/examples/neuron/bench_neuron.py b/examples/neuron/bench_neuron.py index 8a750c8..d426336 100644 --- a/examples/neuron/bench_neuron.py +++ b/examples/neuron/bench_neuron.py @@ -46,6 +46,14 @@ parser.add_argument("--max_model_len", default=16384, type=int) parser.add_argument("--max_num_batched_tokens", default=128, type=int) parser.add_argument("--max_num_seqs", default=1, type=int) parser.add_argument("--nextn", default=0, type=int, help="DSpark drafts per step (0 = plain decode)") +parser.add_argument("--dataset", default="random", choices=["random", "sonnet", "vllm-sonnet"], + help="random token ids (#71's first method) or a contiguous window of real text (--dataset_path, " + "e.g. vllm-neuron#71's sonnet.txt): random prompts make greedy outputs chaotic, so the " + "acceptance changes with any change that moves bf16 rounding") +parser.add_argument("--dataset_path", default="") +parser.add_argument("--chat_tokenizer", default="", + help="vllm-sonnet: a tokenizer directory with a chat template (the checkpoint has none)") +parser.add_argument("--dump", default="", help="write measured requests' prompt / output token ids here (JSON)") parser.add_argument("--out", default="", help="write the results as JSON here (rank 0)") args = parser.parse_args() dp_size = args.ep // args.cp @@ -63,6 +71,47 @@ def _fixed_kv_pool(num_blocks: int) -> None: WorkerBase.profile_num_available_blocks = profile_num_available_blocks +_TEXT_IDS: list[int] = [] + + +def _text_prompt(length: int, seed: int, tokenizer) -> list[int]: # noqa: ANN001 + """``length`` consecutive token ids of --dataset_path, starting at a seed-chosen offset.""" + import numpy as np + + if not _TEXT_IDS: + _TEXT_IDS.extend(tokenizer.encode(Path(args.dataset_path).read_text(), add_special_tokens=False)) + assert len(_TEXT_IDS) > length, f"{args.dataset_path} has {len(_TEXT_IDS)} tokens, fewer than {length}" + start = int(np.random.default_rng(seed).integers(0, len(_TEXT_IDS) - length)) + return _TEXT_IDS[start : start + length] + + +def _vllm_sonnet_prompt(length: int, seed: int) -> list[int]: + """vLLM's SonnetDataset as `vllm bench serve --dataset-name sonnet --sonnet-prefix-len 0` builds it for the + completions backend (vllm-neuron#71's bench): "Pick as many lines as you can from these poem lines:" plus random + poem lines up to ``length`` tokens, chat-formatted.""" + import random + + from transformers import AutoTokenizer + + tok = _CHAT_TOK.get("tok") + if tok is None: + tok = _CHAT_TOK["tok"] = AutoTokenizer.from_pretrained(args.chat_tokenizer) + lines = Path(args.dataset_path).read_text(encoding="utf-8").splitlines(keepends=True) + avg = sum(len(tok(line).input_ids) for line in lines) / len(lines) + base = "Pick as many lines as you can from these poem lines:\n" + fmt = lambda text: tok.apply_chat_template([{"role": "user", "content": text}], add_generation_prompt=True, tokenize=False) # noqa: E731 + base_offset = len(tok(fmt(base)).input_ids) + num_lines = max(round((length - base_offset) / avg), 1) + rng = random.Random(seed) + while True: + ids = tok(fmt(base + "".join(rng.choices(lines, k=num_lines)))).input_ids + if len(ids) <= length: + return list(ids) + + +_CHAT_TOK: dict = {} + + def _random_prompt(length: int, seed: int, vocab: int) -> list[int]: import numpy as np @@ -70,10 +119,16 @@ def _random_prompt(length: int, seed: int, vocab: int) -> list[int]: return rng.integers(1000, vocab - 1000, size=length, dtype=np.int64).astype(np.int32).tolist() +_OUT_IDS: list[int] = [] +#: --dump: every measured request's prompt and output token ids (to replay them on another engine). +_DUMP: list[dict] = [] + + def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int] | None: # noqa: ANN001 """Run the engine until it drains. With a prompt, returns (ttft_s, latency_s, output_tokens, steps_with_output) for it.""" engine = llm.llm_engine start = time.perf_counter() + _OUT_IDS.clear() if prompt is not None: llm._add_request(None, params, prompt) first = last = None @@ -85,12 +140,18 @@ def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int] steps += len(outputs) > 0 for out in outputs: tokens += 1 + if args.dump: + from deepinfer._C.engine import RequestOutput # pyright: ignore[reportMissingModuleSource] + + _OUT_IDS.append(int(RequestOutput.clone(out).token_id)) first = first if first is not None else now last = now if not unfinished: break if prompt is None: return None + if args.dump: + _DUMP.append({"prompt": list(prompt), "output": list(_OUT_IDS)}) return first - start, last - start, tokens, steps @@ -119,18 +180,24 @@ def main(dp_size: int, dp_rank: int) -> None: params = SamplingParams(temperature=0.0, max_tokens=args.output_len, ignore_eos=True) vocab = llm.llm_engine.tokenizer.vocab_size if hasattr(llm.llm_engine.tokenizer, "vocab_size") else 128000 + tokenizer = llm.llm_engine.tokenizer + + def make_prompt(length: int, seed: int) -> list[int]: + if args.dataset == "vllm-sonnet": + return _vllm_sonnet_prompt(length, seed) + return _text_prompt(length, seed, tokenizer) if args.dataset == "sonnet" else _random_prompt(length, seed, vocab) busy = True lens = [int(x) for x in args.input_lens.split(",")] for w in range(args.warmup): - _run_one(llm, _random_prompt(lens[0], 10_000 + w + 7919 * dp_rank, vocab) if busy else None, params) + _run_one(llm, make_prompt(lens[0], 10_000 + w + 7919 * dp_rank) if busy else None, params) results = [] for length in lens: # One untimed request a cell: a cell's first request pays for its own context bucket's lazy compile. - _run_one(llm, _random_prompt(length, 20_000 + length + 7919 * dp_rank, vocab) if busy else None, params) + _run_one(llm, make_prompt(length, 20_000 + length + 7919 * dp_rank) if busy else None, params) cell = [] for r in range(args.repeats): - got = _run_one(llm, _random_prompt(length, 1000 * length + r + 7919 * dp_rank, vocab) if busy else None, params) + got = _run_one(llm, make_prompt(length, 1000 * length + r + 7919 * dp_rank) if busy else None, params) if got is not None: ttft, latency, n, steps = got tpot = (latency - ttft) / max(n - 1, 1) @@ -152,6 +219,8 @@ def main(dp_size: int, dp_rank: int) -> None: print(f"| {r['input']} | {r['output']} | {r['ttft_p50_ms']:.0f} | {r['tpot_p50_ms']:.1f} | {r['decode_tok_s_per_request']:.2f} |") if args.out: Path(args.out).write_text(json.dumps(results, indent=2)) + if args.dump: + Path(args.dump).write_text(json.dumps(_DUMP)) llm.close() diff --git a/examples/neuron/reference_dspark.py b/examples/neuron/reference_dspark.py new file mode 100644 index 0000000..747406a --- /dev/null +++ b/examples/neuron/reference_dspark.py @@ -0,0 +1,254 @@ +"""DSpark draft ground truth: the OFFICIAL DeepSeek-V4.1 reference (inference/model.py forward_spec) on the host CPU. + +Teacher-forced on reference_golden.py tokens; records each decode step\'s 5 drafts. Shares the loading code of +reference_golden.py (copied below the docstring). + +Original reference_golden.py docstring follows. + +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("--golden-dir", default="~/di/golden", help="reference_golden.py outputs (tokens to feed)") + ap.add_argument("--out", default="~/di/golden_dspark") + ap.add_argument("--prompts", default="short,long") + ap.add_argument("--threads", type=int, default=16) + ap.add_argument("--max-seq-len", type=int, default=2048) + args = ap.parse_args() + for k in ("hf_ckpt", "ref_dir", "converted", "golden_dir", "out"): + setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k)))) + os.makedirs(args.out, 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)" + 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 + + 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) + model_args = ref_model.ModelArgs(**margs) + tokenizer = AutoTokenizer.from_pretrained(args.converted) + state = load_file(ckpt_file) + dropped = sorted(k for k in state if k.startswith(("vision.", "aligner.", "image_")) or k.endswith(".bias_vl")) + state = {k: v for k, v in state.items() if k not in set(dropped)} + torch.set_default_dtype(torch.bfloat16) + 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"]) + torch.manual_seed(0) + with torch.device("cpu"): + model = ref_model.Transformer(model_args, tokenizer) + missing, unexpected = model.load_state_dict(state, strict=False) + # the MTP stages share the backbone's embed / head (assigned in Transformer.__init__) + missing = [k for k in missing if not (k.startswith("mtp.") and k.endswith((".embed.weight", ".head.weight")))] + assert not missing and not unexpected, (missing, unexpected) + model.eval() + x = torch.tensor([ids], dtype=torch.long) + out_ids, logits, main_hidden = model(x, 0) + model.forward_spec(out_ids, main_hidden, 0) # seeds the DSpark window KV + pos = len(ids) + drafts, draft_logits = [], [] + # decode step k is fed golden token k (teacher forcing); its drafts follow the token it predicts + for k in range(len(gtok) - 1): + x = torch.tensor([[gtok[k]]], dtype=torch.long) + out_ids, logits, main_hidden = model(x, pos) + res = model.forward_spec(out_ids, main_hidden, pos) + pos += 1 + spec_ids, spec_logits, conf = res + drafts.append(spec_ids[0, 1:].tolist()) + draft_logits.append(spec_logits[0].float().clone()) + print(f"[dspark] {name} step {k}: fed {gtok[k]}, main argmax {int(out_ids[0])} (golden {gtok[k + 1]}), " + f"drafts {drafts[-1]}", flush=True) + torch.save({"name": name, "prompt_ids": ids, "fed": gtok[:-1], "drafts": drafts, "draft_logits": draft_logits}, + os.path.join(args.out, f"{name}.pt")) + + +if __name__ == "__main__": + main() diff --git a/examples/neuron/run_deepinfer_tiny.py b/examples/neuron/run_deepinfer_tiny.py index ff432f1..44a7af4 100644 --- a/examples/neuron/run_deepinfer_tiny.py +++ b/examples/neuron/run_deepinfer_tiny.py @@ -271,6 +271,20 @@ def run(args, rank: int = 0): rec = Recorder() install_logits_hook(rec) + if os.environ.get("RUN_TINY_DRAFTS"): + # Print every DSpark draft block (to compare with reference_dspark.py). + from deepinfer.model_executor.models.mtp.dflash import DFlash + + orig_sample = DFlash.sample_drafts + + def sample_drafts(self, hidden_states_before_norm, driver_input): # noqa: ANN001, ANN202 + pred = orig_sample(self, hidden_states_before_norm, driver_input) + print(f"[drafts] seed {driver_input.golden_tokens.tolist()} lvo {driver_input.last_valid_offset.tolist()} " + f"-> {pred.token_ids.tolist()}", flush=True) + return pred + + DFlash.sample_drafts = sample_drafts + from deepinfer.entrypoints.llm import LLM from deepinfer.sampling_params import SamplingParams