Prepare NVFP4 draft and Markov heads behind independent flags (#65)

DeepSeek V4 and V4.1 can prepare separate NVFP4 copies of the draft
vocabulary head and the Markov transition projection. Each is selected by its
own opt-in flag and recipe capability; omitting both keeps checkpoint
precision, and the target verification head always keeps its checkpoint
precision. Quantized heads run with BF16 activations, and the Markov kernel's
BF16 output is widened to FP32 before bias addition.

The rebuild record gains optional fields 8 and 9 so a rebuild reproduces the
selected head precision instead of silently falling back to the checkpoint.
Fields are written only when selected, so existing BF16 manifests remain
byte-identical. The export path rejects unknown fields, so the schema must
carry every key the manifest writes; the Python bindings are regenerated with
protoc 35.1, the version that produced the checked-in file.

Tests cover the four draft/Markov combinations through the export schema, the
independent flag and plan validation, the Markov projection's BF16-to-FP32
path, and a record written before the fields existed. Focused suites pass;
six Kimi-K3 failures pre-exist because the checkpoint fixture is absent. CUDA
tests were not run in this environment; they require FLARE_RUN_CUDA_TESTS=1.
A paired A/B on four DGX Sparks measured +4.0% coding throughput and flat
acceptance (64.6% to 63.8%) for both heads together versus checkpoint
precision, with the NVFP4 package noisier across rounds.


Claude-Session: https://claude.ai/code/session_01PFTkgqizAoWiQUS1Z3uR9v
Design-Signed-Off-By: ds41; design-review=6aeb

Co-authored-by: ds41 <noreply@invalid>

Author: original-el8

docs/model-preparation.md

@@ -452,6 +452,16 @@ The three preceding tokens live in continuation state, including reused prefixes
 DSpark proposals, target verification, and accepted-state commit use the shared
 native speculative contract; Helios contains no DSpark model implementation.
 
+The optional preparation flags `--draft-head-quantization nvfp4` and
+`--markov-head-quantization nvfp4` select independent NVFP4 copies of the draft
+vocabulary head and Markov output projection. Both use BF16 activations; the
+Markov kernel's BF16 output is widened to FP32 before bias addition. Target
+verification retains the checkpoint head. Omitting these flags preserves the
+checkpoint precision. Each selection requires tuning, capture, and package export;
+it cannot change an existing package at serving time. The same options apply to
+the DeepSeek V4 Flash recipe. Quantization changes proposal probabilities and can
+change acceptance and throughput; measure both separately and together.
+
 ## DeepSeek V4 Flash
 
 DeepSeek V4 Flash uses `local-inference-lab/DeepSeek-V4-Flash-Vision-Exp-lossless-CSF`, including its

helios/proto/runtime_package.proto

@@ -107,6 +107,8 @@ message RebuildSpeculation {
   uint64 token_capacity = 5;
   bool packed_verification = 6;
   string draft_prefix_selection = 7;
+  optional string draft_head_quantization = 8;
+  optional string markov_head_quantization = 9;
 }
 // Views of retained prepared weight storage, with no additional payload copies.
 // Names are paths in the recipe's weight objects, not original checkpoint keys.

ignition/cli/__init__.py

@@ -165,6 +165,12 @@ def _preparation_arguments(running):
                                " excludes mandatory width-one and ordinary graphs. Requires packed adaptive"
                                " learned speculation (default: see MODEL --help)")
     advanced = running.add_argument_group("advanced options")
+    advanced.add_argument("--draft-head-quantization", choices=("nvfp4",),
+                          help="prepare a separate NVFP4 vocabulary head for learned draft proposals;"
+                               " target verification retains its precision (DeepSeek V4/V4.1)")
+    advanced.add_argument("--markov-head-quantization", choices=("nvfp4",),
+                          help="prepare the draft Markov transition projection as NVFP4"
+                               " (DeepSeek V4/V4.1; default: checkpoint precision)")
     advanced.add_argument("--mla-compute-mode", choices=("auto", "bf16"),
                           help="constrain MLA arithmetic while retaining packed KV storage and compatible"
                                " kernel tuning (default: see MODEL --help)")

ignition/cli/display.py

