Add forward/backward integration tests

[fynnsu]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

More integration tests calling the model's forward and backwards and testing edge cases.

Tests

Added by this pr.

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: 65740481-0418-4cf6-952f-741f2d354a9b

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

This PR adds comprehensive CUDA-gated integration tests for three draft model architectures (Eagle3, DFlash, PEagle). It establishes shared test infrastructure—tiny model configs, factories, and synthetic data generators—then provides training-mode validation, multi-batch statefulness tests, and vocab boundary coverage for each model independently.

Changes

Draft Model Integration Tests

Layer / File(s)	Summary
Shared test fixtures and model factories tests/integration/conftest.py	Tiny Llama and Qwen3 configs, NaN-weight sanitization, model factories for Eagle3/DFlash/PEagle with deterministic initialization, synthetic sample generation (make_sample for per-sequence tensors, make_batch for collation), and pytest fixtures with CUDA cleanup.
DFlash forward-pass integration tests tests/integration/models/test_dflash_forward.py	Training-mode forward/backward tests with block_size and max_anchors sweeps, multi-batch stress tests across varying batch compositions, and vocab boundary tests with draft-vocab mapping validation.
Eagle3 forward-pass integration tests tests/integration/models/test_eagle3_forward.py	Sample generation with optional verifier states, training tests across sequence lengths and ttt_steps, multi-batch tests with torch.compiler resets, and vocab boundary validation ensuring draft token IDs stay within mapped ranges.
PEagle forward-pass integration tests tests/integration/models/test_peagle_forward.py	Training tests validating forward/backward with num_depths and down_sample_ratio sweeps, multi-batch tests exercising cache behavior, and vocab boundary tests using CUDA-resident draft-vocab mappings.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

enhancement, two-reviews

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% 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 'Add forward/backward integration tests' accurately captures the main change: adding new integration test modules for forward/backward passes across three draft models.
Description check	✅ Passed	The description clearly relates to the changeset by explaining the purpose (more integration tests for forward/backward paths and edge cases) and confirming tests were added.
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.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch add_integration_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.}

[coderabbitai]

Actionable comments posted: 2

🤖 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 `@tests/integration/models/test_dflash_forward.py`:
- Around line 159-165: The test_boundary_tokens in
tests/integration/models/test_dflash_forward.py currently never asserts that
draft_tokens respect the restricted vocab range; update the test to verify that
the generated draft_tokens from dflash_draft_vocab_model are all within the
provided boundary_token_ids range (0..31) after calling
dflash_draft_vocab_model(**batch). Locate the test_boundary_tokens function and
add an assertion that every value in draft_tokens (or
draft_tokens.detach().cpu().numpy() if tensors) is >= 0 and <= 31 (or uses the
boundary_token_ids variable) so regressions in the vocab mapping will fail the
test.

In `@tests/integration/models/test_peagle_forward.py`:
- Around line 158-164: The test_boundary_tokens should also assert that produced
draft_tokens respect the restricted vocab bounds; update the test (function
test_boundary_tokens) to check draft_tokens values are within the inclusive
range given by boundary_token_ids (e.g., min >= 0 and max <= 31) after calling
peagle_draft_vocab_model(**batch), using the draft_tokens tensor/array returned
to compute min/max and fail the test if any token falls outside [0,31]; keep
existing loss/gradient checks intact.

🪄 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: 6424fe24-2ed7-4cc1-becc-e8b677da30f5

📥 Commits

Reviewing files that changed from the base of the PR and between 1bc3788 and 695a2a6.

📒 Files selected for processing (4)

• tests/integration/conftest.py
• tests/integration/models/test_dflash_forward.py
• tests/integration/models/test_eagle3_forward.py
• tests/integration/models/test_peagle_forward.py

[shanjiaz]

Actually could this be merged to one file, seem to overlap a lot. Thanks!

[shanjiaz]

Looks good! looks like the ttt test just failed tho.