feat(peagle): add random anchor subsampling (max_anchors)

[orestis-z]

Closes #567.

Summary

• Add max_anchors config to P-EAGLE's COD sampling that randomly subsamples chain starting points at depth 1+, capping attention mask size for longer sequences while keeping depth 0 as the full sequence
• Change --max-anchors CLI default from 256 to None so P-EAGLE uses all positions by default (DFlash retains its 3072 hardcoded fallback)
• Add unit tests verifying depth-0 preservation, chain capping, sorted order, and backward compatibility

Motivation

P-EAGLE's COD sampling starts a chain from every loss_mask=1 position, producing attention masks that OOM at 8K+ sequences. The ablation study (#567 (comment)) showed that randomly subsampling with max_anchors=1024 actually improves d0 accuracy by +3.1pp while dramatically reducing memory, enabling training at longer sequence lengths.

Supersedes #590 (contiguous window approach, which degraded quality) and #683.

Test plan

• 6 new unit tests in tests/unit/models/test_peagle_data.py — all pass
• 38 existing P-EAGLE integration tests — all pass
• ruff check clean on all changed files
• CI green

🤖 Generated with Claude Code

[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: 37fc3497-9114-41ca-957f-a4bd237ca05f

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

B{max_anchors set and candidates exceed cap?}

B -- Yes --> C[Random permutation subset of all_valid_indices to max_anchors]
C --> D[Sort selected anchors ascending]
B -- No --> D
D --> E[Existing multi-depth COD sampling logic]

</diagram>
</layer>
<layer id="model_wiring" title="PEagleDraftModel forward and from_training_args wiring" depends_on="data_and_config">
<summary>forward passes config.max_anchors into generate_cod_sample_indices; from_training_args forwards the max_anchors kwarg into PEagleSpeculatorConfig construction.</summary>
<ranges>
range_25fa92804844
range_bf2c00990afc
</ranges>
</layer>
<layer id="cli_and_dflash" title="CLI default change and DFlash fallback adjustment" depends_on="data_and_config">
<summary>The --max-anchors CLI argument default changes from 256 to None with updated help text. DFlash from_training_args switches to or-based fallback so None propagates correctly to the 3072 default.</summary>
<ranges>
range_18a526e0e910
range_1748325f566c
</ranges>
</layer>
<layer id="tests" title="Unit tests for max_anchors in generate_cod_sample_indices" depends_on="model_wiring">
<summary>New TestMaxAnchors suite verifies depth-0 preservation, anchor count capping at depth 1, None equality with default, below-cap passthrough, and sorted output ordering.</summary>
<ranges>
range_db455ca48ff3
</ranges>
</layer>
</cohort>
<unassigned_ranges>
</unassigned_ranges>
-->

## Walkthrough

Adds `max_anchors` support to P-EAGLE COD sampling: a new optional config field in `PEagleSpeculatorConfig`, random capping logic in `generate_cod_sample_indices`, wiring through `PEagleDraftModel.forward` and `from_training_args`, a CLI default change from `256` to `None`, a DFlash fallback adjustment, and new unit tests.

## Changes

**P-EAGLE max\_anchors COD sampling cap**

| Layer / File(s) | Summary |
|---|---|
| **`generate_cod_sample_indices` cap + config field** <br> `src/speculators/models/peagle/data.py`, `src/speculators/models/peagle/config.py` | `generate_cod_sample_indices` gains `max_anchors: int \| None = None`; when set and candidates exceed the cap, a random permutation subset is drawn and sorted before multi-depth sampling. `PEagleSpeculatorConfig` adds the matching optional `max_anchors` field with `ge=1` validation. |
| **Model wiring** <br> `src/speculators/models/peagle/core.py` | `forward` passes `max_anchors=self.config.max_anchors` into `generate_cod_sample_indices`; `from_training_args` forwards the `max_anchors` kwarg into `PEagleSpeculatorConfig`. |
| **CLI default and DFlash fallback** <br> `scripts/train.py`, `src/speculators/models/dflash/core.py` | `--max-anchors` default changes from `256` to `None` with updated help text. DFlash `from_training_args` switches from `kwargs.get("max_anchors", 3072)` to `kwargs.get("max_anchors") or 3072` so a `None` CLI value still resolves to `3072` for DFlash. |
| **Unit tests** <br> `tests/unit/models/test_peagle_data.py` | New `TestMaxAnchors` class asserts depth-0 preservation, depth-1 anchor count capping, `None` equality with default, below-cap passthrough, and sorted anchor ordering. |

## Possibly related issues

- **vllm-project/speculators#567**: This PR directly implements the RFC's proposed `max_anchors` cap on COD anchor selection via `generate_cod_sample_indices` and `PEagleSpeculatorConfig`.

</details>

<!-- walkthrough_end -->
<!-- pre_merge_checks_walkthrough_start -->

<details>
<summary>🚥 Pre-merge checks | ✅ 5</summary>

<details>
<summary>✅ Passed checks (5 passed)</summary>

|         Check name         | Status   | Explanation                                                                                                                               |
| :------------------------: | :------- | :---------------------------------------------------------------------------------------------------------------------------------------- |
|         Title check        | ✅ Passed | The title clearly describes the main change: adding random anchor subsampling via max_anchors for P-EAGLE.                                |
|     Docstring Coverage     | ✅ Passed | Docstring coverage is 83.33% 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 code changes: adding P-EAGLE max_anchors support, updating the CLI default, and adding unit tests. |

</details>

</details>

<!-- pre_merge_checks_walkthrough_end -->
<!-- finishing_touch_checkbox_start -->

<details>
<summary>✨ Finishing Touches</summary>

<details>
<summary>🧪 Generate unit tests (beta)</summary>

- [ ] <!-- {"checkboxId": "f47ac10b-58cc-4372-a567-0e02b2c3d479", "radioGroupId": "utg-output-choice-group-unknown_comment_id"} -->   Create PR with unit tests
- [ ] <!-- {"checkboxId": "6ba7b810-9dad-11d1-80b4-00c04fd430c8", "radioGroupId": "utg-output-choice-group-unknown_comment_id"} -->   Commit unit tests in branch `peagle-max-anchors`

</details>

</details>

<!-- finishing_touch_checkbox_end -->
<!-- tips_start -->

---

Thanks for using [CodeRabbit](https://coderabbit.ai?utm_source=oss&utm_medium=github&utm_campaign=vllm-project/speculators&utm_content=687)! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

<details>
<summary>❤️ Share</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai)
- [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)

</details>


<sub>Comment `@coderabbitai help` to get the list of available commands.</sub>

<!-- tips_end -->

[coderabbitai]

Actionable comments posted: 2

Caution

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

⚠️ Outside diff range comments (1)

src/speculators/models/dflash/core.py (1)

158-163: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Default on None, not on any falsy value.

kwargs.get("max_anchors") or 3072 makes an explicit 0 indistinguishable from an omitted value. That means --max-anchors 0 silently becomes 3072 for DFlash, while the same value is rejected by the new P-EAGLE config. Only fall back when the argument is actually None.

💡 Suggested fix

-            max_anchors=kwargs.get("max_anchors") or 3072,
+            max_anchors=(
+                3072
+                if kwargs.get("max_anchors") is None
+                else kwargs["max_anchors"]
+            ),

🤖 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/dflash/core.py` around lines 158 - 163, The
max_anchors defaulting in DFlashSpeculatorConfig is using a falsy check, so an
explicit 0 gets replaced with 3072 instead of being preserved or handled
consistently. Update the logic in the DFlash config निर्माण path to default only
when kwargs.get("max_anchors") is None, matching the behavior of the related
config handling and keeping 0 distinct from a missing value.

🧹 Nitpick comments (2)

tests/unit/models/test_peagle_data.py (2)

26-38: 🧹 Nitpick | 🔵 Trivial | 💤 Low value

Consider extending chain cap verification to depth 2+.

The test only asserts depth1_count <= max_anchors, but the docstring says "depth-1+ chains." Depth 2+ counts derive from depth 1's already-capped pool, so they're transitively bounded, but explicit assertions would strengthen the contract. Not a blocker.

The anchor_pos unused variable is acceptable test clarity; ignore RUF059.

🤖 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/unit/models/test_peagle_data.py` around lines 26 - 38, The max_anchors
test only checks depth 1, but the contract covers depth-1+ chains; update
test_max_anchors_caps_chains in test_peagle_data.py to explicitly verify that
depth 2 and deeper counts are also bounded by the capped depth-1 pool, using
generate_cod_sample_indices and the depth tensor. Keep anchor_pos unused as-is
and add assertions for the deeper depths so the chain-cap behavior is tested
end-to-end.

---

49-68: 📐 Maintainability & Code Quality | 🔵 Trivial

Remove unnecessary torch.manual_seed calls.

With max_anchors=None, generate_cod_sample_indices does not invoke any random sampling (no torch.randperm path is taken), so the manual seeding has no effect on the output. The equality assertion holds deterministically without it. The seed calls may mislead readers into thinking the default path is stochastic.

     def test_max_anchors_none_uses_all(self):
         """max_anchors=None should use all valid positions (default behavior)."""
         seq_len = 32
-        torch.manual_seed(42)
         anchor_pos_none, depth_none = generate_cod_sample_indices(
             seq_length=seq_len,
             loss_mask=_loss_mask(seq_len),
             num_depths=4,
             max_anchors=None,
         )
-        torch.manual_seed(42)
         anchor_pos_default, depth_default = generate_cod_sample_indices(
             seq_length=seq_len,
             loss_mask=_loss_mask(seq_len),
             num_depths=4,
         )
         assert torch.equal(anchor_pos_none, anchor_pos_default)
         assert torch.equal(depth_none, depth_default)

🤖 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/unit/models/test_peagle_data.py` around lines 49 - 68, Remove the
unnecessary torch.manual_seed calls from test_max_anchors_none_uses_all in
test_peagle_data.py, since generate_cod_sample_indices with max_anchors=None
follows a deterministic path and does not use random sampling. Keep the
assertions comparing anchor_pos_none/depth_none to
anchor_pos_default/depth_default, and leave the test focused on verifying the
default behavior of generate_cod_sample_indices without implying stochasticity.

🤖 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/peagle/data.py`:
- Around line 40-43: The capped sampling in the anchor-selection logic is
mutating all_valid_indices so it no longer represents the full loss-mask-valid
set, which later breaks the torch.isin(next_candidates, all_valid_indices) chain
check. Update the data flow in the relevant function in data.py so the full
valid-position collection stays separate from the capped seed subset, and use
distinct symbols for the complete valid indices versus the max_anchors-limited
anchors when building and extending chains.

In `@tests/unit/models/test_peagle_data.py`:
- Around line 69-81: The test in test_max_anchors_fewer_valid_than_cap is too
weak because anchor_pos.shape[0] > 0 does not prove the below-cap behavior in
generate_cod_sample_indices. Strengthen the assertion by checking that the full
set of valid positions from loss_mask is represented in the sampled indices (or
that the unique sampled valid indices across depths matches the valid count,
accounting for depth shifts), so the test verifies max_anchors is not truncating
when valid positions are fewer than the cap.

---

Outside diff comments:
In `@src/speculators/models/dflash/core.py`:
- Around line 158-163: The max_anchors defaulting in DFlashSpeculatorConfig is
using a falsy check, so an explicit 0 gets replaced with 3072 instead of being
preserved or handled consistently. Update the logic in the DFlash config निर्माण
path to default only when kwargs.get("max_anchors") is None, matching the
behavior of the related config handling and keeping 0 distinct from a missing
value.

---

Nitpick comments:
In `@tests/unit/models/test_peagle_data.py`:
- Around line 26-38: The max_anchors test only checks depth 1, but the contract
covers depth-1+ chains; update test_max_anchors_caps_chains in
test_peagle_data.py to explicitly verify that depth 2 and deeper counts are also
bounded by the capped depth-1 pool, using generate_cod_sample_indices and the
depth tensor. Keep anchor_pos unused as-is and add assertions for the deeper
depths so the chain-cap behavior is tested end-to-end.
- Around line 49-68: Remove the unnecessary torch.manual_seed calls from
test_max_anchors_none_uses_all in test_peagle_data.py, since
generate_cod_sample_indices with max_anchors=None follows a deterministic path
and does not use random sampling. Keep the assertions comparing
anchor_pos_none/depth_none to anchor_pos_default/depth_default, and leave the
test focused on verifying the default behavior of generate_cod_sample_indices
without implying stochasticity.

🪄 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: 20db3352-98f6-47bb-98dd-5f9e6428e227

📥 Commits

Reviewing files that changed from the base of the PR and between c5e7b8a and c612fb0.

📒 Files selected for processing (6)

• scripts/train.py
• src/speculators/models/dflash/core.py
• src/speculators/models/peagle/config.py
• src/speculators/models/peagle/core.py
• src/speculators/models/peagle/data.py
• tests/unit/models/test_peagle_data.py

[mergify]

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

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

[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]

Looks good, but please update the max_anchor piping to not use the model's config.

[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

[shanjiaz]

Thanks for the detailed ablation. Glad this works out!

[fynnsu]

Looks good! One optional nit