Add D-PACE loss implementation for D-Flash training

[weifanjiang]

Purpose

This PR implements D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting. D-PACE proposes an alternative per-position loss weighting for training block-parallel draft models. Instead of applying a fixed exponential decay over block positions, D-PACE dynamically adjusts the weights as training progresses. This is based on the intuition that although earlier positions are naturally more important (i.e., an early rejection forces all subsequent draft positions to be discarded), when they are sufficiently stabilized, the training priority should shift to later positions to improve overall acceptance rate.

Two arguments are added to train.py to support D-PACE loss:

• --per-position-loss-weight ["fixed-exp-decay", "dpace"] (default fixed-exp-decay). This flag enables D-PACE loss; while missing, it defaults to the original fixed exponential decay.
• --dpace-alpha (default 0.5). This flag is only effective when D-PACE is enabled. It smoothes per-position confidence according to equation 7 of the paper. The paper uses 0.5 as the default value for their evaluation and justifies this choice in sensitivity analysis.

Note that D-PACE requires cross-entropy as the training loss (--loss-fn ce), allowing per-position confidence score computation to be piggybacked on the loss computation. Applying D-PACE to other losses is not supported and will result in ValueError.

D-PACE is currently enabled for D-Flash and D-Spark draft models. Theoretically, it is compatible with all block-parallel drafters.

Tests

D-Flash + D-PACE evaluation. Multiple D-Flash checkpoints are trained with different loss function configurations:

• --loss-fn kld: KL Divergence (current default)
• --loss-fn ce: Cross Entropy with fixed exponential decay
• --per-position-loss-weight dpace --dpace-alpha ALPHA: D-PACE with ALPHA = 0.3, 0.5, 0.7.

The training set is 100K random samples from Magpie+Ultrachat. The verifier model is Qwen3-8B. Each checkpoint uses 1 A100 for hidden state generation and 7 A100s for training. Training configurations are identical for all checkpoints:

--target-layer-ids 1 9 17 25 34  --num-layers 5 --draft-vocab-size 32000 --max-anchors 3072 --block-size 8 --sliding-window 2048 --sliding-window-indices 0 1 2 3 4 --scheduler-type cosine --lr 0.0006 --epochs 1 --noise-std 0.0

Metric: since D-PACE is a pure training-time optimization and does not change draft model's forward pass complexity, its gains are captured entirely by acceptance rates.

Evaluation results: Mean-Accepted-Length (MAL) on Speculator Benchmarks:

Dataset	KL Divergence	CE	D-PACE $\alpha=0.3$	D-PACE $\alpha=0.5$	D-PACE $\alpha=0.7$
HumanEval	2.977	2.956	3.047	3.169	3.139
math_reasoning	3.275	3.322	3.541	3.532	3.531
qa	2.279	2.343	2.455	2.495	2.431
question	2.585	2.567	2.685	2.680	2.708
rag	2.302	2.269	2.393	2.372	2.336
summarization	1.906	1.918	1.981	1.980	1.991
tool_call	2.372	2.327	2.480	2.454	2.467
translation	2.096	2.072	2.176	2.183	2.191
writing	2.565	2.586	2.683	2.684	2.718
Mean	2.484	2.485	2.605	2.617	2.612

D-Spark + D-PACE evaluation. D-Spark checkpoints use Markov rank 256 and vanilla Markov head type. The confidence head is disabled. Remaining training configurations match above D-Flash experiments.

• --loss-fn ce: Vanilla fixed-decay cross-entropy
• --position-loss-weight dpace --dpace-alpha 0.5: D-PACE with $\alpha=0.5$

MALs:

Dataset	CE	D-PACE $\alpha=0.5$
HumanEval	3.145	3.302
math_reasoning	3.457	3.615
qa	2.352	2.492
question	2.671	2.770
rag	2.276	2.384
summarization	1.916	1.996
tool_call	2.451	2.538
translation	2.082	2.202
writing	2.686	2.776
Mean	2.560	2.675

D-PACE improves upon fixed-decay loss for both D-Flash and D-Spark drafters.

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds D-PACE (per-position confidence-based) loss weighting as an alternative to fixed-exponential decay. Introduces CLI flags, extends DFlashDraftModel's forward and get_trainer_kwargs to propagate the new configuration, adds branching logic in compute_metrics, and adds a new dpace_loss_weight helper plus optional precomputed CE support in compound_loss and loss_function.

Changes

D-PACE loss weighting

Layer / File(s)	Summary
dpace_loss_weight helper and loss function plumbing src/speculators/models/metrics.py	Adds dpace_loss_weight to compute masked per-position weights from per-token NLL, and extends compound_loss/loss_function with an optional dpace_precomputed_ce parameter to reuse precomputed elementwise CE.
dflash compute_metrics D-PACE branch src/speculators/models/dflash/metrics.py	Adds per_position_loss_weight and dpace_alpha parameters to compute_metrics, imports ce_loss/dpace_loss_weight, and branches decay_fn selection between the new D-PACE weighting and the existing fixed-exp-decay behavior.
DFlashDraftModel forward/trainer kwargs propagation src/speculators/models/dflash/core.py	Extends forward with per_position_loss_weight and dpace_alpha parameters forwarded into compute_metrics, and updates get_trainer_kwargs to propagate these values into train/validation kwargs.
CLI flags for D-PACE configuration scripts/train.py	Adds --per-position-loss-weight and --dpace-alpha CLI arguments to parse_args().

Possibly related PRs

• vllm-project/speculators#542: Also extends DFlashDraftModel.forward and dflash/metrics.compute_metrics signatures for configurable loss computation, overlapping the same call sites modified here.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: adding D-PACE loss support for D-Flash training.
Description check	✅ Passed	The description is directly related to the changeset and explains the D-PACE training update and tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (3)

