fix(dflash): gate flex_attention BACKEND=TRITON on torch>=2.11 (fixes ROCm) - #773
Open
zhentaocc wants to merge 1 commit into
Open
fix(dflash): gate flex_attention BACKEND=TRITON on torch>=2.11 (fixes ROCm)#773zhentaocc wants to merge 1 commit into
zhentaocc wants to merge 1 commit into
Conversation
zhentaocc
marked this pull request as draft
August 19, 2026 09:01
zhentaocc
marked this pull request as ready for review
August 20, 2026 07:48
… ROCm)
OnlineDFlashModel._forward_draft_blocks (shared by DFlash / Domino / DSpark)
pins the flex_attention kernel option {"BACKEND": "TRITON"} to keep every
DFlash-family batch on the general Triton Flex Attention kernel (AUTO can
otherwise route short-query batches to the flex-decoding kernel, whose config
set can be empty for DFlash's sparse BlockMask).
However, the "BACKEND" kernel option is only understood by the inductor
flex-attention lowering on torch>=2.11, where it is sanitized out of the
generated Triton constexprs. On older builds the string leaks into the kernel
as a bare identifier and compilation fails with NameError: 'TRITON' is not
defined. This affects any pre-2.11 torch, in particular current torch ROCm
wheels (2.9.x-based).
Gate the kernel option on the torch version instead of the vendor:
- torch >= 2.11 -> {"BACKEND": "TRITON"} (unchanged CUDA behaviour)
- older -> {"FORCE_USE_FLEX_ATTENTION": True}
FORCE_USE_FLEX_ATTENTION has been supported since torch 2.5 and selects the
same general Triton Flex Attention kernel (verified by dumping the generated
inductor code on torch 2.11: both options produce a byte-identical
triton_tem_fused_flex_attention kernel, no flex-decoding, no FLASH/CuteDSL).
This fixes every pre-2.11 build, ROCm included, while leaving the CUDA path
identical.
Update test_online_wrapper_forces_standard_triton_flex_backend to assert the
version-appropriate kernel options.
zhentaocc
force-pushed
the
fix/rocm-dflash-flex-backend
branch
from
August 20, 2026 07:49
bae7d28 to
50c3407
Compare
3 tasks
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.
Problem
OnlineDFlashModel._forward_draft_blocks(shared by DFlash / Domino / DSpark)pins the flex_attention kernel option
{"BACKEND": "TRITON"}to keep everyDFlash-family batch on the general Triton Flex Attention kernel (AUTO can
otherwise route short-query batches to the flex-decoding kernel, whose config
set can be empty for DFlash's sparse
BlockMask).However, the
"BACKEND"kernel option is only understood by the inductorflex-attention lowering on torch >= 2.11, where it is sanitized out of the
generated Triton constexprs. On older builds the string leaks into the
generated kernel as a bare identifier and compilation fails with:
This affects any pre-2.11 torch, and in particular current torch ROCm
wheels, which ship a 2.9.x-based build.
Fix
Gate the kernel option on the torch version rather than on the vendor:
{"BACKEND": "TRITON"}(unchanged CUDA behaviour){"FORCE_USE_FLEX_ATTENTION": True}FORCE_USE_FLEX_ATTENTIONhas been supported since torch 2.5 and selects theexact same general Triton Flex Attention kernel — verified by dumping the
generated inductor code on torch 2.11 and confirming both options produce a
byte-identical
triton_tem_fused_flex_attentionkernel (no flex-decoding, noFLASH/CuteDSL template). This fixes every pre-2.11 build, ROCm included, while
leaving the CUDA path identical.
test_online_wrapper_forces_standard_triton_flex_backendis updated to assertthe version-appropriate kernel options.