Repository navigation
[DFlash] Keep grouped-conv taps inside each request's block so NaN/Inf cannot leak across requests - #41736
Merged
kpham-sgl merged 5 commits intoOct 1, 2026
Conversation
…ct#148) * dflash: keep grouped-conv taps inside each request's block _grouped_conv shifts rows across the flattened [bs * block_size] token dimension and masked the taps that cross a block boundary by multiplying with a 0/1 mask. NaN * 0 and Inf * 0 are NaN, so a non-finite value in request i's last block rows reached request i+1's first rows, and the DFlash2 draft stack (attention_conv and mlp_conv, prepare and finish, in every layer) carried it one request further per layer. Select the cross-block taps to exact zeros with torch.where before the multiply. Finite outputs are unchanged apart from the sign of exact zeros. * dflash: test the compiled grouped conv at the block boundary The engine calls _grouped_conv through torch.compile, but every boundary test forced the eager original with set_stance("force_eager"), so CI never ran the compiled function on non-finite input. Run the boundary check and the DFlashGroupedConv prepare/finish check both eagerly and through the compiled function (inductor on the CPU runner). With the multiply-by-mask formulation restored in dflash.py, the compiled checks fail at the same boundary rows as the eager ones (taps 2 row 7, taps 3 rows 6-7, for NaN, +Inf and -Inf) and in prepare/finish. The first CPU compile takes about 20 s, so est_time goes to 30. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
rodamani
marked this pull request as ready for review
September 29, 2026 21:15
Contributor
Author
|
/rerun-test -c test_dflash_logits.py test_basic_sanity_dflash.py test_basic_sanity_dspark.py |
Contributor
|
Results for 🚀 🚀 🚀 ⛔ |
Contributor
Author
|
/tag-and-rerun-ci |
3 of 5 tasks
kpham-sgl
approved these changes
Oct 1, 2026
kpham-sgl
left a comment
kpham-sgl
left a comment
Collaborator
There was a problem hiding this comment.
Please address the comments then ping me again
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Contributor
Author
|
@kpham-sgl comments addressed! |
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
3 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.
Motivation
_grouped_convinmodels/dflash.pyshifts rows across the flattened[bs * block_size]token dimension and masks taps that cross a request's block boundary by multiplying with a 0/1 mask.NaN * 0andInf * 0are NaN, so a non-finite value in request i's last rows leaks into request i+1's first rows, and the DFlash2 draft stack (attention_conv and mlp_conv, prepare and finish, every layer) carries it one request further per layer. One request's bad activations can corrupt its neighbours' drafts.Modifications
torch.wherebefore the multiply. Finite outputs are unchanged apart from the sign of exact zeros.DSpark
DSparkDraftModelsubclassesDFlashDraftModeland reuses its decoder layers and_grouped_conv, so DSpark drafts withconv_kernel_sizeset get this fix too. Not run on a DSpark checkpoint.Accuracy Tests
Covered by the existing DFlash CI tests (test_dflash_logits.py, test_basic_sanity_dflash.py, test_basic_sanity_dspark.py).
Speed Tests and Profiling
No hot-path change beyond the fix itself; not separately benchmarked.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ⏳ Run #36806531824
Latest PR Test (Extra): ❌ Run #36806531535
Latest PR Test (AMD ROCm 10): ⏳ Run #36806532282