add DFlash2 npu training recipe - #789
Merged
Merged
Conversation
curnane-lab
force-pushed
the
dflash2_npu
branch
from
August 29, 2026 01:47
a943f19 to
b0c0d1a
Compare
curnane-lab
marked this pull request as ready for review
August 30, 2026 11:08
curnane-lab
requested review from
FlamingoPg,
FrankLeeeee,
shuaills and
sleepcoo
as code owners
August 30, 2026 11:08
jiapingW
approved these changes
Sep 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Follow-up to #772: make DFlash 2 trainable on Ascend NPU by adding the missing recipe layer — one draft config, one external-server recipe, and a tutorial section. DFlash 2 reuses the DFlash capture path and the NPU disaggregated plumbing from #722 unchanged. Mirrors the Domino/DSpark NPU recipes in #726 / #727.
Modifications
configs/qwen3.5-4b-dflash2.json:DFlash2DraftModelconfig for Qwen3.5-4B — 5 full-attention layers on the same target layers as the 4B DFlash config ([1, 8, 15, 22, 29]), 16,conv_kernel_size: 2,conv_group_size: 16,selector_rank: 256,selector_top_k: 16`.examples/configs/online/disaggregated/external/qwen3.5-4b-dflash2-online-npu.yaml: field-aligned with the existing DFlash NPU recipe, exceptattention_backend: sdpa(flex_attention is unavailable on Ascend) anddflash2_selector_loss_alpha: 1.0.docs/basic_usage/Ascend/ascend_npu.md: a "Training DFlash 2 drafts" section — capture setup unchanged from DFlash;sdpamandatory; the convolutions and selector are stock torch ops; serving needs SGLang with DFlash 2 ([Spec] DFlash2: local convolution + candidate selector sglang#35371, merged).Related Issues
Follow-up to #772; builds on #722; serving side is sgl-project/sglang#35371.
Accuracy Test
Config / recipe / docs only; no model-side code. End-to-end validation on Ascend NPU with the checked-in recipe (Qwen3.5-4B target + DFlash 2 draft, ~1.75k optimizer steps):
train/loss~6.9 → ~1.25,train/acc~0.09 → ~0.71,selector_loss2.05 → ~0.25,selector_accuracy~0.36 → ~0.93,selector_coverage~0.44 → ~0.90; boundedgrad_norm, no NaN or stall.Benchmark & Profiling
N/A (no performance-affecting code changes).
Checklist