Skip to content

fix(dflash): gate flex_attention BACKEND=TRITON on torch>=2.11 (fixes ROCm) - #773

Open
zhentaocc wants to merge 1 commit into
sgl-project:mainfrom
zhentaocc:fix/rocm-dflash-flex-backend
Open

fix(dflash): gate flex_attention BACKEND=TRITON on torch>=2.11 (fixes ROCm)#773
zhentaocc wants to merge 1 commit into
sgl-project:mainfrom
zhentaocc:fix/rocm-dflash-flex-backend

Conversation

@zhentaocc

@zhentaocc zhentaocc commented Aug 19, 2026

Copy link
Copy Markdown

Problem

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
generated kernel as a bare identifier and compilation fails with:

NameError: 'TRITON' is not defined

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:

  • 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
exact 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_attention kernel (no flex-decoding, no
FLASH/CuteDSL template). This fixes every pre-2.11 build, ROCm included, while
leaving the CUDA path identical.

test_online_wrapper_forces_standard_triton_flex_backend is updated to assert
the version-appropriate kernel options.

@zhentaocc
zhentaocc marked this pull request as draft August 19, 2026 09:01
@zhentaocc zhentaocc changed the title fix(rocm): don't pin flex_attention BACKEND=TRITON for DFlash-family on ROCm fix(dflash): gate flex_attention BACKEND=TRITON on torch>=2.11 (fixes ROCm) Aug 20, 2026
@zhentaocc
zhentaocc marked this pull request as ready for review August 20, 2026 07:48
@zhentaocc
zhentaocc requested a review from FrankLeeeee as a code owner 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant