From d93fd4909f05a0ad8e5d7fbdaf2caa2a6e95ec3b Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 17 Aug 2026 15:39:58 -0500 Subject: [PATCH 1/5] [TRTLLM-14818][test] Port Kimi K3 DFlash/DSpark eval helpers and KDA FP8 prefill test to main These files already have consumers on main that reference them but were left behind in the feat/kimi_k3 -> main mergeback, leaving dangling paths: - make_synthetic_dflash_drafter.py: referenced by run_eval_kimi_k3.sbatch, a modeling_kimi_linear.py comment, and test_kimi_k3_dflash_scaffold.py (which importlib-loads it and currently skips its schema tests when the file is absent; this activates them). - measure_dspark_acceptance.py: exec'd by run_dspark_acceptance.sbatch. - eval_extra_llm_options_dflash.yaml: selected by run_eval_kimi_k3.sbatch --dflash. Also ports the standalone KDA FP8 packed-prefill parity unit test (test_kimi_kda_fp8_packed_prefill.py). It is skip-guarded (pytest.importorskip("fla") + skipif for SM100/SM103), so it collects cleanly and is not wired into any L0 list in this change. No production code or L0 test-list changes. Signed-off-by: Brian Nguyen --- .../eval_extra_llm_options_dflash.yaml | 39 ++ .../kimi_k3/make_synthetic_dflash_drafter.py | 363 ++++++++++++++++++ examples/kimi_k3/measure_dspark_acceptance.py | 308 +++++++++++++++ .../test_kimi_kda_fp8_packed_prefill.py | 202 ++++++++++ 4 files changed, 912 insertions(+) create mode 100644 examples/kimi_k3/eval_extra_llm_options_dflash.yaml create mode 100644 examples/kimi_k3/make_synthetic_dflash_drafter.py create mode 100644 examples/kimi_k3/measure_dspark_acceptance.py create mode 100644 tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py diff --git a/examples/kimi_k3/eval_extra_llm_options_dflash.yaml b/examples/kimi_k3/eval_extra_llm_options_dflash.yaml new file mode 100644 index 000000000000..50653723c930 --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_dflash.yaml @@ -0,0 +1,39 @@ +# Extra LLM-API options for Kimi K3 GSM8K/eval runs WITH DFlash speculative +# decoding. STATUS: the real K3 drafter is a DSpark drafter (schema known +# from the training team's dummy-dspark0724 checkpoint; weights still +# training). The drafter-forward dspark semantics (Markov intra-block +# bias, SWA, shift_label) ARE implemented; confidence-scheduled +# verification is the remaining gap (weights loaded, unused — see +# docs/kimi_k3_dflash_scaffold.md). Point speculative_model at the real +# drafter when it drops, or at the output of +# make_synthetic_dflash_drafter.py to exercise the wiring only (random +# weights draft gibberish; acceptance ~0, so that validates plumbing, not +# speedup). Deployment mirrors the SA config +# (DEP16: attention-DP + MoE EP16, eager, max_batch_size <= 8) since that +# is the memory-viable GB300 deployment; DFlash x attention-DP parity has +# NOT been certified the way SA was — rerun the parity harness before +# trusting outputs. +enable_attention_dp: true +moe_expert_parallel_size: 16 +disable_overlap_scheduler: true +# Eager like certified SA: the K3 CUDA-graphs regime has a known +# verify/accept parity bug under investigation (SA+graphs probe); job +# 2659946 accidentally ran DFlash with graphs on and scored 90.11 vs the +# >=96.4 baseline. Keep null until the graphs regime is certified. +cuda_graph_config: null +enable_chunked_prefill: false +kv_cache_config: + enable_block_reuse: false + # 0.20 (vs SA's 0.25): the drafter weights, hidden-state capture buffer, + # and DFlash context-KV slots eat into the SA config's already-thin + # warmup headroom on GB300 (0.25 OOMs there). + free_gpu_memory_fraction: 0.20 + tokens_per_block: 64 +speculative_config: + decoding_type: DFlash + # block_size in the drafter config governs tokens drafted per step; the + # K3 dspark drafter uses 8 (same as K2.7). + max_draft_len: 7 + # Path to the DFlash drafter checkpoint (real, or synthetic from + # make_synthetic_dflash_drafter.py). + speculative_model: /path/to/kimi-k3-dflash-drafter diff --git a/examples/kimi_k3/make_synthetic_dflash_drafter.py b/examples/kimi_k3/make_synthetic_dflash_drafter.py new file mode 100644 index 000000000000..884f0f4c0fce --- /dev/null +++ b/examples/kimi_k3/make_synthetic_dflash_drafter.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Emit a synthetic (random-weight) Kimi K3 DFlash/DSpark drafter checkpoint. + +The real K3 drafter (training in progress) is a DSpark drafter — DeepSeek's +DFlash follow-up (arXiv 2607.05147): a dense Qwen3-style parallel block +backbone (q/k-norm attention, SiLU MLP) plus the DFlash pooling projection, +extended with a low-rank Markov head (token-conditioned intra-block logit +bias) and a confidence head (per-position acceptance prediction). This +generator materializes that schema — verified against the training team's +dummy-weight checkpoint (dummy-dspark0724, 73 tensors) — so the structural +wiring (DFlashForCausalLM construction via the model_type=qwen3 fallback, +load_weights key remapping, target-layer hidden-state capture in +KimiLinearModel.forward) can be exercised end-to-end before real weights +drop. Outputs are gibberish by construction. + +Checkpoint contents (no embed_tokens / lm_head; shared with the target): + +* fc.weight [hidden, hidden * len(target_layer_ids)] +* hidden_norm.weight [hidden] +* layers.{i}.self_attn.{q,k,v,o}_proj.weight, {q,k}_norm.weight +* layers.{i}.mlp.{gate,up,down}_proj.weight +* layers.{i}.{input,post_attention}_layernorm.weight +* norm.weight [hidden] +* markov_w1.weight, markov_w2.weight [vocab, markov_rank] (dspark) +* confidence_proj.weight [1, hidden + markov_rank] (dspark) +* confidence_proj.bias [1] (dspark) + +Three modes: + +* --config: adopt a REAL drafter config.json verbatim (authoritative). +* --ckpt-dir: read hidden_size / vocab_size / num_hidden_layers from a + REAL K3 target checkpoint's config.json; drafter dims default to the + dummy-dspark0724 drafter's. +* --tiny: minimal dims, no checkpoint access — for unit tests. + +target_layer_ids defaults to len-6 even spacing over the target stack — +confirmed by the real config: [1, 19, 37, 54, 72, 90] over K3's 93 layers +(same convention as K2.7's [1, 12, 24, 35, 47, 58] over 61). + +Usage: + python make_synthetic_dflash_drafter.py --config --out + python make_synthetic_dflash_drafter.py --ckpt-dir --out + python make_synthetic_dflash_drafter.py --tiny --out +""" + +import argparse +import json +import os + +# torch/safetensors are imported lazily in main() so the schema helpers +# (drafter_tensor_plan, drafter_config, even_target_layer_ids) stay +# importable on hosts without the container venv. + +# nvidia/Kimi-K2.7-Code-DFlash dims (plain DFlash, no dspark heads): kept +# as the schema-compat reference the unit tests pin against. +K27_DRAFTER = dict( + num_hidden_layers=6, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + intermediate_size=18432, + block_size=8, +) + +# Drafter defaults = the K3 DSpark drafter (dummy-dspark0724 config from +# the training team, 2026-07-24). +K3_DRAFTER = dict( + num_hidden_layers=6, + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + intermediate_size=12288, + block_size=8, + markov_rank=256, + use_confidence_head=True, + swa_window_size=1024, +) + +TINY = dict( + hidden_size=64, + vocab_size=512, + num_target_layers=8, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + intermediate_size=128, + block_size=4, + markov_rank=8, + use_confidence_head=True, + swa_window_size=32, +) + + +def even_target_layer_ids(num_target_layers: int, k: int = 6): + """Evenly spaced capture layers, K2.7 convention (first=1, last=L-3).""" + lo, hi = 1, max(1, num_target_layers - 3) + if k == 1: + return [lo] + return sorted({round(lo + i * (hi - lo) / (k - 1)) for i in range(k)}) + + +def drafter_tensor_plan( + hidden, cfg, num_capture, vocab=None, markov_rank=None, use_confidence_head=False +): + """Return {key: shape} for the drafter checkpoint. + + Base keys follow the K2.7 DFlash schema; the dspark heads (markov_w1/w2 + and confidence_proj) are added when requested. The confidence head reads + the concatenation [hidden, markov_features], hence its in-dim. + """ + heads, kv = cfg["num_attention_heads"], cfg["num_key_value_heads"] + hd, inter = cfg["head_dim"], cfg["intermediate_size"] + plan = { + "fc.weight": (hidden, hidden * num_capture), + "hidden_norm.weight": (hidden,), + "norm.weight": (hidden,), + } + if markov_rank: + assert vocab is not None, "markov head tensors need vocab_size" + plan["markov_w1.weight"] = (vocab, markov_rank) + plan["markov_w2.weight"] = (vocab, markov_rank) + if use_confidence_head: + plan["confidence_proj.weight"] = (1, hidden + (markov_rank or 0)) + plan["confidence_proj.bias"] = (1,) + for i in range(cfg["num_hidden_layers"]): + p = f"layers.{i}." + plan.update( + { + p + "self_attn.q_proj.weight": (heads * hd, hidden), + p + "self_attn.k_proj.weight": (kv * hd, hidden), + p + "self_attn.v_proj.weight": (kv * hd, hidden), + p + "self_attn.o_proj.weight": (hidden, heads * hd), + p + "self_attn.q_norm.weight": (hd,), + p + "self_attn.k_norm.weight": (hd,), + p + "mlp.gate_proj.weight": (inter, hidden), + p + "mlp.up_proj.weight": (inter, hidden), + p + "mlp.down_proj.weight": (hidden, inter), + p + "input_layernorm.weight": (hidden,), + p + "post_attention_layernorm.weight": (hidden,), + } + ) + return plan + + +def drafter_config(hidden, vocab, num_target_layers, target_layer_ids, mask_token_id, cfg): + dflash_cfg = { + "mask_token_id": mask_token_id, + "target_layer_ids": target_layer_ids, + } + if cfg.get("markov_rank"): + # DSpark extras, mirroring the dummy-dspark0724 config. + dflash_cfg.update( + { + "use_swa": True, + "swa_window_size": cfg["swa_window_size"], + "causal": False, + "projector_type": "dspark", + "shift_label": True, + "markov_rank": cfg["markov_rank"], + "markov_head_type": "vanilla", + "use_confidence_head": cfg.get("use_confidence_head", False), + } + ) + is_dspark = bool(cfg.get("markov_rank")) + return { + "architectures": ["DFlashDraftModel"], + # model_type drives DFlashForCausalLM's backbone fallback + # (Qwen3ForCausalLM), matching the K2.7 and K3 drafters. + "model_type": "qwen3", + "block_size": cfg["block_size"], + "dflash_config": dflash_cfg, + "hidden_size": hidden, + "num_hidden_layers": cfg["num_hidden_layers"], + "num_attention_heads": cfg["num_attention_heads"], + "num_key_value_heads": cfg["num_key_value_heads"], + "head_dim": cfg["head_dim"], + "intermediate_size": cfg["intermediate_size"], + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "vocab_size": vocab, + "max_position_embeddings": 1048576 if is_dspark else 262144, + "initializer_range": 0.02, + "attention_bias": False, + "attention_dropout": 0.0, + # RoPE per the dummy-dspark0724 config: plain 1e4 theta, no scaling. + "rope_theta": 10000.0, + "rope_scaling": None, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "num_target_layers": num_target_layers, + "layer_types": (["sliding_attention"] if is_dspark else ["full_attention"]) + * cfg["num_hidden_layers"], + **({"sliding_window": cfg["swa_window_size"]} if is_dspark else {}), + "synthetic_random_weights": True, + } + + +def target_dims_from_ckpt(ckpt_dir): + with open(os.path.join(ckpt_dir, "config.json")) as f: + cfg = json.load(f) + text = cfg.get("text_config", cfg) + return (text["hidden_size"], text["vocab_size"], text["num_hidden_layers"]) + + +def drafter_cfg_from_real_config(path): + """--config mode: adopt the REAL drafter config.json verbatim. + + Random weights, exact real module structure — no schema guessing. + Returns (real_cfg_dict, drafter_dims_dict). Errors clearly on fields + the tensor plan needs; anything else in the file is passed through + untouched so TRT-LLM sees exactly what the trained checkpoint will + ship. + """ + with open(path) as f: + real = json.load(f) + required = ( + "hidden_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "intermediate_size", + "vocab_size", + ) + missing = [k for k in required if k not in real] + if missing: + raise KeyError( + f"real drafter config is missing {missing}; " + "ask the training team for the full HF config.json" + ) + dflash_cfg = real.get("dflash_config") or {} + if "target_layer_ids" not in dflash_cfg: + raise KeyError( + "real drafter config has no " + "dflash_config.target_layer_ids — the capture wiring " + "cannot be derived without it" + ) + dims = dict( + num_hidden_layers=real["num_hidden_layers"], + num_attention_heads=real["num_attention_heads"], + num_key_value_heads=real["num_key_value_heads"], + head_dim=real.get("head_dim", real["hidden_size"] // real["num_attention_heads"]), + intermediate_size=real["intermediate_size"], + block_size=real.get("block_size", 8), + # DSpark heads: emitted only if the real config declares them. + markov_rank=dflash_cfg.get("markov_rank"), + use_confidence_head=dflash_cfg.get("use_confidence_head", False), + ) + archs = real.get("architectures", []) + if not any("Laguna" in a for a in archs) and real.get("model_type") not in ("qwen3", "llama"): + print( + f"WARNING: architectures={archs} model_type=" + f"{real.get('model_type')} may not resolve through the " + "generic DFlashForCausalLM fallback — drafter-side code " + "changes may be needed. Generating anyway." + ) + return real, dims + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + mode = ap.add_mutually_exclusive_group(required=True) + mode.add_argument("--ckpt-dir", help="real K3 target checkpoint dir (reads config.json)") + mode.add_argument( + "--config", + help="REAL drafter config.json from the training team: " + "adopt it verbatim and emit random weights matching " + "its exact module structure (no schema guessing)", + ) + mode.add_argument( + "--tiny", action="store_true", help="minimal dims for unit tests; no checkpoint access" + ) + ap.add_argument("--out", required=True, help="output drafter dir") + ap.add_argument( + "--target-layer-ids", + type=int, + nargs="+", + default=None, + help="override capture layers (default: even spacing)", + ) + ap.add_argument( + "--mask-token-id", + type=int, + default=None, + help="default: the real config's value in --config mode, " + "else vocab_size - 2 (NB: the real K3 drafter uses " + "163606, NOT vocab-2)", + ) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + import torch + from safetensors.torch import save_file + + real_cfg = None + if args.config: + real_cfg, cfg = drafter_cfg_from_real_config(args.config) + hidden, vocab = real_cfg["hidden_size"], real_cfg["vocab_size"] + target_layer_ids = args.target_layer_ids or real_cfg["dflash_config"]["target_layer_ids"] + mask_token_id = ( + args.mask_token_id + if args.mask_token_id is not None + else real_cfg["dflash_config"].get("mask_token_id", vocab - 2) + ) + else: + if args.tiny: + cfg = dict(TINY) + hidden, vocab = cfg["hidden_size"], cfg["vocab_size"] + num_target_layers = cfg["num_target_layers"] + default_k = 2 + else: + hidden, vocab, num_target_layers = target_dims_from_ckpt(args.ckpt_dir) + cfg = dict(K3_DRAFTER) + default_k = 6 + target_layer_ids = args.target_layer_ids or even_target_layer_ids( + num_target_layers, default_k + ) + assert all(0 <= t < num_target_layers for t in target_layer_ids), ( + f"target_layer_ids {target_layer_ids} out of range [0, {num_target_layers})" + ) + mask_token_id = args.mask_token_id if args.mask_token_id is not None else vocab - 2 + + torch.manual_seed(args.seed) + plan = drafter_tensor_plan( + hidden, + cfg, + len(target_layer_ids), + vocab=vocab, + markov_rank=cfg.get("markov_rank"), + use_confidence_head=cfg.get("use_confidence_head", False), + ) + tensors = { + k: (torch.randn(s, dtype=torch.float32) * 0.02).to(torch.bfloat16) for k, s in plan.items() + } + + if real_cfg is not None: + out_cfg = dict(real_cfg) + out_cfg["dflash_config"] = { + **real_cfg.get("dflash_config", {}), + "mask_token_id": mask_token_id, + "target_layer_ids": target_layer_ids, + } + out_cfg["synthetic_random_weights"] = True + else: + out_cfg = drafter_config( + hidden, vocab, num_target_layers, target_layer_ids, mask_token_id, cfg + ) + + os.makedirs(args.out, exist_ok=True) + save_file(tensors, os.path.join(args.out, "model.safetensors")) + with open(os.path.join(args.out, "config.json"), "w") as f: + json.dump(out_cfg, f, indent=2) + total = sum(t.numel() for t in tensors.values()) + print( + f"wrote {len(tensors)} tensors ({total / 1e6:.1f}M params) to " + f"{args.out} (target_layer_ids={target_layer_ids}, " + f"mask_token_id={mask_token_id}) — SYNTHETIC RANDOM WEIGHTS" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k3/measure_dspark_acceptance.py b/examples/kimi_k3/measure_dspark_acceptance.py new file mode 100644 index 000000000000..a55d1b16b862 --- /dev/null +++ b/examples/kimi_k3/measure_dspark_acceptance.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Kimi K3 DSpark acceptance / speedup readiness harness. + +Drives generation over a GSM8K prompt subset with the DSpark (DFlash) +drafter and reports the weights-drop-day figures of merit: + +- AL: mean accepted tokens per target verify step (bonus token included), + per request (from RequestOutput.avg_decoded_tokens_per_iter) and + aggregate (from the runtime accept-site histogram), +- per-position acceptance-rate curve (position 1..max_draft_len), +- accepted-draft-count histogram, +- TPOT / decode-throughput numbers for the no-spec A/B (run once with + --spec-off, once with --drafter; see run_dspark_acceptance.sbatch which + fires both in one submission), +- confidence-vs-acceptance calibration table, for calibrating the DSpark + confidence_threshold and Sequential Temperature Scaling. Rows are only + collected when the confidence provider + (``dspark_confidence``) is in the tree and the drafter ships a + confidence head; otherwise "confidence_calibration" is null. For + UNBIASED calibration leave --confidence-threshold unset (0): trimmed + positions are forced rejects. + +The per-position / calibration counters come from the opt-in recorder in +tensorrt_llm/_torch/speculative/accept_stats.py, enabled here by setting +TLLM_DFLASH_ACCEPT_STATS_DIR before the LLM is built (attention-DP ranks +hold disjoint request sets, so all per-rank files are merged). The +recorder syncs a few scalars per step: keep it OFF (--no-accept-stats) +for the TPOT-reference legs of an A/B. + +Example (inside the container, see run_dspark_acceptance.sbatch): + trtllm-llmapi-launch python3 examples/kimi_k3/measure_dspark_acceptance.py \ + --model /path/to/kimi-k3 --drafter /path/to/dspark-drafter \ + --tp-size 16 --num-prompts 64 --max-tokens 256 \ + --stats-dir /tmp/dspark-stats --output-json results_spec.json +""" + +import argparse +import json +import os +import time + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Path to the Kimi K3 target checkpoint") + parser.add_argument( + "--drafter", + default=None, + help="Path to the DSpark/DFlash drafter checkpoint. " + "Omit (or pass --spec-off) for the no-spec reference.", + ) + parser.add_argument( + "--spec-off", action="store_true", help="Force speculation off (no-spec TPOT reference)." + ) + parser.add_argument("--tp-size", type=int, default=16) + parser.add_argument("--max-batch-size", type=int, default=8) + parser.add_argument("--max-seq-len", type=int, default=8192) + parser.add_argument("--max-num-tokens", type=int, default=4096) + parser.add_argument("--kv-frac", type=float, default=0.20) + parser.add_argument( + "--max-draft-len", + type=int, + default=7, + help="K = drafter block_size - 1 (K3 dspark: 8-1=7).", + ) + parser.add_argument( + "--confidence-threshold", + type=float, + default=None, + help="Enable confidence-scheduled verification (leave unset for calibration runs).", + ) + parser.add_argument( + "--confidence-policy", default="first_below", choices=["first_below", "cumulative"] + ) + parser.add_argument( + "--num-prompts", type=int, default=64, help="GSM8K test-split prompts to run." + ) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument( + "--prompt-file", + default=None, + help="Optional JSONL with a 'prompt' field per line; overrides the GSM8K subset.", + ) + parser.add_argument( + "--stats-dir", + default=None, + help="Directory for the runtime accept-stats JSONs (default: alongside --output-json).", + ) + parser.add_argument( + "--no-accept-stats", + action="store_true", + help="Do not enable the runtime recorder (use for " + "TPOT-reference legs; per-position stats are lost).", + ) + parser.add_argument("--output-json", default="dspark_acceptance.json") + return parser.parse_args() + + +def load_prompts(args) -> list: + if args.prompt_file: + prompts = [] + with open(args.prompt_file) as f: + for line in f: + line = line.strip() + if line: + prompts.append(json.loads(line)["prompt"]) + return prompts[: args.num_prompts] + from datasets import load_dataset + + ds = load_dataset("openai/gsm8k", "main", split="test") + return [ + f"Question: {row['question']}\nAnswer:" + for row in ds.select(range(min(args.num_prompts, len(ds)))) + ] + + +def build_llm(args): + from tensorrt_llm import LLM + from tensorrt_llm.llmapi import KvCacheConfig + + spec_on = args.drafter is not None and not args.spec_off + llm_kwargs = dict( + model=args.model, + tensor_parallel_size=args.tp_size, + enable_attention_dp=True, + moe_expert_parallel_size=args.tp_size, + trust_remote_code=True, + max_batch_size=args.max_batch_size, + max_seq_len=args.max_seq_len, + max_num_tokens=args.max_num_tokens, + # Mirror eval_extra_llm_options_dflash.yaml: eager (K3 CUDA-graphs + # verify/accept regime not certified; the recorder also skips + # graph batches), no overlap scheduler, no chunked prefill. + cuda_graph_config=None, + disable_overlap_scheduler=True, + enable_chunked_prefill=False, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=args.kv_frac, + tokens_per_block=64, + ), + ) + if spec_on: + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig + + spec_kwargs = dict( + max_draft_len=args.max_draft_len, + speculative_model=args.drafter, + ) + if args.confidence_threshold is not None: + # Shipped by the DSpark confidence-scheduled verification MR; + # feature-detect so this harness also runs on trees without it + # (where only threshold-off calibration-free measurement works). + if "confidence_threshold" not in DFlashDecodingConfig.model_fields: + raise SystemExit( + "--confidence-threshold requires the DSpark " + "confidence-scheduled verification support in the tree " + "(DFlashDecodingConfig has no confidence_threshold " + "field). Leave it unset for AL/AR measurement." + ) + spec_kwargs["confidence_threshold"] = args.confidence_threshold + spec_kwargs["confidence_policy"] = args.confidence_policy + llm_kwargs["speculative_config"] = DFlashDecodingConfig(**spec_kwargs) + return LLM(**llm_kwargs), spec_on + + +def main() -> None: + args = parse_arguments() + + spec_on = args.drafter is not None and not args.spec_off + stats_dir = None + if spec_on and not args.no_accept_stats: + stats_dir = args.stats_dir or ( + os.path.splitext(os.path.abspath(args.output_json))[0] + ".stats" + ) + # Must be set before the LLM (and its worker processes) is built. + os.environ["TLLM_DFLASH_ACCEPT_STATS_DIR"] = stats_dir + + prompts = load_prompts(args) + llm, spec_on = build_llm(args) + + from tensorrt_llm import SamplingParams + + sampling_params = SamplingParams(max_tokens=args.max_tokens, temperature=0.0) + + result = { + "mode": "dspark" if spec_on else "spec_off", + "model": args.model, + "drafter": args.drafter if spec_on else None, + "confidence_threshold": args.confidence_threshold if spec_on else None, + "confidence_policy": args.confidence_policy if spec_on else None, + "num_prompts": len(prompts), + "max_tokens": args.max_tokens, + "max_batch_size": args.max_batch_size, + } + + try: + # Warmup (excluded from timing): one short batch to pay JIT and + # allocator costs so the timed section measures steady state. + llm.generate(prompts[: min(2, len(prompts))], SamplingParams(max_tokens=8, temperature=0.0)) + + t0 = time.monotonic() + outputs = llm.generate(prompts, sampling_params) + wall_s = time.monotonic() - t0 + + total_out_tokens = 0 + per_request_al = [] + for out in outputs: + n = len(out.outputs[0].token_ids) + total_out_tokens += n + al = getattr(out, "avg_decoded_tokens_per_iter", None) + if al is not None: + per_request_al.append(float(al)) + + result.update( + wall_s=wall_s, + total_output_tokens=total_out_tokens, + output_tokens_per_s=total_out_tokens / wall_s if wall_s else None, + # Batched-decode proxy for TPOT: seconds per generated token. + # A/B speedup = spec_off.tpot_proxy_ms / dspark.tpot_proxy_ms + # at identical prompt set / batch size / max_tokens. + tpot_proxy_ms=(wall_s / total_out_tokens * 1e3) if total_out_tokens else None, + per_request_al={ + "mean": (sum(per_request_al) / len(per_request_al)) if per_request_al else None, + "min": min(per_request_al) if per_request_al else None, + "max": max(per_request_al) if per_request_al else None, + "n": len(per_request_al), + }, + ) + finally: + llm.shutdown() + + if stats_dir: + from tensorrt_llm._torch.speculative.accept_stats import ( + calibration_table, + load_rank_snapshots, + merge_snapshots, + summarize_hist, + ) + + snaps = load_rank_snapshots(stats_dir) if os.path.isdir(stats_dir) else [] + if not snaps: + # Env changes made here don't reach MPI worker ranks that + # trtllm-llmapi-launch pre-spawned; multi-rank runs must export + # TLLM_DFLASH_ACCEPT_STATS_DIR in the launching shell (the + # sbatch runner does this for the dspark leg). + print( + f"WARNING: accept-stats requested but no " + f"dflash_accept_stats_rank*.json found in {stats_dir}; " + f"per-position AR / calibration unavailable. Export " + f"TLLM_DFLASH_ACCEPT_STATS_DIR before trtllm-llmapi-launch " + f"for multi-rank runs." + ) + else: + merged = merge_snapshots(snaps) + summary = summarize_hist(merged["accepted_draft_hist"]) + cc = merged["confidence_calibration"] + has_calib = any(any(row) for row in cc["attempts"]) + result["accept_stats"] = { + "stats_dir": stats_dir, + "num_rank_files": len(snaps), + "accepted_draft_hist": merged["accepted_draft_hist"], + **summary, + "confidence_calibration": calibration_table(cc["attempts"], cc["accepted"]) + if has_calib + else None, + } + + with open(args.output_json, "w") as f: + json.dump(result, f, indent=2) + + print("=" * 72) + print( + f"mode={result['mode']} prompts={result['num_prompts']} " + f"wall={result.get('wall_s', 0):.1f}s " + f"out_tokens={result.get('total_output_tokens')} " + f"tok/s={result.get('output_tokens_per_s') or 0:.1f} " + f"tpot_proxy={result.get('tpot_proxy_ms') or 0:.2f} ms" + ) + if result.get("per_request_al", {}).get("mean") is not None: + al = result["per_request_al"] + print( + f"AL (per-request avg_decoded_tokens_per_iter): " + f"mean={al['mean']:.3f} min={al['min']:.3f} max={al['max']:.3f}" + ) + stats = result.get("accept_stats") + if stats: + print( + f"AL (accept-site, {stats['num_steps']} steps): " + f"{stats['al']:.3f} hist={stats['accepted_draft_hist']}" + ) + curve = " ".join(f"p{k + 1}={v:.3f}" for k, v in enumerate(stats["ar_per_position"])) + print(f"AR per position: {curve}") + if stats["confidence_calibration"]: + print( + "confidence calibration: per-position ECE = " + + " ".join( + f"p{k + 1}={p['ece']:.3f}" if p["ece"] is not None else f"p{k + 1}=n/a" + for k, p in enumerate(stats["confidence_calibration"]["per_position"]) + ) + ) + print(f"results written to {os.path.abspath(args.output_json)}") + print("=" * 72) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py new file mode 100644 index 000000000000..302de7781399 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime parity for the Kimi K3 FP8 packed q/k/v prefill projection.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +pytest.importorskip("fla") + +from tensorrt_llm._torch.models.modeling_kimi_linear import ( + KimiKDARuntime, + _convert_kda_projections_to_fp8_weight_read, +) + + +class _Cfg: + hidden_size = 768 + rms_norm_eps = 1e-6 + linear_attn_config = { + "num_heads": 6, + "head_dim": 128, + "short_conv_kernel_size": 4, + "use_full_rank_gate": True, + "gate_lower_bound": -5.0, + } + + +class _Layer(nn.Module): + def __init__(self, runtime: KimiKDARuntime) -> None: + super().__init__() + self.is_kda = True + self.self_attn = runtime + + +class _Model(nn.Module): + def __init__(self, runtime: KimiKDARuntime) -> None: + super().__init__() + self.layers = nn.ModuleList([_Layer(runtime)]) + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="Kimi K3 FP8 projection is supported only on SM100/SM103", +) + + +def _make_runtime() -> KimiKDARuntime: + runtime = KimiKDARuntime(_Cfg(), layer_idx=0).to("cuda") + assert _convert_kda_projections_to_fp8_weight_read(_Model(runtime)) == 5 + return runtime + + +def _assert_numerically_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + actual_float = actual.float().flatten() + expected_float = expected.float().flatten() + cosine = torch.nn.functional.cosine_similarity(actual_float, expected_float, dim=0).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + assert cosine > 0.999 + assert relative_l2 < 3e-2 + + +@torch.no_grad() +def test_fp8_packed_qkv_projection_matches_separate_views() -> None: + torch.manual_seed(0) + runtime = _make_runtime() + mixer = runtime.mixer + hidden = torch.randn(1, 193, _Cfg.hidden_size, device="cuda", dtype=torch.bfloat16) * 0.05 + + packed = mixer.qkvg_proj(hidden)[..., : 3 * runtime.proj_size] + actual = packed.split(runtime.proj_size, dim=-1) + expected = (mixer.q_proj(hidden), mixer.k_proj(hidden), mixer.v_proj(hidden)) + + for packed_part, separate_part in zip(actual, expected): + _assert_numerically_close(packed_part, separate_part) + + +@pytest.mark.parametrize( + "sequence_lengths,use_initial_states,has_initial_states", + [ + ([17, 31, 64], False, [False, False, False]), + ([1, 129], True, [True, False]), + ], +) +@torch.no_grad() +def test_fp8_packed_qkv_prefill_matches_separate_path_and_updates_state( + sequence_lengths, use_initial_states, has_initial_states +) -> None: + torch.manual_seed(1) + runtime = _make_runtime() + mixer = runtime.mixer + num_prefills = len(sequence_lengths) + num_tokens = sum(sequence_lengths) + slots = num_prefills + 3 + d = runtime.proj_size + h = _Cfg.linear_attn_config["num_heads"] + head_dim = _Cfg.linear_attn_config["head_dim"] + conv_size = _Cfg.linear_attn_config["short_conv_kernel_size"] + slot_indices = torch.arange(2, 2 + num_prefills, device="cuda", dtype=torch.long) + cu_seqlens = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], device="cuda", dtype=torch.long + ) + metadata = SimpleNamespace( + use_initial_states=use_initial_states, + has_initial_states=torch.tensor(has_initial_states, device="cuda", dtype=torch.bool), + ) + hidden = torch.randn(num_tokens, _Cfg.hidden_size, device="cuda", dtype=torch.bfloat16) * 0.05 + hidden_pristine = hidden.clone() + conv_seed = torch.randn(slots, 3 * d, conv_size, device="cuda", dtype=torch.bfloat16) * 0.02 + state_seed = ( + torch.randn(slots, h, head_dim, head_dim, device="cuda", dtype=torch.float32) * 0.01 + ) + + calls = {"qkvg": 0, "q": 0, "k": 0, "v": 0} + + def _count(name): + def _hook(_module, _inputs, _output): + calls[name] += 1 + + return _hook + + handles = [ + mixer.qkvg_proj.register_forward_hook(_count("qkvg")), + mixer.q_proj.register_forward_hook(_count("q")), + mixer.k_proj.register_forward_hook(_count("k")), + mixer.v_proj.register_forward_hook(_count("v")), + ] + fused_qkvg = mixer.qkvg_proj + try: + mixer.qkvg_proj = None + ref_conv = conv_seed.clone() + ref_state = state_seed.clone() + expected = runtime._forward_prefill( + hidden, + cu_seqlens, + metadata, + num_prefills, + ref_conv, + ref_state, + slot_indices, + ) + assert calls == {"qkvg": 0, "q": 1, "k": 1, "v": 1} + + calls.update(qkvg=0, q=0, k=0, v=0) + mixer.qkvg_proj = fused_qkvg + actual_conv = conv_seed.clone() + actual_state = state_seed.clone() + actual = runtime._forward_prefill( + hidden, + cu_seqlens, + metadata, + num_prefills, + actual_conv, + actual_state, + slot_indices, + ) + assert calls == {"qkvg": 1, "q": 0, "k": 0, "v": 0} + finally: + mixer.qkvg_proj = fused_qkvg + for handle in handles: + handle.remove() + + torch.testing.assert_close(hidden, hidden_pristine, rtol=0, atol=0) + _assert_numerically_close(actual, expected) + _assert_numerically_close( + actual_conv.index_select(0, slot_indices), ref_conv.index_select(0, slot_indices) + ) + _assert_numerically_close( + actual_state.index_select(0, slot_indices), ref_state.index_select(0, slot_indices) + ) + + untouched = torch.tensor([0, 1, slots - 1], device="cuda", dtype=torch.long) + torch.testing.assert_close( + actual_conv.index_select(0, untouched), conv_seed.index_select(0, untouched), rtol=0, atol=0 + ) + torch.testing.assert_close( + actual_state.index_select(0, untouched), + state_seed.index_select(0, untouched), + rtol=0, + atol=0, + ) + + repeat_conv = conv_seed.clone() + repeat_state = state_seed.clone() + repeated = runtime._forward_prefill( + hidden, + cu_seqlens, + metadata, + num_prefills, + repeat_conv, + repeat_state, + slot_indices, + ) + torch.testing.assert_close(repeated, actual, rtol=0, atol=0) + torch.testing.assert_close(repeat_conv, actual_conv, rtol=0, atol=0) + torch.testing.assert_close(repeat_state, actual_state, rtol=0, atol=0) From acf1ea0bfb7785b0c0d133b57056dc6c8054960d Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 17 Aug 2026 19:30:41 -0500 Subject: [PATCH 2/5] [None][chore] Address review: type hints, layer-id validation, accept-stats ordering in kimi_k3 examples Signed-off-by: Brian Nguyen --- .../kimi_k3/make_synthetic_dflash_drafter.py | 58 +++++++++++++++---- examples/kimi_k3/measure_dspark_acceptance.py | 16 ++++- .../test_kimi_kda_fp8_packed_prefill.py | 7 ++- 3 files changed, 65 insertions(+), 16 deletions(-) diff --git a/examples/kimi_k3/make_synthetic_dflash_drafter.py b/examples/kimi_k3/make_synthetic_dflash_drafter.py index 884f0f4c0fce..45412580d570 100644 --- a/examples/kimi_k3/make_synthetic_dflash_drafter.py +++ b/examples/kimi_k3/make_synthetic_dflash_drafter.py @@ -44,6 +44,8 @@ python make_synthetic_dflash_drafter.py --tiny --out """ +from __future__ import annotations + import argparse import json import os @@ -93,7 +95,7 @@ ) -def even_target_layer_ids(num_target_layers: int, k: int = 6): +def even_target_layer_ids(num_target_layers: int, k: int = 6) -> list[int]: """Evenly spaced capture layers, K2.7 convention (first=1, last=L-3).""" lo, hi = 1, max(1, num_target_layers - 3) if k == 1: @@ -101,9 +103,37 @@ def even_target_layer_ids(num_target_layers: int, k: int = 6): return sorted({round(lo + i * (hi - lo) / (k - 1)) for i in range(k)}) +def validate_target_layer_ids( + target_layer_ids: list[int], num_target_layers: int | None = None +) -> None: + """Validate capture-layer ids in every mode. + + Uses real exceptions rather than ``assert`` so the checks survive + ``python -O`` (which strips asserts). ``num_target_layers`` is the + target stack depth when known (unavailable in --config mode, where the + target checkpoint is not read); the upper-bound check is skipped when it + is None. + """ + if not target_layer_ids: + raise ValueError("target_layer_ids must be a non-empty list") + if len(set(target_layer_ids)) != len(target_layer_ids): + raise ValueError(f"target_layer_ids {target_layer_ids} contains duplicates") + if min(target_layer_ids) < 0: + raise ValueError(f"target_layer_ids {target_layer_ids} contains negative ids") + if num_target_layers is not None and max(target_layer_ids) >= num_target_layers: + raise ValueError( + f"target_layer_ids {target_layer_ids} out of range [0, {num_target_layers})" + ) + + def drafter_tensor_plan( - hidden, cfg, num_capture, vocab=None, markov_rank=None, use_confidence_head=False -): + hidden: int, + cfg: dict, + num_capture: int, + vocab: int | None = None, + markov_rank: int | None = None, + use_confidence_head: bool = False, +) -> dict[str, tuple[int, ...]]: """Return {key: shape} for the drafter checkpoint. Base keys follow the K2.7 DFlash schema; the dspark heads (markov_w1/w2 @@ -144,7 +174,14 @@ def drafter_tensor_plan( return plan -def drafter_config(hidden, vocab, num_target_layers, target_layer_ids, mask_token_id, cfg): +def drafter_config( + hidden: int, + vocab: int, + num_target_layers: int, + target_layer_ids: list[int], + mask_token_id: int, + cfg: dict, +) -> dict: dflash_cfg = { "mask_token_id": mask_token_id, "target_layer_ids": target_layer_ids, @@ -197,14 +234,14 @@ def drafter_config(hidden, vocab, num_target_layers, target_layer_ids, mask_toke } -def target_dims_from_ckpt(ckpt_dir): +def target_dims_from_ckpt(ckpt_dir: str) -> tuple[int, int, int]: with open(os.path.join(ckpt_dir, "config.json")) as f: cfg = json.load(f) text = cfg.get("text_config", cfg) return (text["hidden_size"], text["vocab_size"], text["num_hidden_layers"]) -def drafter_cfg_from_real_config(path): +def drafter_cfg_from_real_config(path: str) -> tuple[dict, dict]: """--config mode: adopt the REAL drafter config.json verbatim. Random weights, exact real module structure — no schema guessing. @@ -258,7 +295,7 @@ def drafter_cfg_from_real_config(path): return real, dims -def main(): +def main() -> None: ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) mode = ap.add_mutually_exclusive_group(required=True) mode.add_argument("--ckpt-dir", help="real K3 target checkpoint dir (reads config.json)") @@ -298,6 +335,9 @@ def main(): real_cfg, cfg = drafter_cfg_from_real_config(args.config) hidden, vocab = real_cfg["hidden_size"], real_cfg["vocab_size"] target_layer_ids = args.target_layer_ids or real_cfg["dflash_config"]["target_layer_ids"] + # No target checkpoint is read in --config mode, so the stack depth + # is only known if the real config carries num_target_layers. + validate_target_layer_ids(target_layer_ids, real_cfg.get("num_target_layers")) mask_token_id = ( args.mask_token_id if args.mask_token_id is not None @@ -316,9 +356,7 @@ def main(): target_layer_ids = args.target_layer_ids or even_target_layer_ids( num_target_layers, default_k ) - assert all(0 <= t < num_target_layers for t in target_layer_ids), ( - f"target_layer_ids {target_layer_ids} out of range [0, {num_target_layers})" - ) + validate_target_layer_ids(target_layer_ids, num_target_layers) mask_token_id = args.mask_token_id if args.mask_token_id is not None else vocab - 2 torch.manual_seed(args.seed) diff --git a/examples/kimi_k3/measure_dspark_acceptance.py b/examples/kimi_k3/measure_dspark_acceptance.py index a55d1b16b862..9528073703c6 100644 --- a/examples/kimi_k3/measure_dspark_acceptance.py +++ b/examples/kimi_k3/measure_dspark_acceptance.py @@ -28,17 +28,27 @@ recorder syncs a few scalars per step: keep it OFF (--no-accept-stats) for the TPOT-reference legs of an A/B. -Example (inside the container, see run_dspark_acceptance.sbatch): +Example (inside the container, see run_dspark_acceptance.sbatch). The +recorder dir MUST be exported before trtllm-llmapi-launch: it pre-spawns the +MPI worker ranks (where DFlashWorker lives), so setting os.environ inside this +script never reaches them and the accept-stats come back empty for TP>1: + export TLLM_DFLASH_ACCEPT_STATS_DIR=/tmp/dspark-stats trtllm-llmapi-launch python3 examples/kimi_k3/measure_dspark_acceptance.py \ --model /path/to/kimi-k3 --drafter /path/to/dspark-drafter \ --tp-size 16 --num-prompts 64 --max-tokens 256 \ --stats-dir /tmp/dspark-stats --output-json results_spec.json """ +from __future__ import annotations + import argparse import json import os import time +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from tensorrt_llm import LLM def parse_arguments() -> argparse.Namespace: @@ -97,7 +107,7 @@ def parse_arguments() -> argparse.Namespace: return parser.parse_args() -def load_prompts(args) -> list: +def load_prompts(args: argparse.Namespace) -> list[str]: if args.prompt_file: prompts = [] with open(args.prompt_file) as f: @@ -115,7 +125,7 @@ def load_prompts(args) -> list: ] -def build_llm(args): +def build_llm(args: argparse.Namespace) -> tuple[LLM, bool]: from tensorrt_llm import LLM from tensorrt_llm.llmapi import KvCacheConfig diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py index 302de7781399..4bd8dd4efced 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Runtime parity for the Kimi K3 FP8 packed q/k/v prefill projection.""" +from collections.abc import Callable from types import SimpleNamespace import pytest @@ -90,7 +91,7 @@ def test_fp8_packed_qkv_projection_matches_separate_views() -> None: ) @torch.no_grad() def test_fp8_packed_qkv_prefill_matches_separate_path_and_updates_state( - sequence_lengths, use_initial_states, has_initial_states + sequence_lengths: list[int], use_initial_states: bool, has_initial_states: list[bool] ) -> None: torch.manual_seed(1) runtime = _make_runtime() @@ -119,8 +120,8 @@ def test_fp8_packed_qkv_prefill_matches_separate_path_and_updates_state( calls = {"qkvg": 0, "q": 0, "k": 0, "v": 0} - def _count(name): - def _hook(_module, _inputs, _output): + def _count(name: str) -> Callable[[nn.Module, tuple, object], None]: + def _hook(_module: nn.Module, _inputs: tuple, _output: object) -> None: calls[name] += 1 return _hook From 9f480b0c1545fd32d52f786961fd22aa228f2fa9 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 17 Aug 2026 17:53:24 -0700 Subject: [PATCH 3/5] Address trivial review comments Signed-off-by: Brian Nguyen --- examples/kimi_k3/make_synthetic_dflash_drafter.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/kimi_k3/make_synthetic_dflash_drafter.py b/examples/kimi_k3/make_synthetic_dflash_drafter.py index 45412580d570..e4bec2ecfec9 100644 --- a/examples/kimi_k3/make_synthetic_dflash_drafter.py +++ b/examples/kimi_k3/make_synthetic_dflash_drafter.py @@ -148,7 +148,8 @@ def drafter_tensor_plan( "norm.weight": (hidden,), } if markov_rank: - assert vocab is not None, "markov head tensors need vocab_size" + if vocab is None: + raise ValueError("markov head tensors need vocab_size") plan["markov_w1.weight"] = (vocab, markov_rank) plan["markov_w2.weight"] = (vocab, markov_rank) if use_confidence_head: From 5d3a472ac15d19accb6251fe533cb52b52eccda2 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 19 Aug 2026 14:32:56 -0700 Subject: [PATCH 4/5] [TRTLLM-14818][test] Wire Kimi K3 FP8 packed-prefill test into l0_b200 Add test_kimi_kda_fp8_packed_prefill.py to the SM100 pre-merge list so CI exercises the FP8 weight-read + fused qkvg projection path on B200. Eventual removal tracked in TRTLLM-15633. Signed-off-by: Brian Nguyen --- tests/integration/test_lists/test-db/l0_b200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e1acb231f60a..9a6b4a454c29 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -97,6 +97,7 @@ l0_b200: # ------------- Kimi K3 (KimiLinear) unit tests --------------- - unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py - unittest/_torch/modeling/test_kimi_kda_verify_parity.py + - unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py - unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py From 13c2a58287623841253bfe73c9b82e06e815129d Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 20 Aug 2026 21:36:50 -0500 Subject: [PATCH 5/5] [TRTLLM-14818][fix] Exclude warmup from DSpark accept-stats aggregates The DFlash accept-stats recorder accumulates every eager step for the life of the worker process, so the harness warmup batch was folded into the reported AL/AR, accepted-draft histogram, and confidence calibration. The recorder lives in the pre-spawned MPI worker ranks and can't be reset across the launcher process boundary, so snapshot its post-warmup counts on the driver and subtract them from the final totals. Force TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY=1 on the stats leg (in-script for the single-process case, exported by run_dspark_acceptance.sbatch for TP>1) so the short warmup batch is flushed to disk before timing. The stats leg is a measurement run, not a TPOT reference (--no-accept-stats), so the per-step flush is acceptable. Signed-off-by: Brian Nguyen --- examples/kimi_k3/measure_dspark_acceptance.py | 65 +++++++++++++++++++ examples/kimi_k3/run_dspark_acceptance.sbatch | 2 +- 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/examples/kimi_k3/measure_dspark_acceptance.py b/examples/kimi_k3/measure_dspark_acceptance.py index 9528073703c6..643369e87187 100644 --- a/examples/kimi_k3/measure_dspark_acceptance.py +++ b/examples/kimi_k3/measure_dspark_acceptance.py @@ -28,6 +28,14 @@ recorder syncs a few scalars per step: keep it OFF (--no-accept-stats) for the TPOT-reference legs of an A/B. +The recorder accumulates every eager DFlash step, including the warmup +batch, and can't be reset across the launcher process boundary; the +harness instead snapshots the post-warmup counts and subtracts them from +the final totals so the reported figures cover only the timed run. That +subtraction needs the warmup steps flushed to disk before timing, so the +stats leg forces TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY=1 (export it before +trtllm-llmapi-launch for TP>1, as run_dspark_acceptance.sbatch does). + Example (inside the container, see run_dspark_acceptance.sbatch). The recorder dir MUST be exported before trtllm-llmapi-launch: it pre-spawns the MPI worker ranks (where DFlashWorker lives), so setting os.environ inside this @@ -175,6 +183,26 @@ def build_llm(args: argparse.Namespace) -> tuple[LLM, bool]: return LLM(**llm_kwargs), spec_on +def _remove_warmup(merged: dict, warmup: dict) -> None: + """Subtract warmup-phase counts from the merged accept-stats, in place. + + ``merged`` and ``warmup`` are both outputs of ``merge_snapshots``; after + this the reported histogram, AL/AR, and calibration reflect only the timed + generation. Only the fields the harness reports (the accepted-draft + histogram and the confidence-calibration counts) are adjusted. Counts are + monotonic, so the differences are non-negative. + """ + hist = merged["accepted_draft_hist"] + for i, n in enumerate(warmup["accepted_draft_hist"]): + hist[i] -= n + merged["num_steps"] -= warmup["num_steps"] + mcc, wcc = merged["confidence_calibration"], warmup["confidence_calibration"] + for field in ("attempts", "accepted"): + for k, row in enumerate(wcc[field]): + for b, n in enumerate(row): + mcc[field][k][b] -= n + + def main() -> None: args = parse_arguments() @@ -186,6 +214,16 @@ def main() -> None: ) # Must be set before the LLM (and its worker processes) is built. os.environ["TLLM_DFLASH_ACCEPT_STATS_DIR"] = stats_dir + # Flush after every verify step so the short warmup batch (a handful + # of steps, well under the default flush period) is on disk before the + # timed run starts — the harness reads it back as a baseline and + # subtracts it below. This env, like the stats dir, only reaches + # pre-spawned MPI workers when exported before trtllm-llmapi-launch + # (run_dspark_acceptance.sbatch does this for the TP>1 dspark leg); + # the in-script set covers the single-process case. The per-step flush + # writes a tiny JSON and only runs on the stats leg, which is a + # measurement run, not a TPOT reference (--no-accept-stats). + os.environ.setdefault("TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY", "1") prompts = load_prompts(args) llm, spec_on = build_llm(args) @@ -210,6 +248,31 @@ def main() -> None: # allocator costs so the timed section measures steady state. llm.generate(prompts[: min(2, len(prompts))], SamplingParams(max_tokens=8, temperature=0.0)) + # The accept-stats recorder lives in the worker ranks and accumulates + # every eager DFlash step, warmup included — so the warmup batch would + # otherwise bias the aggregate AL/AR, histogram, and calibration. We + # can't reset the recorder across the launcher process boundary, so + # snapshot its post-warmup counts here and subtract them from the final + # totals below (see _remove_warmup). Relies on the per-step flush set + # above so warmup is already on disk. + warmup_stats = None + if stats_dir and os.path.isdir(stats_dir): + from tensorrt_llm._torch.speculative.accept_stats import ( + load_rank_snapshots, + merge_snapshots, + ) + + base_snaps = load_rank_snapshots(stats_dir) + if base_snaps: + warmup_stats = merge_snapshots(base_snaps) + else: + print( + "WARNING: could not snapshot warmup accept-stats " + "(no rank files yet); aggregate AL/AR will include the " + "warmup batch. Export TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY=1 " + "before trtllm-llmapi-launch for multi-rank runs." + ) + t0 = time.monotonic() outputs = llm.generate(prompts, sampling_params) wall_s = time.monotonic() - t0 @@ -264,6 +327,8 @@ def main() -> None: ) else: merged = merge_snapshots(snaps) + if warmup_stats is not None: + _remove_warmup(merged, warmup_stats) summary = summarize_hist(merged["accepted_draft_hist"]) cc = merged["confidence_calibration"] has_calib = any(any(row) for row in cc["attempts"]) diff --git a/examples/kimi_k3/run_dspark_acceptance.sbatch b/examples/kimi_k3/run_dspark_acceptance.sbatch index 83510ab1c0de..30dfbcddf53e 100644 --- a/examples/kimi_k3/run_dspark_acceptance.sbatch +++ b/examples/kimi_k3/run_dspark_acceptance.sbatch @@ -167,7 +167,7 @@ if [[ "$SKIP_BASELINE" -eq 0 ]]; then fi echo "=== leg (b): DSpark-on measurement ===" -LEG_ENV_EXPORT="export TLLM_DFLASH_ACCEPT_STATS_DIR='$OUTDIR/accept-stats'" \ +LEG_ENV_EXPORT="export TLLM_DFLASH_ACCEPT_STATS_DIR='$OUTDIR/accept-stats' TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY=1" \ run_leg dspark --drafter "$DRAFTER" $CONF_ARGS \ --stats-dir "$OUTDIR/accept-stats" \ --output-json "$OUTDIR/results_dspark.json"