feat: add DFlash2 training support

[abatilo]

Summary

• Add DFlash2 configuration, grouped dynamic causal convolution, candidate selection, and a thin wrapper around the shared DFlash training path.
• Export the public DFlash2DraftModel configuration and checkpoint schema used by SGLang and vLLM.
• Add fail-fast checks for unsupported configuration values that would change training-to-serving behavior.
• Add the Qwen3-8B recipe and focused tests for configuration, masking, loss, gradients, selector behavior, distributed export, and artifact schema.

Public inference contract

TorchSpec follows the public DFlash2 architecture and serving behavior described by the Inco overview, z-lab/dflash, SGLang PR #35371, and vLLM PR #52816:

• Apply grouped dynamic causal convolution around attention and the multilayer perceptron.
• Select the strict unary top-K candidates at each position.
• Apply the output multiplier and optional softcap.
• Add the predecessor-conditioned bilinear transition score.
• Use the verified anchor for the first predecessor, followed by the selected candidate from the previous position.
• Export the public configuration and checkpoint parameter names.

Training objective boundary

The public sources describe the selector architecture and inference computation. They do not publish selector targets, a loss, negative construction, initialization, or a training schedule.

TorchSpec therefore defines a teacher-forced cross-entropy objective. Training uses the gold predecessor and inserts the gold successor when it is outside the unary top-K candidates. This insertion is training-only. Serving retains the strict public top-K set.

This pull request does not claim parity with an unpublished upstream training recipe.

Validation

• Focused tests compare the grouped convolution and selector behavior with the public implementation and cover configuration, masks, positions, loss, gradients, distributed export, and checkpoint schema.
• Repository pre-commit checks pass.
• Lint and the CPU-safe test suite pass at the current pull-request head. The CPU-safe suite reports 580 passed and 44 skipped.
• The Qwen3-8B recipe was also exercised through training, export, and SGLang serving in a private test environment. Those artifacts are not linked here and do not replace the public repository checks.

Limitations

• For an artifact with block_size=B, configure vLLM with num_speculative_tokens=B-1.
• The public sources do not disclose the upstream selector-training recipe or a production acceptance target. The included recipe and private integration exercise are not evidence of generalization or production throughput.
• The current configuration path rejects mixed full/sliding attention layouts and non-default RoPE scaling instead of exporting a model with different behavior.
• SGLang compatibility was exercised against its public DFlash2 implementation. The vLLM integration was source-reviewed but was not runtime-validated for this change.

Related work

• DFlash: feat: DFlash (block-parallel) draft model training #64
• DSpark: feat: DeepSpec - DSpark trainer support #129
• Shared DFlash and DSpark training path: fix: share forward pass / loss in DFlash / DSpark #134

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5b543e126e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Remove a normalized zero softcap from the nested config

When an official-style dflash_config supplies final_logit_softcapping: 0.0, these lines disable it for training by setting the attribute to None, but the original 0.0 remains in nested. dflash2_config_for_serving() then preserves that nested value, so the exported artifact advertises a zero softcap while training used no softcap, potentially producing invalid serving logits or a training/serving mismatch. Remove the nested key or overwrite it with None when zero is normalized.

Useful? React with 👍 / 👎.

[yubofredwang]

Thanks for the contribution!

Here is some ongoing training with the implementation indicating the convergence

[yuyangxie96]

@yubofredwang Can you successfully reproduce the results of dflash2 using this PR?

[abatilo]

@yuyangxie96 I didn't try reproducing the exact open source models but I did try a 6 way comparison to train Qwen3-8b with a regenerated dataset based on these prompts: https://huggingface.co/datasets/lightseekorg/kimi-mtp-dataset

DFlash, DFlash2, DSpark each with either the default cross entropy loss or D-PACE based loss. DFlash2 + D-PACE won 13 of the 16 categories in my eval set after 1 epoch.