Add D-PACE loss implementation for D-Flash training (#736)

<!-- markdownlint-disable -->

<!-- PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS
(AT THE BOTTOM) HAVE BEEN CONSIDERED. -->

## Purpose

This PR implements [D-PACE: Dynamic Position-Aware Cross-Entropy for
Parallel Speculative Drafting](https://arxiv.org/abs/2605.18810). 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 are
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](https://huggingface.co/datasets/RedHatAI/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:

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

---------

Signed-off-by: Weifan Jiang <weijiang@redhat.com>
Co-authored-by: Weifan Jiang <weifanjiang@a100-08.nemg-001.lab.rdu2.dc.redhat.com>
Co-authored-by: Weifan Jiang <weijiang@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@a100-04.nemg-001.lab.rdu2.dc.redhat.com>

Author: 3 people

scripts/train.py

@@ -961,6 +961,20 @@ def parse_args():
         default=4.0,
         help="Decay gamma for DFlash/DSpark loss weighting (default: 4.0)",
     )
+    # D-Pace specific arguments (loss weight option + smoothing)
+    parser.add_argument(
+        "--per-position-loss-weight",
+        choices=["fixed-exp-decay", "dpace"],
+        default="fixed-exp-decay",
+        help="Per-position loss weight option for D-PACE support"
+        "default: fixed-exp-decay",
+    )
+    parser.add_argument(
+        "--dpace-alpha",
+        type=float,
+        default=0.5,
+        help="Smoothing constant for D-PACE loss (default: 0.5)",
+    )
     # DSpark-specific arguments (sequential Markov head + confidence head).
     parser.add_argument(
         "--markov-rank",
@@ -1132,6 +1146,13 @@ def parse_args():
     provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS)
     validate_draft_init_args(parser, args, provided)
     resolve_loss_config(args.loss_fn)
+
+    if args.per_position_loss_weight == "dpace":
+        if args.loss_fn != "ce":
+            parser.error("--per-position-loss-weight=dpace requires --loss-fn=ce")
+        if not 0.0 < args.dpace_alpha <= 1.0:
+            raise ValueError(f"alpha must be in (0, 1], got {args.dpace_alpha}")
+
     return args
 
 

src/speculators/models/dflash/core.py

@@ -207,10 +207,16 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         loss_config = resolve_loss_config(kwargs["loss_fn"])
         gamma = kwargs.get("dflash_decay_gamma", 4.0)
         max_anchors = kwargs.get("max_anchors", 3072)
+        per_position_loss_weight = kwargs.get(
+            "per_position_loss_weight", "fixed-exp-decay"
+        )
+        dpace_alpha = kwargs.get("dpace_alpha", 0.5)
         shared = {
             "loss_config": loss_config,
             "gamma": gamma,
             "max_anchors": max_anchors,
+            "per_position_loss_weight": per_position_loss_weight,
+            "dpace_alpha": dpace_alpha,
         }
         return dict(shared), dict(shared)
 
@@ -393,6 +399,8 @@ def forward(
         loss_config: LossConfig | None = None,
         gamma: float = 4.0,
         max_anchors: int = 3072,
+        per_position_loss_weight: str = "fixed-exp-decay",
+        dpace_alpha: float = 0.5,
         **kwargs,
     ):
         _, logits, targets, aligned_loss_mask, _ = self._backbone_forward(
@@ -412,6 +420,8 @@ def forward(
             self.block_size,
             gamma=gamma,
             loss_config=loss_config,
+            per_position_loss_weight=per_position_loss_weight,
+            dpace_alpha=dpace_alpha,
         )
         draft_tokens = torch.argmax(logits, dim=-1)
 

src/speculators/models/dflash/metrics.py

@@ -10,6 +10,7 @@
     compound_loss,
     compute_accuracy_multi_step,
     dflash_loss_decay,
+    dpace_loss_decay,
     kl_div_loss,
 )
 
@@ -23,6 +24,8 @@ def compute_metrics(
     block_size: int = 1,
     gamma: float = 4.0,
     loss_config: LossConfig | None = None,
+    per_position_loss_weight: str = "fixed-exp-decay",
+    dpace_alpha: float = 0.5,
 ) -> tuple[torch.Tensor, dict]:
     """Compute loss and accuracy metrics for draft model predictions.
 
@@ -33,6 +36,8 @@ def compute_metrics(
         block_size: Block size for per-position metrics
         gamma: Temperature for exponential decay in loss weighting
         loss_config: Mapping of ``{name: (loss_fn, weight)}``
+        per_position_loss_weight: Weighting option for per-position block-drafting loss
+        dpace_alpha: Smoothing constant for D-Pace loss weighting
 
     Returns:
         Tuple of (loss, metrics_dict) where metrics_dict contains:
@@ -47,13 +52,23 @@ def compute_metrics(
     pos_idx = torch.arange(seq_len, device=logits.device) % block_size
     pos_idx = pos_idx.unsqueeze(0)  # shape: [1, T]
 
+    if per_position_loss_weight == "dpace":
+        decay_fn = partial(
+            dpace_loss_decay,
+            loss_mask=loss_mask,
+            block_size=block_size,
+            dpace_alpha=dpace_alpha,
+        )
+    else:
+        decay_fn = partial(dflash_loss_decay, gamma=gamma)
+
     loss, term_losses = compound_loss(
         logits,
         targets,
         loss_mask,
         pos_idx,
         loss_config=loss_config,
-        decay_fn=partial(dflash_loss_decay, gamma=gamma),
+        decay_fn=decay_fn,
     )
 
     pred_ids = torch.argmax(logits, dim=-1)

src/speculators/models/dspark/core.py

@@ -84,11 +84,17 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         gamma = kwargs.get("dflash_decay_gamma", 4.0)
         max_anchors = kwargs.get("max_anchors", 3072)
         confidence_head_alpha = kwargs.get("confidence_head_alpha", 1.0)
+        per_position_loss_weight = kwargs.get(
+            "per_position_loss_weight", "fixed-exp-decay"
+        )
+        dpace_alpha = kwargs.get("dpace_alpha", 0.5)
         shared = {
             "loss_config": loss_config,
             "gamma": gamma,
             "max_anchors": max_anchors,
             "confidence_head_alpha": confidence_head_alpha,
+            "per_position_loss_weight": per_position_loss_weight,
+            "dpace_alpha": dpace_alpha,
         }
         return dict(shared), dict(shared)
 
@@ -105,6 +111,8 @@ def forward(
         gamma: float = 4.0,
         max_anchors: int = 3072,
         confidence_head_alpha: float = 1.0,
+        per_position_loss_weight: str = "fixed-exp-decay",
+        dpace_alpha: float = 0.5,
         **kwargs,
     ):
         hidden, logits, targets, aligned_loss_mask, anchored_block_indices = (
@@ -167,6 +175,8 @@ def forward(
             loss_config=loss_config or _DEFAULT_LOSS_CONFIG,
             gamma=gamma,
             confidence_head_alpha=confidence_head_alpha,
+            per_position_loss_weight=per_position_loss_weight,
+            dpace_alpha=dpace_alpha,
         )
         draft_tokens = torch.argmax(logits, dim=-1)
         return draft_tokens, loss, metrics

src/speculators/models/dspark/metrics.py

@@ -18,6 +18,7 @@
     compound_loss,
     compute_accuracy_multi_step,
     dflash_loss_decay,
+    dpace_loss_decay,
 )
 
 __all__ = [
@@ -31,13 +32,15 @@ def _masked_decayed_mean(
     elementwise: torch.Tensor,  # [1, T]
     loss_mask: torch.Tensor,  # [1, T]
     pos_idx: torch.Tensor,  # [1, T]
-    decay_fn: Callable[[torch.Tensor], torch.Tensor] | None,
+    decay_fn: Callable[..., torch.Tensor] | None,
 ) -> torch.Tensor:
     """Masked, optionally position-decayed mean of a precomputed per-position term."""
     loss_mask = loss_mask.to(elementwise.dtype)
     weighted = elementwise * loss_mask
     if decay_fn is not None:
-        weighted = weighted * decay_fn(pos_idx.to(weighted.dtype))
+        weighted = weighted * decay_fn(
+            pos_idx.to(weighted.dtype), elementwise_loss=elementwise
+        )
     denominator = loss_mask.sum(dim=1) + _EPS
     return (weighted.sum(dim=1) / denominator).mean()
 
@@ -51,13 +54,23 @@ def compute_metrics(
     loss_config: LossConfig,
     gamma: float = 4.0,
     confidence_head_alpha: float = 1.0,
+    per_position_loss_weight: str = "fixed-exp-decay",
+    dpace_alpha: float = 0.5,
 ) -> tuple[torch.Tensor, dict]:
     """Compute the DSpark loss and a metrics dict (``*_sum``/``*_total`` pairs)."""
 
     device = logits.device
     seq_len = logits.shape[1]
     pos_idx = (torch.arange(seq_len, device=device) % block_size).unsqueeze(0)
-    decay_fn = partial(dflash_loss_decay, gamma=gamma)
+    if per_position_loss_weight == "dpace":
+        decay_fn = partial(
+            dpace_loss_decay,
+            loss_mask=loss_mask,
+            block_size=block_size,
+            dpace_alpha=dpace_alpha,
+        )
+    else:
+        decay_fn = partial(dflash_loss_decay, gamma=gamma)
 
     loss, term_losses = compound_loss(
         logits, targets, loss_mask, pos_idx, loss_config=loss_config, decay_fn=decay_fn

src/speculators/models/metrics.py

@@ -244,7 +244,7 @@ def lk_hybrid_loss(
     return elementwise_loss  # noqa: RET504
 
 
-def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float):
+def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float, **_kwargs):
     """Compute DFlash-style exponential decay weights per position.
 
     Position 0 gets weight 0, position 1 gets weight 1, and subsequent positions
@@ -265,7 +265,7 @@ def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float):
     return decay_mult  # noqa: RET504
 
 
