Skip to content

Add max_anchors random chain sampling to P-EAGLE - #682

Closed
orestis-z wants to merge 2 commits into
vllm-project:mainfrom
orestis-z:peagle-max-anchors
Closed

Add max_anchors random chain sampling to P-EAGLE#682
orestis-z wants to merge 2 commits into
vllm-project:mainfrom
orestis-z:peagle-max-anchors

Conversation

@orestis-z

@orestis-z orestis-z commented Jun 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stacked on #683 — review/merge that first. After #683 merges, this PR's diff will show only the max_anchors changes.

Summary

  • Adds max_anchors parameter to P-EAGLE's COD sampling to randomly subsample chain starting points at depth 1+
  • Depth 0 always retains the full sequence (no contiguity issues)
  • Combined with streaming attention (Add StreamingLLM attention to P-EAGLE #683), this reduces both attention cost and token count for long sequences
  • DFlash fallback default preserved at 3072

Changes

  • data.py: random subsampling of all_valid_indices when max_anchors is set, with .sort()[0] to maintain causal order
  • config.py: new max_anchors: int | None field
  • core.py: wire max_anchors through forward and from_training_args
  • train.py: --max-anchors default changed from 256 to None
  • dflash/core.py: handle None default with or 3072 fallback

Test plan

🤖 Generated with Claude Code

@orestis-z orestis-z self-assigned this Jun 29, 2026
@coderabbitai

coderabbitai Bot commented Jun 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

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: d465bca1-f89a-4f34-98d8-08c2efe0f687

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

Comment @coderabbitai help to get the list of available commands.

orestis-z and others added 2 commits June 29, 2026 12:55
@orestis-z @claude
Add StreamingLLM (sliding window + sink) attention to P-EAGLE
68bf3ea 
Adds streaming attention support to P-EAGLE's FlexAttention mask,
enabling O(seq_len × W) depth-0 attention cost instead of O(seq_len²).
Two new config fields (sink_size, max_context_window) control the
attention pattern per document in packed sequences.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>
@orestis-z @claude
Add max_anchors random chain sampling to P-EAGLE
cc7940d 
Adds a `max_anchors` parameter that randomly subsamples COD chain
starting points at depth 1+ while keeping depth-0 as the full sequence.
Combined with StreamingLLM attention (sliding window + sink), this
enables long-context P-EAGLE training without the context blindness
that affected the contiguous-window approach.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>
@orestis-z
orestis-z force-pushed the peagle-max-anchors branch from 551d857 to cc7940d Compare June 29, 2026 12:57
@mergify

mergify Bot commented Jun 29, 2026

Copy link
Copy Markdown

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 mergify Bot added the needs-rebase label Jun 29, 2026
@orestis-z

orestis-z commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator Author

Moved to orestis-z#2 — stacked on #683, will retarget to upstream after streaming PR merges.

@orestis-z orestis-z closed this Jun 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@fynnsu fynnsu Awaiting requested review from fynnsu fynnsu will be requested when the pull request is marked ready for review fynnsu is a code owner
@shanjiaz shanjiaz Awaiting requested review from shanjiaz shanjiaz will be requested when the pull request is marked ready for review shanjiaz is a code owner
@rahul-tuli rahul-tuli Awaiting requested review from rahul-tuli rahul-tuli will be requested when the pull request is marked ready for review rahul-tuli is a code owner

Assignees

@orestis-z orestis-z

Labels

needs-rebase

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@orestis-z