support dflash2

[jiapingW]

Motivation

Add end-to-end support for training DFlash 2 draft models in SpecForge.

DFlash 2 reuses the existing DFlash capture and training pipeline while extending the draft architecture with grouped dynamic convolutions and candidate-path selection. This PR
also ensures that exported checkpoints follow the public SGLang DFlash 2 parameter and configuration contract.

• Add and register DFlash2DraftModel as a DFlash-compatible architecture.
• Add grouped dynamic depthwise convolutions around each attention and MLP sublayer.
• Add a low-rank candidate selector for reranking top-k token candidates using predecessor-token transitions.
• Support DFlash 2 unary-logit transformations, including output_multiplier and final_logit_softcapping.
• Extend the DFlash objective with a configurable selector loss:
  • add training.dflash2_selector_loss_alpha;
  • insert the gold token during training when it is missing from the top-k candidates;
  • report selector loss, accuracy, and original top-k coverage.
• Persist DFlash 2 architecture and selector settings in the resume contract.
• Allow the DFlash provider and draft registry to load both DFlashDraftModel and DFlash2DraftModel.
• Support block_size from either the top-level draft config or dflash_config, with conflict validation during export.
• Preserve the DFlash2DraftModel architecture during Hugging Face export and validate all required DFlash 2 fields.
• Add a Qwen3.6-27B DFlash 2 draft config and a managed two-GPU disaggregated training recipe.
• Document DFlash 2 training, configuration, export, and serving behavior.
• Add unit tests covering:
  • model registration and construction;
  • grouped-convolution behavior and gradients;
  • candidate scoring and lattice construction;
  • selector supervision and gradient propagation;
  • resume metadata;
  • export normalization and validation;
  • configuration schema and example topology.

Modifications

Related Issues

Accuracy Test

Benchmark & Profiling

{
  "passed": true,
  "input_format": "training_jsonl",
  "spec_accept_length": 8.0,
  "target_prefix_match_tokens": 24,
  "generated_tokens": 24,
  "target_tokens": 3785,
  "clean_block_tokens": 8,
  "errors": [],
  "result_path": "/disk3/wjp/projects/codex-run/SpecForge/outputs/Qwen3.6-27B-Dflash-overfit/serving-gate.json"
}

Checklist

• [] Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[maocheng23]

would this introduce Nan denominator?

Add a clamp like 1e-10?

[maocheng23]

move clamp to L691

[maocheng23]

Actually, I think we can add a clamp here, to avoid 0/0 situation.

[maocheng23]

Would it be better to separate the selector part into a sub-function, e.g. _selector_chunk_terms(...) -> SelectorTerms (a small NamedTuple), with _dflash_objective_chunk_terms splicing its fields into the flat return tuple?

Beyond readability this pays off twice:

• the six new_zeros(()) initializations collapse into one SelectorTerms.zeros(...) — which also exposes that selector_tv_num is never assigned anywhere (a permanently-zero dead slot in the 12-tuple; it can just be dropped);
• tests currently reach into terms[6] / terms[10] / terms[11]; named fields make them self-describing.

Two constraints the helper must keep (worth stating in its docstring):

1. The chunk-function boundary must remain a flat tuple of additive tensors (checkpointed_chunk_reduce's contract) — use the NamedTuple inside, flatten at the return.
2. Gate the helper only on static conditions (candidate_selector is not None, configured self.selector_loss_alpha > 0) — never on the per-step effective alpha. During selector warmup the effective alpha is 0 and forward adds 0.0 * selector_ce_num; that 0-scaled term is load-bearing: it keeps the selector parameters in the autograd graph on every step, which the DDP path (NO_SHARD → find_unused_parameters=False) requires. An if effective_alpha == 0: return zeros early-return would fail multi-GPU runs on the first warmup step with DDP's unused-parameters error. For the same reason, keep the effective_selector_alpha * scaling in forward after chunk reduction — the ratio metrics also need the unscaled selector numerator.

[maocheng23]

Suggestion: mirror the parent's explicit signature instead of **kwargs + a hand-maintained exclusion set.

This works today because DFlashDraftModel.forward happens to call layers with exactly these keyword names, but it is fragile:

• if a caller (or a future parent change) passes a kwarg that is explicitly re-passed here but missing from the exclusion set below (e.g. past_key_values vs past_key_value), it gets forwarded twice → TypeError: got multiple values for keyword argument;
• positional calls break, while the parent accepts them;
• the past_key_value → past_key_values rename is hidden inside a .get().

Mirroring Qwen3DFlashDecoderLayer.forward's signature removes the exclusion set entirely, and the diff vs the parent becomes exactly the four conv-sandwich lines — which is the real content of this subclass:

def forward(
    self,
    target_hidden: Optional[torch.Tensor] = None,
    hidden_states: Optional[torch.Tensor] = None,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Cache] = None,
    output_attentions: Optional[bool] = False,
    use_cache: Optional[bool] = False,
    cache_position: Optional[torch.LongTensor] = None,
    position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    **kwargs: Unpack[FlashAttentionKwargs],
) -> torch.Tensor:
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    hidden_states, attention_kernel = self.attention_conv.prepare(hidden_states)
    hidden_states = self.self_attn(
        hidden_states=hidden_states,
        target_hidden=target_hidden,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_value,
        output_attentions=output_attentions,
        use_cache=use_cache,
        cache_position=cache_position,
        position_embeddings=position_embeddings,
        **kwargs,
    )[0]
    hidden_states = self.attention_conv.finish(hidden_states, attention_kernel)
    hidden_states = residual + hidden_states

    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states, mlp_kernel = self.mlp_conv.prepare(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = self.mlp_conv.finish(hidden_states, mlp_kernel)
    return residual + hidden_states

Imports needed (same sources dflash.py already uses): Optional from typing; Tuple, Unpack from typing_extensions; Cache from transformers.cache_utils; FlashAttentionKwargs from transformers.models.qwen3.modeling_qwen3.

[maocheng23]

The identity-init story is only half-implemented: kernel_projection keeps PyTorch's default (kaiming-uniform) init, so a fresh conv is not a no-op.

The effective kernel is base + Δ with Δ = kernel_projection(x). With hidden_size=5120 the default init bound is ≈ 1/√5120 ≈ 0.014, and x is an O(1) RMSNorm output, so each Δ entry is a ~5120-term dot product of magnitude ~O(0.5) — the same order as the identity coefficient 1.0. At init the conv substantially perturbs every sublayer's input and output.

That contradicts the class docstring ("makes enabling the module a stable extension of a DFlash backbone"), and test_identity_initialization_preserves_inputs has to manually zero this weight before its identity assertion holds. It also defeats DFlash → DFlash2 warm-starts: the selector side gets this right (zero-init successor_codebook → exact unary no-op), but the conv side scrambles the loaded backbone at step 0.

Zero-init has no dead-gradient risk here because Δ enters additively: ∂L/∂W = ∂L/∂Δ ⊗ x, independent of W (same math as LoRA's zero-init B).

Suggested change

	self.kernel_projection = nn.Linear(
	int(hidden_size),
	2 * self.taps * self.num_groups,
	bias=False,
	)
	self.kernel_projection = nn.Linear(
	int(hidden_size),
	2 * self.taps * self.num_groups,
	bias=False,
	)
	nn.init.zeros_(self.kernel_projection.weight)

If the reference DFlash2 recipe deliberately random-inits this projection, the docstring should say so instead — but currently the docstring, the test, and the selector's symmetric zero-init all point to making this an exact identity.