-def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
+def exp_loss_decay(pos_idx: torch.Tensor, gamma: float, **_kwargs):
     """Compute simple exponential decay weights as gamma^pos_idx.
 
     Args:
@@ -278,6 +278,57 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
     return gamma**pos_idx
 
 
+def dpace_loss_decay(
+    pos_idx: torch.Tensor,  # noqa: ARG001
+    loss_mask: torch.Tensor,
+    block_size: int,
+    dpace_alpha: float,
+    elementwise_loss: torch.Tensor,
+    **_kwargs,
+):
+    """
+    Per-position block-drafting loss weight based on D-PACE
+
+    Args:
+        elementwise_loss: requires to be cross-entropy loss, negative log-likelihood
+            of per-position confidence
+        dpace_alpha: confidence smoothing constant
+
+    Returns:
+        Decay multiplier tensor with same shape as pos_idx.
+    """
+    with torch.no_grad():
+        # convert CE to per-position confidence
+        q = torch.exp(-elementwise_loss).float()
+
+        # reshape loss to [num_anchors, block_size]
+        # for intra-block cumulative multiplication
+        if q.shape[1] % block_size != 0:
+            raise ValueError(
+                f"q.shape[1] ({q.shape[1]}) must be divisible by "
+                f"block_size ({block_size})"
+            )
+        num_anchors = q.shape[1] // block_size
+        q = q.reshape(num_anchors, block_size)
+        mask = loss_mask.reshape(num_anchors, block_size).to(q.dtype)
+
+        # smoothed confidence for numerical stability
+        smooth = (1.0 - dpace_alpha) * q + dpace_alpha
+        smooth = torch.where(mask > 0, smooth, torch.ones_like(smooth))
+
+        # prefix cumulative production
+        prefix = torch.cumprod(smooth, dim=-1)
+
+        # suffix summation: flip -> cumsum -> flip
+        weight = torch.flip(
+            torch.cumsum(torch.flip(prefix * mask, dims=[-1]), dim=-1), dims=[-1]
+        )
+        weight = weight * mask
+
+    # reshape weight
+    return weight.reshape(1, -1)
+
+
 _LOSS_FN_MAP: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {
     "kl_div": kl_div_loss,
     "rkl": reverse_kl_div_loss,
@@ -360,7 +411,7 @@ def compound_loss(
     loss_mask: torch.Tensor,
     pos_idx: torch.Tensor,
     loss_config: LossConfig,
-    decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    decay_fn: Callable[..., torch.Tensor] | None = None,
 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
     """Compute a weighted sum of loss terms.
 
