[hootandy321]

Part of #16. This PR covers the model/loss prerequisite. I am continuing to iterate and test the remaining work locally and plan to submit follow-up PRs once it is more complete.

Main changes

Following the scope suggested in the issue, this PR adds a standalone Qwen3 P-EAGLE model and forward-KL loss module as components for parallel draft training.

The model computes multiple prediction depths in one forward pass, using shared projection and decoder parameters, a trainable token mask embedding, and a trainable auxiliary hidden mask. It preserves Qwen3 Q/K normalization and makes the alignment between Target features, input tokens, and per-depth distillation targets explicit.

The main changes are:

• Implement the parallel model, document-isolated attention, and forward-KL loss.
• Return per-depth KL sums, valid-position counts, and correct-prediction counts for later aggregation and evaluation.
• Add independent numerical references and regression tests under tests/cpu, where the upstream CI workflow discovers them automatically.
• Document usage, provenance, and support boundaries.

This is still a minimal implementation. It currently uses a deterministic dense anchor/depth layout to validate the model and loss semantics. The supported scope is one packed input row, a full vocabulary, full attention, zero attention dropout, and eager/SDPA execution.

If this direction and scope look reasonable, I would be happy to continue improving it and complete the remaining work through follow-up PRs.

Verification (agent-run tests, human-reviewed)

Checks were run in an isolated Linux venv using Python 3.13.3, Torch 2.13.0+cpu, Transformers 5.12.1, and the repository's CPU test dependencies:

• New tests: 64 passed, no skips. These cover an independent layer-by-layer mathematical reference, forward values and all parameter gradients, target indexing, mask updates, future-token/document/padding isolation, and empty supervision.
• The automatically discovered CPU suite (python -m pytest -q tests/cpu): 711 passed, no skips, including the 64 new tests.
• The dependency-light source suite (python -m pytest -q tests/source): 9 passed, no skips.
• Full-repository Ruff, Python source compilation, and local documentation-link checks passed.
• Python wheel/sdist build and audit passed. The packaged model source matches the tested files.

The branch has been rebased onto upstream fd4dfa2e7aeab30cc447670d927478cb44bec04b. The P-EAGLE tests now live in tests/cpu/test_peagle.py; the CI workflow itself matches upstream.

The first CPU run hit an upstream DMTP strict round-trip assertion with a 2.38e-7 difference. I reproduced the same assertion on unmodified upstream main with seed 0. An unchanged rerun of this PR passed all 711 CPU tests. The initial failure and reproduction logs have been retained.

Component-level comparisons were also rerun against the pinned upstream source:

• All 16 attention-mask cases matched exactly.
• With the same dense outer layout and the original upstream first-layer class substituted, 1/2/4-layer toy models had zero maximum difference in logits and parameter gradients.
• The maximum absolute KL loss difference was approximately 5.96e-8.

Due to time constraints, training with the real Qwen3-8B Target, CUDA numerics, and related runtime measurements have not yet been validated.

This pull request adds a Qwen3-only P-EAGLE model core and a forward-KL loss function. It adds no public trainer, CLI option, checkpoint format, or export path.

Changed files (6 files, 782 lines added, 0 removed):

• aspec/models/peagle.py (+267): new model, layout, mask, and loss code.
• tests/cpu/test_peagle.py (+381): CPU tests and independent scalar reference helpers.
• docs/user-guide/peagle.md (+126): usage, input shape rules, masking semantics, loss definition, and deferred work.
• docs/user-guide/algorithms.md (+6), docs/support.md (+1), docs/framework/REFERENCES.md (+1): support table, guide section, and source attribution.

New public API in aspec/models/peagle.py:

• peagle_dense_layout(document_ids, num_depths): builds anchor and depth positions.
• peagle_attention_mask(anchor_pos, depth, document_ids): builds the attention mask.
• PEagleOutput: dataclass with logits, layout, validity, and depth metadata.
• PEagleDraftModel(config, *, mask_token_id, num_aux_hidden_states=3, num_depths=4): Qwen3 draft model. forward takes input_ids, hidden_states, document_ids, and optional position_ids.
• peagle_kl_loss(output, teacher_logits, loss_mask): returns the loss and per-depth observations.

Behavior and constraints, confirmed in the diff:

• The model accepts only Qwen3 configurations. It requires eager or SDPA attention, full attention, zero attention dropout, and an untied full-vocabulary head. It rejects other inputs with ValueError.
• The model accepts only one packed row with shape [1, S]. It requires int64 token IDs and document IDs. It uses -1 for padding.
• mask_token_id must lie within the Target vocabulary.
• The model uses a learned mask hidden state sized for num_aux_hidden_states Target features.
• The loss requires a binary loss mask and matching teacher and mask shapes and devices. It divides the summed KL by the selected count plus 1e-5.

