Build flash-attn ourselves to fix missing Blackwell support - #6
Open
MarioSieg wants to merge 2 commits into
Open
Build flash-attn ourselves to fix missing Blackwell support#6MarioSieg wants to merge 2 commits into
MarioSieg wants to merge 2 commits into
Conversation
The mjun0812/flash-attention-prebuild-wheels build prime-rl consumed is built against CUDA 12.6, which predates the 12.8 floor flash-attn's setup.py requires before it emits sm_100/sm_120 gencodes. Confirmed via cuobjdump: that wheel only ever contained sm_80/sm_90 cubins and no PTX, hard-failing flash_attention_2 on every Blackwell GPU (the fallback path for workstation Blackwell/SM120 in resolve_auto_attn). Building v2.8.3 from source here, against our CUDA 12.9 container with FLASH_ATTN_CUDA_ARCHS=80;90;100;120, produces a wheel with real sm_80/90/100/120 cubins (verified locally on a B200). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Both arches OOM-killed (exit 137) within 2 minutes. flash-attn's own NinjaBuildExtension throttles job count from actual free memory, but only when MAX_JOBS is unset -- our shared convention of always setting MAX_JOBS=8 disabled that safety check, and with 4 archs baked into every gencode list each nvcc job peaks far above the ~8-9GB flash-attn's own heuristic assumes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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
1 participant
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.
Summary
flash-attnwheel prime-rl consumes (frommjun0812/flash-attention-prebuild-wheels) is built against CUDA 12.6, which predates the CUDA 12.8 floor flash-attn's ownsetup.pyrequires before it emitssm_100/sm_120gencodes.cuobjdump --list-elfon the actual pinned wheel: it only containssm_80/sm_90cubins and no PTX, soflash_attention_2hard-fails on every Blackwell GPU — including workstation Blackwell (SM120), whichresolve_auto_attnin prime-rl explicitly falls back to FA2 for.flash-attnas a 5th kernel in this repo's build matrix (same pattern as deep-ep/deep-gemm/torchao), building the upstreamv2.8.3tag from source against our CUDA 12.9 container withFLASH_ATTN_CUDA_ARCHS=80;90;100;120.Test plan
>= 12.8floor as CI's 12.9) — verified withcuobjdumpthat the resulting wheel contains realsm_80/sm_90/sm_100/sm_120cubins (the old wheel had none of the last two).flash_attn_funccall on an actual B200 (SM100) GPU against the new wheel — succeeded.flash-attnfor bothx86_64/aarch64)🤖 Generated with Claude Code