Build vllm-flash-attn FA2 kernels with SM12x SASS

[ehartford]

Problem

vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so carries only sm_80 code. On an RTX 5090 (sm_120) the driver has to JIT the compute_80 PTX, and the R595 driver (CUDA 13.2) rejects PTX produced by nvcc 13.3:

torch.AcceleratorError: CUDA error: the provided PTX was compiled with an unsupported toolchain.
No device code available for GPU ISA 120

This breaks every FA2 varlen call on SM12x, including the FLASH_ATTN ViT path used by the Qwen3.5/3.6 vision encoders. It is why qwen36-nvfp4-1 and qwen36-fp8-1 (rtx5090) set LD_LIBRARY_PATH=/usr/local/cuda-13.3/compat.

Root cause

The pinned upstream vllm-flash-attn (ed4b7342) hardcodes the FA2 kernels (csrc/flash_attn/src/flash_fwd_*.cu) to "8.0+PTX". For a TORCH_CUDA_ARCH_LIST=12.0 build, its loose intersection picks 8.0+PTX, so the extension holds sm_80 SASS plus compute_80 PTX and nothing for SM12x. This is the same class of bug #88 fixed for MARLIN_OTHER_ARCHS / HADACORE_ARCHS. FA3 is Hopper-only (9.0a) and is not affected.

Fix

cmake/external_projects/vllm_flash_attn.cmake now does what MARLIN_OTHER_ARCHS does:

• Before FetchContent_MakeAvailable, it resolves 12.0f (CUDA >= 13) or 12.0a;12.1a (CUDA 12.8+) against CUDA_ARCHS with SlimServe's cuda_archs_loose_intersection. This has to happen first because vllm-flash-attn installs its own older copy of that function over ours, and that copy doesn't understand the a/f suffixes.
• After FetchContent_MakeAvailable, it appends the matching -gencode to the FA2 sources in vllm-flash-attn's directory scope (set_property(SOURCE ... DIRECTORY <fa src>)). The flag shape is the same one set_gencode_flag_for_srcs produces.

There is no patch to the fetched source and the pin is unchanged. The override also applies when VLLM_FLASH_ATTN_SRC_DIR is used. If the upstream FA2 source glob ever comes back empty, the configure stops with a FATAL_ERROR instead of skipping the override without warning. Builds without an SM12x target are unchanged, and the FA3 gencodes are untouched (verified arch=compute_90a,code=sm_90a only).

Why this mechanism: the arch list lives inside the fetched upstream CMakeLists, and SlimServe passes no variable that controls it. A PATCH_COMMAND would need idempotent re-apply logic on top of FetchContent's git update step. Setting the property after the fetch keeps the arch policy in SlimServe next to the Marlin/hadacore lists.

Evidence

Configure, CUDA 13.3, TORCH_CUDA_ARCH_LIST=12.0 (this box):

-- FA2_ARCHS: 8.0+PTX
-- FA3_ARCHS:
-- FA2 SM12x archs: 12.0f
flash_fwd_hdim128_fp16_sm80.cu: -gencode arch=compute_80,code=sm_80 -gencode arch=compute_80,code=compute_80 -gencode arch=compute_120f,code=sm_120f

Configure only, TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0 10.0 12.0": FA2 gets the same three gencodes, FA3_ARCHS: 9.0a, and the FA3 sources still use only compute_90a,code=sm_90a.

cuobjdump of _vllm_fa2_C.abi3.so:

	before	after
--list-elf	76 x sm_80	76 x sm_80, 76 x sm_120
--list-ptx	76 x sm_80	76 x sm_80

All tests below ran on an RTX 5090 with the R595 driver, with LD_LIBRARY_PATH unset. Only the FA2 .so changed; every other extension was the installed build.

• FA2 varlen paged-KV tests (upstream tests/kernels/attention/test_flash_attn.py::test_varlen_with_paged_kv, recovered from history):
  • old .so: 160 failed, 480 skipped. Every failure is cudaErrorUnsupportedPtxVersion.
  • new .so: 160 passed, 480 skipped. The skips are FA3 (unsupported on sm_120) and FP8-Q with FA2.
• Qwen3.5-0.8B vision, using vllm.LLM with an image request, seed 1234, temperature 1.0 / top_p 0.95 / top_k 20, thinking on, gpu_memory_utilization=0.3:
  • old .so: engine init fails in the encoder profiling run (vit_flash_attn_wrapper -> flash_attn_varlen_func) with "the provided PTX was compiled with an unsupported toolchain".
  • new .so: the engine comes up with FLASH_ATTN for vit attention. The image request answers "There is a red circle and a blue square." (synthetic red circle + blue rectangle), and the text request completes.

What remains

• The LD_LIBRARY_PATH=/usr/local/cuda-13.3/compat env on qwen36-nvfp4-1 and qwen36-fp8-1 (rtx5090) should become removable once installs are rebuilt with this change. That needs a live profile serving check (text and image) without the compat path first, so this PR leaves the env in place. The operator-local affine-king-nvfp4-1 record carries the same env and needs the same check.
• The CUDA 12.8–12.9 branch (12.0a;12.1a) mirrors Fix test failures on main #88 but was not built here, because only CUDA 13.x toolkits are installed.

Summary by CodeRabbit

• New Features
  • Added support for building vLLM Flash Attention 2 kernels for applicable SM12x GPU architectures with supported CUDA toolkits.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: aea779a4-ed93-40f2-9766-f999ed858e8b

📥 Commits

Reviewing files that changed from the base of the PR and between 470a782 and 9b696a3.

📒 Files selected for processing: 1

• cmake/external_projects/vllm_flash_attn.cmake

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

📝 Walkthrough

The CMake configuration now selects SM12x architectures based on the CUDA version. When applicable, it finds FA2 forward-kernel sources and adds matching SASS gencode options. Configuration fails if no matching sources are found. FA3 is unchanged.

Changes: FA2 SM12x architecture support

Layer / File(s)	Summary
Select and apply SM12x architectures cmake/external_projects/vllm_flash_attn.cmake	CUDA 13.0 and newer select 12.0f; CUDA 12.8 through 12.x select 12.0a;12.1a. When an architecture is selected and _vllm_fa2_C exists, the configuration finds matching FA2 forward-kernel sources and adds SASS gencode options. It fails if no matching sources are found.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 9b696

FA2 kernels built for RTX 50-series (SM12x) GPUs will now include native machine code for those GPUs instead of relying on runtime compilation. No concrete defects remain in the build configuration change, so it appears ready to merge. CUDA 12.8–12.9 builds were not tested.

Pre-merge checks: 5

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: building vllm-flash-attn FA2 kernels with SM12x SASS.
Docstring Coverage		No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}