Verification:

• The PR reports 64 new tests and 654 CPU tests passing. These results come from the PR description. I did not run them.
• The PR reports exact agreement with pinned upstream references for 16 attention-mask cases. It also reports zero logit and gradient difference on 1-, 2-, and 4-layer toy models, and a maximum KL difference of about 5.96e-8. These are author-reported results.
• Not validated: training with the real Qwen3-8B Target, CUDA numerics, runtime measurements, and any public training, COD sampling, checkpoint, export, or decoding path.

Compatibility: The change adds new modules and documentation only. It changes no existing public declarations.

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-12T07:22:36.336814Z	97af8a0	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[zhouaihui]

Thanks @hootandy321, and welcome to ASpec! 🎉 Really happy to see our first community PR, and to hear you're already trying ASpec for your own experiments.

Thanks for putting so much care into the tests and docs, too. Starting with the Qwen3 model/loss core sounds like a good first step, and building this up through smaller PRs makes sense.

Let's work through the review here and keep discussing the next steps in #16. Feel free to share questions or anything you get stuck on. We're happy to help!

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

Adds a Qwen3 P-EAGLE draft model with dense packed-document layout, depth-aware attention masking, and a forward-KL loss. CPU tests check numerical behavior and validation. Documentation describes the API and identifies training, sampling, recovery, export, and consumer validation as pending.

Changes: P-EAGLE Core

Layer / File(s)	Summary
Dense layout and attention rules aspec/models/peagle.py, tests/cpu/test_peagle.py	Adds document segmentation, dense anchor/depth slots, and attention rules for causal base positions and depth chains. Tests compare layout and masks with scalar reference rules.
Qwen3 draft model aspec/models/peagle.py, tests/cpu/test_peagle.py, docs/user-guide/peagle.md	Adds Qwen3 configuration checks, first-layer fusion, model initialization, and the forward path. Tests cover reference outputs and gradients, masking, positions, state-dict loading, bfloat16 execution, and invalid inputs. The guide describes the model API and its input and attention semantics.
Forward-KL loss and API documentation aspec/models/peagle.py, tests/cpu/test_peagle.py, docs/framework/REFERENCES.md, docs/support.md, docs/user-guide/algorithms.md	Adds the forward-KL loss and per-depth metrics. Tests cover target alignment, supervision masks, empty supervision, and invalid loss inputs. The documentation records the API scope and source attribution.

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Possibly related issues

• inclusionAI/ASpec#16 — The PR adds the standalone Qwen3 P-EAGLE model and forward-KL loss described as prerequisite scope; it does not add the recipe, export, or resume outcomes.

Merge Risk: Moderate · up to d91cc

This PR adds a new P-EAGLE draft model and KL training loss that are not yet connected to the training pipeline. CPU numerical checks pass according to the author, but real-model and GPU behavior are unvalidated. The loss and tensor contracts should be confirmed on the target setup before downstream training depends on them.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title names the main change: adding a Qwen3 P-EAGLE parallel model and KL loss core. The raw summary and objectives both describe the same model and loss module, so the title matches the changeset…
Description check		The description follows the template in substance. It covers the outcome (standalone model and forward-KL loss, with the scope boundaries stated), verification (commands, test counts, pass and skip st…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Title check

Explanation

The title names the main change: adding a Qwen3 P-EAGLE parallel model and KL loss core. The raw summary and objectives both describe the same model and loss module, so the title matches the changeset. It is concise and specific.

Full details: Description check

Explanation

The description follows the template in substance. It covers the outcome (standalone model and forward-KL loss, with the scope boundaries stated), verification (commands, test counts, pass and skip status, and unvalidated GPU and Qwen3-8B training), and provenance (the upstream revision and rebase). The template's compatibility section asks for public config, checkpoint, and export implications. The description does not state these explicitly, although it says the ASpec training CLI does not yet select this algorithm and that export is deferred. The verification section also omits the 654 versus 711 test-count discrepancy with the objectives, which should be reconciled.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

Recent review info

Run configuration

• Source mode: Auto
• Repository: inclusionAI/ASpec
• Branch: main
• File: .coderabbit.yaml
• Commit: fd4dfa2
• Layers: Repository YAML
• Review profile: ASSERTIVE
• Plan: Advanced
• Run ID: f0692e13-0d70-4c46-8ae0-2f7f2994f563

Commits

Reviewing files that changed from the base of the PR and between fd4dfa2 and d91cc80.

Files selected for processing: 6

• aspec/models/peagle.py
• docs/framework/REFERENCES.md
• docs/support.md
• docs/user-guide/algorithms.md
• docs/user-guide/peagle.md
• tests/cpu/test_peagle.py

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.