@@ -29,6 +29,9 @@ def speculation(settings, *, brief: bool = False) -> str:
     drafts = settings.draft_tokens
     depth = f"{'up to' if settings.adaptive else 'exactly'} {drafts} draft token{'' if drafts == 1 else 's'}"
     selection = " · GPU prefix selection" if settings.draft_prefix_selection == "gpu" else ""
+    heads = [f"{head} head {getattr(settings, head + '_head_quantization').upper()}"
+             for head in ("draft", "markov") if getattr(settings, head + '_head_quantization') is not None]
+    precision = " · " + ", ".join(heads) if heads else ""
     if brief:
-        return depth + selection
-    return f"{depth} per round · MTP module" + selection
+        return depth + selection + precision
+    return f"{depth} per round · MTP module" + selection + precision

ignition/cli/preparation.py

@@ -140,7 +140,9 @@ def speculation(args, model: ResolvedModel, configuration: Configuration | None
     given = [flag for flag, value in (
         ("--mtp-draft-tokens", args.mtp_draft_tokens),
         ("--speculative-token-capacity", args.speculative_token_capacity),
-        ("--no-mtp-adaptive-verification", args.no_mtp_adaptive_verification or None)) if value is not None]
+        ("--no-mtp-adaptive-verification", args.no_mtp_adaptive_verification or None),
+        ("--draft-head-quantization", getattr(args, "draft_head_quantization", None)),
+        ("--markov-head-quantization", getattr(args, "markov_head_quantization", None))) if value is not None]
     if explicit_graph_budget:
         given.append("--speculative-graph-budget")
     if default is None:
@@ -172,8 +174,13 @@ def speculation(args, model: ResolvedModel, configuration: Configuration | None
     budget = graph_budget.value if packed else None
     if budget is not None and (type(budget) is not int or budget < 1):
         raise UsageError("--speculative-graph-budget must be a positive integer")
+    draft_head = getattr(args, "draft_head_quantization", None)
+    markov_head = getattr(args, "markov_head_quantization", None)
+    if (draft_head or markov_head) and drafts < 1:
+        raise UsageError("draft head quantization requires positive learned draft depth")
     return Speculation(drafts, adaptive, token_capacity=capacity, packed_verification=packed,
-                       draft_prefix_selection=prefix_selection, graph_budget=budget)
+                       draft_prefix_selection=prefix_selection, graph_budget=budget,
+                       draft_head_quantization=draft_head, markov_head_quantization=markov_head)
 
 
 def fetch(model: ResolvedModel, hub_cache, *, progress_bars: bool = True) -> str:

ignition/components/dspark.py

@@ -225,25 +225,30 @@ class PreparedMarkov:
 
     A replicated head runs independently on each rank without a collective.
     A sharded head uses the supplied runtime to gather its FP32 bits unchanged.
