Repository navigation
[2/3] feat(dflash): share DFlash masks with inference, default masks for spec_generate - #23
Closed
maocheng23 wants to merge 1 commit into
Closed
maocheng23 wants to merge 1 commit into
maocheng23 wants to merge 1 commit into
Conversation
Move the DFlash mask builders (and resolve_dflash_is_causal) from algorithms/common/dflash_family_model.py to modeling/draft/dflash_mask.py so the draft model can use them; dflash_family_model re-imports them, so existing import paths and test patches keep working. DFlashDraftModel.forward now builds per-layer-type masks when called without one (spec_generate): one draft block after the cached plus new context, following config.is_causal and the sliding window like SGLang. Bidirectional full layers get no mask. flash_attention_2 gets no mask and relies on the attention module's is_causal (now resolved from the config instead of hard-coded False) plus its native window. FA2's window is symmetric, so for a bidirectional sliding layer with sliding_window < block_size it also bounds the right side, unlike the explicit masks. spec_generate no longer forces is_causal=False. Split out of sgl-project#905. Co-authored-by: Inhyuk Cho <ihcho@lgresearch.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
maocheng23
force-pushed
the
dflash-inference-masks
branch
from
October 6, 2026 02:48
9c404db to
e0db52f
Compare
maocheng23
force-pushed
the
dflash-is-causal-min
branch
from
October 6, 2026 02:48
b99ebd4 to
4ba5c24
Compare
Owner
Author
|
Superseded by sgl-project#935. |
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
1 participant
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.
Split of sgl-project#905, part 2 of 3. Stacked on #22.
Changes
resolve_dflash_is_causalfromalgorithms/common/dflash_family_model.pyto a newmodeling/draft/dflash_mask.py. Names are kept, anddflash_family_modelre-imports them, so existing imports, mock patches andbenchmark_dflash2_conv.pykeep working. Fix DFlash attention masking sgl-project/SpecForge#905 renamed them tobuild_*andflex_block_sizetoBLOCK_SIZE; that rename is what broke the benchmark that later landed on main.Qwen3DFlashAttentionBase.is_causalis resolved from the config viaresolve_dflash_is_causal(is_causal, config.layer_types[layer_idx])instead of being hard-codedFalse.DFlashDraftModel._prepare_attention_mask: whenforwardgets no mask (spec_generate), it builds per-layer-type masks for one draft block after cached plus new context.module.is_causalplus FA's native window with bottom-right causal alignment. Fix DFlash attention masking sgl-project/SpecForge#905 raisedNotImplementedErrorfor FA2 here. FA2's window is symmetric(W-1, W-1), so for a bidirectional sliding layer withsliding_window < block_sizeit also bounds the right side, unlike the explicit masks (left bound only, see [1/3] fix(dflash): honor config.is_causal in DFlash training masks #22). This is an edge case onlyspec_generatecan reach.spec_generateno longer forcesis_causal=False.Tests
Ported from sgl-project#905, adapted to the kept names:
test_dflash_sliding.py(TestDFlashGenerationMasks),test_dflash_mla.py(spec_generate mask check across eager/sdpa × is_causal × window), andtest_dflash_eager_attention.py. The FA2 test now checks that no mask is built, and the reference masks use the left-only in-block window.CPU run: all
tests/test_modelingpass (124 tests, 14 skipped).tests/test_utilsgives the same 20 failures as main on this machine (Eagle3 flex tests and 3 modules that can't import here); this PR adds none.Co-authored with @cih9088 (original sgl-project#905).
🤖 Generated with Claude Code