DeepInfer/examples/neuron/bench_neuron.py

"""TTFT / TPOT benchmark of the full DeepSeek-V4.1-Flash on one trn2.48xlarge, the vllm-neuron#71 method offline.

vllm-neuron#71 (``contrib/models/deepseek_v41/test/performance/run_bench.py``) serves the model and drives it with
``vllm bench serve --dataset-name random``: random token ids fix the input length exactly and share no prefix,
output 128 tokens, and per cell it reports the median TTFT, the median TPOT (decode only:
``(latency - ttft) / (output_tokens - 1)``) and the decode tokens per second. DeepInfer has no server, so this drives
the engine directly with the same requests: random prompt token ids of the cell's input length, ``ignore_eos`` and
``max_tokens`` 128, timed per request from submission.

Topology: one engine process per DP rank (EP=64, DP-attention), and every rank runs its own request of the cell at
the same time: batch 1 per rank, 64 requests in flight (the engine steps all DP ranks together). #71's bs1 cells are
one request for the whole machine; per request the two loads cost the same here, since the static MoE computes
every rank's row either way.

    export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0 DEEPINFER_NEURON_PREFILL_BUCKET=128 \\
        DEEPINFER_MAX_BACKTRACKING_TOKENS=128
    python examples/neuron/bench_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from pathlib import Path

EXAMPLES = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(EXAMPLES))

from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices  # noqa: E402
from _launch import launch_dp  # noqa: E402

parser = argparse.ArgumentParser(description="V4.1 TTFT/TPOT benchmark on Neuron (vllm-neuron#71 method)")
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
parser.add_argument("--ep", default=64, type=int)
parser.add_argument("--cp", default=1, type=int, help="TP group size: attention heads sharded over cp ranks running one batch (issue #28)")
parser.add_argument("--input_lens", default="512,1024,2048,8192", help="comma-separated prompt lengths (tokens)")
parser.add_argument("--output_len", default=128, type=int)
parser.add_argument("--repeats", default=3, type=int, help="requests per cell (medians are reported)")
parser.add_argument("--warmup", default=1, type=int, help="untimed requests before the first cell")
parser.add_argument("--num_kv_blocks", default=400, type=int)
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


def _fixed_kv_pool(num_blocks: int) -> None:
    os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_blocks)  # CP followers read it too
    from deepinfer.worker.base import WorkerBase

    def profile_num_available_blocks(self):  # noqa: ANN001, ANN202
        self.model_runner.set_block_size(self.cache_config.block_size)
        self.cache_config.num_gpu_blocks = num_blocks
        return num_blocks

    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

    rng = np.random.default_rng(seed)
    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
    tokens = steps = 0
    while True:
        unfinished = engine.has_unfinished_requests()
        outputs = engine.step() if unfinished else engine.flush_inflight_reqs()
        now = time.perf_counter()
        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


def main(dp_size: int, dp_rank: int) -> None:
    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = str(args.nextn)
    os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
    setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
    from deep_cluster.hardware import get_actual_num_gpus

    setup_cuda_visible_devices(dp_rank, args.cp, get_actual_num_gpus())
    _fixed_kv_pool(args.num_kv_blocks)
    llm = create_llm(
        args.model_name,
        dp_rank,
        ep_size=args.ep,
        cp_size=args.cp,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": "100"},
        data_parallel_size=dp_size,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_model_len=args.max_model_len,
        max_num_seqs=args.max_num_seqs,
        gpu_memory_utilization=0.85,
        lookahead_tokens=1 if args.nextn else 0,
    )
    from deepinfer.sampling_params import SamplingParams

    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, 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, 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, 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)
                accept = (n - 1) / max(steps - 1, 1)  # tokens per decode step (the acceptance length with DSpark)
                cell.append((ttft, tpot, n, accept))
                if dp_rank == 0:
                    print(f"BENCH input={length} rep={r} ttft_ms={ttft * 1e3:.1f} tpot_ms={tpot * 1e3:.1f} tokens={n} "
                          f"tokens_per_step={accept:.2f}", flush=True)
        if dp_rank == 0 and cell:
            ttft = statistics.median(c[0] for c in cell)
            tpot = statistics.median(c[1] for c in cell)
            results.append({"input": length, "output": args.output_len, "ttft_p50_ms": ttft * 1e3, "tpot_p50_ms": tpot * 1e3,
                            "decode_tok_s_per_request": 1.0 / tpot,
                            "tokens_per_step": statistics.median(c[3] for c in cell)})
    if dp_rank == 0:
        print("BENCH_RESULTS " + json.dumps(results), flush=True)
        print(f"| input | output | TTFT p50 (ms) | TPOT p50 (ms) | decode tok/s per request |   (batch 1 per DP rank)")
        for r in results:
            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()


if __name__ == "__main__":
    launch_dp(main, dp_size)