Add max_anchors and max_context_window support to P-EAGLE

[orestis-z]

Closes #567

Purpose

Add max_anchors and max_context_window support to P-EAGLE to cap memory usage during training on long sequences.

P-EAGLE's COD sampling predicts next tokens for every starting position, producing attention masks of size (seq_len × 2.73)². When max_anchors is set, generate_cod_sample_indices selects a random contiguous window containing at most max_anchors valid (loss_mask=1) positions, shrinking the effective attention footprint while preserving context for the selected anchors. max_context_window (default 4096) hard-caps the window size to prevent sparse loss masks from reinflating it.

Changes:

• Add max_anchors and max_context_window fields to PEagleSpeculatorConfig
• Wire the new params through PEagleDraftModel and generate_cod_sample_indices
• Add --max-context-window CLI arg; make --max-anchors default to None (was 256), defaulting per-model (256 for DFlash, unlimited for P-EAGLE)
• Fix DFlash max_anchors default from 3072 to 256
• Fix DFlash max_anchors falsy-value handling: use explicit is None check instead of or coercion
• Sparse sequence edge case: Ensured the max_context_window clamp fires even when len(valid_indices) < max_anchors. Previously, sparse sequences would fall back to the full sequence length, inadvertently causing the exact OOMs the window was designed to prevent.

Tests

• New integration test test_varying_max_anchors parametrized over [4, 16, None]
• Integration test test_max_context_window with small window size (32) to exercise the cap branch
• Unit tests (test_peagle_data.py) directly testing generate_cod_sample_indices for the sparse bypass bug, normal cap path, and no-windowing case
• Verifies loss is finite and backward pass succeeds for each setting

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]

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: 20a9f14d-80f6-449f-bcc0-3df1725ea3b9

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

Two new training-time limits—max_anchors and max_context_window—are added to P-EAGLE. The CLI gains --max-context-window and updates --max-anchors to default None. PEagleSpeculatorConfig receives both fields, PEagleDraftModel stores and forwards them, and generate_cod_sample_indices implements contiguous windowed subsampling. DFlashDraftModel adopts an or 3072 fallback to handle the now-None CLI default.

Changes

P-EAGLE anchor/context window limits

Layer / File(s)	Summary
PEagleSpeculatorConfig fields and COD windowed-sampling algorithm src/speculators/models/peagle/config.py, src/speculators/models/peagle/data.py	PEagleSpeculatorConfig gains max_anchors: int | None and max_context_window: int Pydantic fields (both ge=1). generate_cod_sample_indices adds matching parameters and implements contiguous-window selection: when valid positions exceed max_anchors, a random contiguous segment is chosen from valid indices, then optionally truncated to max_context_window, and the result replaces the full-sequence candidate pool.
PEagleDraftModel init, forward, and factory wiring src/speculators/models/peagle/core.py	__init__ stores config.max_anchors and config.max_context_window as instance attributes; forward passes them into generate_cod_sample_indices; from_training_args includes both in the constructed PEagleSpeculatorConfig, defaulting max_context_window to 4096.
CLI arguments and DFlash falsy-value fix scripts/train.py, src/speculators/models/dflash/core.py	--max-anchors is changed to default=None with updated help text describing per-model defaults; --max-context-window is added with default=4096. DFlashDraftModel.from_training_args changes kwargs.get("max_anchors", 3072) to kwargs.get("max_anchors") or 3072 so a None from the CLI is replaced with 3072.
Integration test fixture and parameterization tests/integration/conftest.py, tests/integration/models/test_model_forward.py	make_peagle_model accepts and forwards an optional max_anchors argument into PEagleSpeculatorConfig. test_varying_max_anchors updates its parameter set from [2, 8, 16] to [4, 16, None] to cover the unlimited (None) path.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 55.56% 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 and specifically summarizes the main changes: adding max_anchors and max_context_window support to P-EAGLE, which aligns with the primary objectives across all modified files.
Linked Issues check	✅ Passed	The PR fully implements the requirements from issue #567: adds max_anchors parameter to P-EAGLE config [#567], wires it through generate_cod_sample_indices [#567], and achieves the asymmetric attention mask optimization [#567].
Out of Scope Changes check	✅ Passed	All changes are directly in-scope: CLI args match the new parameters, test updates validate the feature, and DFlash default fix aligns with making max_anchors default to None with per-model defaults.
Description check	✅ Passed	The pull request description clearly explains the purpose (memory optimization for P-EAGLE), outlines specific changes across multiple files, and provides detailed context about the feature including the sparse sequence edge case fix.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 and usage tips.}

[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

[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

[coderabbitai]

Actionable comments posted: 3

🤖 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/core.py`:
- Line 159: The line using `kwargs.get("max_anchors") or 3072` silently converts
any falsy value (including 0) to 3072, which masks explicit user input and can
cause unexpected behavior. Replace this logic with an explicit None check
instead of relying on truthiness, so that only when max_anchors is not provided
(None) should it default to 3072, allowing valid falsy values like 0 to be
respected.

In `@src/speculators/models/peagle/data.py`:
- Around line 50-70: The max_context_window constraint is only being applied in
the if block when all_valid_indices.shape[0] exceeds max_anchors. In the else
block (lines 68-69), the code falls back to using the full seq_length without
applying any windowing, which bypasses the max_context_window guard. Modify the
else block to also apply the max_context_window constraint by capping the
window_end to window_start plus max_context_window, ensuring that sample_indices
respects the max_context_window limit regardless of whether max_anchors was
exceeded.

In `@tests/integration/models/test_model_forward.py`:
- Around line 341-349: The test_varying_max_anchors method does not cover the
max_context_window branch because with MAX_LEN=128 and the default
max_context_window=4096, the condition window_end - window_start >
max_context_window is never triggered. Add an additional test case (either by
extending this test or creating a new parametrized case) where you pass a small
max_context_window value to make_peagle_model or use a sparse loss_mask in the
batch so that the window size exceeds the max_context_window limit and the COD
sampling branch is executed and validated.

🪄 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: 5bbfeff8-7feb-4bd0-b556-eb4091de0003

📥 Commits

Reviewing files that changed from the base of the PR and between 6f46d97 and fa168ed.

📒 Files selected for processing (7)

• 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/integration/conftest.py
• tests/integration/models/test_model_forward.py

[shanjiaz]

I'm a little bit worried that this bit causing more graph breaks but otherwise this diff looks great!

[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

[orestis-z]

ON HOLD, see #567 (comment)