feat(dflash): remap fused QKV weights for Laguna warm-start by orestis-z · Pull Request #922 · vllm-project/speculators

orestis-z · 2026-08-04T12:12:23Z

Summary

Adds _remap_weights() to DFlashConverter that splits Laguna-style fused qkv_proj into separate q_proj/k_proj/v_proj, drops incompatible keys (g_proj, aux_hidden_norms), and slices fc.weight when source has more target layers than needed.
Enables warm-starting DSpark training from published Laguna DFlash checkpoints (e.g. poolside/Laguna-S-2.1-DFlash) instead of training from scratch (~13.5 day ETA → significantly faster convergence).
No-op when checkpoint already uses separate Q/K/V projections (Qwen3-style).

Test plan

5 new unit tests covering QKV splitting, key dropping, fc slicing, passthrough, and key preservation
All 57 converter tests pass
End-to-end conversion of poolside/Laguna-S-2.1-DFlash succeeds (6 fused qkv_proj split, 12 keys dropped, fc sliced 18432→15360)

🤖 Generated with Claude Code

…rm-start Laguna-style DFlash checkpoints use fused qkv_proj + per-head g_proj gating, while DFlashDraftModel expects separate q/k/v_proj (Qwen3-style). Add a _remap_weights step that splits fused QKV, drops incompatible keys (g_proj, aux_hidden_norms), and slices fc.weight when the source has more target layers than the model expects. Enables warm-starting DSpark training from published Laguna DFlash checkpoints. Signed-off-by: Claude <noreply@anthropic.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis-z@users.noreply.github.com>

coderabbitai · 2026-08-04T12:12:44Z

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 Plus

Run ID: 23cad7a7-1853-473e-8b4f-48af9c15c2cd

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

DFlash weight conversion

Layer / File(s)	Summary
Implement checkpoint weight remapping `src/speculators/convert/dflash/converter.py`	`_remap_weights` splits fused QKV tensors, removes unsupported weights, slices oversized `fc.weight` tensors, and runs before draft-model loading in `_save`.
Validate remapped weight cases `tests/unit/convert/test_dflash_converter.py`	Tests cover QKV splitting, unsupported-weight removal, FC slicing, passthrough behavior, and preservation of unrelated parameters.

Possibly related PRs

vllm-project/speculators#617: This PR extends the DFlashConverter implementation from that PR with weight remapping and related tests.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly describes the main change: adding weight remapping functionality for fused QKV weights from Laguna-style DFlash checkpoints.
Description check	✅ Passed	The description is comprehensive and directly related to the changeset, covering the purpose, implementation details, test coverage, and validation results.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch feat/dflash-converter-fused-qkv-remap

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

mergify · 2026-08-04T12:13:01Z

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 Require approval from approved reviewers list

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

mergify · 2026-08-04T12:14:07Z

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

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 `@src/speculators/convert/dflash/converter.py`:
- Around line 169-173: Reject mixed fused and separate projection formats in the
remapping logic around the qkv_proj handling: before inserting any q_proj,
k_proj, or v_proj destination, check whether that destination already exists in
remapped and raise ValueError on a collision. Add a mixed-format checkpoint test
in tests/unit/convert/test_dflash_converter.py covering the q_proj/k_proj/v_proj
destinations and asserting ValueError.

🪄 Autofix

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 Plus

Run ID: e6f9da45-7fbd-498b-8579-3129bc204862

📥 Commits

Reviewing files that changed from the base of the PR and between c667ee1 and b1ce58c.

📒 Files selected for processing (2)

src/speculators/convert/dflash/converter.py
tests/unit/convert/test_dflash_converter.py

shanjiaz

Looks good pending quality!

Laguna models use nested per-layer-type rope_parameters (with `full_attention` and `sliding_attention` sub-dicts), which Qwen3RotaryEmbedding does not understand — it expects a flat dict with `rope_theta` at the top level. Since all DFlash draft layers use sliding-window attention, flatten to the `sliding_attention` variant before constructing the rotary embedding. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis-z@users.noreply.github.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis-z@users.noreply.github.com>

