fix(dspark): loss mask off-by-one

[WindChimeRan]

Purpose

Affects sample_from_anchor=True, DSpark's default. DFlash defaults to False and is unchanged.

A draft slot was gated by the loss mask at its own position instead of at the token it predicts. A hidden state at position i yields the distribution over token i+1, so the mask has to be read one position ahead.

Assistant turn at positions 8-39, block_size=4, anchor a=37:

slot j             0    1    2    3
position a+j      37   38   39   40
predicts token    38   39   40   41
is it content?   yes  yes   no   no     <- turn ends at 39

before           [1,   1,   1,   0]     <- loss_mask[a+j]
after            [1,   1,   0,   0]     <- loss_mask[a+j+1]
                            ^ trained on the turn separator

# before
aligned_loss_mask = loss_mask.clone()[:, anchored_block_indices]

# after
mask_indices = (
    anchored_block_indices + 1
    if self.config.sample_from_anchor
    else anchored_block_indices
)
aligned_loss_mask = loss_mask.clone()[:, mask_indices]

The two indices only disagree where the mask changes, so this is confined to turn boundaries — 0.47% of supervised slots at 200-token turns, 5.7% at 20-token turns, always over-supervision.

Tests

One test: every slot the model marks trainable must be predicting a supervised token. Red with the fix reverted, green with it — the failure is the bug stated directly, loss_mask sampled at each trained slot's predicted token, which should be all ones.

E  assert tensor(False)
E   + where ... = tensor([1., 1., 1., ..., 1., 1., 0.]).all

pytest tests/integration/models/test_model_forward.py::TestDFlashParams -q   19 passed
pytest tests/unit/models -q                                                 137 passed

ruff check and ruff format --check pass.

Note for reviewers: this mask also feeds full_acc, eal, accept_rate and the confidence-head target, so acceptance numbers will shift slightly. That is the correction, not a regression.

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

Changes

DFlash loss-mask alignment

Layer / File(s)	Summary
Anchor-aware loss-mask indexing src/speculators/models/dflash/core.py	aligned_loss_mask now indexes loss_mask using anchor indices shifted by one when sample_from_anchor is enabled.
Integration configuration and validation tests/integration/conftest.py, tests/integration/models/test_model_forward.py	The test model factory forwards sample_from_anchor, and an integration test verifies supervision for the predicted token slots.

Possibly related PRs

• vllm-project/speculators#550: Extends the same DFlash integration-test factory and coverage.
• vllm-project/speculators#760: Introduces related sample_from_anchor target and loss alignment behavior.
• vllm-project/speculators#832: Modifies related anchored token index selection in _backbone_forward.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: fixing the DSpark loss-mask off-by-one bug.
Description check	✅ Passed	The description is directly about the loss-mask off-by-one fix and the added test coverage.
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.

✨ 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.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz

[coderabbitai]

Caution

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

⚠️ Outside diff range comments (1)

tests/integration/conftest.py (1)

138-150: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

Keep the proposal count aligned with sample_from_anchor.

This helper forwards sample_from_anchor=True, but the proposal configuration below still hard-codes speculative_tokens=block_size - 1. Production configuration uses block_size when sampling from the anchor, so tests can exercise inconsistent proposal-length metadata.

As per path instructions, the integration helper should propagate the new code path consistently across its configuration.

Suggested fix

             proposal_methods=[
-                GreedyTokenProposalConfig(speculative_tokens=block_size - 1)
+                GreedyTokenProposalConfig(
+                    speculative_tokens=(
+                        block_size if sample_from_anchor else block_size - 1
+                    )
+                )
             ],

🤖 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 `@tests/integration/conftest.py` around lines 138 - 150, Update the
DFlashSpeculatorConfig setup in the integration model helper so
speculative_tokens is derived from sample_from_anchor: use block_size when
sampling from the anchor and block_size - 1 otherwise. Keep the helper’s
proposal-length metadata consistent with production configuration.

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.

Outside diff comments:
In `@tests/integration/conftest.py`:
- Around line 138-150: Update the DFlashSpeculatorConfig setup in the
integration model helper so speculative_tokens is derived from
sample_from_anchor: use block_size when sampling from the anchor and block_size
- 1 otherwise. Keep the helper’s proposal-length metadata consistent with
production configuration.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 536716cb-7006-4b7b-87c4-caff5a67e70a

📥 Commits

Reviewing files that changed from the base of the PR and between f7ec341 and 9e499cc.

📒 Files selected for processing (3)

• src/speculators/models/dflash/core.py
• tests/integration/conftest.py
• tests/integration/models/test_model_forward.py