feat(dflash): share DFlash masks with inference, default masks for spec_generate [2/3] - #935
Open
maocheng23 wants to merge 1 commit into
Open
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 #905. Co-authored-by: Inhyuk Cho <ihcho@lgresearch.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
maocheng23
requested review from
FlamingoPg,
FrankLeeeee,
shuaills and
sleepcoo
as code owners
October 6, 2026 04:52
cih9088
reviewed
Oct 6, 2026
| mask_builder = create_dflash_block_mask | ||
| if flex_attention_backend() == "FLASH": | ||
| # FLASH requires a minimum of this block size. | ||
| mask_args["flex_block_size"] = (256, 128) |
Contributor
There was a problem hiding this comment.
Suggested change
| mask_args["flex_block_size"] = (256, 128) | |
| mask_args["BLOCK_SIZE"] = (256, 128) |
how about just passing additional kwargs to the block mask creation function directly, instead of introducing merely a wrapper argument?
Comment on lines
+137
to
+139
| kwargs = {} | ||
| if flex_block_size is not None: | ||
| kwargs["BLOCK_SIZE"] = flex_block_size |
Contributor
There was a problem hiding this comment.
Suggested change
| kwargs = {} | |
| if flex_block_size is not None: | |
| kwargs["BLOCK_SIZE"] = flex_block_size |
how about just passing additional kwargs to the block mask creation function directly, instead of introducing merely a wrapper argument?
Comment on lines
+90
to
+92
| flex_block_size=None, | ||
| sliding_window: Optional[int] = None, | ||
| is_causal: Optional[bool] = None, |
Contributor
There was a problem hiding this comment.
Suggested change
| flex_block_size=None, | |
| sliding_window: Optional[int] = None, | |
| is_causal: Optional[bool] = None, | |
| sliding_window: Optional[int] = None, | |
| is_causal: Optional[bool] = None, | |
| **kwargs, |
how about just passing additional kwargs to the block mask creation function directly, instead of introducing merely a wrapper argument?
This branch has not been deployed
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
2 participants
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 #905, part 2 of 3. Stacked on #934.
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 #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 #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 fix(dflash): honor config.is_causal in DFlash training masks [1/3] #934). This is an edge case onlyspec_generatecan reach.spec_generateno longer forcesis_causal=False.Tests
Ported from #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-by @cih9088 (original #905); every commit carries a
Co-authored-by: Inhyuk Chotrailer.🤖 Generated with Claude Code