flash-attn-hopper: bump to 00756db for TE 2.17 window kwargs - #10
Merged
Conversation
transformer_engine 2.17 moved use_flash_attn_3 out of the window_size-tuple branch in cp_p2p_fwd/bwd_flash_attn, so it now passes window_size_left/right to FA3. The pinned revision builds 3.0.0b1, whose _flash_attn_forward only accepts a single window_size tuple, so every context-parallel path raises TypeError. TE has no FA3-only kill switch -- NVTE_FLASH_ATTN disables FA2 too. 00756db builds 3.0.0, whose forward and backward are torch.library custom ops with SymInt window_size_left/right in their schemas, and which returns (out, softmax_lse, out_accum, softmax_lse_accum) -- indices 0 and 1, exactly what TE reads. Verified on an H200 against the miles image: TE selects 'FlashAttention backend (version 3.0.0)' and matches an SDPA reference to 2e-3. Same sm_80 + sm_90a coverage as before; hopper/setup.py hardcodes those from source-file suffixes and ignores TORCH_CUDA_ARCH_LIST. Also force a source build, since setup.py otherwise falls back to downloading a prebuilt release wheel built against a different CUDA/torch, and drop the separate flash_attn_interface.py fetch: the wheel now ships it top-level plus a flash_attn_3.flash_attn_interface re-export shim, and overwriting that shim with the full module double-registers the flash_attn_3:: custom ops.
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
miles bumped transformer_engine to 2.17.0 (radixark/miles#1781). TE 2.17 moved
use_flash_attn_3out of thewindow_size-tuple branch incp_p2p_fwd_flash_attn/cp_p2p_bwd_flash_attn:The pinned revision
fbf24f67buildsflash_attn_3 3.0.0b1, whose_flash_attn_forwardonly accepts a singlewindow_sizetuple, so everycontext-parallel path raises
TypeError. There is no FA3-only kill switch inTE —
NVTE_FLASH_ATTNdisables FA2 and FA3 together.FA2 in the same image (
2.7.4.post1) already takes the new kwargs, so only theFA3 path is affected.
Change
hopper/pin to00756db, which builds3.0.0.setup.pyprintsGuessing wheel URL: .../v3.0.0/flash_attn_3-3.0.0+cu122torch2.11...whlanddownloads a prebuilt wheel built against a different CUDA/torch than the image.
flash_attn_interface.pyfetch intest_wheels.py. It was asecond pin that can drift from the one the extension module was built at. The
wheel now ships the interface top-level and a
flash_attn_3.flash_attn_interfacere-export shim, so nothing extra is needed — and copying the full module over
that shim would re-run
@torch.library.custom_op("flash_attn_3::...")anddouble-register the ops.
Verification
Built and installed on an h200-sci devbox running the
radixark/miles:pr-1795image itself (torch 2.11.0+cu130, CUDA 13.0), so the wheel is compiled against
the toolchain it will ship with.
Notes on things checked so they are not rediscovered later:
PackageNotFoundErroraround the version lookup, not
ImportErroraround the imports themselves, soa wheel missing any of
flash_attn_func,flash_attn_varlen_func,flash_attn_with_kvcache,_flash_attn_forward,_flash_attn_backwardmakesimport transformer_engine.pytorchraise rather than degrade to cuDNN.3.0.0b1->3.0.0is benign. It flipsFlashAttentionUtils.v3_0_0_betatoFalse; the only consumer(
backends.py:1230) uses it to append an "update your flash-attn v3 beta"hint to an already-raised
TypeError. No functional branch depends on it.sm_80+sm_90a, verified withcuobjdumpagainst both the old and new
_C.abi3.so.hopper/setup.pyderives arch fromsource-file suffixes (
_sm90.cu->sm_90a) and ignores theTORCH_CUDA_ARCH_LISTthat_setup_envsets, so that variable has never hadany effect on FA3. FA3 has no Blackwell kernels here either way.
3.0.0returns(out, softmax_lse, out_accum, softmax_lse_accum); TE's CP code reads indices[0]and[1]underuse_flash_attn_3, which line up.