Skip to content

[Build] Bump vllm-flash-attn to compress fatbins with --compress-mode=size - #65518

Open
mgoin wants to merge 2 commits into
vllm-project:mainfrom
mgoin:mgoin/bump-fa-compress
Open

mgoin wants to merge 2 commits into
vllm-project:mainfrom
mgoin:mgoin/bump-fa-compress

Conversation

@mgoin

@mgoin mgoin commented Nov 3, 2026 •

Copy link
Copy Markdown
Member

Overview

Bump vllm-flash-attn to d8dca8f to pick up vllm-project/flash-attention#216, which compresses FA's fatbins with --compress-mode=size on CUDA >= 13 as vLLM's own extensions already do. The x86_64 wheel shrinks by about 82 MB (about 25%).

Claims

  • The wheel is about 82 MB smaller (335 MB → about 253 MB for v0.33.0 x86_64), and installed size drops by about 347 MB.
  • Kernel code is unchanged: every FA2/FA3 SASS kernel is byte-identical to the v0.33.0 release.
  • The only other commit picked up is Add SM90 FP8 MMA forward for head dimension 512 flash-attention#199 (FA4 CuteDSL: SM90 FP8 MMA forward for hdim 512). It is Python-only and doesn't change the compiled .so files.

Validation

I built _vllm_fa2_C/_vllm_fa3_C from this branch's FA commit (CUDA 13.0, TORCH_CUDA_ARCH_LIST='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0') and compared against the .so files in the v0.33.0 x86_64 wheel on PyPI (FA b273a02, also nvcc 13.0.88):

installed .nv_fatbin deflated (wheel cost)
_vllm_fa2_C v0.33.0 308.0 MB 297.2 MB 129.6 MB
_vllm_fa2_C this PR 77.1 MB 63.3 MB 65.0 MB
_vllm_fa3_C v0.33.0 172.6 MB 133.7 MB 42.6 MB
_vllm_fa3_C this PR 56.9 MB 17.2 MB 24.8 MB
  • Same kernels. I extracted every image with cuobjdump -xelf/-xptx and hashed each kernel's .text.* section. FA2 has 2,904/2,904 SASS kernels identical to v0.33.0, and FA3 has 1,333/1,333. The 76 FA2 PTX files differ only in nvcc's per-TU _INTERNAL_<hash> symbol names.
  • Runtime on B300. FA2 on B300 (sm_103) JITs from the compressed compute_80 PTX. A causal varlen call (bf16, hdim 128, GQA 32/8) matches fp32 scaled_dot_product_attention with a max abs error of 0.0118.
  • Startup cost. I timed cuModuleLoadData on B300 for one FA2 TU compiled to sm_100 SASS (4.4 MB uncompressed). It takes 12.7 ms uncompressed and 14.6 ms with --compress-mode=size, so about 2 ms per TU (about 0.45 ms per MB of SASS). With lazy module loading, only the TUs whose kernels run are loaded, once per process. Loading every image for one arch would cost about 100 ms for FA2 (sm_80) and about 60 ms for FA3 (sm_90a).
  • Fetch. FetchContent checks out d8dca8f from vllm-project/flash-attention.
  • Not run on Hopper or Ampere. I have no such hardware, so FA3 and the FA2 sm_80 SASS path weren't run. Coverage there rests on the identical-SASS check above, plus CI.

Details

  • VLLM_FA_GPU_FLAGS doesn't inherit VLLM_GPU_FLAGS, so FA never got the --compress-mode=size that CMakeLists.txt sets for vLLM's own targets. By default fatbinary compresses only PTX and debug images, which left all FA SASS stored raw.
  • -Xfatbin=-compress-all (FlashInfer's approach) adds nothing on top of --compress-mode=size and gives larger output when used alone. Details are in the FA PR.
  • Add support for H100 #199 compatibility. It adds new FA4 SM90 fp8_kv_dequant guards (for example, it now requires head_dim 512). vLLM already restricts SM90 FA4 FP8-KV to head size 512 (fa_utils.py flash_attn_supports_fp8), so no supported vLLM path changes.
  • Duplicate check. [Do Not Merge][Attention] Sync FA with upstream (54b044c8) #65203 also moves the FA pin, but it's a do-not-merge upstream sync. No other open PR bumps the FA pin.

Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

  • AI assistance was used during the creation of this PR.

  • Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

  • Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

  • Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

  • Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

…=size

Picks up vllm-project/flash-attention#216, which passes
--compress-mode=size for CUDA >= 13 like vLLM's own extensions, plus
vllm-project/flash-attention#199 (SM90 FP8 forward for head dim 512).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
@mergify mergify Bot added the ci/build label Nov 3, 2026
@mgoin
mgoin marked this pull request as ready for review November 3, 2026 23:46

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mgoin

mgoin commented Nov 3, 2026

Copy link
Copy Markdown
Member Author

/ci run

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99744 for commit bc196dfbb0b2.

ℹ️ This PR is 7 commits behind upstream main. Merged upstream main into this branch as bc196dfbb0b2 before starting CI. Pull the branch before pushing more changes.

@mgoin

mgoin commented Nov 4, 2026

Copy link
Copy Markdown
Member Author

Blocked on #60372 landing first

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants