From 2ec20d0f63c1c2e1675ef30886a942db431e0118 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 29 Jul 2026 02:11:09 +0000 Subject: [PATCH 01/15] add script Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../scripts/quantize_drafter.py | 216 ++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 examples/speculative_decoding/scripts/quantize_drafter.py diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py new file mode 100644 index 00000000000..541668154d8 --- /dev/null +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Calibration-free W4A16 weight-only PTQ for a speculative-decoding drafter. + +Block-wise ``max`` quantization derives every scale from the weight tensor itself, so +this needs neither a dataset nor a forward pass -- only the drafter's safetensors file. +That in turn means we never have to import the drafter's modeling code: each 2-D weight +is wrapped in a throwaway ``nn.Linear`` under its checkpoint name, and ModelOpt's normal +``quantizer_name`` patterns select over those names exactly as they would on the real +module tree. Works for any drafter layout (DSpark / DFlash / EAGLE3 / Medusa). + +The format is fixed to ModelOpt's ``w4a16_nvfp4`` preset (E2M1 values, block 16, FP8 E4M3 +scales). It is calibration-free because the block scales are dynamic; AWQ formats are +deliberately not offered, since ``awq_lite`` silently degrades to plain RTN when no +``forward_loop`` is supplied. + +Example: + python quantize_drafter_w4a16.py \ + --drafter_path nvidia/MiniMax-M3-DSpark \ + --export_path ./MiniMax-M3-DSpark-W4A16 +""" + +import argparse +import copy +import json +import shutil +from pathlib import Path + +import torch +import torch.nn as nn +from safetensors.torch import load_file, save_file + +import modelopt.torch.quantization as mtq +from modelopt.recipe.presets import QUANT_CFG_CHOICES +from modelopt.torch.export.quant_utils import ( + get_quant_config, + get_quantization_format, + get_weight_block_size, + get_weight_scaling_factor, + get_weight_scaling_factor_2, + to_quantized_weight, +) +from modelopt.torch.quantization.utils import is_quantized_linear + +# NVFP4 weight-only: block-16 E2M1 values with dynamic FP8 E4M3 scales, so every scale is +# derived from the weight block itself and no calibration data is involved. +QFORMAT = "w4a16_nvfp4" + +# The Markov head writes straight into the draft logits and is only a few percent of the +# drafter, so the bits it would save are not worth the acceptance-rate risk. It is also not +# a plain GEMM in the real module tree -- markov_w1 is an nn.Embedding, which ModelOpt's +# stock presets already exclude via `parent_class: nn.Embedding`; the flat-linear view here +# has lost the original module classes, so the exclusion has to be restated by name. +# lm_head is excluded by the preset itself (see --quantize_lm_head). +DEFAULT_EXCLUDE = ["*markov_head*"] + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--drafter_path", required=True, help="HF repo id or local dir of the drafter checkpoint." + ) + parser.add_argument("--export_path", required=True, help="Output directory.") + parser.add_argument( + "--dtype", + default="bfloat16", + choices=["bfloat16", "float16", "float32"], + help="Compute dtype the weights are cast to before quantizing.", + ) + parser.add_argument( + "--exclude", + nargs="*", + default=[], + metavar="PATTERN", + help="Extra fnmatch patterns to leave unquantized, in `quantizer_name` form. " + f"Appended to the defaults ({' '.join(DEFAULT_EXCLUDE)}), which always apply.", + ) + parser.add_argument( + "--quantize_lm_head", + action="store_true", + help="Also quantize lm_head. It is the single largest drafter tensor, but it feeds " + "the acceptance test directly -- measure AL before shipping this.", + ) + return parser.parse_args() + + +def load_drafter(drafter_path: str) -> tuple[Path, dict[str, torch.Tensor]]: + """Resolve a local dir or HF repo id to (dir, state_dict).""" + local_dir = Path(drafter_path) + if not local_dir.is_dir(): + from huggingface_hub import snapshot_download + + local_dir = Path(snapshot_download(drafter_path)) + + shards = sorted(local_dir.glob("*.safetensors")) + assert shards, f"No .safetensors found under {local_dir}" + state_dict: dict[str, torch.Tensor] = {} + for shard in shards: + state_dict.update(load_file(shard)) + return local_dir, state_dict + + +def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) -> nn.Module: + """Expose every 2-D weight as an nn.Linear whose module name is its checkpoint key. + + Nested ModuleDicts are used so that ``named_modules()`` reproduces the dotted keys + (``layers.0.self_attn.q_proj``), which is what the preset's ``quantizer_name`` + patterns match against. + """ + root = nn.ModuleDict() + for key, weight in state_dict.items(): + if weight.dim() != 2 or not key.endswith(".weight"): + continue + *parents, leaf = key[: -len(".weight")].split(".") + node = root + for part in parents: + if part not in node: + node[part] = nn.ModuleDict() + node = node[part] + out_features, in_features = weight.shape + linear = nn.Linear(in_features, out_features, bias=False, dtype=dtype) + linear.weight.data = weight.to(dtype) + node[leaf] = linear + return root + + +def build_quant_cfg(exclude: list[str], quantize_lm_head: bool) -> dict: + """Take the shipped preset and layer the drafter-specific exclusions on top.""" + quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[QFORMAT]) + if quantize_lm_head: + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*lm_head*weight_quantizer", "enable": True} + ) + for pattern in DEFAULT_EXCLUDE + exclude: + quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) + return quant_cfg + + +def export_quantized_state_dict( + root: nn.Module, state_dict: dict[str, torch.Tensor], dtype: torch.dtype +) -> dict[str, torch.Tensor]: + """Pack each quantized weight and emit it alongside its scales. + + Follows the unified-HF naming convention: ``w.weight`` / ``w.weight_scale`` / + ``w.weight_scale_2``. Untouched tensors are carried through in ``dtype``. + """ + export_sd = {k: v.to(dtype) for k, v in state_dict.items()} + for name, module in root.named_modules(): + if not is_quantized_linear(module) or not module.weight_quantizer.is_enabled: + continue + quantization = get_quantization_format(module) + assert quantization is not None, f"{name}: enabled quantizer resolved to no format" + weight_scale = get_weight_scaling_factor(module) + weight_scale_2 = get_weight_scaling_factor_2(module) + export_sd[f"{name}.weight"] = to_quantized_weight( + module.weight, + weight_scale, + quantization, + weight_scale_2, + get_weight_block_size(module), + ) + export_sd[f"{name}.weight_scale"] = weight_scale + if weight_scale_2 is not None: + export_sd[f"{name}.weight_scale_2"] = weight_scale_2 + return export_sd + + +def main(): + args = parse_args() + dtype = getattr(torch, args.dtype) + + source_dir, state_dict = load_drafter(args.drafter_path) + root = build_linear_view(state_dict, dtype) + + quant_cfg = build_quant_cfg(args.exclude, args.quantize_lm_head) + mtq.quantize(root, quant_cfg) # no forward_loop: scales come from the weights + mtq.print_quant_summary(root) + + export_sd = export_quantized_state_dict(root, state_dict, dtype) + + export_dir = Path(args.export_path) + export_dir.mkdir(parents=True, exist_ok=True) + save_file(export_sd, export_dir / "model.safetensors", metadata={"format": "pt"}) + + config = json.loads((source_dir / "config.json").read_text()) + hf_quant_config = get_quant_config(root) + config["quantization_config"] = hf_quant_config["quantization"] + config["torch_dtype"] = args.dtype + (export_dir / "config.json").write_text(json.dumps(config, indent=2)) + (export_dir / "hf_quant_config.json").write_text(json.dumps(hf_quant_config, indent=2)) + + for extra in ("tokenizer.json", "tokenizer_config.json", "generation_config.json"): + if (source_dir / extra).is_file(): + shutil.copy2(source_dir / extra, export_dir / extra) + + before = sum(v.numel() * v.element_size() for v in state_dict.values()) + after = sum(v.numel() * v.element_size() for v in export_sd.values()) + print(f"\n{QFORMAT}: {before / 2**30:.2f} GiB -> {after / 2**30:.2f} GiB") + print(f"Exported to {export_dir}") + + +if __name__ == "__main__": + main() From 0bea10e584aa6deb18526f77dea33acbfc40d74d Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 07:53:26 +0000 Subject: [PATCH 02/15] example(specdec): FP8/NVFP4 weight+activation PTQ for drafters Extends quantize_drafter.py from W4A16-only to five formats, and adds a calibration-free way to set the static activation scale that the weight+activation formats need. Formats: w4a16_nvfp4, nvfp4, fp8, fp8_pc_pt, fp8_pb_wo -- the ModelOpt formats vLLM can actually serve. need_calibration() decides which need a static activation amax. For those, --act_scale_heuristic fixed (the default) applies one amax to every layer, set by --act_scale_amax (default 448, i.e. input_scale 1.0 for FP8). A single fixed scale sounds crude but measures well, because acceptance length is governed almost entirely by clipping rather than by resolution. Sweeping input_scale over three decades on Qwen3-8B, both formats fall off a cliff below ~0.03 -- where the declared range is far under the activations' true magnitude and most of the tensor is clipped -- and both sit on a flat plateau from ~0.3 to 4.0 with no drop-off at the top. NVFP4 trails FP8 by a roughly constant 3.5% across the plateau; that gap is the 4-bit resolution cost and no choice of scale recovers it. Two weight-derived estimators are kept for reference but land 1-2 orders of magnitude below the activation range (-31% to -46% AL). Deploying a quantized drafter needed four export fixes and one caller fix: - emit quant_method (modelopt_fp4 / modelopt); vLLM reads it, not quant_algo - emit the exclusion list under `ignore` too, the key read from the flat quantization_config in config.json - add * wildcards so exclusions match a runtime's nested module prefix - add *qkv_proj / *gate_up_proj aliases for layers a runtime fuses - specdec_bench: pass the draft's own quantization into speculative_config. vLLM otherwise copies the target's onto the draft, so a quantized drafter under a bf16 target is built as bf16 and dies on the packed weights. Also excludes embed_tokens (an Embedding the drafter inherits from the target, not a GEMM) and confidence_head (a [1, H] projection whose per-channel scale collapses to 0-dim), plus a generic 0-dim guard, and writes .input_scale without which every activation scale was silently absent from the export. Testing: Qwen3-8B + deepseek-ai/dspark_qwen3_8b_block7, MT-Bench 80q, greedy. bf16 AL 3.1423. fp8 @ input_scale 1.0 -> 3.1457 (+0.11%); fp8_pc_pt dynamic per-token -> 3.1228 (-0.62%); w4a16_nvfp4 -> 3.0392 (-3.28%); nvfp4 @ 1.0 -> 3.0193 (-3.91%). Full 20-point sweep of both formats in the module docstring. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- examples/specdec_bench/run.py | 9 + .../specdec_bench/models/vllm.py | 20 ++ .../scripts/quantize_drafter.py | 317 +++++++++++++++++- 3 files changed, 330 insertions(+), 16 deletions(-) diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index ca2f9908966..342dffb9ad8 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -188,6 +188,7 @@ def run_simple(args): sampling_kwargs=sampling_kwargs, speculative_algorithm=args.speculative_algorithm, draft_model_dir=args.draft_model_dir, + draft_quantization=args.draft_quantization, speculative_num_steps=args.draft_length, speculative_num_draft_tokens=args.block_size, tensor_parallel_size=args.tp_size, @@ -316,6 +317,14 @@ def run_simple(args): default=None, help="Path to the draft model directory", ) + parser.add_argument( + "--draft_quantization", + type=str, + required=False, + default=None, + help="Quantization method of the draft checkpoint (e.g. modelopt, modelopt_fp4). " + "Read from the draft's config.json when omitted; set it only to override.", + ) parser.add_argument( "--runtime_params", type=str, diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index 24062399cb8..e0801a4585a 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/examples/specdec_bench/specdec_bench/models/vllm.py @@ -14,6 +14,8 @@ # limitations under the License. import asyncio +import json +import os import time from .base import Model @@ -117,6 +119,24 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs elif kwargs.get("speculative_algorithm") == "NONE": specdec = None + # A quantized draft checkpoint has to declare its own quantization. vLLM otherwise + # copies the target's onto the draft (config/speculative.py, "Align the + # quantization of draft model"), so a quantized drafter under a bf16 target is + # built as if it were bf16 and dies loading the packed weights. The drafter's + # config.json already records the format, so read it back rather than making the + # caller repeat it; --draft_quantization overrides. + if specdec is not None and specdec.get("model"): + draft_quantization = kwargs.get("draft_quantization") + if draft_quantization is None: + draft_config = os.path.join(specdec["model"], "config.json") + if os.path.isfile(draft_config): + with open(draft_config) as f: + quant_config = json.load(f).get("quantization_config") or {} + draft_quantization = quant_config.get("quant_method") + if draft_quantization: + specdec["quantization"] = draft_quantization + print(f"Draft model quantization: {draft_quantization}") + if specdec is None: num_speculative_tokens = 1 else: diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index 541668154d8..c2547ee6320 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Calibration-free W4A16 weight-only PTQ for a speculative-decoding drafter. +"""Calibration-free PTQ for a speculative-decoding drafter. Block-wise ``max`` quantization derives every scale from the weight tensor itself, so this needs neither a dataset nor a forward pass -- only the drafter's safetensors file. @@ -22,15 +22,53 @@ ``quantizer_name`` patterns select over those names exactly as they would on the real module tree. Works for any drafter layout (DSpark / DFlash / EAGLE3 / Medusa). -The format is fixed to ModelOpt's ``w4a16_nvfp4`` preset (E2M1 values, block 16, FP8 E4M3 -scales). It is calibration-free because the block scales are dynamic; AWQ formats are -deliberately not offered, since ``awq_lite`` silently degrades to plain RTN when no -``forward_loop`` is supplied. +Avoiding the modeling code is what makes this work at all for DFlash-family drafters. An +exported drafter declares ``architectures: ["DFlashDraftModel"]`` but ships no importable +class, so ``AutoModelForCausalLM`` silently loads it as a plain Qwen3 and drops ``fc`` / +``hidden_norm`` / the DSpark heads; and the real draft module's ``forward`` takes +``(noise_embedding, target_hidden, ...)``, not ``input_ids``, so the stock ``hf_ptq.py`` +calibration loop cannot drive it either. The flat-linear view sidesteps both problems. + +Two families of format are supported: + +* **Weight-only / dynamic-activation** (``w4a16_nvfp4``, ``fp8_pc_pt``, ``fp8_pb_wo``) -- + every scale is either derived from the weight or computed at runtime per token. +* **Static-activation** (``fp8``, ``nvfp4``) -- these need an ``input_quantizer`` amax, + normally measured with calibration data. Instead one fixed amax is applied to every + layer (``--act_scale_amax``, default 448 = input_scale 1.0 for FP8). + +A fixed activation scale sounds crude but measures well, because acceptance length is +governed almost entirely by *clipping* rather than by resolution. Sweeping input_scale over +three decades on Qwen3-8B + a DSpark drafter (MT-Bench, 80 questions): + +=========== ============== ============== +input_scale FP8 AL NVFP4 AL +=========== ============== ============== +0.003 2.220 (-29.3%) 2.208 (-29.8%) +0.03 2.975 (-5.3%) 2.926 (-6.9%) +0.3 3.137 (-0.2%) 3.022 (-3.8%) +1.0 3.146 (+0.1%) 3.019 (-3.9%) +4.0 3.125 (-0.6%) 3.003 (-4.4%) +=========== ============== ============== + +(bf16 baseline 3.142.) Both formats fall off a cliff below ~0.03, where the declared range +is far under the activations' true magnitude and most of the tensor is clipped, and both +sit on a flat plateau from ~0.3 to 4.0 with no drop-off at the top -- so the scale needs to +be big enough, and little else. NVFP4 trails FP8 by a roughly constant 3.5% across the +plateau: that gap is the 4-bit resolution cost itself and no choice of scale recovers it. + +Two weight-derived estimators (``weight_amax``, ``weight_rms``) are kept for reference. +They land 1-2 orders of magnitude below the activation range and measure -31% to -46% AL; +they are not recommended. + +AWQ formats are deliberately not offered: ``awq_lite`` silently degrades to plain RTN when +no ``forward_loop`` is supplied. Example: - python quantize_drafter_w4a16.py \ + python quantize_drafter.py \ --drafter_path nvidia/MiniMax-M3-DSpark \ - --export_path ./MiniMax-M3-DSpark-W4A16 + --qformat fp8 \ + --export_path ./MiniMax-M3-DSpark-FP8 """ import argparse @@ -46,6 +84,7 @@ import modelopt.torch.quantization as mtq from modelopt.recipe.presets import QUANT_CFG_CHOICES from modelopt.torch.export.quant_utils import ( + get_activation_scaling_factor, get_quant_config, get_quantization_format, get_weight_block_size, @@ -53,19 +92,54 @@ get_weight_scaling_factor_2, to_quantized_weight, ) +from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.utils import is_quantized_linear -# NVFP4 weight-only: block-16 E2M1 values with dynamic FP8 E4M3 scales, so every scale is -# derived from the weight block itself and no calibration data is involved. -QFORMAT = "w4a16_nvfp4" +# Formats this script offers, in the order they are most likely to be wanted. +# w4a16_nvfp4 -- block-16 E2M1 weights, dynamic FP8 E4M3 scales. Smallest, weight-only. +# nvfp4 -- same weights plus NVFP4 activations (static amax -> needs a scale). +# fp8 -- E4M3 weights and activations, per-tensor (static amax -> needs a scale). +# fp8_pc_pt -- E4M3 per-channel weights, DYNAMIC per-token activations. The +# calibration-free way to get FP8 activations. +# fp8_pb_wo -- E4M3 per-block weight-only fallback. +# INT8/INT4 weight-only formats are deliberately absent: vLLM's ModelOpt backend accepts +# only FP8, FP8_PER_CHANNEL_PER_TOKEN, FP8_PB_WO, NVFP4, W4A16_NVFP4, MXFP8 and +# MIXED_PRECISION, so an INT checkpoint quantizes cleanly but cannot be served. +SUPPORTED_QFORMATS = [ + "w4a16_nvfp4", + "nvfp4", + "fp8", + "fp8_pc_pt", + "fp8_pb_wo", +] # The Markov head writes straight into the draft logits and is only a few percent of the # drafter, so the bits it would save are not worth the acceptance-rate risk. It is also not # a plain GEMM in the real module tree -- markov_w1 is an nn.Embedding, which ModelOpt's # stock presets already exclude via `parent_class: nn.Embedding`; the flat-linear view here # has lost the original module classes, so the exclusion has to be restated by name. +# +# The confidence head is excluded for the same reason plus a mechanical one: it projects to +# a single output (``[1, hidden]``), so a per-channel scale collapses to a 0-dim tensor and +# the per-channel export path (fp8_pc_pt) indexes it as ``scale[:, None]`` and raises. It is +# one row of weights, so there is nothing to gain by quantizing it. +# +# ``fc`` is deliberately NOT excluded here: it is the single largest non-lm_head tensor and +# the formats' own presets decide its fate. Exclude it explicitly with --exclude '*fc*' when +# comparing acceptance length. # lm_head is excluded by the preset itself (see --quantize_lm_head). -DEFAULT_EXCLUDE = ["*markov_head*"] +# +# ``embed_tokens`` must be excluded too. It is 2-D, so the flat-linear view happily treats +# it as a GEMM, but it is an ``nn.Embedding``: a row lookup, not a matmul. ModelOpt's stock +# presets skip embeddings via ``parent_class: nn.Embedding``, which the flat view cannot +# see. Quantizing it also breaks deployment -- a drafter inherits ``embed_tokens`` from the +# target, and vLLM's loader then fails with ``KeyError: 'embed_tokens.weight_scale'``. +DEFAULT_EXCLUDE = ["*markov_head*", "*confidence_head*", "*embed_tokens*"] + +# Largest value FP8 E4M3 can represent, and the default static activation amax. The +# exported input_scale is amax/448 for FP8 but amax/(6*448) for NVFP4 -- the amax is what +# both formats share, which is why the CLI takes it rather than an input_scale. +FP8_E4M3_MAX = 448.0 def parse_args(): @@ -74,6 +148,43 @@ def parse_args(): "--drafter_path", required=True, help="HF repo id or local dir of the drafter checkpoint." ) parser.add_argument("--export_path", required=True, help="Output directory.") + parser.add_argument( + "--qformat", + default="w4a16_nvfp4", + choices=SUPPORTED_QFORMATS, + help="Quantization format. Weight-only and dynamic-activation formats are fully " + "calibration-free; static-activation formats (fp8, nvfp4) additionally need " + "--act_scale_heuristic.", + ) + parser.add_argument( + "--act_scale_heuristic", + default="fixed", + choices=["fixed", "none", "weight_amax", "weight_rms"], + help="How to set the static activation amax for formats that need one, without " + "calibration data. 'fixed' (default) applies --act_scale_amax to every " + "layer. 'none' refuses to guess and errors out. The two weight-derived estimators " + "are kept for reference only -- they were measured at -31%% to -46%% acceptance " + "length; see estimate_activation_amax.", + ) + parser.add_argument( + "--act_scale_amax", + type=float, + default=FP8_E4M3_MAX, + help="Static activation amax applied to every layer with --act_scale_heuristic " + "fixed, i.e. the largest activation the quantizer will represent without clipping. " + "Default 448 corresponds to input_scale 1.0 for FP8; it is lossless for FP8 and " + "within 4%% for NVFP4 on Qwen3-8B, and anything from ~134 (scale 0.3) upward sits " + "on the same plateau. Below ~13 the activations get clipped and acceptance length " + "falls off a cliff. Note the exported input_scale is amax/448 for FP8 but " + "amax/(6*448) for NVFP4, so the same amax yields different input_scale values.", + ) + parser.add_argument( + "--act_scale_multiplier", + type=float, + default=1.0, + help="Extra factor on the estimated activation amax. >1 clips less and quantizes " + "coarser; <1 the reverse. Only used with --act_scale_heuristic.", + ) parser.add_argument( "--dtype", default="bfloat16", @@ -137,9 +248,84 @@ def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) - return root -def build_quant_cfg(exclude: list[str], quantize_lm_head: bool) -> dict: +def estimate_activation_amax( + module: nn.Linear, method: str, multiplier: float, fixed_value: float = 1.0 +) -> torch.Tensor | None: + """Estimate a static activation amax for ``module`` from its weight alone. + + Static-activation formats (fp8, nvfp4) need an ``input_quantizer`` amax, which is + normally measured by pushing calibration data through the model. A drafter cannot be + driven that way here (no importable modeling code, and its real forward takes hidden + states rather than token ids), so these estimators stand in for the measurement. + + Both assume the layer's input is roughly unit-scale -- true for a transformer's linear + inputs, which are RMSNorm outputs -- and use the weight only to set the dynamic range: + + ``weight_amax`` + amax = max |W|. An upper-ish bound: activations feeding a layer with large weights + are assumed to span a comparable range. Errs toward clipping less, quantizing + coarser. + ``weight_rms`` + amax = 4 * rms(W), i.e. a 4-sigma range for a roughly Gaussian weight. Tighter than + weight_amax on layers with a few outlier weights, so it usually preserves more + resolution -- at the cost of clipping genuine activation outliers. + + ``fixed`` + amax = ``--act_scale_amax``, the same value for every layer, + ignoring the weight + entirely. Measured on Qwen3-8B the weight-derived estimators land 1-2 orders of + magnitude below the activations they are meant to bound (max|W| averages 0.79 while + a RMSNorm'd activation is O(1) with outlier channels in the tens), so most values + are clipped. A fixed amax lets the range be set directly on the scale the + activations actually live on. + + None is a substitute for calibration. All are deliberately crude, and which one wins is + model-dependent; measure acceptance length rather than trusting any of them. + """ + weight = module.weight.detach().float() + if method == "weight_amax": + amax = weight.abs().max() + elif method == "weight_rms": + amax = 4.0 * weight.pow(2).mean().sqrt() + elif method == "fixed": + amax = torch.tensor(fixed_value, dtype=torch.float32) + else: + return None + return (amax * multiplier).to(module.weight.dtype) + + +def apply_activation_scale_heuristic( + root: nn.Module, method: str, multiplier: float, fixed_value: float = 1.0 +) -> int: + """Set every enabled ``input_quantizer``'s amax from the weight heuristic. + + Returns the number of quantizers that were given a scale. Only quantizers that are + both enabled and still missing an amax are touched, so dynamic-activation formats + (whose scales are computed per token at runtime) are left alone. + """ + count = 0 + for _, module in root.named_modules(): + if not is_quantized_linear(module): + continue + input_quantizer = getattr(module, "input_quantizer", None) + if input_quantizer is None or not input_quantizer.is_enabled: + continue + # A dynamic quantizer derives its scale at runtime; leave it untouched. + if getattr(input_quantizer, "_dynamic", False): + continue + if getattr(input_quantizer, "amax", None) is not None: + continue + amax = estimate_activation_amax(module, method, multiplier, fixed_value) + if amax is None: + continue + input_quantizer.amax = amax + count += 1 + return count + + +def build_quant_cfg(qformat: str, exclude: list[str], quantize_lm_head: bool) -> dict: """Take the shipped preset and layer the drafter-specific exclusions on top.""" - quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[QFORMAT]) + quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[qformat]) if quantize_lm_head: quant_cfg["quant_cfg"].append( {"quantizer_name": "*lm_head*weight_quantizer", "enable": True} @@ -165,6 +351,12 @@ def export_quantized_state_dict( assert quantization is not None, f"{name}: enabled quantizer resolved to no format" weight_scale = get_weight_scaling_factor(module) weight_scale_2 = get_weight_scaling_factor_2(module) + # A per-channel scale over a single-output projection (``[1, hidden]``) collapses to + # a 0-dim tensor that the packing helpers index as ``scale[:, None]``. Such a layer + # is one row of weights, so skip it rather than crash -- it stays in ``dtype``. + if weight_scale is not None and weight_scale.dim() == 0 and module.weight.shape[0] == 1: + print(f"Skipping {name}: single-output projection, per-channel scale is scalar") + continue export_sd[f"{name}.weight"] = to_quantized_weight( module.weight, weight_scale, @@ -175,6 +367,12 @@ def export_quantized_state_dict( export_sd[f"{name}.weight_scale"] = weight_scale if weight_scale_2 is not None: export_sd[f"{name}.weight_scale_2"] = weight_scale_2 + # Static-activation formats also need the input scale in the checkpoint. Without + # it the runtime has no activation scale to apply, so every heuristic produces a + # byte-identical export and the format silently degrades. + activation_scale = get_activation_scaling_factor(module) + if activation_scale is not None: + export_sd[f"{name}.input_scale"] = activation_scale return export_sd @@ -185,8 +383,36 @@ def main(): source_dir, state_dict = load_drafter(args.drafter_path) root = build_linear_view(state_dict, dtype) - quant_cfg = build_quant_cfg(args.exclude, args.quantize_lm_head) + quant_cfg = build_quant_cfg(args.qformat, args.exclude, args.quantize_lm_head) + needs_static_act = need_calibration(quant_cfg) + if needs_static_act and args.act_scale_heuristic == "none": + raise SystemExit( + f"--qformat {args.qformat} quantizes activations with a static scale, which is " + "normally measured with calibration data. This script has none, so pass " + "--act_scale_heuristic {weight_amax,weight_rms} to estimate it from the weights " + "(approximate -- verify acceptance length), or pick a calibration-free format " + "such as w4a16_nvfp4 (weight-only) or fp8_pc_pt (dynamic per-token activations)." + ) + mtq.quantize(root, quant_cfg) # no forward_loop: scales come from the weights + + if needs_static_act: + n = apply_activation_scale_heuristic( + root, + args.act_scale_heuristic, + args.act_scale_multiplier, + args.act_scale_amax, + ) + detail = ( + f"amax={args.act_scale_amax:g}" + if args.act_scale_heuristic == "fixed" + else f"derived from weights (x{args.act_scale_multiplier:g})" + ) + print( + f"Set {n} static activation amax values via '{args.act_scale_heuristic}' " + f"({detail}) -- not calibrated." + ) + mtq.print_quant_summary(root) export_sd = export_quantized_state_dict(root, state_dict, dtype) @@ -197,7 +423,66 @@ def main(): config = json.loads((source_dir / "config.json").read_text()) hf_quant_config = get_quant_config(root) - config["quantization_config"] = hf_quant_config["quantization"] + # ``get_quant_config`` only knows about the modules in the linear view, so anything the + # view never saw -- embeddings, norms, and any 1-D weight -- is absent from + # ``exclude_modules``. A loader that walks the checkpoint (vLLM does) then expects a + # ``weight_scale`` for those too and dies with e.g. KeyError: 'embed_tokens.weight_scale'. + # List every weight that was not quantized so the exclusion set is complete. + quantized = { + name + for name, module in root.named_modules() + if is_quantized_linear(module) + and f"{name}.weight" in export_sd + and f"{name}.weight_scale" in export_sd + } + unquantized = sorted( + key[: -len(".weight")] + for key in state_dict + if key.endswith(".weight") and key[: -len(".weight")] not in quantized + ) + exclude_modules = hf_quant_config["quantization"].get("exclude_modules", []) + for name in unquantized: + if name not in exclude_modules: + exclude_modules.append(name) + # A runtime matches this list against its own module prefix, which is usually + # nested relative to the checkpoint key (vLLM builds the draft's ``fc`` at + # ``model.fc`` via ``maybe_prefix``). A bare ``fc`` then fails to match, the layer + # is built quantized, and loading the bf16 weight into a packed parameter raises a + # size assertion. Add a suffix wildcard so the exclusion matches at any depth. + wildcard = f"*{name}" + if wildcard not in exclude_modules: + exclude_modules.append(wildcard) + # Runtimes fuse sibling projections into one layer whose name appears in no checkpoint + # key: q/k/v -> ``qkv_proj``, gate/up -> ``gate_up_proj``. Excluding only the individual + # names leaves the fused layer quantized, and its merged shard shapes then disagree with + # the unpacked weights being loaded into it. Emit the fused aliases whenever every + # component of a fusion group was excluded. + for fused, parts in ( + ("qkv_proj", ("q_proj", "k_proj", "v_proj")), + ("gate_up_proj", ("gate_proj", "up_proj")), + ): + if all(any(p in name for name in exclude_modules) for p in parts): + alias = f"*{fused}" + if alias not in exclude_modules: + exclude_modules.append(alias) + hf_quant_config["quantization"]["exclude_modules"] = exclude_modules + config["quantization_config"] = dict(hf_quant_config["quantization"]) + # ModelOpt names the format ``quant_algo``; vLLM's ModelConfig reads + # ``quant_cfg["quant_method"]`` and treats a config without that key as unquantized, + # then dies on the packed weight shapes. Emit both so either loader is satisfied. + # vLLM splits ModelOpt checkpoints into two backends: ``modelopt_fp4`` for the + # block-scaled NVFP4 layouts and ``modelopt`` for the per-tensor/per-channel ones. + quant_algo = str(hf_quant_config["quantization"].get("quant_algo") or "") + config["quantization_config"].setdefault( + "quant_method", "modelopt_fp4" if "NVFP4" in quant_algo.upper() else "modelopt" + ) + # The exclusion list is read under two different keys. ModelOpt's own + # ``hf_quant_config.json`` nests it under ``quantization.exclude_modules``, but when a + # runtime parses the flat ``quantization_config`` block inside ``config.json`` it looks + # for ``ignore`` (vLLM's ModelOptQuantConfigBase.from_config). Emitting only + # ``exclude_modules`` there yields an empty exclusion set, every layer is built + # quantized, and loading an untouched bf16 weight into a packed parameter raises. + config["quantization_config"]["ignore"] = list(exclude_modules) config["torch_dtype"] = args.dtype (export_dir / "config.json").write_text(json.dumps(config, indent=2)) (export_dir / "hf_quant_config.json").write_text(json.dumps(hf_quant_config, indent=2)) @@ -208,7 +493,7 @@ def main(): before = sum(v.numel() * v.element_size() for v in state_dict.values()) after = sum(v.numel() * v.element_size() for v in export_sd.values()) - print(f"\n{QFORMAT}: {before / 2**30:.2f} GiB -> {after / 2**30:.2f} GiB") + print(f"\n{args.qformat}: {before / 2**30:.2f} GiB -> {after / 2**30:.2f} GiB") print(f"Exported to {export_dir}") From cfdf9f510c7e2b7604dade189e94ed28b4f6808e Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 08:33:20 +0000 Subject: [PATCH 03/15] example(specdec): fix the drafter PTQ activation scale, drop the heuristic flags The activation-scale knobs existed to run a sweep that is now finished, and its answer was a single value: input_scale 1.0 measures +0.11% AL on FP8, and the usable plateau spans ~0.3 to 4.0 with no drop-off at the top, so there is nothing for a caller to tune. Hardcode it and remove --act_scale_heuristic, --act_scale_amax and --act_scale_multiplier. Two of the four heuristic choices (weight_amax, weight_rms) were documented in the script as not recommended -- they measured -31% to -46% AL. Keeping known-bad options in a CLI only creates a way to pick one; the numbers belong in the docs, which is where they now live. Route the remaining logic through resolve_activation_scales(), the one place that decides where a static activation amax comes from. Real calibration slots in there as a second source ahead of the fixed fallback without changing the CLI or the call site, since set_static_activation_amax() already skips quantizers that have an amax and so composes as a fallback rather than an overwrite. Exports are byte-identical to before this change for all five formats: fp8 still records input_scale 1.0, nvfp4 0.1667 (the 6*448 divisor), and the three calibration-free formats none. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../scripts/quantize_drafter.py | 177 ++++++------------ 1 file changed, 54 insertions(+), 123 deletions(-) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index c2547ee6320..d8bc79432e4 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -34,8 +34,8 @@ * **Weight-only / dynamic-activation** (``w4a16_nvfp4``, ``fp8_pc_pt``, ``fp8_pb_wo``) -- every scale is either derived from the weight or computed at runtime per token. * **Static-activation** (``fp8``, ``nvfp4``) -- these need an ``input_quantizer`` amax, - normally measured with calibration data. Instead one fixed amax is applied to every - layer (``--act_scale_amax``, default 448 = input_scale 1.0 for FP8). + normally measured with calibration data. Instead a fixed ``input_scale`` of 1.0 is + applied to every layer. A fixed activation scale sounds crude but measures well, because acceptance length is governed almost entirely by *clipping* rather than by resolution. Sweeping input_scale over @@ -54,12 +54,13 @@ (bf16 baseline 3.142.) Both formats fall off a cliff below ~0.03, where the declared range is far under the activations' true magnitude and most of the tensor is clipped, and both sit on a flat plateau from ~0.3 to 4.0 with no drop-off at the top -- so the scale needs to -be big enough, and little else. NVFP4 trails FP8 by a roughly constant 3.5% across the +be big enough, and little else. 1.0 sits in the middle of that plateau, which is why it is +fixed rather than exposed as a knob. NVFP4 trails FP8 by a roughly constant 3.5% across the plateau: that gap is the 4-bit resolution cost itself and no choice of scale recovers it. -Two weight-derived estimators (``weight_amax``, ``weight_rms``) are kept for reference. -They land 1-2 orders of magnitude below the activation range and measure -31% to -46% AL; -they are not recommended. +Estimating the amax from the weights instead was tried and does not work: max|W| averages +0.79 while a RMSNorm'd activation is O(1) with outlier channels in the tens, so the range +lands 1-2 orders of magnitude low and clips. That measured -31% to -46% AL. AWQ formats are deliberately not offered: ``awq_lite`` silently degrades to plain RTN when no ``forward_loop`` is supplied. @@ -136,10 +137,14 @@ # target, and vLLM's loader then fails with ``KeyError: 'embed_tokens.weight_scale'``. DEFAULT_EXCLUDE = ["*markov_head*", "*confidence_head*", "*embed_tokens*"] -# Largest value FP8 E4M3 can represent, and the default static activation amax. The -# exported input_scale is amax/448 for FP8 but amax/(6*448) for NVFP4 -- the amax is what -# both formats share, which is why the CLI takes it rather than an input_scale. +# Static activation amax applied to every layer of a static-activation format, expressed as +# the amax that yields input_scale 1.0. The exported input_scale is amax/448 for FP8 but +# amax/(6*448) for NVFP4, so this is the FP8 convention; an NVFP4 checkpoint built from the +# same amax records input_scale 0.1667. See the module docstring for why 1.0 and why this is +# not a CLI knob: the usable plateau spans ~0.3 to 4.0, so the value only has to be big +# enough, and 1.0 sits in the middle of it. FP8_E4M3_MAX = 448.0 +STATIC_ACT_AMAX = FP8_E4M3_MAX def parse_args(): @@ -152,38 +157,10 @@ def parse_args(): "--qformat", default="w4a16_nvfp4", choices=SUPPORTED_QFORMATS, - help="Quantization format. Weight-only and dynamic-activation formats are fully " - "calibration-free; static-activation formats (fp8, nvfp4) additionally need " - "--act_scale_heuristic.", - ) - parser.add_argument( - "--act_scale_heuristic", - default="fixed", - choices=["fixed", "none", "weight_amax", "weight_rms"], - help="How to set the static activation amax for formats that need one, without " - "calibration data. 'fixed' (default) applies --act_scale_amax to every " - "layer. 'none' refuses to guess and errors out. The two weight-derived estimators " - "are kept for reference only -- they were measured at -31%% to -46%% acceptance " - "length; see estimate_activation_amax.", - ) - parser.add_argument( - "--act_scale_amax", - type=float, - default=FP8_E4M3_MAX, - help="Static activation amax applied to every layer with --act_scale_heuristic " - "fixed, i.e. the largest activation the quantizer will represent without clipping. " - "Default 448 corresponds to input_scale 1.0 for FP8; it is lossless for FP8 and " - "within 4%% for NVFP4 on Qwen3-8B, and anything from ~134 (scale 0.3) upward sits " - "on the same plateau. Below ~13 the activations get clipped and acceptance length " - "falls off a cliff. Note the exported input_scale is amax/448 for FP8 but " - "amax/(6*448) for NVFP4, so the same amax yields different input_scale values.", - ) - parser.add_argument( - "--act_scale_multiplier", - type=float, - default=1.0, - help="Extra factor on the estimated activation amax. >1 clips less and quantizes " - "coarser; <1 the reverse. Only used with --act_scale_heuristic.", + help="Quantization format. All are calibration-free: weight-only and " + "dynamic-activation formats derive every scale from the weights or per token at " + "runtime, and the static-activation formats (fp8, nvfp4) use a fixed input_scale " + "of 1.0.", ) parser.add_argument( "--dtype", @@ -248,60 +225,19 @@ def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) - return root -def estimate_activation_amax( - module: nn.Linear, method: str, multiplier: float, fixed_value: float = 1.0 -) -> torch.Tensor | None: - """Estimate a static activation amax for ``module`` from its weight alone. +def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) -> int: + """Give every static ``input_quantizer`` the same fixed amax. Returns how many were set. Static-activation formats (fp8, nvfp4) need an ``input_quantizer`` amax, which is - normally measured by pushing calibration data through the model. A drafter cannot be - driven that way here (no importable modeling code, and its real forward takes hidden - states rather than token ids), so these estimators stand in for the measurement. - - Both assume the layer's input is roughly unit-scale -- true for a transformer's linear - inputs, which are RMSNorm outputs -- and use the weight only to set the dynamic range: - - ``weight_amax`` - amax = max |W|. An upper-ish bound: activations feeding a layer with large weights - are assumed to span a comparable range. Errs toward clipping less, quantizing - coarser. - ``weight_rms`` - amax = 4 * rms(W), i.e. a 4-sigma range for a roughly Gaussian weight. Tighter than - weight_amax on layers with a few outlier weights, so it usually preserves more - resolution -- at the cost of clipping genuine activation outliers. - - ``fixed`` - amax = ``--act_scale_amax``, the same value for every layer, - ignoring the weight - entirely. Measured on Qwen3-8B the weight-derived estimators land 1-2 orders of - magnitude below the activations they are meant to bound (max|W| averages 0.79 while - a RMSNorm'd activation is O(1) with outlier channels in the tens), so most values - are clipped. A fixed amax lets the range be set directly on the scale the - activations actually live on. - - None is a substitute for calibration. All are deliberately crude, and which one wins is - model-dependent; measure acceptance length rather than trusting any of them. - """ - weight = module.weight.detach().float() - if method == "weight_amax": - amax = weight.abs().max() - elif method == "weight_rms": - amax = 4.0 * weight.pow(2).mean().sqrt() - elif method == "fixed": - amax = torch.tensor(fixed_value, dtype=torch.float32) - else: - return None - return (amax * multiplier).to(module.weight.dtype) - - -def apply_activation_scale_heuristic( - root: nn.Module, method: str, multiplier: float, fixed_value: float = 1.0 -) -> int: - """Set every enabled ``input_quantizer``'s amax from the weight heuristic. - - Returns the number of quantizers that were given a scale. Only quantizers that are - both enabled and still missing an amax are touched, so dynamic-activation formats - (whose scales are computed per token at runtime) are left alone. + normally *measured* by pushing calibration data through the model. A drafter cannot be + driven that way here: it has no importable modeling code, and its real forward takes + hidden states rather than token ids. See the module docstring for why one fixed value + stands in for that measurement, and why 1.0. + + Only quantizers that are enabled and still missing an amax are touched, so + dynamic-activation formats -- whose scales are computed per token at runtime -- are left + alone, and any amax already established (by a future calibration pass, see + ``resolve_activation_scales``) wins over the fixed default. """ count = 0 for _, module in root.named_modules(): @@ -315,14 +251,34 @@ def apply_activation_scale_heuristic( continue if getattr(input_quantizer, "amax", None) is not None: continue - amax = estimate_activation_amax(module, method, multiplier, fixed_value) - if amax is None: - continue - input_quantizer.amax = amax + input_quantizer.amax = torch.tensor(amax, dtype=torch.float32).to(module.weight.dtype) count += 1 return count +def resolve_activation_scales(root: nn.Module, quant_cfg: dict) -> None: + """Establish activation scales for a format that quantizes activations statically. + + The single place that decides *where* a static activation amax comes from. Today there + is one source -- a fixed value applied uniformly -- because the drafter cannot be run + forward without its modeling code. Real calibration would slot in here as a second + source ahead of the fixed fallback:: + + if calib_forward_loop is not None: + mtq.calibrate(root, quant_cfg["algorithm"], forward_loop=calib_forward_loop) + set_static_activation_amax(root) # fills in whatever calibration did not reach + + ``set_static_activation_amax`` deliberately skips quantizers that already have an amax, + so it composes as a fallback rather than overwriting measured values, and the callers + below do not change when that day comes. + """ + if not need_calibration(quant_cfg): + # Weight-only, or activations scaled per token at runtime -- nothing to establish. + return + n = set_static_activation_amax(root) + print(f"Set {n} static activation amax values (fixed, input_scale 1.0) -- not calibrated.") + + def build_quant_cfg(qformat: str, exclude: list[str], quantize_lm_head: bool) -> dict: """Take the shipped preset and layer the drafter-specific exclusions on top.""" quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[qformat]) @@ -384,34 +340,9 @@ def main(): root = build_linear_view(state_dict, dtype) quant_cfg = build_quant_cfg(args.qformat, args.exclude, args.quantize_lm_head) - needs_static_act = need_calibration(quant_cfg) - if needs_static_act and args.act_scale_heuristic == "none": - raise SystemExit( - f"--qformat {args.qformat} quantizes activations with a static scale, which is " - "normally measured with calibration data. This script has none, so pass " - "--act_scale_heuristic {weight_amax,weight_rms} to estimate it from the weights " - "(approximate -- verify acceptance length), or pick a calibration-free format " - "such as w4a16_nvfp4 (weight-only) or fp8_pc_pt (dynamic per-token activations)." - ) mtq.quantize(root, quant_cfg) # no forward_loop: scales come from the weights - - if needs_static_act: - n = apply_activation_scale_heuristic( - root, - args.act_scale_heuristic, - args.act_scale_multiplier, - args.act_scale_amax, - ) - detail = ( - f"amax={args.act_scale_amax:g}" - if args.act_scale_heuristic == "fixed" - else f"derived from weights (x{args.act_scale_multiplier:g})" - ) - print( - f"Set {n} static activation amax values via '{args.act_scale_heuristic}' " - f"({detail}) -- not calibrated." - ) + resolve_activation_scales(root, quant_cfg) mtq.print_quant_summary(root) From 0da04f969beda72b88a8cd1e6cbd6fc2d595d4bd Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:19:17 +0000 Subject: [PATCH 04/15] example(specdec): trim the drafter PTQ comments Cut commentary that restated the code, enumerated formats already named in the list below it, or narrated failure modes reached while debugging. Kept the non-obvious reasons a later reader would otherwise have to rediscover: why embeddings must be excluded by name under the flat-linear view, why exclusions need suffix wildcards and fused aliases, and why both quant_method/quant_algo and ignore/exclude_modules are written. Exports are unchanged (byte-identical for all five formats). Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../scripts/quantize_drafter.py | 184 +++++------------- 1 file changed, 54 insertions(+), 130 deletions(-) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index d8bc79432e4..a35653a1199 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -15,52 +15,26 @@ """Calibration-free PTQ for a speculative-decoding drafter. -Block-wise ``max`` quantization derives every scale from the weight tensor itself, so -this needs neither a dataset nor a forward pass -- only the drafter's safetensors file. -That in turn means we never have to import the drafter's modeling code: each 2-D weight -is wrapped in a throwaway ``nn.Linear`` under its checkpoint name, and ModelOpt's normal -``quantizer_name`` patterns select over those names exactly as they would on the real -module tree. Works for any drafter layout (DSpark / DFlash / EAGLE3 / Medusa). - -Avoiding the modeling code is what makes this work at all for DFlash-family drafters. An +Every scale is derived from the weights, so this needs no dataset and no forward pass -- +only the drafter's safetensors file. That also means the drafter's modeling code is never +imported: each 2-D weight is wrapped in a throwaway ``nn.Linear`` under its checkpoint +name, and ModelOpt's usual ``quantizer_name`` patterns select over those names. Works for +any drafter layout (DSpark / DFlash / EAGLE3 / Medusa). + +Avoiding the modeling code is what makes this work for DFlash-family drafters at all: an exported drafter declares ``architectures: ["DFlashDraftModel"]`` but ships no importable -class, so ``AutoModelForCausalLM`` silently loads it as a plain Qwen3 and drops ``fc`` / -``hidden_norm`` / the DSpark heads; and the real draft module's ``forward`` takes -``(noise_embedding, target_hidden, ...)``, not ``input_ids``, so the stock ``hf_ptq.py`` -calibration loop cannot drive it either. The flat-linear view sidesteps both problems. - -Two families of format are supported: - -* **Weight-only / dynamic-activation** (``w4a16_nvfp4``, ``fp8_pc_pt``, ``fp8_pb_wo``) -- - every scale is either derived from the weight or computed at runtime per token. -* **Static-activation** (``fp8``, ``nvfp4``) -- these need an ``input_quantizer`` amax, - normally measured with calibration data. Instead a fixed ``input_scale`` of 1.0 is - applied to every layer. - -A fixed activation scale sounds crude but measures well, because acceptance length is -governed almost entirely by *clipping* rather than by resolution. Sweeping input_scale over -three decades on Qwen3-8B + a DSpark drafter (MT-Bench, 80 questions): - -=========== ============== ============== -input_scale FP8 AL NVFP4 AL -=========== ============== ============== -0.003 2.220 (-29.3%) 2.208 (-29.8%) -0.03 2.975 (-5.3%) 2.926 (-6.9%) -0.3 3.137 (-0.2%) 3.022 (-3.8%) -1.0 3.146 (+0.1%) 3.019 (-3.9%) -4.0 3.125 (-0.6%) 3.003 (-4.4%) -=========== ============== ============== - -(bf16 baseline 3.142.) Both formats fall off a cliff below ~0.03, where the declared range -is far under the activations' true magnitude and most of the tensor is clipped, and both -sit on a flat plateau from ~0.3 to 4.0 with no drop-off at the top -- so the scale needs to -be big enough, and little else. 1.0 sits in the middle of that plateau, which is why it is -fixed rather than exposed as a knob. NVFP4 trails FP8 by a roughly constant 3.5% across the -plateau: that gap is the 4-bit resolution cost itself and no choice of scale recovers it. - -Estimating the amax from the weights instead was tried and does not work: max|W| averages -0.79 while a RMSNorm'd activation is O(1) with outlier channels in the tens, so the range -lands 1-2 orders of magnitude low and clips. That measured -31% to -46% AL. +class, so ``AutoModelForCausalLM`` silently loads it as a plain Qwen3 and drops the draft +tensors, and the real draft ``forward`` takes hidden states rather than ``input_ids``, so +the stock ``hf_ptq.py`` calibration loop cannot drive it either. + +``fp8`` and ``nvfp4`` quantize activations against a static amax that is normally measured +on calibration data; a fixed ``input_scale`` of 1.0 is applied instead. That holds up +because acceptance length is governed by clipping rather than resolution: on Qwen3-8B + +a DSpark drafter (MT-Bench), AL falls off a cliff below input_scale ~0.03 but sits on a +flat plateau from ~0.3 to 4.0 with no drop-off at the top, so the scale only has to be big +enough. 1.0 is the middle of that plateau, and measures +0.1% AL for FP8 and -3.9% for +NVFP4 (that gap is the 4-bit resolution cost; no scale recovers it). Deriving the amax +from the weights instead clips badly and measures -31% to -46%. AWQ formats are deliberately not offered: ``awq_lite`` silently degrades to plain RTN when no ``forward_loop`` is supplied. @@ -96,13 +70,6 @@ from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.utils import is_quantized_linear -# Formats this script offers, in the order they are most likely to be wanted. -# w4a16_nvfp4 -- block-16 E2M1 weights, dynamic FP8 E4M3 scales. Smallest, weight-only. -# nvfp4 -- same weights plus NVFP4 activations (static amax -> needs a scale). -# fp8 -- E4M3 weights and activations, per-tensor (static amax -> needs a scale). -# fp8_pc_pt -- E4M3 per-channel weights, DYNAMIC per-token activations. The -# calibration-free way to get FP8 activations. -# fp8_pb_wo -- E4M3 per-block weight-only fallback. # INT8/INT4 weight-only formats are deliberately absent: vLLM's ModelOpt backend accepts # only FP8, FP8_PER_CHANNEL_PER_TOKEN, FP8_PB_WO, NVFP4, W4A16_NVFP4, MXFP8 and # MIXED_PRECISION, so an INT checkpoint quantizes cleanly but cannot be served. @@ -114,35 +81,19 @@ "fp8_pb_wo", ] -# The Markov head writes straight into the draft logits and is only a few percent of the -# drafter, so the bits it would save are not worth the acceptance-rate risk. It is also not -# a plain GEMM in the real module tree -- markov_w1 is an nn.Embedding, which ModelOpt's -# stock presets already exclude via `parent_class: nn.Embedding`; the flat-linear view here -# has lost the original module classes, so the exclusion has to be restated by name. -# -# The confidence head is excluded for the same reason plus a mechanical one: it projects to -# a single output (``[1, hidden]``), so a per-channel scale collapses to a 0-dim tensor and -# the per-channel export path (fp8_pc_pt) indexes it as ``scale[:, None]`` and raises. It is -# one row of weights, so there is nothing to gain by quantizing it. -# -# ``fc`` is deliberately NOT excluded here: it is the single largest non-lm_head tensor and -# the formats' own presets decide its fate. Exclude it explicitly with --exclude '*fc*' when -# comparing acceptance length. -# lm_head is excluded by the preset itself (see --quantize_lm_head). -# -# ``embed_tokens`` must be excluded too. It is 2-D, so the flat-linear view happily treats -# it as a GEMM, but it is an ``nn.Embedding``: a row lookup, not a matmul. ModelOpt's stock -# presets skip embeddings via ``parent_class: nn.Embedding``, which the flat view cannot -# see. Quantizing it also breaks deployment -- a drafter inherits ``embed_tokens`` from the -# target, and vLLM's loader then fails with ``KeyError: 'embed_tokens.weight_scale'``. +# These are 2-D, so the flat-linear view treats them as GEMMs, but none of them is one: +# markov_w1 and embed_tokens are embeddings (row lookups), which ModelOpt's presets +# normally skip via `parent_class: nn.Embedding` -- a class the flat view cannot see, so +# the exclusion is restated by name. Quantizing embed_tokens also breaks deployment, since +# a drafter inherits it from the target (vLLM: KeyError: 'embed_tokens.weight_scale'). +# confidence_head projects to a single output, whose per-channel scale collapses to a +# 0-dim tensor. All are tiny; nothing is lost by leaving them alone. +# Not excluded: `fc` (the presets decide; use --exclude '*fc*' to compare) and `lm_head` +# (the preset already excludes it -- see --quantize_lm_head). DEFAULT_EXCLUDE = ["*markov_head*", "*confidence_head*", "*embed_tokens*"] -# Static activation amax applied to every layer of a static-activation format, expressed as -# the amax that yields input_scale 1.0. The exported input_scale is amax/448 for FP8 but -# amax/(6*448) for NVFP4, so this is the FP8 convention; an NVFP4 checkpoint built from the -# same amax records input_scale 0.1667. See the module docstring for why 1.0 and why this is -# not a CLI knob: the usable plateau spans ~0.3 to 4.0, so the value only has to be big -# enough, and 1.0 sits in the middle of it. +# The amax that yields input_scale 1.0 for FP8. NVFP4 divides by 6*448 instead, so the same +# amax records as input_scale 0.1667 there; both mean the same activation range. FP8_E4M3_MAX = 448.0 STATIC_ACT_AMAX = FP8_E4M3_MAX @@ -228,16 +179,8 @@ def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) - def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) -> int: """Give every static ``input_quantizer`` the same fixed amax. Returns how many were set. - Static-activation formats (fp8, nvfp4) need an ``input_quantizer`` amax, which is - normally *measured* by pushing calibration data through the model. A drafter cannot be - driven that way here: it has no importable modeling code, and its real forward takes - hidden states rather than token ids. See the module docstring for why one fixed value - stands in for that measurement, and why 1.0. - - Only quantizers that are enabled and still missing an amax are touched, so - dynamic-activation formats -- whose scales are computed per token at runtime -- are left - alone, and any amax already established (by a future calibration pass, see - ``resolve_activation_scales``) wins over the fixed default. + Skips quantizers that are dynamic (scale computed per token at runtime) or that already + have an amax, so this composes as a fallback rather than an overwrite. """ count = 0 for _, module in root.named_modules(): @@ -259,21 +202,14 @@ def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) - def resolve_activation_scales(root: nn.Module, quant_cfg: dict) -> None: """Establish activation scales for a format that quantizes activations statically. - The single place that decides *where* a static activation amax comes from. Today there - is one source -- a fixed value applied uniformly -- because the drafter cannot be run - forward without its modeling code. Real calibration would slot in here as a second - source ahead of the fixed fallback:: + The single place deciding where a static amax comes from. Real calibration would slot + in here ahead of the fixed fallback, leaving the CLI and call site unchanged:: if calib_forward_loop is not None: mtq.calibrate(root, quant_cfg["algorithm"], forward_loop=calib_forward_loop) - set_static_activation_amax(root) # fills in whatever calibration did not reach - - ``set_static_activation_amax`` deliberately skips quantizers that already have an amax, - so it composes as a fallback rather than overwriting measured values, and the callers - below do not change when that day comes. + set_static_activation_amax(root) # fills in what calibration did not reach """ if not need_calibration(quant_cfg): - # Weight-only, or activations scaled per token at runtime -- nothing to establish. return n = set_static_activation_amax(root) print(f"Set {n} static activation amax values (fixed, input_scale 1.0) -- not calibrated.") @@ -307,9 +243,9 @@ def export_quantized_state_dict( assert quantization is not None, f"{name}: enabled quantizer resolved to no format" weight_scale = get_weight_scaling_factor(module) weight_scale_2 = get_weight_scaling_factor_2(module) - # A per-channel scale over a single-output projection (``[1, hidden]``) collapses to - # a 0-dim tensor that the packing helpers index as ``scale[:, None]``. Such a layer - # is one row of weights, so skip it rather than crash -- it stays in ``dtype``. + # A per-channel scale over a single-output projection collapses to a 0-dim tensor + # that the packing helpers index as ``scale[:, None]``. Skip rather than crash -- + # it is one row of weights and stays in ``dtype``. if weight_scale is not None and weight_scale.dim() == 0 and module.weight.shape[0] == 1: print(f"Skipping {name}: single-output projection, per-channel scale is scalar") continue @@ -323,9 +259,8 @@ def export_quantized_state_dict( export_sd[f"{name}.weight_scale"] = weight_scale if weight_scale_2 is not None: export_sd[f"{name}.weight_scale_2"] = weight_scale_2 - # Static-activation formats also need the input scale in the checkpoint. Without - # it the runtime has no activation scale to apply, so every heuristic produces a - # byte-identical export and the format silently degrades. + # Static-activation formats need the input scale in the checkpoint too; without it + # the runtime has no activation scale to apply and the format silently degrades. activation_scale = get_activation_scaling_factor(module) if activation_scale is not None: export_sd[f"{name}.input_scale"] = activation_scale @@ -354,11 +289,9 @@ def main(): config = json.loads((source_dir / "config.json").read_text()) hf_quant_config = get_quant_config(root) - # ``get_quant_config`` only knows about the modules in the linear view, so anything the - # view never saw -- embeddings, norms, and any 1-D weight -- is absent from - # ``exclude_modules``. A loader that walks the checkpoint (vLLM does) then expects a - # ``weight_scale`` for those too and dies with e.g. KeyError: 'embed_tokens.weight_scale'. - # List every weight that was not quantized so the exclusion set is complete. + # ``get_quant_config`` only knows the modules in the linear view, so anything it never + # saw (norms, 1-D weights) is missing from ``exclude_modules`` and a loader walking the + # checkpoint expects a ``weight_scale`` for it. List every unquantized weight instead. quantized = { name for name, module in root.named_modules() @@ -375,19 +308,15 @@ def main(): for name in unquantized: if name not in exclude_modules: exclude_modules.append(name) - # A runtime matches this list against its own module prefix, which is usually - # nested relative to the checkpoint key (vLLM builds the draft's ``fc`` at - # ``model.fc`` via ``maybe_prefix``). A bare ``fc`` then fails to match, the layer - # is built quantized, and loading the bf16 weight into a packed parameter raises a - # size assertion. Add a suffix wildcard so the exclusion matches at any depth. + # A runtime matches this against its own module prefix, which is nested relative to + # the checkpoint key (vLLM builds the draft's ``fc`` at ``model.fc``). Add a suffix + # wildcard so the exclusion matches at any depth. wildcard = f"*{name}" if wildcard not in exclude_modules: exclude_modules.append(wildcard) # Runtimes fuse sibling projections into one layer whose name appears in no checkpoint - # key: q/k/v -> ``qkv_proj``, gate/up -> ``gate_up_proj``. Excluding only the individual - # names leaves the fused layer quantized, and its merged shard shapes then disagree with - # the unpacked weights being loaded into it. Emit the fused aliases whenever every - # component of a fusion group was excluded. + # key (q/k/v -> ``qkv_proj``, gate/up -> ``gate_up_proj``), so excluding only the parts + # leaves the fused layer quantized. Emit the alias once every component is excluded. for fused, parts in ( ("qkv_proj", ("q_proj", "k_proj", "v_proj")), ("gate_up_proj", ("gate_proj", "up_proj")), @@ -398,21 +327,16 @@ def main(): exclude_modules.append(alias) hf_quant_config["quantization"]["exclude_modules"] = exclude_modules config["quantization_config"] = dict(hf_quant_config["quantization"]) - # ModelOpt names the format ``quant_algo``; vLLM's ModelConfig reads - # ``quant_cfg["quant_method"]`` and treats a config without that key as unquantized, - # then dies on the packed weight shapes. Emit both so either loader is satisfied. - # vLLM splits ModelOpt checkpoints into two backends: ``modelopt_fp4`` for the - # block-scaled NVFP4 layouts and ``modelopt`` for the per-tensor/per-channel ones. + # ModelOpt names the format ``quant_algo``; vLLM reads ``quant_method`` and treats its + # absence as unquantized. Emit both. vLLM splits ModelOpt checkpoints across two + # backends: ``modelopt_fp4`` for block-scaled NVFP4, ``modelopt`` for the rest. quant_algo = str(hf_quant_config["quantization"].get("quant_algo") or "") config["quantization_config"].setdefault( "quant_method", "modelopt_fp4" if "NVFP4" in quant_algo.upper() else "modelopt" ) - # The exclusion list is read under two different keys. ModelOpt's own - # ``hf_quant_config.json`` nests it under ``quantization.exclude_modules``, but when a - # runtime parses the flat ``quantization_config`` block inside ``config.json`` it looks - # for ``ignore`` (vLLM's ModelOptQuantConfigBase.from_config). Emitting only - # ``exclude_modules`` there yields an empty exclusion set, every layer is built - # quantized, and loading an untouched bf16 weight into a packed parameter raises. + # The exclusion list is read under two different keys: ModelOpt nests it under + # ``quantization.exclude_modules``, but a runtime parsing the flat + # ``quantization_config`` in config.json looks for ``ignore``. Emit both. config["quantization_config"]["ignore"] = list(exclude_modules) config["torch_dtype"] = args.dtype (export_dir / "config.json").write_text(json.dumps(config, indent=2)) From 851d4a25aa78ed0cbe24c78c1e6e2bcdc48e1fca Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:29:53 +0000 Subject: [PATCH 05/15] example(specdec): tighten comments and docstrings Address review: drop the hf_ptq background paragraph from the module docstring and cut the draft-quantization comment in specdec_bench to two lines. Same pass over the rest of the script -- remove what the code already says, keep only the non-obvious reasons. No behaviour change; exports are byte-identical for all five formats. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../specdec_bench/models/vllm.py | 9 +- .../scripts/quantize_drafter.py | 119 +++++++----------- 2 files changed, 46 insertions(+), 82 deletions(-) diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index e0801a4585a..35d60e039f2 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/examples/specdec_bench/specdec_bench/models/vllm.py @@ -119,12 +119,9 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs elif kwargs.get("speculative_algorithm") == "NONE": specdec = None - # A quantized draft checkpoint has to declare its own quantization. vLLM otherwise - # copies the target's onto the draft (config/speculative.py, "Align the - # quantization of draft model"), so a quantized drafter under a bf16 target is - # built as if it were bf16 and dies loading the packed weights. The drafter's - # config.json already records the format, so read it back rather than making the - # caller repeat it; --draft_quantization overrides. + # vLLM copies the target's quantization onto the draft, so a quantized drafter under + # a bf16 target is built as bf16 and dies loading the packed weights. Read the format + # from the drafter's own config.json; --draft_quantization overrides. if specdec is not None and specdec.get("model"): draft_quantization = kwargs.get("draft_quantization") if draft_quantization is None: diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index a35653a1199..269716b9a0d 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -15,29 +15,18 @@ """Calibration-free PTQ for a speculative-decoding drafter. -Every scale is derived from the weights, so this needs no dataset and no forward pass -- -only the drafter's safetensors file. That also means the drafter's modeling code is never -imported: each 2-D weight is wrapped in a throwaway ``nn.Linear`` under its checkpoint -name, and ModelOpt's usual ``quantizer_name`` patterns select over those names. Works for -any drafter layout (DSpark / DFlash / EAGLE3 / Medusa). - -Avoiding the modeling code is what makes this work for DFlash-family drafters at all: an -exported drafter declares ``architectures: ["DFlashDraftModel"]`` but ships no importable -class, so ``AutoModelForCausalLM`` silently loads it as a plain Qwen3 and drops the draft -tensors, and the real draft ``forward`` takes hidden states rather than ``input_ids``, so -the stock ``hf_ptq.py`` calibration loop cannot drive it either. - -``fp8`` and ``nvfp4`` quantize activations against a static amax that is normally measured -on calibration data; a fixed ``input_scale`` of 1.0 is applied instead. That holds up -because acceptance length is governed by clipping rather than resolution: on Qwen3-8B + -a DSpark drafter (MT-Bench), AL falls off a cliff below input_scale ~0.03 but sits on a -flat plateau from ~0.3 to 4.0 with no drop-off at the top, so the scale only has to be big -enough. 1.0 is the middle of that plateau, and measures +0.1% AL for FP8 and -3.9% for -NVFP4 (that gap is the 4-bit resolution cost; no scale recovers it). Deriving the amax -from the weights instead clips badly and measures -31% to -46%. - -AWQ formats are deliberately not offered: ``awq_lite`` silently degrades to plain RTN when -no ``forward_loop`` is supplied. +Every scale is derived from the weights, so this needs no dataset and no forward pass, and +never imports the drafter's modeling code: each 2-D weight is wrapped in a throwaway +``nn.Linear`` under its checkpoint name, and ModelOpt's usual ``quantizer_name`` patterns +select over those names. Works for any drafter layout (DSpark / DFlash / EAGLE3 / Medusa), +including exported ones that ship no importable model class. + +``fp8`` and ``nvfp4`` need a static activation amax, normally measured on calibration data; +a fixed ``input_scale`` of 1.0 is applied instead. Acceptance length is governed by +clipping rather than resolution, and AL sits on a flat plateau from input_scale ~0.3 to 4.0 +(Qwen3-8B + DSpark, MT-Bench: +0.1% for FP8, -3.9% for NVFP4), so the scale only has to be +big enough. AWQ is not offered: ``awq_lite`` silently degrades to RTN without a +``forward_loop``. Example: python quantize_drafter.py \ @@ -70,9 +59,8 @@ from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.utils import is_quantized_linear -# INT8/INT4 weight-only formats are deliberately absent: vLLM's ModelOpt backend accepts -# only FP8, FP8_PER_CHANNEL_PER_TOKEN, FP8_PB_WO, NVFP4, W4A16_NVFP4, MXFP8 and -# MIXED_PRECISION, so an INT checkpoint quantizes cleanly but cannot be served. +# INT8/INT4 are absent on purpose: they quantize cleanly but vLLM's ModelOpt backend +# cannot serve them. SUPPORTED_QFORMATS = [ "w4a16_nvfp4", "nvfp4", @@ -81,19 +69,14 @@ "fp8_pb_wo", ] -# These are 2-D, so the flat-linear view treats them as GEMMs, but none of them is one: -# markov_w1 and embed_tokens are embeddings (row lookups), which ModelOpt's presets -# normally skip via `parent_class: nn.Embedding` -- a class the flat view cannot see, so -# the exclusion is restated by name. Quantizing embed_tokens also breaks deployment, since -# a drafter inherits it from the target (vLLM: KeyError: 'embed_tokens.weight_scale'). -# confidence_head projects to a single output, whose per-channel scale collapses to a -# 0-dim tensor. All are tiny; nothing is lost by leaving them alone. -# Not excluded: `fc` (the presets decide; use --exclude '*fc*' to compare) and `lm_head` -# (the preset already excludes it -- see --quantize_lm_head). +# All 2-D, so the flat view treats them as GEMMs, but none is one: markov_w1/embed_tokens +# are embeddings (ModelOpt's presets skip these via `parent_class`, which the flat view +# cannot see), and confidence_head has a single output whose per-channel scale is 0-dim. +# All tiny. `fc` is left to the presets; `lm_head` is excluded by the preset itself. DEFAULT_EXCLUDE = ["*markov_head*", "*confidence_head*", "*embed_tokens*"] -# The amax that yields input_scale 1.0 for FP8. NVFP4 divides by 6*448 instead, so the same -# amax records as input_scale 0.1667 there; both mean the same activation range. +# The amax that yields input_scale 1.0 for FP8. NVFP4 divides by 6*448, so the same amax +# records as 0.1667 there; both mean the same activation range. FP8_E4M3_MAX = 448.0 STATIC_ACT_AMAX = FP8_E4M3_MAX @@ -108,10 +91,8 @@ def parse_args(): "--qformat", default="w4a16_nvfp4", choices=SUPPORTED_QFORMATS, - help="Quantization format. All are calibration-free: weight-only and " - "dynamic-activation formats derive every scale from the weights or per token at " - "runtime, and the static-activation formats (fp8, nvfp4) use a fixed input_scale " - "of 1.0.", + help="Quantization format. All are calibration-free; fp8 and nvfp4 additionally " + "quantize activations, using a fixed input_scale of 1.0.", ) parser.add_argument( "--dtype", @@ -130,8 +111,8 @@ def parse_args(): parser.add_argument( "--quantize_lm_head", action="store_true", - help="Also quantize lm_head. It is the single largest drafter tensor, but it feeds " - "the acceptance test directly -- measure AL before shipping this.", + help="Also quantize lm_head -- the largest drafter tensor, but it feeds the " + "acceptance test directly, so measure AL first.", ) return parser.parse_args() @@ -155,9 +136,8 @@ def load_drafter(drafter_path: str) -> tuple[Path, dict[str, torch.Tensor]]: def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) -> nn.Module: """Expose every 2-D weight as an nn.Linear whose module name is its checkpoint key. - Nested ModuleDicts are used so that ``named_modules()`` reproduces the dotted keys - (``layers.0.self_attn.q_proj``), which is what the preset's ``quantizer_name`` - patterns match against. + Nested ModuleDicts so ``named_modules()`` reproduces the dotted checkpoint keys, which + is what ``quantizer_name`` patterns match against. """ root = nn.ModuleDict() for key, weight in state_dict.items(): @@ -179,8 +159,8 @@ def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) - def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) -> int: """Give every static ``input_quantizer`` the same fixed amax. Returns how many were set. - Skips quantizers that are dynamic (scale computed per token at runtime) or that already - have an amax, so this composes as a fallback rather than an overwrite. + Skips dynamic quantizers and any that already have an amax, so it composes as a + fallback rather than an overwrite. """ count = 0 for _, module in root.named_modules(): @@ -189,7 +169,6 @@ def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) - input_quantizer = getattr(module, "input_quantizer", None) if input_quantizer is None or not input_quantizer.is_enabled: continue - # A dynamic quantizer derives its scale at runtime; leave it untouched. if getattr(input_quantizer, "_dynamic", False): continue if getattr(input_quantizer, "amax", None) is not None: @@ -202,12 +181,8 @@ def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) - def resolve_activation_scales(root: nn.Module, quant_cfg: dict) -> None: """Establish activation scales for a format that quantizes activations statically. - The single place deciding where a static amax comes from. Real calibration would slot - in here ahead of the fixed fallback, leaving the CLI and call site unchanged:: - - if calib_forward_loop is not None: - mtq.calibrate(root, quant_cfg["algorithm"], forward_loop=calib_forward_loop) - set_static_activation_amax(root) # fills in what calibration did not reach + The single place deciding where a static amax comes from: real calibration would call + ``mtq.calibrate`` here, ahead of the fixed fallback. """ if not need_calibration(quant_cfg): return @@ -232,8 +207,7 @@ def export_quantized_state_dict( ) -> dict[str, torch.Tensor]: """Pack each quantized weight and emit it alongside its scales. - Follows the unified-HF naming convention: ``w.weight`` / ``w.weight_scale`` / - ``w.weight_scale_2``. Untouched tensors are carried through in ``dtype``. + Unified-HF naming (``w.weight_scale`` etc). Untouched tensors carry through in ``dtype``. """ export_sd = {k: v.to(dtype) for k, v in state_dict.items()} for name, module in root.named_modules(): @@ -243,9 +217,8 @@ def export_quantized_state_dict( assert quantization is not None, f"{name}: enabled quantizer resolved to no format" weight_scale = get_weight_scaling_factor(module) weight_scale_2 = get_weight_scaling_factor_2(module) - # A per-channel scale over a single-output projection collapses to a 0-dim tensor - # that the packing helpers index as ``scale[:, None]``. Skip rather than crash -- - # it is one row of weights and stays in ``dtype``. + # The packing helpers index the scale as ``scale[:, None]``, which a 0-dim scale + # cannot satisfy. One row of weights, so leave it in ``dtype``. if weight_scale is not None and weight_scale.dim() == 0 and module.weight.shape[0] == 1: print(f"Skipping {name}: single-output projection, per-channel scale is scalar") continue @@ -259,8 +232,7 @@ def export_quantized_state_dict( export_sd[f"{name}.weight_scale"] = weight_scale if weight_scale_2 is not None: export_sd[f"{name}.weight_scale_2"] = weight_scale_2 - # Static-activation formats need the input scale in the checkpoint too; without it - # the runtime has no activation scale to apply and the format silently degrades. + # Without this the runtime has no activation scale and the format silently degrades. activation_scale = get_activation_scaling_factor(module) if activation_scale is not None: export_sd[f"{name}.input_scale"] = activation_scale @@ -289,9 +261,8 @@ def main(): config = json.loads((source_dir / "config.json").read_text()) hf_quant_config = get_quant_config(root) - # ``get_quant_config`` only knows the modules in the linear view, so anything it never - # saw (norms, 1-D weights) is missing from ``exclude_modules`` and a loader walking the - # checkpoint expects a ``weight_scale`` for it. List every unquantized weight instead. + # ``get_quant_config`` only knows the linear view, so tensors it never saw (norms, 1-D + # weights) are missing and a loader walking the checkpoint expects a scale for them. quantized = { name for name, module in root.named_modules() @@ -308,15 +279,13 @@ def main(): for name in unquantized: if name not in exclude_modules: exclude_modules.append(name) - # A runtime matches this against its own module prefix, which is nested relative to - # the checkpoint key (vLLM builds the draft's ``fc`` at ``model.fc``). Add a suffix - # wildcard so the exclusion matches at any depth. + # Runtimes match against their own module prefix, which is nested relative to the + # checkpoint key (vLLM builds the draft's ``fc`` at ``model.fc``). wildcard = f"*{name}" if wildcard not in exclude_modules: exclude_modules.append(wildcard) - # Runtimes fuse sibling projections into one layer whose name appears in no checkpoint - # key (q/k/v -> ``qkv_proj``, gate/up -> ``gate_up_proj``), so excluding only the parts - # leaves the fused layer quantized. Emit the alias once every component is excluded. + # Runtimes fuse sibling projections into one layer whose name is in no checkpoint key, + # so excluding only the parts would leave the fused layer quantized. for fused, parts in ( ("qkv_proj", ("q_proj", "k_proj", "v_proj")), ("gate_up_proj", ("gate_proj", "up_proj")), @@ -328,15 +297,13 @@ def main(): hf_quant_config["quantization"]["exclude_modules"] = exclude_modules config["quantization_config"] = dict(hf_quant_config["quantization"]) # ModelOpt names the format ``quant_algo``; vLLM reads ``quant_method`` and treats its - # absence as unquantized. Emit both. vLLM splits ModelOpt checkpoints across two - # backends: ``modelopt_fp4`` for block-scaled NVFP4, ``modelopt`` for the rest. + # absence as unquantized, splitting NVFP4 off into its own backend. Emit both. quant_algo = str(hf_quant_config["quantization"].get("quant_algo") or "") config["quantization_config"].setdefault( "quant_method", "modelopt_fp4" if "NVFP4" in quant_algo.upper() else "modelopt" ) - # The exclusion list is read under two different keys: ModelOpt nests it under - # ``quantization.exclude_modules``, but a runtime parsing the flat - # ``quantization_config`` in config.json looks for ``ignore``. Emit both. + # Same list, second key: the flat ``quantization_config`` in config.json is read for + # ``ignore``, not ``exclude_modules``. config["quantization_config"]["ignore"] = list(exclude_modules) config["torch_dtype"] = args.dtype (export_dir / "config.json").write_text(json.dumps(config, indent=2)) From ac2631f03768ef42a5c8e5e8293f75a90405dbce Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:48:38 +0000 Subject: [PATCH 06/15] example(launcher): add a DSpark NVFP4 PTQ pipeline for Qwen3-8B Quantizes an exported DSpark drafter weight-only to NVFP4 and measures acceptance length, so the cost of quantizing is visible in the same run. Needs no calibration data. Adds the thin common/specdec/quantize_drafter.sh wrapper, which also resolves a training output_dir to its newest exported-checkpoint-. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../common/specdec/quantize_drafter.sh | 43 +++++++++++ .../Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml | 72 +++++++++++++++++++ 2 files changed, 115 insertions(+) create mode 100644 tools/launcher/common/specdec/quantize_drafter.sh create mode 100644 tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml diff --git a/tools/launcher/common/specdec/quantize_drafter.sh b/tools/launcher/common/specdec/quantize_drafter.sh new file mode 100644 index 00000000000..15db02f627e --- /dev/null +++ b/tools/launcher/common/specdec/quantize_drafter.sh @@ -0,0 +1,43 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Calibration-free PTQ for an exported speculative-decoding drafter. +# +# Required env vars: +# DRAFTER_CKPT — exported drafter path, or a training output_dir to auto-detect under + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source ${SCRIPT_DIR}/../service_utils.sh + +trap 'error_handler $0 $LINENO' ERR + +################################################################################################### + +DRAFTER=${DRAFTER_CKPT} +# Training writes exported-checkpoint-/ under output_dir; take the newest. +if [ ! -f "${DRAFTER}/config.json" ]; then + DRAFTER=$(ls -d ${DRAFTER_CKPT}/exported-checkpoint-* 2>/dev/null | sort -t- -k3 -n | tail -1) + if [ -z "${DRAFTER}" ]; then + echo "ERROR: no drafter checkpoint at ${DRAFTER_CKPT}" + exit 1 + fi + echo "Auto-detected drafter: ${DRAFTER}" +fi + +python modules/Model-Optimizer/examples/speculative_decoding/scripts/quantize_drafter.py \ + --drafter_path ${DRAFTER} \ + ${@} diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml new file mode 100644 index 00000000000..00dedac8053 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml @@ -0,0 +1,72 @@ +# Calibration-free NVFP4 PTQ of a DSpark drafter for Qwen3-8B. +# +# Takes an exported drafter (train one with hf_streaming_dspark.yaml, or point +# drafter at a published checkpoint) and quantizes it weight-only to NVFP4, then +# measures acceptance length so the cost is visible. No calibration data needed: +# every scale comes from the weights. +# +# 2-step pipeline: +# task_0: Quantize the drafter (CPU-only, ~1 min for an 8B-class draft) +# task_1: Benchmark acceptance length on MT-Bench via vLLM +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml --yes + +job_name: Qwen3-8B_DSpark_PTQ_nvfp4 +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + # Exported drafter to quantize. /scratchspace/export is where the streaming + # examples leave theirs; a plain checkpoint dir or HF repo id also works. + draft_model: /scratchspace/export + + # Step 1: Quantize. w4a16_nvfp4 keeps activations in bf16; use --qformat nvfp4 + # for weight+activation (fixed input_scale 1.0, ~3.9% AL on Qwen3-8B). + # DFlash-family drafters read qkv_proj.weight raw for their fused context-KV + # projection, so those layers must stay unquantized; o_proj and the MLP still + # quantize. + task_0: + script: common/specdec/quantize_drafter.sh + args: + - --qformat w4a16_nvfp4 + - --export_path /scratchspace/export_quantized + - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' '*fc*' + environment: + - DRAFTER_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 2: Acceptance length on MT-Bench. Compare against the same run with + # --draft_model_dir <> to see what quantization cost. + task_1: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export_quantized + # DFLASH reads --block_size, not --draft_length; must match the drafter's + # dflash_block_size. + - --block_size 16 + - --output_length 4096 + - --engine VLLM + - --tp_size 1 + - --ep_size 1 + # DSpark runs through the DFLASH path; the drafter's own config carries the + # DSpark heads. + - --speculative_algorithm DFLASH + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:nightly From c564e254df5cf5a37e1e44e0f715f7e024305688 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 12:07:47 +0000 Subject: [PATCH 07/15] specdec_bench: add DSPARK speculative algorithm The launcher PTQ example benchmarks a DSpark drafter, but specdec_bench had no DSPARK path -- an exported Qwen3DSparkModel would have gone through DFLASH and been built with vLLM method="dflash". Adds the branch: vLLM method="dspark", and a draft_sample_method matched to the target's verify mode (greedy target + probabilistic draft, or the reverse, crushes acceptance at temp > 0). DSpark also runs eager, since its block-parallel draft can outgrow the workspace during CUDA-graph capture; acceptance length is unaffected by graph capture. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- examples/specdec_bench/run.py | 2 +- .../specdec_bench/models/vllm.py | 21 ++++++++++++++++++- .../Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml | 8 +++---- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index 342dffb9ad8..b398de589f6 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -306,7 +306,7 @@ def run_simple(args): type=str, required=False, default="EAGLE3", - choices=["EAGLE3", "EAGLE", "DRAFT_TARGET", "NGRAM", "MTP", "DFLASH", "NONE"], + choices=["EAGLE3", "EAGLE", "DRAFT_TARGET", "NGRAM", "MTP", "DFLASH", "DSPARK", "NONE"], help="Speculative algorithm to use", ) parser.add_argument("--model_dir", type=str, required=True, help="Path to the model directory") diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index 35d60e039f2..db352c117e2 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/examples/specdec_bench/specdec_bench/models/vllm.py @@ -116,6 +116,17 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs "model": kwargs.get("draft_model_dir"), "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 8), } + elif kwargs.get("speculative_algorithm") == "DSPARK": + # Match draft sampling to the target's verify mode: a greedy target with a + # probabilistic draft (or the reverse) crushes acceptance at temp > 0. + temperature = sampling_kwargs.get("temperature", 1.0) + specdec = { + "method": "dspark", + "model": kwargs.get("draft_model_dir"), + "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 7), + "draft_sample_method": kwargs.get("dspark_draft_sample_method") + or ("greedy" if temperature == 0 else "probabilistic"), + } elif kwargs.get("speculative_algorithm") == "NONE": specdec = None @@ -139,6 +150,14 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs else: num_speculative_tokens = specdec.get("num_speculative_tokens", 3) + # DSpark's block-parallel draft can outgrow the pre-allocated workspace during + # CUDA-graph capture; acceptance length is unaffected by graph capture, so skip it. + # SPECDEC_ENFORCE_EAGER=1 forces the same for other algorithms. + enforce_eager = ( + kwargs.get("speculative_algorithm") == "DSPARK" + or os.environ.get("SPECDEC_ENFORCE_EAGER") == "1" + ) + engine_args = AsyncEngineArgs( model=model_dir, tokenizer=kwargs.get("tokenizer_path"), @@ -150,7 +169,7 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs max_num_seqs=max_concurrent_requests * num_speculative_tokens, skip_tokenizer_init=False, async_scheduling=kwargs.get("async_scheduling", True), - enforce_eager=False, + enforce_eager=enforce_eager, max_model_len=kwargs.get("max_model_len"), ) self.engine_args = engine_args diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml index 00dedac8053..01d4e3b7716 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml @@ -50,16 +50,14 @@ pipeline: script: common/specdec_bench/quick_check.sh args: - --draft_model_dir /scratchspace/export_quantized - # DFLASH reads --block_size, not --draft_length; must match the drafter's - # dflash_block_size. + # DSPARK/DFLASH read --block_size, not --draft_length; must match the + # drafter's dflash_block_size. - --block_size 16 - --output_length 4096 - --engine VLLM - --tp_size 1 - --ep_size 1 - # DSpark runs through the DFLASH path; the drafter's own config carries the - # DSpark heads. - - --speculative_algorithm DFLASH + - --speculative_algorithm DSPARK - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - --concurrency 32 environment: From 907dd45eabf09c0282b6fe95307c39661db0a633 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:02:13 +0000 Subject: [PATCH 08/15] example(launcher): quantize fc by default in the DSpark PTQ example Measured on Qwen3-8B: quantizing fc saves ~12% more size (1.6 -> 1.4 GiB) for ~0.7 points of acceptance length (3.0392 -> 3.0186). Worth taking by default, and the comment now separates it from the q/k/v exclusions, which are mandatory rather than a tuning choice. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml index 01d4e3b7716..759792bcc11 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml @@ -26,15 +26,18 @@ pipeline: # Step 1: Quantize. w4a16_nvfp4 keeps activations in bf16; use --qformat nvfp4 # for weight+activation (fixed input_scale 1.0, ~3.9% AL on Qwen3-8B). - # DFlash-family drafters read qkv_proj.weight raw for their fused context-KV - # projection, so those layers must stay unquantized; o_proj and the MLP still - # quantize. + # + # The q/k/v exclusions are mandatory, not a tuning choice: DFlash-family + # drafters build their fused context-KV projection by reading qkv_proj.weight + # raw, which cannot be a packed tensor. o_proj and the MLP still quantize. + # `fc` is optional -- quantizing it saves ~12% more size for ~0.7% AL on + # Qwen3-8B (3.0186 vs 3.0392); add '*fc*' to keep it in bf16. task_0: script: common/specdec/quantize_drafter.sh args: - --qformat w4a16_nvfp4 - --export_path /scratchspace/export_quantized - - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' '*fc*' + - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' environment: - DRAFTER_CKPT: <> slurm_config: From dbedb9c1a68b1df00e7352663b1b38180c6dad8a Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:10:11 +0000 Subject: [PATCH 09/15] example(specdec): fix --quantize_lm_head on W+A formats, drop fp8_pb_wo Review findings: --quantize_lm_head re-enabled only lm_head's weight quantizer, but the preset disables all of them. On fp8/nvfp4 that exported lm_head with weight_scale and no input_scale, while the config still advertised it as fully quantized and left it out of the exclusion list -- a load error on the one layer the flag exists to quantize. Re-enable the input quantizer too. Verified: lm_head.input_scale is now written. fp8_pb_wo is the one offered format with no canonical quant_algo mapping (it falls through to a raw lowercase "fp8_pb_wo" where the others emit "FP8" / "FP8_PER_CHANNEL_PER_TOKEN") and no measurement behind it, so the claim that every offered format is servable did not hold for it. Dropped. Also keep the static amax in fp32 rather than casting to the weight dtype; exact for 448.0 but bf16 would round an arbitrary measured amax by ~0.3%. Exports for the remaining formats are byte-identical. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../scripts/quantize_drafter.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index 269716b9a0d..0343d375918 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -66,7 +66,6 @@ "nvfp4", "fp8", "fp8_pc_pt", - "fp8_pb_wo", ] # All 2-D, so the flat view treats them as GEMMs, but none is one: markov_w1/embed_tokens @@ -173,7 +172,9 @@ def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) - continue if getattr(input_quantizer, "amax", None) is not None: continue - input_quantizer.amax = torch.tensor(amax, dtype=torch.float32).to(module.weight.dtype) + # Keep amax in fp32, as ModelOpt does everywhere else -- casting to the weight + # dtype would round a measured amax through bf16's 8-bit mantissa. + input_quantizer.amax = torch.tensor(amax, dtype=torch.float32) count += 1 return count @@ -194,9 +195,13 @@ def build_quant_cfg(qformat: str, exclude: list[str], quantize_lm_head: bool) -> """Take the shipped preset and layer the drafter-specific exclusions on top.""" quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[qformat]) if quantize_lm_head: - quant_cfg["quant_cfg"].append( - {"quantizer_name": "*lm_head*weight_quantizer", "enable": True} - ) + # The preset disables *all* of lm_head's quantizers. Re-enabling only the weight + # one would leave a W+A format exporting lm_head with no input_scale while the + # config still advertises it as fully quantized, which a runtime fails to load. + for quantizer in ("weight_quantizer", "input_quantizer"): + quant_cfg["quant_cfg"].append( + {"quantizer_name": f"*lm_head*{quantizer}", "enable": True} + ) for pattern in DEFAULT_EXCLUDE + exclude: quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) return quant_cfg From 63ce3b5c848f5da620693039bb8f3f8f7bfed5f7 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:23:59 +0000 Subject: [PATCH 10/15] example(specdec): address review findings on the wrapper and export Wrapper: quote "$DRAFTER" and "$@". Unquoted ${@} globbed the wildcard exclusions the launcher YAML passes -- verified: with a file matching *q_proj* in the CWD, '*q_proj*' expanded to that filename and the drafter would have been quantized with the wrong exclusion set. Also restrict checkpoint auto-detection to an existing local directory so an HF repo id passes through instead of erroring, and sort with -V so exported-checkpoint-1000 beats -900. Export: copy the modeling files an auto_map points at. The export carries the source config verbatim, so a drafter shipping custom modeling code would otherwise reference .py files that are not there. (DFlash/DSpark exports have no auto_map; this is for the ones that do.) Download: pass allow_patterns to snapshot_download rather than pulling whole repos for the weights, config and three sidecars this reads. Exports are byte-identical for all four formats. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../scripts/quantize_drafter.py | 21 +++++++++++++++++-- .../common/specdec/quantize_drafter.sh | 20 +++++++++++------- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index 0343d375918..bb527c0544c 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -74,6 +74,10 @@ # All tiny. `fc` is left to the presets; `lm_head` is excluded by the preset itself. DEFAULT_EXCLUDE = ["*markov_head*", "*confidence_head*", "*embed_tokens*"] +# Sidecars carried over to the export, and -- with the weights and config.json -- the only +# files fetched when --drafter_path is a repo id rather than a local directory. +SIDECAR_FILES = ("tokenizer.json", "tokenizer_config.json", "generation_config.json") + # The amax that yields input_scale 1.0 for FP8. NVFP4 divides by 6*448, so the same amax # records as 0.1667 there; both mean the same activation range. FP8_E4M3_MAX = 448.0 @@ -122,7 +126,11 @@ def load_drafter(drafter_path: str) -> tuple[Path, dict[str, torch.Tensor]]: if not local_dir.is_dir(): from huggingface_hub import snapshot_download - local_dir = Path(snapshot_download(drafter_path)) + local_dir = Path( + snapshot_download( + drafter_path, allow_patterns=["*.safetensors", "config.json", *SIDECAR_FILES] + ) + ) shards = sorted(local_dir.glob("*.safetensors")) assert shards, f"No .safetensors found under {local_dir}" @@ -314,10 +322,19 @@ def main(): (export_dir / "config.json").write_text(json.dumps(config, indent=2)) (export_dir / "hf_quant_config.json").write_text(json.dumps(hf_quant_config, indent=2)) - for extra in ("tokenizer.json", "tokenizer_config.json", "generation_config.json"): + for extra in SIDECAR_FILES: if (source_dir / extra).is_file(): shutil.copy2(source_dir / extra, export_dir / extra) + # A drafter that ships custom modeling code points at it from auto_map; the export + # carries that config verbatim, so the .py files have to come along or the reference + # dangles. (The DFlash/DSpark exports have no auto_map -- this is for the ones that do.) + if config.get("auto_map"): + for module in {ref.split("--")[-1].split(".")[0] for ref in config["auto_map"].values()}: + source_py = source_dir / f"{module}.py" + if source_py.is_file(): + shutil.copy2(source_py, export_dir / source_py.name) + before = sum(v.numel() * v.element_size() for v in state_dict.values()) after = sum(v.numel() * v.element_size() for v in export_sd.values()) print(f"\n{args.qformat}: {before / 2**30:.2f} GiB -> {after / 2**30:.2f} GiB") diff --git a/tools/launcher/common/specdec/quantize_drafter.sh b/tools/launcher/common/specdec/quantize_drafter.sh index 15db02f627e..ddaac78ed00 100644 --- a/tools/launcher/common/specdec/quantize_drafter.sh +++ b/tools/launcher/common/specdec/quantize_drafter.sh @@ -27,17 +27,21 @@ trap 'error_handler $0 $LINENO' ERR ################################################################################################### -DRAFTER=${DRAFTER_CKPT} -# Training writes exported-checkpoint-/ under output_dir; take the newest. -if [ ! -f "${DRAFTER}/config.json" ]; then - DRAFTER=$(ls -d ${DRAFTER_CKPT}/exported-checkpoint-* 2>/dev/null | sort -t- -k3 -n | tail -1) - if [ -z "${DRAFTER}" ]; then - echo "ERROR: no drafter checkpoint at ${DRAFTER_CKPT}" +DRAFTER="${DRAFTER_CKPT}" +# Training writes exported-checkpoint-/ under output_dir; take the newest. Only for a +# local directory -- anything else (an HF repo id) is passed through for the script to +# resolve. -V sorts numerically, so checkpoint-1000 beats checkpoint-900. +if [ -d "${DRAFTER}" ] && [ ! -f "${DRAFTER}/config.json" ]; then + latest=$(find "${DRAFTER}" -maxdepth 1 -mindepth 1 -type d \ + -name 'exported-checkpoint-*' -printf '%f\n' 2>/dev/null | sort -V | tail -1) + if [ -z "${latest}" ]; then + echo "ERROR: ${DRAFTER} is not a checkpoint and holds no exported-checkpoint-* directory." exit 1 fi + DRAFTER="${DRAFTER}/${latest}" echo "Auto-detected drafter: ${DRAFTER}" fi python modules/Model-Optimizer/examples/speculative_decoding/scripts/quantize_drafter.py \ - --drafter_path ${DRAFTER} \ - ${@} + --drafter_path "${DRAFTER}" \ + "$@" From 7a79da146e007b755d1cd352ccdba351c3caf474 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:27:54 +0000 Subject: [PATCH 11/15] example(specdec): build the linear view on the meta device nn.Linear was allocating and running reset_parameters() over a full-size weight that the next line replaced, costing a whole model's worth of RNG plus a transient allocation per layer. Constructing on meta skips both: ~4.6 s -> ~0.3 s of setup for a 30-layer model, and peak host memory drops by the largest layer. Exports are byte-identical. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- examples/speculative_decoding/scripts/quantize_drafter.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index bb527c0544c..72277d60761 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -157,8 +157,11 @@ def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) - node[part] = nn.ModuleDict() node = node[part] out_features, in_features = weight.shape - linear = nn.Linear(in_features, out_features, bias=False, dtype=dtype) - linear.weight.data = weight.to(dtype) + # On meta, so nn.Linear skips allocating and randomly initializing a weight that + # the next line replaces anyway. + with torch.device("meta"): + linear = nn.Linear(in_features, out_features, bias=False, dtype=dtype) + linear.weight = nn.Parameter(weight.to(dtype), requires_grad=False) node[leaf] = linear return root From bb61f9747c3ef8774ddd68c0e6400a19963b5f5c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:47:14 +0000 Subject: [PATCH 12/15] example(launcher): DSpark PTQ example for Nemotron-3.5-Lightning-30B-A3B Same two-step pipeline as the Qwen3-8B example against a hybrid Mamba-MoE target and the published DSpark drafter. All four formats quantize this checkpoint with no script changes. The target differs from Qwen3 in ways the YAML has to reflect: TP8, and a startup dominated by Mamba2 kernel warmup rather than weight loading. The drafter's block_size is 8, and it sets has_lm_head=false (shares the target's), so --quantize_lm_head does not apply. Also fixes --block_size in the Qwen3 example: it said 16, carried over from the DFlash streaming recipe's dflash_block_size, but dspark_qwen3_8b_block7 is block_size 7 -- which is what every AL number in the PR description was measured at. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml | 4 +- .../hf_dspark_ptq_nvfp4.yaml | 78 +++++++++++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml index 759792bcc11..14404a01dd1 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml @@ -54,8 +54,8 @@ pipeline: args: - --draft_model_dir /scratchspace/export_quantized # DSPARK/DFLASH read --block_size, not --draft_length; must match the - # drafter's dflash_block_size. - - --block_size 16 + # drafter's block_size (7 for dspark_qwen3_8b_block7). + - --block_size 7 - --output_length 4096 - --engine VLLM - --tp_size 1 diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml new file mode 100644 index 00000000000..defd30bee1d --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml @@ -0,0 +1,78 @@ +# Calibration-free NVFP4 PTQ of the DSpark drafter for Nemotron-3.5-Lightning-30B-A3B. +# +# Same pipeline as examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml, against a hybrid +# Mamba-MoE target and the published DSpark drafter. Quantizes weight-only to NVFP4, +# then measures acceptance length so the cost is visible. No calibration data needed: +# every scale comes from the weights. +# +# 2-step pipeline: +# task_0: Quantize the drafter (CPU-only, well under a minute for this 0.97B draft) +# task_1: Benchmark acceptance length on MT-Bench via vLLM +# +# Usage: +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml --yes + +job_name: Nemotron-3.5-Lightning-30B-A3B_DSpark_PTQ_nvfp4 +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + draft_model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark + + # Step 1: Quantize. w4a16_nvfp4 keeps activations in bf16; use --qformat nvfp4 + # for weight+activation (fixed input_scale 1.0). + # + # The q/k/v exclusions are mandatory, not a tuning choice: DFlash-family + # drafters build their fused context-KV projection by reading qkv_proj.weight + # raw, which cannot be a packed tensor. o_proj and the MLP still quantize. + # `fc` is optional -- add '*fc*' to keep it in bf16 at the cost of size. + # + # This drafter has has_lm_head=false (it shares the target's), so + # --quantize_lm_head does not apply. embed_tokens is 37% of the checkpoint but + # is excluded by default: it is an nn.Embedding the drafter inherits. + task_0: + script: common/specdec/quantize_drafter.sh + args: + - --qformat w4a16_nvfp4 + - --export_path /scratchspace/export_quantized + - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' + environment: + - DRAFTER_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 2: Acceptance length on MT-Bench. Compare against the same run with + # --draft_model_dir <> to see what quantization cost. + # + # The target is a hybrid Mamba-MoE model, which changes two things versus the + # Qwen3 example: it needs TP8, and startup is dominated by Mamba2 kernel warmup + # (~6 min before the engine is ready) rather than weight loading. + task_1: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export_quantized + # DSPARK reads --block_size, not --draft_length; must match the drafter's + # block_size, which is 8 for this checkpoint. + - --block_size 8 + - --output_length 4096 + - --engine VLLM + - --tp_size 8 + - --ep_size 1 + - --speculative_algorithm DSPARK + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 8 + container: vllm/vllm-openai:nightly From a48db7c46cbab952cef8a4d0da8be66371b3af03 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:12:32 +0000 Subject: [PATCH 13/15] example(launcher): correct the fc size figures with measured values The Qwen3 comment claimed quantizing fc saves ~12% of checkpoint size. That came from du -sh on export directories holding extra files; the weight files are 3.293 -> 3.181 GiB, i.e. ~3%. fc is only 3.5% of that drafter's parameters (embed_tokens is 26% and is excluded by default), so ~12% was never plausible. Fills in the same figures for Nemotron, now that they are measured: 1.316 -> 1.258 GiB (~4%) for 4.2899 -> 4.2334 AL. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml | 2 +- .../hf_dspark_ptq_nvfp4.yaml | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml index 14404a01dd1..0e87e97f3b7 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml @@ -30,7 +30,7 @@ pipeline: # The q/k/v exclusions are mandatory, not a tuning choice: DFlash-family # drafters build their fused context-KV projection by reading qkv_proj.weight # raw, which cannot be a packed tensor. o_proj and the MLP still quantize. - # `fc` is optional -- quantizing it saves ~12% more size for ~0.7% AL on + # `fc` is optional -- quantizing it saves ~3% more size for ~0.7% AL on # Qwen3-8B (3.0186 vs 3.0392); add '*fc*' to keep it in bf16. task_0: script: common/specdec/quantize_drafter.sh diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml index defd30bee1d..1526b585dd3 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml @@ -28,7 +28,8 @@ pipeline: # The q/k/v exclusions are mandatory, not a tuning choice: DFlash-family # drafters build their fused context-KV projection by reading qkv_proj.weight # raw, which cannot be a packed tensor. o_proj and the MLP still quantize. - # `fc` is optional -- add '*fc*' to keep it in bf16 at the cost of size. + # `fc` is optional -- quantizing it saves ~4% more size for ~1.3% AL on this + # model (4.2334 vs 4.2899); add '*fc*' to keep it in bf16. # # This drafter has has_lm_head=false (it shares the target's), so # --quantize_lm_head does not apply. embed_tokens is 37% of the checkpoint but From dc15f02fcb3a685a710836544ba58df327c25b63 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:19:57 +0000 Subject: [PATCH 14/15] specdec_bench: forward model-card engine args; make the Nemotron example runnable The Nemotron example passed --runtime_params, but AsyncEngineArgs is built from an explicit allow-list, so the mamba settings were silently dropped. That is not a cosmetic gap: with vLLM's defaults the first draft token is rejected ~88% of the time and acceptance length falls from ~4.3 to ~1.5, while every later position stays normal -- it reads as a bad drafter rather than a misconfigured engine. Forwards the five mamba keys when the caller sets them, and ships the settings the model card pins as engine_args.json next to the YAML. Every key in that file is now actually plumbed. Also aligns the example's task_1 with the configuration these numbers were measured at: ntasks_per_node 1 (vLLM owns TP internally; 8 tasks would launch 8 duplicate benchmarks), --trust_remote_code, --temperature 0, concurrency 8, and VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 as the card sets. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../specdec_bench/models/vllm.py | 17 +++++++++++++++++ .../engine_args.json | 8 ++++++++ .../hf_dspark_ptq_nvfp4.yaml | 19 ++++++++++++++----- 3 files changed, 39 insertions(+), 5 deletions(-) create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index db352c117e2..0aa76de491b 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/examples/specdec_bench/specdec_bench/models/vllm.py @@ -171,6 +171,23 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs async_scheduling=kwargs.get("async_scheduling", True), enforce_eager=enforce_eager, max_model_len=kwargs.get("max_model_len"), + # Engine knobs a model card may pin, passed through from + # --runtime_params engine_args.. Only keys the caller actually set are + # forwarded, so vLLM keeps its own defaults otherwise. Hybrid Mamba models + # need these: on Nemotron-3.5-Lightning the SSM-cache settings decide whether + # the first draft token is accepted, and leaving them at vLLM's defaults costs + # ~65% of acceptance length while every later position looks normal. + **{ + key: kwargs[key] + for key in ( + "mamba_backend", + "mamba_ssm_cache_dtype", + "mamba_cache_mode", + "mamba_cache_philox_rounds", + "enable_mamba_cache_stochastic_rounding", + ) + if kwargs.get(key) is not None + }, ) self.engine_args = engine_args self.model = AsyncLLM.from_engine_args(engine_args) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json new file mode 100644 index 00000000000..b4d22256dd7 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json @@ -0,0 +1,8 @@ +{ + "engine_args": { + "mamba_backend": "flashinfer", + "mamba_ssm_cache_dtype": "float16", + "enable_mamba_cache_stochastic_rounding": true, + "mamba_cache_philox_rounds": 5 + } +} diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml index 1526b585dd3..76275297cd8 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml @@ -52,9 +52,13 @@ pipeline: # Step 2: Acceptance length on MT-Bench. Compare against the same run with # --draft_model_dir <> to see what quantization cost. # - # The target is a hybrid Mamba-MoE model, which changes two things versus the - # Qwen3 example: it needs TP8, and startup is dominated by Mamba2 kernel warmup - # (~6 min before the engine is ready) rather than weight loading. + # This target is a hybrid Mamba-MoE model, so it differs from the Qwen3 example + # in three ways. It needs TP8. Startup is dominated by Mamba2 kernel warmup + # (~6 min before the engine is ready), not weight loading. And it needs the + # mamba engine settings from the model card (runtime_params below): without + # them the first draft token is rejected ~88% of the time and acceptance length + # collapses from ~4.3 to ~1.5, while every later position stays normal -- so it + # looks like a bad drafter rather than a serving misconfiguration. task_1: script: common/specdec_bench/quick_check.sh args: @@ -67,13 +71,18 @@ pipeline: - --tp_size 8 - --ep_size 1 - --speculative_algorithm DSPARK + - --trust_remote_code + - --temperature 0 + - --runtime_params modules/Model-Optimizer/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 32 + - --concurrency 8 environment: - HF_MODEL_CKPT: <> + # The model card sets this on every Nemotron-3.5 serve command. + - VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" slurm_config: _factory_: "slurm_factory" nodes: 1 - ntasks_per_node: 8 + ntasks_per_node: 1 gpus_per_node: 8 container: vllm/vllm-openai:nightly From 2e20994e33de0daadd64b6e74759623327874c32 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:58:43 +0000 Subject: [PATCH 15/15] example(specdec): fix auto_map handling for list values and Hub downloads Review findings, both in the auto_map support added earlier in this PR. auto_map values are not always strings: HF writes AutoTokenizer as a list whose entries may be null, e.g. [null, "tokenization_x.XTokenizerFast"]. Calling .split() on that raised AttributeError, so any drafter with a tokenizer auto_map crashed at export. Extraction now lives in auto_map_modules(), which handles strings, lists, nulls and repo-- prefixes. The snapshot_download allow-list also excluded *.py, so for a Hub drafter the modules were never fetched and the copy step had nothing to copy -- leaving the exported auto_map pointing at files that do not exist. Added *.py. Verified against all six auto_map shapes plus an end-to-end export with a list-valued auto_map; exports for the four formats are byte-identical. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../scripts/quantize_drafter.py | 30 +++++++++++++++---- 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py index 72277d60761..6bb41e7a5a1 100644 --- a/examples/speculative_decoding/scripts/quantize_drafter.py +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -120,15 +120,34 @@ def parse_args(): return parser.parse_args() +def auto_map_modules(config: dict) -> set[str]: + """Module basenames referenced by a config's ``auto_map``, e.g. {"modeling_x"}. + + Values are either ``"modeling_x.XModel"`` or, for tokenizers, a list whose entries may + be null (``[null, "tokenization_x.XTokenizerFast"]``); a ``repo--`` prefix points at + another repository and is not a local file. + """ + modules = set() + for value in (config.get("auto_map") or {}).values(): + for ref in value if isinstance(value, list) else [value]: + if isinstance(ref, str) and "." in ref: + modules.add(ref.split("--")[-1].rsplit(".", 1)[0]) + return modules + + def load_drafter(drafter_path: str) -> tuple[Path, dict[str, torch.Tensor]]: """Resolve a local dir or HF repo id to (dir, state_dict).""" local_dir = Path(drafter_path) if not local_dir.is_dir(): from huggingface_hub import snapshot_download + # A drafter that ships custom modeling code references it from auto_map, and the + # export carries that config verbatim -- so those .py files have to be fetched too + # or the exported references dangle. local_dir = Path( snapshot_download( - drafter_path, allow_patterns=["*.safetensors", "config.json", *SIDECAR_FILES] + drafter_path, + allow_patterns=["*.safetensors", "config.json", "*.py", *SIDECAR_FILES], ) ) @@ -332,11 +351,10 @@ def main(): # A drafter that ships custom modeling code points at it from auto_map; the export # carries that config verbatim, so the .py files have to come along or the reference # dangles. (The DFlash/DSpark exports have no auto_map -- this is for the ones that do.) - if config.get("auto_map"): - for module in {ref.split("--")[-1].split(".")[0] for ref in config["auto_map"].values()}: - source_py = source_dir / f"{module}.py" - if source_py.is_file(): - shutil.copy2(source_py, export_dir / source_py.name) + for module in auto_map_modules(config): + source_py = source_dir / f"{module}.py" + if source_py.is_file(): + shutil.copy2(source_py, export_dir / source_py.name) before = sum(v.numel() * v.element_size() for v in state_dict.values()) after = sum(v.numel() * v.element_size() for v in export_sd.values())