Add optional Liger kernels for DFlash training by jessiewei7 · Pull Request #824 · vllm-project/speculators

jessiewei7 · 2026-07-20T09:41:43Z

Summary

Adds --use-liger-kernel for DFlash training to opt into Liger's fused
Qwen3 RMSNorm and SwiGLU kernels.

The kernels are resolved lazily and injected through per-model DFlash
factories. This keeps the default path native, avoids global Transformers
class changes, and does not depend on import order. The flag is limited to
--speculator-type dflash and is provided through speculators[liger].

Validation

CLI scope and missing-extra handling
From-scratch, pretrained, and config-only initialization
Native/Liger instance isolation and checkpoint compatibility
bf16 DFlash forward/backward numerical smoke

coderabbitai · 2026-07-20T09:42:00Z

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: 912775b4-5cc1-49c8-b5f8-da5176eaca84

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

Changes

The DFlash model now supports optional Liger RMSNorm and MLP kernels through a new dependency extra, kernel factories, CLI selection, model-construction propagation, and unit-test coverage.

DFlash Liger kernel integration

Layer / File(s)	Summary
Kernel factories and DFlash model injection `src/speculators/models/dflash/kernels.py`, `src/speculators/models/dflash/model_definitions.py`, `src/speculators/models/dflash/core.py`	DFlash accepts native or Liger kernel factories and uses them to construct attention norms, decoder norms, MLPs, and model-level norms.
CLI and training-path wiring `pyproject.toml`, `scripts/train.py`	The `liger` optional dependency and `--use-liger-kernel` flag are added; selected kernels are forwarded through all DFlash initialization paths and rejected for non-DFlash models.
Kernel and training integration tests `tests/unit/models/test_dflash_liger.py`, `tests/unit/train/test_cli_args.py`, `tests/unit/train/test_draft_config_init.py`	Tests cover optional dependency handling, instance-local Liger modules, checkpoint compatibility, CLI validation, and propagation across initialization routes.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.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
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.
Title check	✅ Passed	The title clearly summarizes the main change: optional Liger kernel support for DFlash training.
Description check	✅ Passed	The description matches the changeset and accurately describes the new flag, lazy kernel loading, and validation coverage.

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

coderabbitai

Actionable comments posted: 1

🤖 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/unit/models/test_dflash_liger.py`:
- Around line 57-70: Rename the shadowing globals and locals parameters in the
missing_liger __import__ mock within
test_missing_liger_extra_has_actionable_error to non-builtin names, and update
the positional call accordingly without changing the mock’s behavior.

🪄 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: 7bd68d0e-29be-4de4-b284-640a97740a38

📥 Commits

Reviewing files that changed from the base of the PR and between 065ecd2 and fe4fb87.

📒 Files selected for processing (8)

pyproject.toml
scripts/train.py
src/speculators/models/dflash/core.py
src/speculators/models/dflash/kernels.py
src/speculators/models/dflash/model_definitions.py
tests/unit/models/test_dflash_liger.py
tests/unit/train/test_cli_args.py
tests/unit/train/test_draft_config_init.py

speculatorsbot

Clean integration -- the factory-based kernel injection via DFlashKernels avoids the global monkey-patching that Liger's apply_liger_kernel_to_* helpers typically use, preserving instance isolation and checkpoint compatibility. All three initialization paths (from-scratch, config-only, pretrained) are wired and tested. Lazy import keeps liger-kernel truly optional.

One process note: per CONTRIBUTING.md, changes adding CLI/API surface and touching 3+ files should reference an assigned issue.

LGTM with one non-blocking item below. Recommend approving.

🤖 Generated with Claude Code using the /pr-review skill

speculatorsbot · 2026-07-20T12:13:57Z

+    parser.add_argument(
+        "--use-liger-kernel",
+        action="store_true",
+        help=(
+            "Use Liger Qwen3 RMSNorm/SwiGLU kernels for DFlash. Requires "
+            "the optional `speculators[liger]` extra."
+        ),
+    )


Non-blocking: docs/cli/train.md comprehensively documents all training flags. --use-liger-kernel should be added there (Model Arguments section, near --speculator-type), along with a note about the speculators[liger] optional dependency.

Documented under DFlash-specific arguments, including the optional dependency.

WindChimeRan

Please benchmark the performance gain on this.

Add --use-liger-kernel for DFlash training. Resolve Liger Qwen3 RMSNorm/SwiGLU modules lazily and inject them through per-model factories, leaving the native path unchanged by default. Cover CLI scope, initialization paths, instance isolation, and checkpoint compatibility. Signed-off-by: jessiewei7 <jessiewei747@gmail.com>

mergify · 2026-07-24T03:34:07Z

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

approved-reviews-by = dsikka
approved-reviews-by = fynnsu
approved-reviews-by = orestis-z
approved-reviews-by = rahul-tuli
approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

mergify · 2026-07-24T14:03:19Z

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

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

coderabbitai Bot reviewed Jul 20, 2026

View reviewed changes

Comment thread tests/unit/models/test_dflash_liger.py

This comment was marked as duplicate.

Sign in to view

speculatorsbot reviewed Jul 20, 2026

View reviewed changes

WindChimeRan reviewed Jul 20, 2026

View reviewed changes

jessiewei7 force-pushed the feat/dflash-liger-kernels branch from fe4fb87 to 1193a43 Compare July 21, 2026 02:25

mergify Bot added the documentation Improvements or additions to documentation label Jul 21, 2026

jessiewei7 force-pushed the feat/dflash-liger-kernels branch from 1193a43 to b3529db Compare July 21, 2026 02:43

mergify Bot added the needs-rebase label Jul 24, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add optional Liger kernels for DFlash training#824

Add optional Liger kernels for DFlash training#824
jessiewei7 wants to merge 1 commit into
vllm-project:mainfrom
jessiewei7:feat/dflash-liger-kernels

jessiewei7 commented Jul 20, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading

Review skipped

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

This comment was marked as duplicate.

Uh oh!

speculatorsbot left a comment

Uh oh!

speculatorsbot Jul 20, 2026

Uh oh!

jessiewei7 Jul 21, 2026

Uh oh!

WindChimeRan left a comment

Uh oh!

mergify Bot commented Jul 24, 2026

Uh oh!

mergify Bot commented Jul 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

jessiewei7 commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Uh oh!

coderabbitai Bot commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment was marked as duplicate.

Uh oh!

speculatorsbot left a comment

Choose a reason for hiding this comment

Uh oh!

speculatorsbot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

jessiewei7 Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

WindChimeRan left a comment

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Jul 24, 2026

Merge Protections

🔴 Require approval from approved reviewers list

Uh oh!

mergify Bot commented Jul 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jessiewei7 commented Jul 20, 2026 •

edited

Loading

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading