[maocheng23]

Motivation

This PR adds opt-in residual architectures for DFlash-family drafts: Qwen3.8-Flash-Next gated residual (GR) and DeepSeek-V4.1 single-pass mHC. They are selected by config, not by duplicating draft models.

Modifications

• layers/residual/:
  • StandardResidual is the existing pre-norm add, parameter-free.
  • GatedResidual provides per-branch Gemma RMSNorm, a low-rank gated read, a gated write, and a learned final readout.
  • ManifoldHyperConnection implements the lagged pre-mix, FP32 coefficients with Sinkhorn, and a stream-zero first read.
• DecoderLayer and the DFlash2 decoder read and write through attention_residual and ffn_residual. The DFlash model expands streams at the input and collapses them (mHC) or reads them out (GR) at the head.
• draft_residual in the HF draft config selects the kind; omitting it keeps the standard layout.
• Non-standard residuals force non-reentrant gradient checkpointing.
• The DFlash, Domino and DSpark resume contracts record the residual semantics.

Difference from #1023: the EAGLE3 decoder no longer gets always-standard residual modules, since there's no config path for them there.

Reference check (review)

The math matches SGLang:

• GR matches layers/hyperconnection.py:GatedResidual with hc_per_branch_norm=True, as used by qwen4_exp.py.
• mHC matches models/deepseek_v4_mhc.py: the pre-mix computed at one hyper-connection is consumed by the next, the first read uses stream 0, and the head combines with the last pre-mix (hc_pre_from_prev_sublayer).

Open items for discussion

• Not servable yet. Parameter names and layouts differ from SGLang/DeepSeek:

  • GR uses a [count, hidden] norm_weight; SGLang uses a flat hc_norm.weight.
  • mHC uses mix_weight/mix_scale/mix_bias; DeepSeek uses hc_attn_fn/scale/base and the hc_ffn_* equivalents.

  Either align the names or add an exporter before training drafts with these.

• FSDP1 casts the mHC "FP32" parameters to bf16. MixedPrecision(param_dtype=bf16) applies to them, so the FP32 coefficient path runs on bf16-rounded weights. DeepSeek keeps hc_* in FP32.

• Initialization lives outside HF's hooks. ManifoldHyperConnection builds mix_weight with torch.empty and initializes it in __init__, which HF's _init_weights doesn't cover. Loading a checkpoint that lacks these keys through from_pretrained may leave them uninitialized.

Accuracy Test

• test_residual.py (from Refactor draft layers and add GR/mHC residuals #1023): GR/mHC reference forward/backward, save/load, checkpointing, plus CUDA BF16 cases that are skipped without a GPU.
• With the default StandardResidual, the parity test from the previous PR stays bit-exact.
• tests/test_modeling on CPU: 182 tests, 24 skipped. The only errors are the 4 env-only import errors.

Stack

This is part of a split of #1023 (by @jiapingW) into reviewable layers. Each PR targets the branch of the previous one, and only its own commits show in the diff.

1. [1/6] refactor(draft): move draft layer components under modeling/draft/layers #1025: Move draft layer components (pure move + legacy-import shims)
2. [2/6] refactor(draft): extract GQA/MLA attention into reusable layers #1026: GQA/MLA attention extraction (bit-exact parity test)
3. [3/6] feat(draft): residual hooks with Qwen GR and DeepSeek-V4.1 mHC #1027: Residual hooks + GR + mHC 👈 this PR
4. [4/6] feat(attention): unify visibility across dense attention backends #1028: Visibility unified across dense backends (behavior change)
5. [5/6] feat(attention): sparse attention and SWA interfaces #1029: Sparse attention / SWA interfaces
6. [6/6] feat(attention): DSA and QSA draft attention training #1030: DSA and QSA training

The code is @jiapingW's from #1023 (commits are authored by them). The split adds a parity test, legacy import shims, and two fixes in the visibility PR, each noted in the PR it belongs to.

🤖 Generated with Claude Code

[yushengsu-thu]

To support the DeepSeek-V4 MoE DSpark draft, we need the mHC variant (own pre-mix, learned hc_head). Refer to #1068 if you want to support the DeepSeek-V4 MoE DSpark draft as well now.