@@ -377,7 +428,12 @@ def compound_loss(
     multi = len(loss_config) > 1
     for name, (fn, weight) in loss_config.items():
         term = loss_function(
-            logits, targets, loss_mask, pos_idx, loss_fn=fn, decay_fn=decay_fn
+            logits,
+            targets,
+            loss_mask,
+            pos_idx,
+            loss_fn=fn,
+            decay_fn=decay_fn,
         )
         if multi:
             term_losses[f"{name}_loss"] = term.detach()
@@ -391,7 +447,7 @@ def loss_function(
     loss_mask: torch.Tensor,  # shape: [1, seq_len]
     pos_idx: torch.Tensor,  # shape: [1, seq_len]
     loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = kl_div_loss,
-    decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    decay_fn: Callable[..., torch.Tensor] | None = None,
 ):
     """Compute masked, optionally position-decayed training loss.
 
@@ -412,7 +468,9 @@ def loss_function(
     elementwise_loss = elementwise_loss * loss_mask
 
     if decay_fn is not None:
-        decay_mult = decay_fn(pos_idx.to(elementwise_loss.dtype))
+        decay_mult = decay_fn(
+            pos_idx.to(elementwise_loss.dtype), elementwise_loss=elementwise_loss
+        )
         elementwise_loss = elementwise_loss * decay_mult
 
     denominator = loss_mask.sum(dim=1) + _EPS