[Example]: Calibration-free FP8/NVFP4 PTQ for speculative-decoding drafters (#2027)

### What does this PR do? Type of change: new example Adds `examples/speculative_decoding/scripts/quantize_drafter.py`, a CLI that quantizes an exported speculative-decoding drafter to FP8 or NVFP4 — weight-only or weight+activation — with no calibration data. It needs no modeling code either. Exported drafters such as [`nvidia/MiniMax-M3-DSpark`](https://huggingface.co/nvidia/MiniMax-M3-DSpark) have no importable model class, so each 2-D weight is wrapped in a throwaway `nn.Linear` under its checkpoint key and ModelOpt's usual `quantizer_name` patterns select over those names. Works for any drafter layout (DSpark / DFlash / EAGLE3 / Medusa). **Formats:** `w4a16_nvfp4`, `nvfp4`, `fp8`, `fp8_pc_pt` — the ModelOpt formats vLLM's backend can actually serve. AWQ is deliberately not offered, since `awq_lite` silently degrades to plain RTN without a `forward_loop`. **Static activation scales without calibration.** `fp8` and `nvfp4` normally need an activation amax *measured* on calibration data; a fixed `input_scale` of 1.0 is applied instead. That works because acceptance length is governed almost entirely by **clipping**, not resolution: Sweeping the fixed scale over three decades (same setup as the Testing section below; bf16 baseline 3.1423): | `input_scale` | amax | FP8 AL | vs bf16 | NVFP4 AL | vs bf16 | |---|---|---|---|---|---| | 0.003 | 1.3 | 2.2204 | -29.34% | 2.2076 | -29.75% | | 0.01 | 4.5 | 2.6719 | -14.97% | 2.6641 | -15.22% | | 0.03 | 13.4 | 2.9751 | -5.32% | 2.9259 | -6.89% | | 0.1 | 44.8 | 3.1013 | -1.31% | 3.0206 | -3.88% | | 0.2 | 89.6 | 3.1178 | -0.78% | 3.0015 | -4.48% | | 0.3 | 134.4 | 3.1370 | -0.17% | 3.0222 | -3.82% | | 0.5 | 224.0 | 3.1268 | -0.50% | 3.0360 | -3.38% | | **1.0 (default)** | **448.0** | **3.1457** | **+0.11%** | **3.0193** | **-3.91%** | | 2.0 | 896.0 | 3.1354 | -0.22% | 3.0172 | -3.98% | | 4.0 | 1792.0 | 3.1245 | -0.57% | 3.0034 | -4.42% | Both formats fall off a cliff below ~0.03, where the declared range sits far under the activations' true magnitude and most of the tensor is clipped. Both then sit 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, which is why it is hardcoded rather than exposed. 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. Deriving 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, measuring -31% to -46% AL. **Where calibration would go.** All of this sits behind `resolve_activation_scales()`, the single place deciding where a static amax comes from. Real calibration slots in ahead of the fixed fallback with no change to the CLI or the call site, and composes because `set_static_activation_amax()` skips quantizers that already have an amax: ```python 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 ``` **Serving a quantized drafter.** Four things had to be written into the exported checkpoint before vLLM would load one: - emit `quant_method` (`modelopt_fp4` / `modelopt`) — vLLM reads that key, ModelOpt writes only `quant_algo` - emit the exclusion list under `ignore` too — that is the key read from the flat `quantization_config`; `exclude_modules` alone yields an empty exclusion set - add `*<name>` wildcards so exclusions match a runtime's nested module prefix (`model.fc`) rather than the checkpoint key (`fc`) - add `*qkv_proj` / `*gate_up_proj` aliases for layers a runtime fuses, whose names appear in no checkpoint key Nothing is then needed on the caller side. **This closes the open question left in the previous revision of this PR: vLLM does read `quantization_config` off the draft checkpoint.** `ModelConfig._verify_quantization` fills `quantization` in from `quant_method` when it is unset, so once the export declares that key — the first fix above — detection works on its own. Verified on Nemotron-3.5-Lightning passing nothing: `Detected ModelOpt NVFP4 checkpoint (quant_algo=NVFP4)` → `FlashInferCuteDslNvFp4LinearKernel`, AL 4.278 against 4.203 measured earlier. `specdec_bench` also gains a `DSPARK` algorithm, which it did not have: an exported `Qwen3DSparkModel` would otherwise have to go through `DFLASH` and be built with vLLM `method="dflash"`. The branch sets `method="dspark"` and leaves `draft_sample_method` on vLLM's own default of `greedy`. A target whose fused-collective workspace (sized at CUDA-graph capture) overflows at large speculative batches can disable graphs with `--runtime_params '{"engine_args": {"enforce_eager": true}}'`. For DFlash-family drafters, `qwen3_dflash.py` builds its fused context-KV projection by reading `qkv_proj.weight` raw and calling `F.linear`, which cannot consume a packed weight. Keep those layers in bf16 with `--exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*'`; `o_proj` and the MLP — the bulk of the drafter — still quantize. That exclusion is mandatory, not a tuning choice. `fc` (the projection from the target's captured layers into the draft) is the one real knob, and it is a genuine trade rather than a free win — see the Testing section for both models' numbers. The examples quantize it; add `'*fc*'` to the exclude list to keep it in bf16. `embed_tokens`, `markov_head` and `confidence_head` are excluded by default: they are 2-D so the flat view treats them as GEMMs, but they are embeddings or a single-output projection. `lm_head` is excluded by the preset itself — unlike on a base model it is 37% of this drafter's parameters, so `--quantize_lm_head` is a real lever (~1.9 GiB), but measure AL first. The flag re-enables both of `lm_head`'s quantizers; re-enabling only the weight one would ship a W+A checkpoint whose `lm_head` has no `input_scale` while the config still advertises it as quantized. ### Usage ```bash # weight+activation FP8, calibration-free, lossless on both models measured below python scripts/quantize_drafter.py \ --drafter_path deepseek-ai/dspark_qwen3_8b_block7 \ --qformat fp8 \ --export_path ./dspark-qwen3-8b-fp8 \ --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' # smallest: weight-only NVFP4 python scripts/quantize_drafter.py \ --drafter_path nvidia/MiniMax-M3-DSpark \ --qformat w4a16_nvfp4 \ --export_path ./MiniMax-M3-DSpark-W4A16 ``` Or end to end on Slurm — quantize, then measure AL — via the launcher examples added here, one per target: ```bash uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml --yes uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml --yes ``` Serving one, if you are not going through `specdec_bench`: ```python speculative_config = { "method": "dspark", "model": "./dspark-qwen3-8b-fp8", # quantization is read from its config.json "num_speculative_tokens": 7, } ``` ### Testing Two targets with different architectures, so the conclusions are not one model's quirk: * **Qwen3-8B** (dense transformer) + [`deepseek-ai/dspark_qwen3_8b_block7`](https://huggingface.co/deepseek-ai/dspark_qwen3_8b_block7), `block_size` 7, TP1. * **Nemotron-3.5-Lightning-30B-A3B** (hybrid Mamba-MoE) + [`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark), `block_size` 8, TP8, with the mamba engine settings the model card pins (`mamba_backend=flashinfer`, `mamba_ssm_cache_dtype=float16`, stochastic SSM-cache rounding). Both: MT-Bench 80 questions, greedy, one vLLM instance per point. | recipe | activations | Qwen3-8B AL | vs bf16 | Nemotron-3.5 AL | vs bf16 | |---|---|---|---|---|---| | bf16 baseline | — | 3.1423 | — | 4.3296 | — | | **`fp8`** | static, `input_scale` 1.0 | **3.1457** | **+0.11%** | **4.3289** | **-0.02%** | | `fp8_pc_pt` | dynamic per-token | 3.1228 | -0.62% | 4.3411 | +0.26% | | `w4a16_nvfp4`, `fc` in bf16 | bf16 (weight-only) | 3.0392 | -3.28% | 4.2899 | -0.92% | | `w4a16_nvfp4`, `fc` quantized | bf16 (weight-only) | 3.0186 | -3.94% | 4.2334 | -2.22% | | **`nvfp4`** | static, `input_scale` 1.0 | **3.0193** | **-3.91%** | **4.2030** | **-2.92%** | **FP8 weight+activation at the fixed `input_scale` of 1.0 is lossless on both.** +0.11% and -0.02% are both inside run-to-run noise — the Nemotron baseline was measured twice under identical settings and the two runs differ by 0.94% (4.3093 / 4.3499), which sets the resolution of that column. On the same reading, `fp8` and `fp8_pc_pt` are indistinguishable on Nemotron; the dynamic variant only pulls ahead on Qwen3. NVFP4 costs 3-4% on Qwen3 and 2-3% on Nemotron, i.e. the 4-bit weight resolution is the real price and it is model-dependent but bounded. Whether to quantize `fc` is a per-model call rather than a general recommendation — it buys a few percent of size for an AL cost that differs by ~2x between these two drafters: | `fc` bf16 → quantized | Qwen3-8B | Nemotron-3.5 | |---|---|---| | checkpoint size | 3.293 → 3.181 GiB (-3.4%) | 1.316 → 1.258 GiB (-4.4%) | | AL | 3.0392 → 3.0186 (-0.68%) | 4.2899 → 4.2334 (-1.32%) | `fc` itself is only 3.5% (Qwen3) / 4.5% (Nemotron) of drafter parameters; `embed_tokens` is the bulk (26% / 36%) and is excluded by default. The Qwen3 `w4a16_nvfp4` rows were measured in a later session than the rest of that column; the `fc`-in-bf16 run reproduced the original number to four decimals (3.0392), so the column is internally comparable. Also validated on `nvidia/MiniMax-M3-DSpark`: `w4a16_nvfp4` runs in 67 s on CPU, 9.98 GiB (fp32) -> 3.51 GiB; all 43 quantized tensors round-trip within 0.0952 relative error; the 29 untouched tensors are bit-identical to `bf16(source)`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (example-only) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ❌ — validated manually as above. Can add a `tests/examples/speculative_decoding/` test over a small synthetic drafter if wanted before merge. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (example-only) - Did you get Claude approval on this PR?: ❌ (not yet run) ### Additional Information The measurements above are one drafter on one target with one benchmark; the plateau's location and the ~3.5% NVFP4 gap should be re-measured before assuming they carry to a different drafter. Note when reading an exported checkpoint: `input_scale` is `amax/448` for FP8 but `amax/(6*448)` for NVFP4, so the one fixed amax records as 1.0 in an FP8 checkpoint and 0.1667 in an NVFP4 one. Both mean the same activation range. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
2026-10-02 03:14:52 +08:00 · 2026-08-26 22:04:09 +08:00
parent 73d7784223
commit 5db2682519
13 changed files with 715 additions and 7 deletions
@@ -305,7 +305,7 @@ if __name__ == "__main__":
        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")
@@ -93,6 +93,8 @@ def create_auto_deploy_model(model_path: str, max_concurrent_requests: int, kwar
        )
    elif speculative_algorithm == "NONE":
        specdec = None
+    elif speculative_algorithm == "DSPARK":
+        raise NotImplementedError("DSPARK is only supported by --engine VLLM.")

    max_num_tokens = kwargs.get("max_num_tokens", 8192)

@@ -49,6 +49,8 @@ class SGLANGModel(Model):
            speculative_algorithm = "LOOKAHEAD"
        elif speculative_algorithm == "NONE":
            speculative_algorithm = None
+        elif speculative_algorithm == "DSPARK":
+            raise NotImplementedError("DSPARK is only supported by --engine VLLM.")

        engine_kwargs = {
            "model_path": model_dir,
@@ -125,6 +125,8 @@ def create_executor(model_path: str, max_concurrent_requests, kwargs):
        )
    elif kwargs.get("speculative_algorithm", None) == "NONE":
        specdec = None
+    elif kwargs.get("speculative_algorithm", None) == "DSPARK":
+        raise NotImplementedError("DSPARK is only supported by --engine VLLM.")
    else:
        specdec = None

@@ -28,6 +28,16 @@ except ImportError:
    vllm = None


+# Forwarded from ``--runtime_params`` ``engine_args.<key>``; extend as needed.
+PASSTHROUGH_ENGINE_ARGS = (
+    "mamba_backend",
+    "mamba_ssm_cache_dtype",
+    "mamba_cache_mode",
+    "mamba_cache_philox_rounds",
+    "enable_mamba_cache_stochastic_rounding",
+)
+
+
 class VLLMModel(Model):
    # Cross-engine ``--max_seq_len`` (run.py) lands in kwargs under the
    # vLLM-native name ``max_model_len`` (see run.py's ``_MAX_SEQ_LEN_KEY``)
@@ -114,6 +124,13 @@ class VLLMModel(Model):
                "model": kwargs.get("draft_model_dir"),
                "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 8),
            }
+        elif kwargs.get("speculative_algorithm") == "DSPARK":
+            specdec = {
+                "method": "dspark",
+                "model": kwargs.get("draft_model_dir"),
+                "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 7),
+                "draft_sample_method": kwargs.get("draft_sample_method", "greedy"),
+            }
        elif kwargs.get("speculative_algorithm") == "NONE":
            specdec = None

@@ -133,8 +150,9 @@ class VLLMModel(Model):
            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=kwargs.get("enforce_eager", False),
            max_model_len=kwargs.get("max_model_len"),
+            **{key: kwargs[key] for key in PASSTHROUGH_ENGINE_ARGS if kwargs.get(key) is not None},
        )
        self.engine_args = engine_args
        self.model = AsyncLLM.from_engine_args(engine_args)
@@ -0,0 +1,378 @@
+# 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 a speculative-decoding drafter.
+
+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 \
+        --drafter_path nvidia/MiniMax-M3-DSpark \
+        --qformat fp8 \
+        --export_path ./MiniMax-M3-DSpark-FP8
+"""
+
+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_activation_scaling_factor,
+    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.config import need_calibration
+from modelopt.torch.quantization.utils import is_quantized_linear
+
+# INT8/INT4 are absent on purpose: they quantize cleanly but vLLM's ModelOpt backend
+# cannot serve them.
+SUPPORTED_QFORMATS = [
+    "w4a16_nvfp4",
+    "nvfp4",
+    "fp8",
+    "fp8_pc_pt",
+]
+
+# 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*"]
+
+# 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
+STATIC_ACT_AMAX = FP8_E4M3_MAX
+
+
+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(
+        "--qformat",
+        default="w4a16_nvfp4",
+        choices=SUPPORTED_QFORMATS,
+        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",
+        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 -- the largest drafter tensor, but it feeds the "
+        "acceptance test directly, so measure AL first.",
+    )
+    return parser.parse_args()
+
+
+def auto_map_modules(config: dict, source_dir: Path) -> set[Path]:
+    """Module files an ``auto_map`` points at, relative to ``source_dir``.
+
+    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. A value may also carry a package path
+    (``"pkg/modeling_x.XModel"``), which is preserved so the exported reference resolves.
+
+    The config is untrusted -- it may come straight off the Hub -- so references that
+    escape ``source_dir`` are dropped rather than followed.
+    """
+    modules = set()
+    for value in (config.get("auto_map") or {}).values():
+        for ref in value if isinstance(value, list) else [value]:
+            if not isinstance(ref, str) or "." not in ref:
+                continue
+            relative = Path(ref.split("--")[-1].rsplit(".", 1)[0] + ".py")
+            candidate = (source_dir / relative).resolve()
+            if relative.is_absolute() or not candidate.is_relative_to(source_dir.resolve()):
+                print(f"Skipping auto_map entry outside the checkpoint: {ref}")
+                continue
+            modules.add(relative)
+    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", "*.py", *SIDECAR_FILES],
+            )
+        )
+
+    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 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():
+        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
+        # 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
+
+
+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 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():
+        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
+        if getattr(input_quantizer, "_dynamic", False):
+            continue
+        if getattr(input_quantizer, "amax", None) is not None:
+            continue
+        # 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
+
+
+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 call
+    ``mtq.calibrate`` here, ahead of the fixed fallback.
+    """
+    if not need_calibration(quant_cfg):
+        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])
+    if quantize_lm_head:
+        # 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
+
+
+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.
+
+    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():
+        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)
+        # 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
+        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
+        # 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
+    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.qformat, args.exclude, args.quantize_lm_head)
+
+    mtq.quantize(root, quant_cfg)  # no forward_loop: scales come from the weights
+    resolve_activation_scales(root, quant_cfg)
+
+    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)
+    # ``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()
+        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)
+        # 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 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")),
+    ):
+        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 reads ``quant_method`` and treats its
+    # 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"
+    )
+    # 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))
+    (export_dir / "hf_quant_config.json").write_text(json.dumps(hf_quant_config, indent=2))
+
+    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.)
+    for relative in auto_map_modules(config, source_dir):
+        source_py = source_dir / relative
+        if source_py.is_file():
+            target = export_dir / relative
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(source_py, target)
+
+    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")
+    print(f"Exported to {export_dir}")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,47 @@
+#!/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-<step>/ 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}" \
+    "$@"
@@ -0,0 +1,73 @@
+# 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).
+  #
+  # 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 ~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
+    args:
+      - --qformat w4a16_nvfp4
+      - --export_path /scratchspace/export_quantized
+      - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*'
+    environment:
+      - DRAFTER_CKPT: <<global_vars.draft_model>>
+    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 <<global_vars.draft_model>> to see what quantization cost.
+  task_1:
+    script: common/specdec_bench/quick_check.sh
+    args:
+      - --draft_model_dir /scratchspace/export_quantized
+      # DSPARK/DFLASH read --block_size, not --draft_length; must match the
+      # drafter's block_size (7 for dspark_qwen3_8b_block7).
+      - --block_size 7
+      - --output_length 4096
+      - --engine VLLM
+      - --tp_size 1
+      - --ep_size 1
+      - --speculative_algorithm DSPARK
+      - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl
+      - --concurrency 32
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 1
+      container: vllm/vllm-openai:nightly
@@ -32,13 +32,13 @@ pipeline:

  global_vars:
    hf_model: /hf-local/nvidia/Kimi-K2.5-NVFP4