…n_state_layer_ids When converting a z-lab DFlash checkpoint without explicit aux_hidden_state_layer_ids, the converter auto-derives them from dflash_config.target_layer_ids with a +1 offset. This included the last verifier layer (e.g. layer 48 for a 48-layer model), but speculators treats that layer separately as verifier_last_hidden_states — it should not appear in aux_hidden_state_layer_ids. Including it caused the fc layer to be sized for one extra feature, leading to dimension mismatches at inference when vLLM builds the DFlash/DSpark model from the converted checkpoint. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis-z@users.noreply.github.com>

speculatorsbot

The weight remapping logic in _remap_weights is clean and well-tested — the QKV split, key dropping, and fc slicing all make sense for the Laguna warm-start use case. The _build_config change to exclude the last verifier layer from aux_hidden_state_layer_ids is correct and pairs well with the fc slicing.

One concern about the rope_parameters fix in core.py — see inline comment.

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

…warning Reject checkpoints that contain both fused qkv_proj and separate q/k/v_proj keys, and warn when flattening nested rope_parameters to sliding_attention while full-attention layers are present. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Speculators Bot <23146389+orestis-z@users.noreply.github.com>

orestis-z requested review from fynnsu, rahul-tuli and shanjiaz as code owners August 4, 2026 12:12

Merge branch 'main' into feat/dflash-converter-fused-qkv-remap

b1ce58c

mergify Bot added the quality-failed label Aug 4, 2026

coderabbitai Bot reviewed Aug 4, 2026

View reviewed changes

Comment thread src/speculators/convert/dflash/converter.py Outdated

shanjiaz reviewed Aug 5, 2026

View reviewed changes

shotsan mentioned this pull request Aug 5, 2026

fix(dflash): resolve quality failures and PR feedback for Laguna converter #941

Closed

orestis-z and others added 3 commits August 6, 2026 13:22

style: move deepcopy import to module level

706681e

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis-z@users.noreply.github.com>

santosh-agebold mentioned this pull request Aug 6, 2026

fix(dflash): resolve quality failures and PR feedback for Laguna converter (clean) #945

Closed

speculatorsbot reviewed Aug 7, 2026

View reviewed changes

Comment thread src/speculators/models/dflash/core.py Outdated

mergify Bot removed the quality-failed label Aug 12, 2026

Merge branch 'main' into feat/dflash-converter-fused-qkv-remap

071a288

orestis-z self-assigned this Aug 12, 2026

orestis-z added the ready This PR is ready for review label Aug 12, 2026

orestis-z enabled auto-merge (squash) August 12, 2026 13:09

shanjiaz approved these changes Aug 13, 2026

View reviewed changes

Merge branch 'main' into feat/dflash-converter-fused-qkv-remap

d780fcf

orestis-z merged commit 2682131 into main Aug 13, 2026
9 checks passed

orestis-z deleted the feat/dflash-converter-fused-qkv-remap branch August 13, 2026 14:11

Navigation Menu

Uh oh!

feat(dflash): remap fused QKV weights for Laguna warm-start - #922

feat(dflash): remap fused QKV weights for Laguna warm-start#922
orestis-z merged 8 commits into
mainfrom
feat/dflash-converter-fused-qkv-remap

orestis-z commented Aug 4, 2026

Uh oh!

coderabbitai Bot commented Aug 4, 2026 •

edited

Loading

Review skipped

Walkthrough

Changes

Uh oh!

mergify Bot commented Aug 4, 2026 •

edited

Loading

🟢 Require approval from approved reviewers list

Uh oh!

mergify Bot commented Aug 4, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

shanjiaz left a comment

Uh oh!

speculatorsbot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

orestis-z commented Aug 4, 2026

Summary

Test plan

Uh oh!

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

Uh oh!

Review skipped

Walkthrough

Changes

Uh oh!

mergify Bot commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Protections

🟢 Require approval from approved reviewers list

Uh oh!

mergify Bot commented Aug 4, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

shanjiaz left a comment

Choose a reason for hiding this comment

Uh oh!

speculatorsbot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coderabbitai Bot commented Aug 4, 2026 •

edited

Loading

mergify Bot commented Aug 4, 2026 •

edited

Loading