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

Neuron bench: vllm-neuron#71's sonnet prompts, prompt/output dumps; DSpark draft ground truth

Browse Source 
- 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 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-26 03:16:54 +08:00
parent 39f8aa6eed
commit 1dce2c4995
3 changed files with 340 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
examples/neuron/bench_neuron.py
View File

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

254
examples/neuron/reference_dspark.py Normal file
View File

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

14
examples/neuron/run_deepinfer_tiny.py
View File

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