-    # Trained+exported DFLASH draft; override: pipeline.global_vars.draft_model_dir=<path>
-    draft_model_dir: /hf-local/nvidia/Kimi-K2.5-DFlash
+    # Trained+exported DFLASH draft; override: pipeline.global_vars.draft_model=<path>
+    draft_model: /hf-local/nvidia/Kimi-K2.5-DFlash

  task_0:
    script: common/specdec_bench/run.sh
    args:
-      - --draft_model_dir <<global_vars.draft_model_dir>>
+      - --draft_model_dir <<global_vars.draft_model>>
      - --speculative_algorithm DFLASH
      - --engine VLLM
      - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl
@@ -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
+  }
+}
@@ -0,0 +1,90 @@
+# Calibration-free FP8 PTQ of the DSpark drafter for Nemotron-3.5-Lightning-30B-A3B.
+#
+# Same pipeline as hf_dspark_ptq_nvfp4.yaml in this directory, quantizing both
+# weights and activations to FP8 instead of weight-only NVFP4. Activations use a
+# fixed input_scale of 1.0, so no calibration data is needed. This is the
+# lossless recipe: AL 4.3289 vs 4.3296 for bf16 (-0.02%, inside run-to-run
+# noise), against -2.92% for nvfp4.
+#
+# 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_fp8.yaml --yes
+
+job_name: Nemotron-3.5-Lightning-30B-A3B_DSpark_PTQ_fp8
+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. fp8 quantizes weights and activations; use --qformat
+  # fp8_pc_pt for per-channel weights with dynamic per-token activations, which
+  # needs no activation scale at all (4.3411 here, also within noise).
+  #
+  # 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.
+  # Unlike the nvfp4 recipe, `fc` is left quantized: FP8 is lossless here, so
+  # there is no accuracy reason 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 36% 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 fp8
+      - --export_path /scratchspace/export_quantized
+      - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*'
+    environment:
+      - DRAFTER_CKPT: <<global_vars.draft_model>>
+    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 <<global_vars.draft_model>> to see what quantization cost.
+  #
+  # 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:
+      - --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
+      - --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 8
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+      # 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: 1
+      gpus_per_node: 8
+      container: vllm/vllm-openai:nightly
@@ -0,0 +1,88 @@
+# 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 -- 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
+  # 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: <<global_vars.draft_model>>
+    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 <<global_vars.draft_model>> to see what quantization cost.
+  #
+  # 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:
+      - --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
+      - --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 8
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+      # 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: 1
+      gpus_per_node: 8
+      container: vllm/vllm-openai:nightly
@@ -56,7 +56,7 @@ pipeline:
    # config.json layer-type vocabulary already patched for tf5 (see header).
    hf_model: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16
    # The published drafter the warm start continues from.
-    drafter: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark
+    draft_model: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark

  # Build /scratchspace/data/train.jsonl. Point data.data_path at the full
  # Spec-Decoding-Dataset-v2 corpus to reproduce; eagle_utils also accepts a
@@ -80,7 +80,7 @@ pipeline:
      # causal attention and attention sink all come from this recipe — see header.
      - --config modules/Model-Optimizer/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml
      - model.model_name_or_path=<<global_vars.hf_model>>
-      - dflash.dflash_init_checkpoint=<<global_vars.drafter>>
+      - dflash.dflash_init_checkpoint=<<global_vars.draft_model>>
      - data.data_path=/scratchspace/data/train.jsonl
      # The stock Nemotron template has no {% generation %} tags; without a tagged copy
      # answer_only_loss trains on an all-zero mask (see header).