Add hybrid LK loss (adaptive KL/TV blend) for speculative decoding

[GIREESH7963]

Summary

Adds the hybrid LK loss: an adaptive per-position blend of KL and TV,
L = lambda*KL + (1-lambda)*TV with lambda = exp(-eta * sg[alpha]), alpha = overlap.

Motivation

TV optimizes acceptance directly but has a vanishing gradient from a cold start; KL has a
strong early gradient but optimizes a proxy. The hybrid leans on KL when overlap is low and
shifts to TV as the draft aligns, making TV's acceptance-optimal target trainable from
scratch. Source: Samarin et al., arXiv 2602.23881. Paper's best setting: eta=3.

Changes

• lk_hybrid_loss(logits, targets, eta=3.0) in metrics.py, reusing kl_div_loss and the
  TV overlap term; alpha in the blend weight is detached.
• resolve_loss_fn gains lk_loss_eta (default 3.0), binding eta into the "lk_hybrid"
  entry via functools.partial (single binding point; call sites unchanged).
• --lk-loss-eta CLI flag (mirrors --dflash-decay-gamma), threaded through each
  get_trainer_kwargs into resolve_loss_fn; "lk_hybrid" added to --loss-fn choices.
• Tests: eta->0 == KL, eta->inf == TV, shape/finiteness, alpha-detach gradient check,
  resolve binds eta.

Notes

• The adaptive KL/TV blend and the DFlash positional decay are orthogonal weightings.
• tv/nla are not in the --loss-fn choices list either; happy to add them here if desired.
• Loss key name "lk_hybrid" open to preference.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ee0151b4-5507-4ce6-9c75-47d96c1130dd

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The PR adds an lk_hybrid loss variant, introduces a configurable lk_loss_eta parameter, and threads that parameter through the training CLI and draft-model trainer setup. It also adds unit tests for the hybrid loss and resolver binding.

Changes

LK hybrid loss plumbing

Layer / File(s)	Summary
Hybrid loss definition and resolver binding src/speculators/models/metrics.py, tests/unit/models/test_metrics.py	lk_hybrid_loss blends KL divergence and TV loss with a detached overlap weight, resolve_loss_fn accepts lk_loss_eta and maps lk_hybrid through partial, and unit tests cover edge cases, gradients, and binding.
CLI and trainer loss wiring scripts/train.py, src/speculators/models/dflash/core.py, src/speculators/models/eagle3/core.py, src/speculators/models/peagle/core.py	scripts/train.py adds lk_hybrid and --lk-loss-eta, and the draft-model trainer setup forwards lk_loss_eta into resolve_loss_fn.

Possibly related PRs

• vllm-project/speculators#542: Updates the same resolve_loss_fn and get_trainer_kwargs wiring path that this PR extends with lk_hybrid and lk_loss_eta.
• vllm-project/speculators#648: Adds tv_loss, which is one of the components blended by lk_hybrid_loss.
• vllm-project/speculators#657: Modifies the same metrics and test area with acceptance-overlap-based speculative-decoding losses and related unit coverage.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: introducing an adaptive LK hybrid loss for speculative decoding.
Docstring Coverage	✅ Passed	Docstring coverage is 87.50% which is sufficient. The required threshold is 80.00%.
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.
Description check	✅ Passed	The description accurately matches the changeset, covering the new hybrid LK loss, CLI flag, resolver wiring, 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]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

scripts/train.py (1)

966-971: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Update --loss-fn help to document the lk_hybrid choice.

lk_hybrid was added to choices, but the help text still only describes kl_div and ce, so users won't know what the new option does or that --lk-loss-eta applies to it.

📝 Proposed help text update

         help=(
             "Loss function used during draft model training. "
             "'kl_div' = KL divergence (default). "
-            "'ce' = cross-entropy."
+            "'ce' = cross-entropy. "
+            "'lk_hybrid' = adaptive KL/TV blend (see --lk-loss-eta)."
         ),

🤖 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 966 - 971, The `--loss-fn` help text in the
training CLI is missing the newly added `lk_hybrid` option, so update the
argument help near the `choices` definition in `train.py` to describe
`lk_hybrid` alongside `kl_div` and `ce`, and mention that it uses
`--lk-loss-eta` as its mixing parameter. Keep the description consistent with
the existing `loss_fn`/`lk_loss_eta` argument names so users can understand the
available training modes from the CLI help.

🤖 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.

Outside diff comments:
In `@scripts/train.py`:
- Around line 966-971: The `--loss-fn` help text in the training CLI is missing
the newly added `lk_hybrid` option, so update the argument help near the
`choices` definition in `train.py` to describe `lk_hybrid` alongside `kl_div`
and `ce`, and mention that it uses `--lk-loss-eta` as its mixing parameter. Keep
the description consistent with the existing `loss_fn`/`lk_loss_eta` argument
names so users can understand the available training modes from the CLI help.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 80a07b9c-2c60-4de6-8e9a-39ea47b11988

📥 Commits

Reviewing files that changed from the base of the PR and between 5aa4da9 and 2857b7c.

📒 Files selected for processing (6)

• scripts/train.py
• src/speculators/models/dflash/core.py
• src/speculators/models/eagle3/core.py
• src/speculators/models/metrics.py
• src/speculators/models/peagle/core.py
• tests/unit/models/test_metrics.py

[GIREESH7963]

addressed in 03bc603. The --loss-fn help now lists lk_hybrid and points to --lk-loss-eta for its mixing parameter ('lk_hybrid' = adaptive KL/TV blend (see --lk-loss-eta).).

[fynnsu]

Looks good but thing to fix below.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @GIREESH7963.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[GIREESH7963]

Rebased onto main to clear the conflict (the loss refactor in #671/#677 is what conflicted). Summary of what changed since the approval, since it's a bit more than a plain rebase:

• lk_hybrid is now registered in the new module-level _LOSS_FN_MAP, so it works both as a standalone --loss-fn lk_hybrid and inside the JSON weighted-combination spec (e.g. '{"lk_hybrid": 0.5, "ce": 0.5}').
• Re: the suggestion to add tv/nla to the --loss-fn choices list — that choices=[...] list was removed in the refactor (the arg is now a free-form name-or-JSON spec). tv, nla, and lk_hybrid are all listed in the updated free-form help text instead.
• Dropped the --lk-loss-eta flag. The refactored interface no longer has a per-loss hyperparameter channel (losses are bare (logits, targets) callables in _LOSS_FN_MAP), so eta is now baked at the paper's default of 3.0. Happy to restore tunability as a follow-up if you'd prefer.

Net functional diff is now just the lk_hybrid_loss function + one _LOSS_FN_MAP entry + the help-text mention, plus unit tests. Flagging in case you'd like another glance before merge given the --lk-loss-eta removal.