Add hc_mult support to DFlash for DeepSeek-V4-Flash

[rahul-tuli]

Purpose

Enable training DFlash speculators for DeepSeek-V4-Flash (DSv4), which uses Manifold-Constrained Hyper-Connection (mHC) with hc_mult=4. DSv4's hidden states are (N, hc_mult * hidden_size) per layer rather than the standard (N, hidden_size), and its checkpoint uses non-standard weight names (embed.weight, head.weight, norm.weight). Without these changes the DFlash pipeline cannot initialize or train a speculator for DSv4.

Description

• scripts/train.py: Read hc_mult from the verifier config (default 1) and thread it into the draft model's transformer_layer_config. Pass hc_mult * hidden_size as the effective hidden size to both train and val dataloaders so create_collate_fn / create_empty_sample create correctly shaped tensors.
• src/speculators/models/dflash/core.py:
  • __init__: Read self.hc_mult from config. Expand FC layer input dim to len(target_layer_ids) * hc_mult * hidden_size. Register hc_head_fn, hc_head_base, hc_head_scale buffers when hc_mult > 1.
  • forward: Apply hc_head_project() to collapse verifier_last_hidden_states from (N, hc_mult * hidden_size) → (N, hidden_size) before verifier_norm in the loss path (only when hc_mult > 1).
  • load_verifier_weights: New override that handles DSv4's non-standard weight names and loads hc_head parameters. Delegates to super() when hc_mult == 1.
• src/speculators/models/dflash/utils.py: hc_head_project() added in a prior commit (pure-PyTorch port of vLLM's _hc_head_fused_reference).

All changes are backward-compatible: when hc_mult=1 (all non-DSv4 models), every dimension calculation is algebraically identical to the prior code.

Related Issue

Part of the DFlash DSv4 code changes effort (Diff 1: Hidden States from the companion PRD).

Tests

• All 264 existing unit tests pass (python -m pytest tests/unit/ -x -q).
• Verified end-to-end with a smoke test script that initializes a DFlash speculator against the real DSv4 checkpoint, loads weights, and runs a dummy forward pass on CUDA — init, weight loading, and forward all pass.

--- __init__ checks ---
  hc_mult = 4
  FC: Linear(49152, 4096)
  hc_head_fn:    (4, 16384)
  hc_head_base:  (4,)
  hc_head_scale: (1,)
  PASS

--- load_verifier_weights checks ---
  embed_tokens:     (129280, 4096), loaded
  lm_head:          (32000, 4096), loaded
  verifier_lm_head: (32000, 4096), loaded
  verifier_norm:    (4096,), loaded
  hc_head_fn:       (4, 16384), loaded
  hc_head_base:     (4,), loaded
  hc_head_scale:    (1,), loaded
  PASS

--- forward pass checks ---
  hidden_states:                (1, 64, 49152)
  input_ids:                    (1, 64)
  verifier_last_hidden_states:  (1, 64, 16384)
  draft_tokens: (1, 32)
  loss:         29.4409
  metrics keys: ['full_acc', 'loss', 'position 1 acc', ...]
  PASS

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

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: 2564a22d-75c0-441c-9288-f6c14048c1cb

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
• Commit unit tests in branch rtuli/dflash-dsv4-hc-mult

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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

[mergify]

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

[fynnsu]

The concern I have with this approach is that the transformation logic will have to be mirrored in vllm as well. Maybe we should be doing the transformation in vllm before extracting the hidden states?

Also ideally we could generalize some of this. There is already a base implementation of load_verifier_weights in src/speculators/model.py, if modify that implementation then eagle3, p-eagle, and dflash will all be able to use the code. We've also previously discussed a weight mapping layer to handle non-standard weight names (e.g. a user could pass in a dict mapping custom names to the standard ones). Something like that would let us handle the non-standard naming part better.