+    Quantized A16 projections produce BF16 logits, widened before bias addition.
     """
 
     def __init__(self, weights, tokens, *, runtime, coordinator, session, name):
         rows, device = tokens.numel(), tokens.device
         self.embedding = torch.zeros(
             (rows, weights.config.draft.markov_rank), dtype=torch.bfloat16, device=device
         )
-        columns = weights.draft.markov_head.shape[0]
+        head = weights.draft.markov_head
+        columns = head.shape[0] if isinstance(head, torch.Tensor) else head.out_features
         expected_columns = 129280 if runtime is None else 129280 // weights.shard.world_size
         if columns != expected_columns:
             raise ValueError("Markov head vocabulary partition must match its gather runtime")
         self.local = torch.zeros((rows, columns), dtype=torch.float32, device=device)
+        self.projected = self.local if isinstance(head, torch.Tensor) else torch.empty(
+            (rows, columns), dtype=torch.bfloat16, device=device)
         self.logits = self.local if runtime is None else torch.zeros((rows, 129280), dtype=torch.float32, device=device)
         self.embed = PreparedEmbedding(
             weights.draft.markov_embedding, tokens, self.embedding, session=session, name=name + ".embed"
         )
-        self.projection = UnquantizedProjection(
-            self.embedding, weights.draft.markov_head, self.local, session=session, name=name + ".head"
-        )
+        projection = UnquantizedProjection if isinstance(head, torch.Tensor) else QuantizedProjection
+        options = {} if isinstance(head, torch.Tensor) else {"activation_mode": "a16"}
+        self.projection = projection(self.embedding, head, self.projected,
+                                     session=session, name=name + ".head", **options)
         self.gather = None if runtime is None else PreparedAllGather(
             self.local.view(torch.bfloat16),
             self.logits.view(torch.bfloat16),
@@ -255,6 +260,8 @@ def __init__(self, weights, tokens, *, runtime, coordinator, session, name):
 
     def __call__(self):
         self.projection(self.embed())
+        if self.projected is not self.local:
+            self.local.copy_(self.projected)
         if self.gather is not None:
             self.gather(self.local.view(torch.bfloat16))
         return self.logits

ignition/planning.py

@@ -302,6 +302,8 @@ def to_dict(self) -> dict:
                 "draft_tokens": self.speculation.draft_tokens, "adaptive": self.speculation.adaptive,
                 "packed_verification": self.speculation.packed_verification,
                 "draft_prefix_selection": self.speculation.draft_prefix_selection,
+                "draft_head_quantization": self.speculation.draft_head_quantization,
+                "markov_head_quantization": self.speculation.markov_head_quantization,
                 **({"graph_budget": self.speculation.graph_budget}
                    if self.speculation.graph_budget is not None else {})},
             "context": self.context, "checkpoint_context": self.recipe.context,
@@ -341,6 +343,12 @@ def fit(recipe: Recipe, checkpoint: Checkpoint, gpus: tuple[Gpu, ...], *, concur
     if mla_compute_mode not in recipe.mla_compute_modes:
         raise UsageError(f"--mla-compute-mode {mla_compute_mode}: supported values for this recipe are "
                          + ", ".join(recipe.mla_compute_modes))
+    if speculation is not None:
+        for head in ("draft", "markov"):
+            value = getattr(speculation, f"{head}_head_quantization")
+            if value is not None and value not in getattr(recipe, f"{head}_head_quantizations"):
+                raise UsageError(f"--{head}-head-quantization {value}: {checkpoint.alias}"
+                                 " does not prepare this proposal head format")
     if expert_scale_compression is not None and expert_scale_compression not in checkpoint.scale_compressed_device_bytes:
         raise UsageError(f"--expert-scale-compression {expert_scale_compression}: {checkpoint.alias}"
                          " does not prepare this scale encoding")

ignition/preparation/rebuild_manifest.py

@@ -31,6 +31,11 @@ def from_plan(model, plan, *, compiled_forwards):
             adaptive=s.adaptive,
             token_capacity=plan.speculative_token_capacity, packed_verification=s.packed_verification,
             draft_prefix_selection=s.draft_prefix_selection)
+        # Record proposal-head precision only when it differs from the checkpoint
+        # default, so existing BF16 manifests stay byte-identical.
+        for field in ('draft_head_quantization', 'markov_head_quantization'):
+            if (value := getattr(s, field)) is not None:
+                result['speculation'][field] = value
     return result
 
 

ignition/preparation/runtime_package_pb2.py

@@ -108,8 +108,15 @@ class Speculation:
     packed_verification: bool = False
     draft_prefix_selection: str = "cpu"
     graph_budget: int | None = None
+    draft_head_quantization: str | None = None
+    markov_head_quantization: str | None = None
 
     def __post_init__(self):
+        for value in (self.draft_head_quantization, self.markov_head_quantization):
+            if value not in (None, "nvfp4"):
+                raise ValueError("draft head quantization must be omitted or nvfp4")
+            if value is not None and self.draft_tokens < 1:
+                raise ValueError("draft head quantization requires positive learned draft depth")
         if self.graph_budget is not None:
             if type(self.graph_budget) is not int or self.graph_budget < 1:
                 raise ValueError("speculative graph budget must be a positive integer")
@@ -209,6 +216,9 @@ class Recipe:
     mla_compute_mode: str = "auto"
     expert_scale_compression: str | None = None
     export_repository: str = ""
+    # Preparation overrides for vocabulary and transition projections used only by proposals.
+    draft_head_quantizations: tuple[str, ...] = ()
+    markov_head_quantizations: tuple[str, ...] = ()
 
     def mixed_shapes(self, concurrency, chunk, speculation):
         from ignition.components.capacity import mixed_shapes