@@ -18,9 +18,10 @@ budget.
 A working draft-training reference already exists in NeMo AutoModel at
 `nemo_automodel/components/speculative/dspark/`. This document specifies how the
 same training recipe maps onto Megatron-Bridge and Megatron-Core. It
-intentionally does **not** ship the draft module itself: that is a
-Megatron-Core (or ModelOpt) neural module whose parallel/loss integration must be
-validated on multi-GPU runs, so it belongs in later, separately reviewed PRs.
+intentionally does **not** ship the draft module itself: that is a Megatron-Core
+neural module (see [Decision: Megatron-Core native](#decision-megatron-core-native))
+whose parallel/loss integration must be validated on multi-GPU runs, so it belongs
+in later, separately reviewed PRs.
 
 ## Background: the DSpark draft
 
@@ -57,11 +58,24 @@ emphasizing earlier positions):
 
 - `L_ce` (weight 0.1): cross-entropy of the draft logits against the target's
   teacher-forced next tokens.
-- `L_tv` (weight 0.9): total-variation distance `0.5 * ||p_draft - p_target||_1`
+- `L_l1` (weight 0.9): the **raw** probability L1 distance `||p_draft - p_target||_1`
   to the target's next-token distribution. This is the dominant term and a direct
   acceptance proxy.
 - `L_conf` (weight 1.0): binary cross-entropy training the confidence head against
-  the analytical acceptance label `c_k* = 1 - TV(p_draft, p_target)`.
+  the analytical acceptance label `c_k* = 1 - 0.5 * ||p_draft - p_target||_1`.
+
+The factor `1/2` belongs to the acceptance label only. Paper
+[Eq. 10](https://arxiv.org/html/2607.05147#S3.E10) minimizes the weighted **raw**
+L1 distance at coefficient 0.9; the `1/2` that converts L1 into a total-variation
+distance appears in [Eq. 8](https://arxiv.org/html/2607.05147#S3.E8), where the
+analytical acceptance/confidence target is formed. Both the
+[DeepSpec](https://github.com/deepseek-ai/DeepSpec/blob/main/deepspec/modeling/dspark/loss.py)
+and [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/speculative/dspark/loss.py)
+implementations make the same distinction. Training `0.5 * L1` instead would halve
+the dominant term's gradient while leaving the published coefficient in place, so
+the two must not be conflated. A regression test asserts the pair together
+(disjoint distributions give `l1_loss == 2` and `accept_rate == 0`) so the
+implementation and this page cannot drift apart.
 
 The target's teacher distribution is produced by passing the target's last hidden
 state through the (shared, frozen) LM head and is always detached, so no gradient
@@ -70,7 +84,8 @@ flows into the target.
 ### Acceptance and confidence
 
 Training measures acceptance analytically rather than by running rejection
-sampling. Per position, `accept_rate = clamp(1 - TV, 0, 1)`. The expected accepted
+sampling. Per position, `accept_rate = clamp(1 - 0.5 * ||p_draft - p_target||_1, 0, 1)`
+(this is where the `1/2` of Eq. 8 lives). The expected accepted
 prefix length of a block is `tau = sum_k cumprod(accept_rate)_k + 1` (a token
 survives only if every earlier token in its block is accepted; the `+1` counts the
 verified anchor token). `tau` is the training-time speedup proxy. The confidence
@@ -89,73 +104,172 @@ batch plumbing in `training/gpt_step.py`, and per-model HF weight mappings such
 ModelOpt **export** path, not in training.
 
 DSpark's draft (semi-autoregressive block plus Markov and confidence heads) is a
-new neural module. It therefore has to live in one of two places:
-
-1. **A Megatron-Core native module** (as MTP does in
-   `megatron/core/transformer/multi_token_prediction.py`), with tensor- and
-   pipeline-parallel-aware layers and loss folding. This is a Megatron-LM change,
-   not a Bridge change.
-2. **A ModelOpt-injected module** (as EAGLE/Medusa are, via
-   `modelopt.torch.speculative`), in which case Bridge only needs a provider hook,
-   a custom forward step, and export mappings.
-
-### Blueprint: the ModelOpt distillation subsystem
-
-Bridge already contains a self-contained, bolt-on training subsystem that DSpark
-should mirror: **ModelOpt knowledge distillation**. It slots a new training mode
-in without touching the core training loop, via four pieces:
-
-- `models/distillation_provider.py`: a provider that dynamically subclasses the
-  base provider and registers a pre-wrap hook to mutate the already-built,
-  weight-loaded model (there, it attaches the teacher).
-- `training/post_training/distillation.py`: the custom loss (`loss_func_kd`).
-- `training/gpt_step.py`: `forward_step_modelopt`, which selects that loss when the
-  model is a distillation model.
-- `training/distill.py`: a thin entry point, `distill(config) = pretrain(config,
-  forward_step_modelopt)`.
-
-Plus an `examples/distillation/` recipe and a `docs/training/` page. This
-provider-hook plus custom-forward-step plus custom-loss plus thin-entry pattern is
-the established, low-friction way to add a training subsystem here, and it is the
-recommended shape for DSpark.
+new neural module, so the first thing this RFC has to settle is which repository
+owns it end to end. The two candidates lead to genuinely different contracts, and
+picking one is a prerequisite for every section below.
+
+1. **Megatron-Core native** (as MTP does in
+   `megatron/core/transformer/multi_token_prediction.py`): Megatron-Core owns the
+   parallel module, the loss folding, and the sharded checkpoint; Bridge owns
+   configuration and the HF mapping. Requires a Megatron-LM change and a
+   Bridge-defined HF export schema.
+2. **ModelOpt-injected** (as EAGLE/Medusa are, via `modelopt.torch.speculative`):
+   ModelOpt owns the transformation, module, loss, and export; Bridge only
+   configures it. Requires nothing new in Megatron-LM, but ModelOpt has to grow a
+   Megatron-side DSpark plugin first.
+
+### Decision: Megatron-Core native
+
+**This RFC selects option 1.** ModelOpt's DSpark support is currently
+HF-path-only: `modelopt/torch/speculative/plugins/` ships `hf_dspark.py` and
+`modeling_dspark.py`, and `modelopt/torch/export/plugins/hf_spec_export.py` ships
+`DSparkExporter`, but the Megatron-side plugins are `megatron_eagle.py` and
+`megatron_medusa.py` only. There is no `megatron_dspark.py`. Choosing option 2
+would therefore mean landing a Megatron-native DSpark in Model-Optimizer first and
+then waiting on Bridge's ModelOpt pin (currently `nvidia-modelopt==0.46.0rc1`) to
+pick it up, which puts a cross-repo dependency in front of every Bridge-side
+milestone. Option 1 keeps the critical path inside Megatron-LM and Bridge, where
+MTP already demonstrates the pattern end to end.
+
+Consequences of the decision, which the rest of this document then specifies:
+
+- The parallel draft module, its loss folding, and its sharded checkpoint are a
+  **Megatron-Core** contribution, not a Bridge one.
+- Bridge owns the **HF export schema** for trained drafts, because ModelOpt's
+  `DSparkExporter` covers its own HF-path artifacts and not an MCore-native
+  checkpoint.
+- The architecture-agnostic pieces that depend only on `torch` (the sequential and
+  confidence heads, and the CE / L1 / confidence objective) can be developed and
+  unit-tested in Bridge, and move to Megatron-Core when the parallel module lands
+  there. They are not a substitute for that module.
 
 ## Proposed design
 
-Following the distillation blueprint, the Bridge-side surface is:
+The construction path is `ModelConfig` + `ModelBuilder`, **not** a
+`ModelProviderMixin` subclass: `ModelProviderMixin`-based configuration is
+deprecated in current Bridge main (`models/model_provider.py` raises a
+`DeprecationWarning` directing callers to `ModelConfig` + `ModelBuilder`), so a new
+provider subclass could not be the stable API this RFC promises. The DSpark config
+subclasses `GPTModelConfig` and names a `DSparkModelBuilder`, mirroring how
+`MambaModelConfig` names `MambaModelBuilder`.
 
 | Piece | Location (proposed) | Mirrors |
 |---|---|---|
-| Draft provider hook | `models/dspark_provider.py` | `distillation_provider.py` (dynamic subclass + pre-wrap hook that attaches the frozen target's feature capture and the draft heads) |
+| Draft config + builder | `models/dspark/dspark_builder.py` (`DSparkModelConfig`, `DSparkModelBuilder`) | `models/mamba/mamba_builder.py` (`MambaModelConfig` naming `MambaModelBuilder`) |
+| Parallel draft module | Megatron-Core `megatron/core/transformer/` | `multi_token_prediction.py` |
 | Custom forward + loss dispatch | `training/gpt_step.py` (`forward_step_dspark`) | `forward_step_modelopt` |
-| DSpark loss | `training/post_training/dspark.py` | `post_training/distillation.py` (the three-term CE/TV/confidence loss and the `tau` acceptance metric) |
+| DSpark loss | `training/post_training/dspark/` | `post_training/distillation.py` (the three-term CE/L1/confidence loss and the `tau` acceptance metric) |
 | Thin entry point | `training/dspark.py` | `training/distill.py` |
 | HF to/from Megatron weight mapping for `dspark.*` params | `models/<family>/<name>_bridge.py` mapping registry | the `mtp.*` mappings in `qwen3_bridge.py` |
-| Export flag | ModelOpt `export_extra_modules` path | the existing EAGLE/Medusa/MTP export |
 | Recipe + example | `recipes/<family>/…` + `examples/dspark/` | the distillation recipe/example |
 | Docs | this page | `multi-token-prediction.md` |
 
+### Loss return contract
+
+`dspark_loss` follows the Megatron-Core 3-tuple contract used by
+`training/losses.py`: it returns the **unnormalized** micro-batch numerator, this
+micro-batch's `num_tokens`, and a report mapping each metric to a reducible
+`[numerator, denominator]` pair. Normalization happens once, against the globally
+reduced token count.
+
+This requires `calculate_per_token_loss=True` on the model config, and DSpark
+training must fail fast when it is unset. With the default `False`, Megatron-Core
+divides each micro-batch by its own `num_tokens` before averaging over
+micro-batches, which trains a mean of micro-batch means. That is not a rounding
+detail: with two ranks running two micro-batches each at uneven supervised-token
+counts, the per-micro-batch normalization produces a gradient **37% away** from
+the full-batch reference, while the global-token normalization matches it to
+`1.8e-07`.
+
+Reporting pairs exist for the same reason. A logged ratio must be formed as
+`sum(numerator) / sum(denominator)` over the log window and the data-parallel
+group; averaging per-micro-batch ratios reintroduces the same bias in the metrics.
+
+### Memory contract for the full-vocabulary terms
+
+Both the cross-entropy and the L1 term need the draft's normalized distribution
+over the whole vocabulary. At reference scale (`num_blocks=512`, `block_size=7`,
+`vocab≈150k`) one FP32 distribution tensor is 2.03 GiB, and computing the terms
+separately keeps several live at once: the L1 needs the draft probabilities, the
+target probabilities, and their difference, while the cross-entropy saves its own
+`log_softmax` for the backward. Measured on an 80 GB card, an unchunked pass costs
+**10.15 GiB** of forward-plus-backward overhead above its inputs.
+
+The implementation therefore derives both terms from a single FP32 `log_softmax`,
+chunks the flattened rows, and recomputes each chunk in the backward. That brings
+the overhead to **2.06 GiB** (4.9x) at identical values for both terms. Any owner
+of this module inherits the same constraint: either chunk, or state and enforce a
+much smaller anchor-count contract with measurements. Chunking only the L1 is not
+enough, because the cross-entropy's saved `log_softmax` is a full-size tensor of
+the same order.
+
+### Frozen weights: replicas, not aliases
+
+The draft does **not** physically alias the target's parameters. The AutoModel
+reference creates draft-owned `embed_tokens` and `lm_head`, copies the target
+weights into them, freezes the replicas, and controls their trainability
+separately. This RFC adopts the same storage model, because it is what makes the
+remaining contracts well-defined:
+
+- **Pipeline placement.** Replicas let the draft's embedding and LM head sit on
+  whatever pipeline stage hosts the draft, instead of forcing the draft to share a
+  stage with the target's first and last stages.
+- **Checkpointing.** The frozen replicas are ordinary sharded-checkpoint entries.
+  They are saved (not reconstructed) so a resumed run is bit-identical without
+  needing the target checkpoint present.
+- **Trainability.** "Frozen" is a `requires_grad=False` property of draft-owned
+  parameters, not an aliasing side effect, so unfreezing the LM head later is a
+  config change rather than a structural one.
+
 ### Frozen-target supervision
 
 The draft trains against features captured from the frozen target: the selected
 decoder-layer hidden states (fused by `fc` into the draft's context K/V) and the
 final hidden state (which drives the teacher distribution through the shared LM
-head). In Megatron-Core, this capture must respect tensor, pipeline, and context
-parallelism. Two options carry different trade-offs:
-
-- Colocated capture during the same forward (analogous to how MTP consumes extra
-  token IDs), keeping the target and draft in one model.
-- Offline precomputation of target features, streamed to a draft-only training
-  job (as AutoModel's `precompute_dspark.py` does), which sidesteps holding the
-  target in memory but adds a cache format.
-
-### HF to/from Megatron conversion
-
-Because the draft reuses the target's embedding and LM head and adds `fc`,
-`markov_head`, `confidence_head`, and the draft decoder layers, the bridge mapping
-registry must gain `dspark.*` entries (QKV and gated-MLP splits for the draft
-backbone, plain linears for `fc` and the heads), so trained drafts round-trip
-through `AutoBridge` for serving. The `mtp.*` mappings are the closest existing
-template.
+head). Two capture modes are supported, with different schemas:
+
+- **Colocated capture** during the same forward (analogous to how MTP consumes
+  extra token IDs), keeping the target and draft in one model. No cache format;
+  costs target memory.
+- **Offline precomputation** of target features, streamed to a draft-only training
+  job (as AutoModel's `precompute_dspark.py` does). Sidesteps holding the target
+  in memory, and requires a versioned cache schema recording the target checkpoint
+  id, the selected layer indices, the hidden dtype, the anchor sampling seed, and
+  the block size, so a cache can never be silently paired with a different target.
+
+Both modes must respect tensor, pipeline, and context parallelism:
+
+- **Tensor parallel.** The teacher and draft logits are vocab-sharded, so the L1
+  distance must be computed shard-locally and summed across the TP group rather
+  than gathered: `||p_d - p_t||_1` is the sum over TP shards of the local absolute
+  difference, with the softmax denominators reduced across the group first. The
+  exactness of this term is what the acceptance label depends on, so an
+  approximate (top-k, or gathered in lower precision) path is not acceptable here.
+- **Pipeline parallel.** The selected decoder-layer hidden states and the final
+  hidden state must be transported to the draft's stage. When the draft shares the
+  last stage, this is a local read; otherwise the selected hiddens are additional
+  pipeline payload, and their dtype and layout are part of the module contract.
+- **Context parallel.** Anchor blocks must not straddle CP shard boundaries.
+  Anchors are sampled per CP shard from positions whose whole block is local, and
+  the loss mask zeroes any block that would cross a boundary, so no cross-shard
+  attention or gather is needed for the draft.
+
+### Checkpoint and export ownership
+
+- **Training checkpoint.** MCore-native sharded checkpoint, owned by the
+  Megatron-Core module, covering the draft backbone, `fc`, `markov_head`,
+  `confidence_head`, and the frozen embedding and LM-head replicas.
+- **HF export.** Owned by Bridge. ModelOpt's `DSparkExporter` targets ModelOpt's
+  own HF-path artifacts, so an MCore-native checkpoint needs its own HF schema
+  before `AutoBridge` can round-trip it. The bridge mapping registry gains
+  `dspark.*` entries (QKV and gated-MLP splits for the draft backbone, plain
+  linears for `fc` and the heads); the frozen replicas are exported as ordinary
+  weights so a served draft is self-contained. The `mtp.*` mappings in
+  `qwen3_bridge.py` are the closest existing template.
+
+These mappings are specified only once the storage model above is fixed; listing
+them earlier would have been guesswork about whether the shared tensors are stored
+or reconstructed.
 
 ## References