src/speculators/models/metrics.py (2)

281-320: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Validate alpha range for numerical stability.

Confirmed against the D-PACE paper (arXiv:2605.18810) that this implementation matches the published smoothing/prefix-product/suffix-sum weighting scheme. However, the paper notes the smoothing parameter is only meaningful for a lower bound preventing the cumulative product in the weight from vanishing, and specifically that at α = 0, where the cumulative product vanishes and τ stalls near 1.0. There's no validation here that alpha stays within [0, 1]; an out-of-range value (e.g., negative or >1, or the degenerate alpha=0) would silently produce meaningless/unstable weights instead of failing fast.

🛡️ Proposed fix

     with torch.no_grad():
+        if not 0.0 < alpha <= 1.0:
+            raise ValueError(f"alpha must be in (0, 1], got {alpha}")
         # convert CE to per-position confidence
         q = torch.exp(-neg_log_q).float()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/speculators/models/metrics.py` around lines 281 - 320, Validate the
smoothing parameter in dpace_loss_weight before computing the prefix product:
add an explicit check that alpha is within the supported range, rejecting
out-of-range values and the degenerate zero case with a clear failure. Keep the
existing weighting logic unchanged, but ensure the function fails fast for
invalid alpha so the numerical stability of the D-PACE weighting is preserved.

---

300-302: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value

Reshape relies on exact block-size divisibility.

num_anchors = q.shape[1] // block_size uses floor division; if q.shape[1] isn't an exact multiple of block_size, the subsequent .reshape(num_anchors, block_size) calls will raise a runtime error rather than a clear, actionable message. In the current call path (compute_metrics in dflash/metrics.py), seq_len is always num_anchors * block_size by construction, so this shouldn't trigger today, but a defensive assert would make future misuse fail with a clearer error.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/speculators/models/metrics.py` around lines 300 - 302, The reshape logic
in the metrics path assumes q.shape[1] is exactly divisible by block_size, so
add a defensive validation before computing num_anchors in the relevant reshape
block. Update the code around the q.reshape and loss_mask.reshape operations to
assert or explicitly check divisibility and raise a clear error message if the
sequence length is not compatible with block_size, using the existing
compute_metrics / metrics reshape flow as the location to safeguard.

scripts/train.py (1)

964-976: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Consider validating --dpace-alpha at parse time.

Per the D-PACE paper, alpha is only meaningful in [0, 1], with alpha=0 causing the weighting to degenerate. Argparse currently accepts any float unchecked. Adding a lightweight validator here (or in dpace_loss_weight, see comment in metrics.py) would catch a misconfigured training run early rather than after silently training with a degenerate loss weighting.

As per path instructions, "Check that scripts handle argument parsing robustly."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/train.py` around lines 964 - 976, Add parse-time validation for the
D-PACE smoothing argument so invalid values are rejected early. In the argument
setup around parser.add_argument for --dpace-alpha, enforce that the value is
within the valid [0, 1] range and disallow degenerate settings such as 0, either
via an argparse type/validator or a small reusable check used by
dpace_loss_weight. Keep the validation close to the existing D-Pace-specific
arguments so misconfigured runs fail immediately with a clear message.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/models/dflash/metrics.py`:
- Around line 55-64: Replace the inline lambda assigned to decay_fn in the
dflash metrics logic with a named nested function to satisfy Ruff E731 and
improve readability. In the per_position_loss_weight == "dpace" branch, define a
def that returns decay_mult instead of using lambda, and keep the existing
partial(dflash_loss_decay, gamma=gamma) path unchanged. Use a distinct,
descriptive function name near ce_loss, dpace_loss_weight, and decay_fn so the
intent is clear and no parameter name conflicts with the outer pos_idx.

---

Nitpick comments:
In `@scripts/train.py`:
- Around line 964-976: Add parse-time validation for the D-PACE smoothing
argument so invalid values are rejected early. In the argument setup around
parser.add_argument for --dpace-alpha, enforce that the value is within the
valid [0, 1] range and disallow degenerate settings such as 0, either via an
argparse type/validator or a small reusable check used by dpace_loss_weight.
Keep the validation close to the existing D-Pace-specific arguments so
misconfigured runs fail immediately with a clear message.

In `@src/speculators/models/metrics.py`:
- Around line 281-320: Validate the smoothing parameter in dpace_loss_weight
before computing the prefix product: add an explicit check that alpha is within
the supported range, rejecting out-of-range values and the degenerate zero case
with a clear failure. Keep the existing weighting logic unchanged, but ensure
the function fails fast for invalid alpha so the numerical stability of the
D-PACE weighting is preserved.
- Around line 300-302: The reshape logic in the metrics path assumes q.shape[1]
is exactly divisible by block_size, so add a defensive validation before
computing num_anchors in the relevant reshape block. Update the code around the
q.reshape and loss_mask.reshape operations to assert or explicitly check
divisibility and raise a clear error message if the sequence length is not
compatible with block_size, using the existing compute_metrics / metrics reshape
flow as the location to safeguard.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d76c259c-9738-4941-87eb-0fba2af483e2

📥 Commits

Reviewing files that changed from the base of the PR and between 2d97b17 and 7b26944.

📒 Files selected for processing (4)

• scripts/train.py
• src/speculators/models/dflash/core.py
• src/speculators/models/dflash/metrics.py
• src/speculators/models/metrics.py

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[fynnsu]

Hi @weifanjiang thanks for the pr! I think we need to change how we're piping some things through speculators so that it fits within the existing abstractions.