Skip to content

Releases: ussoewwin/flash-attention

Release list

v1.8.1 — CUTLASS v4.6.2 Submodule Sync & Package Version 2.9.2.post2

v1.8.1 — CUTLASS v4.6.2 Submodule Sync & Package Version 2.9.2.post2 Latest
Latest

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 05 Sep 04:38
v1.8.1
6f4ac8d

Overview

Fork release v1.8.1 synchronizes the bundled C++ CUTLASS submodule (csrc/cutlass) from v4.3.4 to the official tag v4.6.2 (commit 6c65a17), bringing the C++ template headers into strict alignment with FlashAttention-4's nvidia-cutlass-dsl>=4.6.2 requirement while ensuring 100% backward compatibility with FlashAttention-2 CUDA/C++ kernels. Package version is incremented to 2.9.2.post2 across flash_attn/__init__.py and setup2.py. Furthermore, legacy compilation artifacts (build/) have been purged to guarantee a clean build baseline under Windows MSVC.

1. CUTLASS Submodule Synchronization to v4.6.2

Background & Motivation

In recent upstream commits (specifically commit ce088ab, Dao-AILab#2798), the Python DSL dependency floor was raised to nvidia-cutlass-dsl>=4.6.2 to support QuACK packed subtraction compatibility and upcoming CUTLASS 4.7 capabilities. However, the Git submodule in csrc/cutlass had remained at v4.3.4 (commit 7127592). This release updates the submodule to tag v4.6.2 (6c65a175668952f09bcbf66cb97a8de1b734b4a0).

Technical Compatibility with FlashAttention-2

FlashAttention-2 (csrc/flash_attn/) does not instantiate CUTLASS 3.x/4.x Collective Mainloops or CuTe C++ GEMM pipelines; its usage of CUTLASS is strictly scoped to fundamental data types and utility headers:

  • <cutlass/cutlass.h>
  • <cutlass/array.h>
  • <cutlass/numeric_types.h>
  • <cutlass/numeric_conversion.h>

The numeric primitives (cutlass::half_t, cutlass::bfloat16_t, cutlass::Array<T, N>) and conversion routines retain identical semantics and binary layouts in CUTLASS 4.6.2. Under MSVC with /std:c++20, the FA2 compilation units build cleanly without symbol collisions or deprecation regressions.

2. Package Version Bump (2.9.2.post2)

Package metadata and runtime indicators have been updated from 2.9.2.post1 to 2.9.2.post2:

  • flash_attn/__init__.py: Default fallback version updated to "2.9.2.post2".
  • setup2.py: public_version = "2.9.2.post2" in get_package_version().

This ensures that binary wheels produced from this revision are unambiguously distinguishable from earlier post1 builds.

3. Build Tree Hygiene & Clean Build Assurance

  • Build Directory Purge: The legacy build/ tree (temp.win-amd64-*, lib.win-amd64-*) was completely wiped. When building with Ninja/MSVC, cached .obj files can bypass recompilation if source file timestamps appear unchanged, risking subtle ABI or header inconsistencies. Removing build/ guarantees that every object file is freshly compiled against CUTLASS 4.6.2 headers.
  • Git Ignore Hygiene: Added explicit rules in .gitignore for build logs (*.log), temporary directories (temp/, tmp/), and backup files (*.bak*, *.tmp), eliminating unintended staging of build side-effects.

4. Preserved Fork Architecture & Windows Optimizations

All fork-specific Windows enhancements remain active and verified:

  • 24 split_align Forward Kernels: Independent translation units (flash_fwd_split_align_*.cu) for num_splits == 1 forward pass prevent ptxas compiler timeouts.
  • MSVC 2GB COFF Linker Limit Bypass (3:3 SM Split):
    • WindowsWhlBuilder_cuda_2.bat: Blackwell targets (100;120;121).
    • WindowsWhlBuilder_cuda_3.bat: Legacy targets (80;89;90), with MAX_JOBS=5.
  • Windows MSVC Compatibility Overrides:
    • Standard flags: /std:c++20 enabled for modern C++ initializers.
    • C10 ArrayRef operator overrides in namespace c10 (resolves MSVC C2666 ambiguity).
    • Clean linking against PyTorch 2.14.0's 6-argument c10::cuda::c10_cuda_check_implementation.
  • PEP 427 Compliant Wheel Naming: Employs .blackwell and .legacy dot-suffixes.
  • CI Workflow Suppression: .github and tools/ci remain permanently suppressed.

5. Pre-built Windows Wheels

Pre-compiled wheels will be published on Hugging Face:
👉 ussoewwin/Flash-Attention-2_for_Windows

Quick Install Command

pip install https://huggingface.co/ussoewwin/Flash-Attention-2_for_Windows/resolve/main/<matching_wheel_file>.whl

6. Verification and Validation Guide

Detailed accuracy matrix, functional sweep (7 API paths), backward pass gradients, and latency benchmarks on NVIDIA GeForce RTX 5060 Ti (sm_120) are fully documented in:
👉 md/2.9.2.post2_COMPLETE_TEST_AND_VALIDATION_GUIDE.md

Assets 2
Loading

Fork Release v1.8 — PyTorch 2.14.0+cu132 Compatibility & Upstream Sync

Fork Release v1.8 — PyTorch 2.14.0+cu132 Compatibility & Upstream Sync

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 04 Sep 05:28
v1.8
9b45d73

Overview

Fork release v1.8 provides verified compatibility for PyTorch 2.14.0+cu132 under Windows MSVC, resolves linker symbol mismatches (LNK2001 / LNK1120), synchronizes the latest official upstream improvements (ce088ab / FA4 CuTe DSL), and strictly maintains all fork-specific architectural features (24 split_align kernels, MSVC 2GB COFF limit bypass via 3:3 SM split, PEP 427 wheel naming, and complete CI workflow suppression).

1. PyTorch 2.14.0+cu132 C++ ABI Compatibility & Linker Fix

Root Cause of MSVC LNK1120 / LNK2001

When building against PyTorch 2.14.0 under MSVC, link.exe could fail with exit status 1120:

flash_api.obj : error LNK2001: unresolved external symbol "__declspec(dllimport) void __cdecl c10::cuda::c10_cuda_check_implementation(int,char const *,char const *,unsigned int,bool)" (__imp_?c10_cuda_check_implementation@cuda@c10@@YAXHPEBD0I_N@Z)
  Hint on matching symbol:
    "__declspec(dllimport) void __cdecl c10::cuda::c10_cuda_check_implementation(int,char const *,char const *,unsigned int,bool,class c10::cuda::CUDAErrorLogCapture *)" (__imp_?c10_cuda_check_implementation@cuda@c10@@YAXHPEBD0I_NPEAVCUDAErrorLogCapture@12@@Z)
build\lib.win-amd64-cpython-314\flash_attn_2_cuda.cp314-win_amd64.pyd : fatal error LNK1120: 1 unresolved externals

Technical Breakdown

  • Signature Expansion in PyTorch 2.14.0:
    Starting in PyTorch 2.14.0, CUDA error handling was enhanced with c10::cuda::CUDAErrorLogCapture. The function signature of c10_cuda_check_implementation in c10/cuda/CUDAException.h expanded from 5 parameters to 6 parameters: 
    // PyTorch 2.14.0+ signature:
C10_CUDA_API void c10_cuda_check_implementation(
    const int32_t err,
    const char* filename,
    const char* function_name,
    const uint32_t line_number,
    const bool include_device_assertions,
    CUDAErrorLogCapture* error_log = nullptr);
  • ABI Mismatch Mechanism:
    PyTorch 2.14.0's c10_cuda.lib exports only the 6-parameter mangled symbol (?c10_cuda_check_implementation...PEAVCUDAErrorLogCapture...). If an existing flash_api.obj compiled under PyTorch 2.13.0 headers remained in the build tree, ninja skipped recompilation because flash_api.cpp had no source diff, causing the linker to request the obsolete 5-parameter symbol.
  • Resolution:
    Purging intermediate build artifacts (build/ directory) ensures flash_api.cpp compiles freshly against PyTorch 2.14.0's header, correctly emitting references to the 6-parameter symbol with default nullptr and resolving cleanly without any external symbol gaps.

2. Upstream Sync (Dao-AILab/flash-attention @ ce088ab)

The latest changes from the authoritative upstream repository were cleanly merged without touching any FlashAttention-2 code:

  • SM100 Scalar Mask Compilation Speedup (Dao-AILab#2819):
    Replaced static Python loop expansion in flash_attn/cute/mask.py with dynamic CuTeDSL unrolling for 128-element scalar mask-mod loops, reducing cold first-call compilation latency by ~31%.
  • CUTLASS DSL Requirement Relaxation (Dao-AILab#2798):
    Updated pyproject.toml dependency floor to nvidia-cutlass-dsl>=4.6.2 to support newer QuACK releases and CUTLASS DSL 4.7.
  • QuACK Packed Subtraction Compatibility (Dao-AILab#2787):
    Adapted flash_attn/cute/flash_bwd_sm100.py and utils.py to invoke CuTe DSL primitives directly following QuACK's deprecation of sub_packed_f32x2.

3. Preserved Fork Architecture & Windows Optimizations

All fork-specific optimizations remain 100% active and untouched:

  • 24 split_align Forward Kernels:
    Independent translation units (flash_fwd_split_align_*.cu) for num_splits == 1 sequences eliminate ptxas compiler timeouts.
  • MSVC 2GB COFF Linker Limit Bypass (3:3 SM Split):
    • WindowsWhlBuilder_cuda_2.bat: Blackwell targets (100;120;121).
    • WindowsWhlBuilder_cuda_3.bat: Legacy targets (80;89;90), with MAX_JOBS=5.
  • PEP 427 Compliant Wheel Naming:
    Uses .blackwell and .legacy dot-suffixes to avoid pip Invalid build number errors.
  • Dynamic Runtime Versioning:
    flash_attn/__init__.py uses importlib.metadata to dynamically reflect package version (2.9.1 vs 2.9.2).
  • CI Workflow Suppression:
    .github and tools/ci remain purged from the repository.

4. Pre-built Windows Wheels

Pre-compiled wheels are published on Hugging Face:
👉 ussoewwin/Flash-Attention-2_for_Windows

Quick Install Command

pip install https://huggingface.co/ussoewwin/Flash-Attention-2_for_Windows/resolve/main/<matching_wheel_file>.whl
Loading

v1.7 / flash_attn 2.9.2.post1

v1.7 / flash_attn 2.9.2.post1

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 12 Aug 03:03
v1.7
6513475

Fork tag: v1.7
Package: flash_attn 2.9.2.post1
Date: 2026-08-12
Tag commit: 65134752eca23e52bf902946a9bb461931c320fd

This is a patch on the 2.9.2 split_align line (fork v1.6). It does not add new public APIs. It changes the FA2 forward softmax rescale skip (A-1) to the FA4 16-bit threshold with exact max-lagging, and it allows the compiler to schedule the A-2 packed FMA (Delta 4).

Windows FA2 remains the production path for this fork (PyTorch 2.10+ / CUDA 13+ / MSVC). The Blackwell / legacy wheel split from v1.6 is unchanged.

1. Why 2.9.2.post1 exists

Layer Meaning
v1.6 / 2.9.2 Official split_align merge: 24 flash_fwd_split_align_*.cu units; bat2 Blackwell (100;120;121) vs bat3 Ampere/Hopper (80;89;90) to stay under MSVC LNK1189 (2 GB linker).
v1.7 / 2.9.2.post1 Same split_align + wheel split. A-1: float RescaleThreshold (call sites 8.0f) instead of bool / -0.01f skip; on skip, restore row_max to the previous max so O, row_sum, and P share one base (exact output and LSE). A-2: fma_f32x2 uses asm( (not asm volatile() so fma.rn.f32x2 may be scheduled.

1.1 Kernel deltas in this release

Delta Change Files
Delta 1 bool Use_rescale_threshold → float RescaleThreshold=0.0f; skip when scaled_diff >= -RescaleThreshold; call sites pass 8.0f csrc/flash_attn/src/softmax.h, csrc/flash_attn/src/flash_fwd_kernel.h
Delta 4 asm volatile( → asm( in fma_f32x2 csrc/flash_attn/src/utils.h
Delta 5 Comment: documents RescaleThreshold scope csrc/flash_attn/src/softmax.h
Delta 6 FA4-style max lagging: on skip, row_max(mi) = scores_max_prev(mi) csrc/flash_attn/src/softmax.h

1.2 A-1 — rescale threshold and max lagging

Aspect 2.9.0 / 2.9.1 (original A-1) 2.9.2.post1 (this release)
Template bool Use_rescale_threshold float RescaleThreshold=0.0f
Threshold kSoftmaxRescaleSkipThreshold = -0.01f 8.0f at Is_first=false call sites (FA4 16-bit default)
Skip condition scaled_diff >= -0.01 scaled_diff >= -8.0
Skip semantics Keep the new max; approximate (~0.7% error at 0.01) Revert row_max to the previous max; running O, row_sum, and P stay on one base (exact)

Why 8.0f: FA4 on Blackwell uses rescale_threshold = 8.0 when Q is 16-bit (q_dtype.width == 16), else 0.0.

Why Delta 6 is required: Skipping rescale but keeping the new max is bounded (~0.7%) at threshold 0.01. At 8.0, the same pattern can overweight stale block contributions by up to ~256×. Max lagging restores row_max(mi) = scores_max_prev(mi) on skip so output and LSE stay exact.

P saturation: With max lagging, P can reach exp2(8) = 256 before fp16/bf16 conversion. fp16 max is 65504; bf16 max is 3.39e38. No overflow at this threshold.

1.3 A-2 — packed FMA scheduling (Delta 4)

On sm_100+, scale_apply_exp2 uses fma_f32x2 (fma.rn.f32x2). Removing volatile lets the compiler schedule that instruction.

On sm_120, ptxas may still lower fma.rn.f32x2 to scalar FFMA (see §4.5). Accuracy and determinism tests still pass.

2. What did not change

  • Public APIs: flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache, window, ALiBi, softcap.
  • v1.6 split_align (24 compilation units) and the bat2 / bat3 arch split.
  • FA3 (hopper/) and FA4 (flash_attn/cute/) Windows runtime: out of scope for this wheel line (same as v1.6).
  • No new FP8 path (A-3). No extra public kwargs.

3. Version and Windows wheels

Package version is 2.9.2.post1 (import flash_attn / pip show flash-attn). flash_attn/__init__.py reads installed metadata and falls back to "2.9.2.post1". setup2.py sets public_version = "2.9.2.post1".

Script Archs Wheel suffix
WindowsWhlBuilder_cuda_2.bat (setup2.py) 100;120;121 .blackwell
WindowsWhlBuilder_cuda_3.bat 80;89;90 .legacy

Validated wheel (this release):

flash_attn-2.9.2.post1+cu132torch2.13.0cxx11abiTRUE.blackwell-cp314-cp314-win_amd64.whl
PYD size 1,313,640,960 bytes (1.3 GB).

4. Validation (measured)

Hardware / software (one environment):

Item Value
GPU NVIDIA GeForce RTX 5060 Ti (sm_120, capability 12.0)
CUDA (build) 13.2
PyTorch 2.13.0+cu132
Python 3.14.6
OS Windows 11 (10.0.26200)
Installed package flash-attn 2.9.2.post1

Reference for accuracy: PyTorch scaled_dot_product_attention in fp32.
Gates: fp16 rel_Linf <= 2e-2, bf16 rel_Linf <= 1e-2, cosine >= 0.9995.

4.1 Seven public API paths (fp16, B=2, S=512, H=8, D=64)

# Test NaN Inf Verdict
1 flash_attn_func (base, non-causal) False False PASS
2 flash_attn_func (causal) False False PASS
3 flash_attn_varlen_func False False PASS
4 flash_attn_with_kvcache False False PASS
5 flash_attn_func window_size=(32,32) False False PASS
6 flash_attn_func alibi_slopes False False PASS
7 flash_attn_func softcap=30.0 False False PASS

flash_attn_triton: not tested (Triton not installed in that venv).
flash_attn.cute: not tested (Windows FA4 runtime out of scope).

4.2 Accuracy — 8-case sweep (S=512–2048, D=64)

dtype S D causal rel_Linf cosine Verdict
fp16 512 64 True 2.95e-04 0.99999994 PASS
fp16 1024 64 True 2.75e-04 1.00000000 PASS
fp16 2048 64 True 5.81e-04 1.00000000 PASS
bf16 512 64 True 2.37e-03 0.99999809 PASS
bf16 1024 64 True 2.48e-03 0.99999815 PASS
bf16 2048 64 True 1.92e-03 0.99999815 PASS
fp16 512 64 False 7.89e-04 0.99999994 PASS
bf16 512 64 False 6.54e-03 0.99999774 PASS

4.3 Accuracy — 10-shape sweep (D=128, D=256, S=8192, B=1)

dtype S D B causal max_abs rel_Linf cosine Verdict
fp16 1024 64 2 True 1.22e-04 2.45e-04 0.99999994 PASS
fp16 1024 64 2 False 7.63e-06 6.23e-04 0.99999988 PASS
fp16 1024 128 2 True 1.22e-04 3.42e-04 0.99999988 PASS
fp16 1024 128 2 False 7.63e-06 7.49e-04 0.99999994 PASS
bf16 4096 64 2 True 9.77e-04 2.58e-03 0.99999809 PASS
bf16 4096 64 2 False 3.05e-05 4.41e-03 0.99999768 PASS
bf16 4096 128 2 True 9.77e-04 3.36e-03 0.99999815 PASS
bf16 4096 128 2 False 3.05e-05 5.88e-03 0.99999785 PASS
fp16 2048 256 1 False 3.81e-06 7.94e-04 1.00000000 PASS
fp16 8192 64 2 True 1.22e-04 3.02e-04 0.99999994 PASS

Combined 18 cases: fp16 12/12, bf16 6/6.
Worst fp16 rel_Linf = 7.94e-04. Worst bf16 rel_Linf = 6.54e-03. Both inside gates.

4.4 Backward

Finiteness (fp16/bf16, S=2048, causal): dQ, dK, dV all finite.

Gradient vs fp32 (fp16, S=1024, causal):

Gradient max_abs_diff vs fp32
dQ 2.10e-05
dK 1.86e-05
dV 1.47e-03

dV is larger because V gradients accumulate softmax-weighted values. All finite; training path is usable on this config.

Varlen backward (fp16): dQ/dK/dV finite (no NaN/Inf). Ranges: dQ [-1.06, 1.44], dK [-1.82, 2.47], dV [0.00, 1.88].

4.5 Latency (fp16 causal, 50 warmup + 50 timing)

S D ms/iter
1024 64 0.079
1024 128 0.151
4096 64 0.839
4096 128 1.725
8192 64 3.162
8192 128 6.504

Scales as expected (quadratic in S, linear in D). No catastrophic regression vs the 2.9.0 sm_120 baseline on this machine.

4.6 Edge / feature probes

Test Result
return_attn_probs=True (LSE) shape [2, 8, 512], finite, range [-0.0244, 6.2399]
Determinism (two identical calls) max abs diff 0.00e+00, bit-identical (Delta 4 did not add nondeterminism)
S=1 finite
S=2 vs fp32 max_abs_diff 1.22e-04
H=1 vs fp32 max_abs_diff 1.22e-04
D=256 bf16 rel_Linf 3.73e-03, cosine 0.99999821, PASS
softcap + causal finite
window_size=(0,0) vs V max_abs_diff 0.00e+00 (output == V)
window_size=(128,128) finite

4.7 SASS on sm_120 (96 cubins)

Pattern Count
FFMA.X2 not found
FFMA (scalar) 5,760
FMUL.FTZ 12,854
MUFU.EX2 6,862

fma.rn.f32x2 PTX was lowered to scalar FFMA on this sm_120 / CUDA 13.2 build. Correctness is covered by the tables above, not by FFMA.X2 presence.

5. Not in this release

Item Status
Triton (flash_attn_triton) Not tested (module not installed)
CuTe / FA4 (flash_attn.cute) Windows runtime out of scope
FP8 path (A-3) Not implemented
Multi-GPU / multi-node Not in this validation
A-2-off latency ablation Not measured (A-2-on binary only)

This is an unofficial fork build. Use at your own risk. The numbers above are from one Windows + RTX 5060 Ti + CUDA 13.2 + torch 2.13.0+cu132 environment.

6. Commits (kernel + docs on the tag)

Commit Role
5b068de A-1: rescale threshold + FA4-style max lagging (Deltas 1/5/6)
9abf358 A-2: remove volatile from fma_f32x2 (Delta 4)
b6dad32 version 2.9.2.post1
`a293...
Read more
Loading

Fork Release v1.6 (fa2.9.2)

Fork Release v1.6 (fa2.9.2)

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 13 Jul 02:31
v1.6
99651e9

Summary: Fork release v1.6 / package flash_attn 2.9.2

  • Official merge of the split_align architecture.
  • Extracted the num_splits == 1 forward pass alignment kernels into 24 independent compilation units (flash_fwd_split_align_*.cu) to resolve NVIDIA ptxas compiler timeouts and serial compilation crashes.
  • To bypass the Windows MSVC 2GB linker limit (LNK1189) caused by compiling these massive object files, the build scripts were physically split into bat2 (Blackwell: 100;120;121) and bat3 (Ampere/Hopper: 80;89;90).
  • Dynamic runtime versioning implemented in __init__.py to correctly branch between 2.9.1 and 2.9.2.

Technical Specification: split_align_kernels_explanation.md

Loading

v1.5 - Explanation of MSVC Compilation Fixes for PyTorch 2.13.0+cu13.2

v1.5 - Explanation of MSVC Compilation Fixes for PyTorch 2.13.0+cu13.2

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 09 Jul 18:44
v1.5
fab466e

1. Error Details

During the build process, two distinct sets of compilation errors were encountered:

① Overload Resolution Ambiguity in c10::HeaderOnlyArrayRef<int64_t>::operator == (C2666)

When comparing tensor sizes (e.g., IntArrayRef equality checks) in PyTorch headers such as torch/nn/functional/activation.h and loss.h, the compiler failed to resolve which operator== to invoke.

D:\USERFILES\fp8e4m3\venv\Lib\site-packages\torch\include\torch\csrc\api\include\torch/nn/functional/activation.h(674): error C2666: 'c10::HeaderOnlyArrayRef<int64_t>::operator ==': overloaded functions have similar conversions
D:\USERFILES\fp8e4m3\venv\Lib\site-packages\torch\include\torch/headeronly/util/HeaderOnlyArrayRef.h(251): note: could be 'bool c10::HeaderOnlyArrayRef<int64_t>::operator ==(c10::HeaderOnlyArrayRef<int64_t>,c10::HeaderOnlyArrayRef<int64_t>)' [found using argument-dependent lookup]
D:\USERFILES\fp8e4m3\venv\Lib\site-packages\torch\include\c10/util/OptionalArrayRef.h(223): note: or 'bool c10::OptionalArrayRef<int64_t>::operator ==(c10::OptionalArrayRef<int64_t>,c10::ArrayRef<int64_t>)' [found using argument-dependent lookup]
D:\USERFILES\fp8e4m3\venv\Lib\site-packages\torch\include\ATen/core/ivalue.h(284): note: or 'bool c10::operator ==(const c10::IValue &,const c10::IValue &)' [found using argument-dependent lookup]

② nvcc (CUDA Compiler) Rejection of C++20 Syntax

The CUDA compiler rejected modern C++20 syntax elements used in PyTorch's headers, such as designated initializers and default member initializers for bitfields.

D:/USERFILES/fp8e4m3/venv/Lib/site-packages/torch/include\c10/util/StringUtil.h(169): error: expected an expression
      return {.function = function, .file = file, .line = line};
              ^

D:/USERFILES/fp8e4m3/venv/Lib/site-packages/torch/include\c10/core/AutogradState.h(89): error: data member initializer is not allowed
    bool view_replay_enabled_ : 1 = false;
                                  ^

2. Root Cause Analysis

① Root Cause of the C2666 Error

Starting with PyTorch 2.13.0, a refactoring was introduced where ArrayRef<T> inherits from HeaderOnlyArrayRef<T>.
However, operator== was only defined in the base class HeaderOnlyArrayRef, leaving the derived class ArrayRef without its own explicit operator== definition.

When comparing two ArrayRef<int64_t> objects, MSVC's ADL (Argument-Dependent Lookup) detects multiple potential overloads:

  1. HeaderOnlyArrayRef's operator== (requiring implicit conversion of both sides to their base class).
  2. OptionalArrayRef's operator== (requiring implicit construction of one side).
  3. IValue's operator== (requiring implicit construction of both sides to IValue).

While compilers like GCC and Clang silently resolve this overload, MSVC strictly enforces C++ standard overload resolution rules. Since all candidates require some form of implicit user-defined conversion of equal rank, MSVC flags the call as ambiguous (C2666).

② Root Cause of the C++20 Syntax Error

PyTorch 2.13.0's C++ headers internally require C++20 features. However, FlashAttention's setup.py explicitly appended -std=c++17 (and /std:c++17 on Windows) to the compiler flags.

Because the -std=c++17 flag was appended after PyTorch's internal -std=c++20 flag, the compiler fallback to C++17 mode. This caused compiling errors when parsing C++20 features:

  • Designated initializers: {.function = ...}
  • Default bitfield member initializers: bool view_replay_enabled_ : 1 = false;

3. Resolution

① Overview

  1. Patching PyTorch Headers (ArrayRef.h):
    Add explicit exact-match operator== and operator!= function templates for c10::ArrayRef<T> inside namespace c10. This eliminates the need for any implicit conversions, resolving MSVC's overload ambiguity.
  2. Updating FlashAttention Build Configuration (setup.py):
    Change the default C++ standard flag from C++17 to C++20 to align with PyTorch's requirements.

② Patched Code (Full Diff)

1) Patch applied to ArrayRef.h

Modified c10/util/ArrayRef.h:L166-L188:

       std::initializer_list<U>) = delete;
 
   /// @}
 };
 
+/// MSVC C2666 workaround: namespace-scope exact-match operator== for ArrayRef.
+/// MSVC fails to find hidden friends in derived class templates via ADL,
+/// causing ambiguity with HeaderOnlyArrayRef, OptionalArrayRef, IValue, and
+/// SymbolicShape overloads. Namespace-scope functions are reliably found via
+/// ADL and, as exact matches, unambiguously win overload resolution.
+template <typename T>
+inline bool operator==(ArrayRef<T> lhs, ArrayRef<T> rhs) {
+  return lhs.equals(rhs);
+}
+template <typename T>
+inline bool operator!=(ArrayRef<T> lhs, ArrayRef<T> rhs) {
+  return !lhs.equals(rhs);
+}
+
 /// Deduction guides for ArrayRef to support CTAD with inherited constructors
 /// These mirror the constructors inherited from HeaderOnlyArrayRef
 /// @{

2) Patch applied to setup.py

Modified setup.py:L320-L355:

     nvcc_flags = [
     "-O3",
-    "-std=c++17",
+    "-std=c++20",
     "-U__CUDA_NO_HALF_OPERATORS__",
     "-U__CUDA_NO_HALF_CONVERSIONS__",
     "-U__CUDA_NO_HALF2_OPERATORS__",
     "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
     "--expt-relaxed-constexpr",
     "--expt-extended-lambda",
     "--use_fast_math",
     # "--ptxas-options=-v",
     # "--ptxas-options=-O2",
     # "-lineinfo",
     # "-DFLASHATTENTION_DISABLE_BACKWARD",
     # "-DFLASHATTENTION_DISABLE_DROPOUT",
     # "-DFLASHATTENTION_DISABLE_ALIBI",
     # "-DFLASHATTENTION_DISABLE_SOFTCAP",
     # "-DFLASHATTENTION_DISABLE_UNEVEN_K",
     # "-DFLASHATTENTION_DISABLE_LOCAL",
     ]
 
-    compiler_c17_flag=["-O3", "-std=c++17"]
+    compiler_c17_flag=["-O3", "-std=c++20"]
     # Add Windows-specific flags
     if sys.platform == "win32" and os.getenv('DISTUTILS_USE_SDK') == '1':
         # CUDA 13.2 CCCL headers require MSVC conforming preprocessor (see md/CUDA_13.0_TO_13.2_BUILD_FIX.md).
         nvcc_flags.extend(
             ["-Xcompiler", "/Zc:__cplusplus", "-Xcompiler", "/Zc:preprocessor"]
         )
         compiler_c17_flag = [
             "-O2",
-            "/std:c++17",
+            "/std:c++20",
             "/Zc:__cplusplus",
             "/Zc:preprocessor",
         ]

③ Technical Explanation

Significance of the ArrayRef.h Patch

In C++ overload resolution rules, an Exact Match (requiring no implicit conversions) always takes absolute precedence over candidates that require implicit user-defined conversions.
By defining the operator== template in the c10 namespace scope, the compiler can directly match comparisons of two ArrayRef<T> objects (such as IntArrayRef) without attempting to cast them to their base class (HeaderOnlyArrayRef) or wrapping them into another type. This prevents MSVC's ADL from flagging the comparison as ambiguous.

Significance of the setup.py Patch

Configuring the host compiler and nvcc to use c++20 standard options ensures that the entire compilation unit runs in C++20 mode. This enables native support for modern features used by PyTorch's backend, such as designated initializers ({.function = ...}) and in-class bitfield initializers, allowing compilation to proceed without syntax errors.

Loading

v1.4 - FlashAttention Fork: 2.9.0 → 2.9.1 Release Notes

v1.4 - FlashAttention Fork: 2.9.0 → 2.9.1 Release Notes

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 23 May 22:57
v1.4
139430a

Package: flash_attn (this fork)
Upstream baseline: FlashAttention-2 feature line at version 2.9.0
This release: 2.9.1 (patch release — same features, corrected packaging metadata)

This document describes the meaning, files changed, full modified source, and behavioral impact of the 2.9.0 → 2.9.1 correction on this fork. It covers the A-1 (rescale-threshold skip) and A-2 (packed fma.rn.f32x2 in scale_apply_exp2) refinements that justify the 2.9.0 feature label, plus the version bump to 2.9.1.

Table of contents

  1. Why 2.9.1 exists
  2. Modified files
  3. Package version (flash_attn/__init__.py)
  4. CUDA: csrc/flash_attn/src/softmax.h
  5. CUDA: csrc/flash_attn/src/utils.h
  6. Tooling: bench/check_sass_gates.py
  7. Call sites (unchanged, for context)
  8. Build, install, and verification
  9. References

1. Why 2.9.1 exists

1.1 Fork versioning vs upstream FlashAttention

Layer Meaning
Upstream FA2 Independent project (e.g. 2.8.x on PyPI upstream).
This fork’s feature line Tracks 2.9.0 when A-1 (rescale-threshold skip in forward softmax) and A-2 (packed FMA in scale_apply_exp2 on sm_100+) are present in the built kernels.
2.9.1 Patch release on that line: no new features, only a version string bump so wheels, pip show, and import flash_attn report 2.9.1 after a rebuild.

The kernel code for A-1/A-2 was already integrated before this bump. 2.9.1 does not change numerical behavior relative to the tree at commit time; it changes __version__ only (unless you rebuild wheels, in which case the binary reflects whatever is currently compiled).

1.2 What A-1 and A-2 do (summary)

Plan Name Purpose
A-1 Rescale-threshold skip When updating online softmax across K blocks, if the change in row max is tiny (scaled_diff >= -0.01 in log₂ scale), skip rescaling O and row_sum — error bounded by ~0.7% on normalized output.
A-2 Packed FMA in scale_apply_exp2 On sm_100+ (Hopper / Blackwell), compute two exp2(scale*t - max_scaled) lanes per step using one fma.rn.f32x2 instead of scalar fmaf + exp2f.

Refinement tasks (R1–R6, A1-R1–R3) are documented in AI/A1_A2_REFINEMENTS_PLAN.md and md/FA2_CHANGES_v1.2.md.

1.3 What 2.9.1 changes vs 2.9.0 (this release)

Item Changed?
flash_attn/__init__.py __version__ Yes → "2.9.1"
setup.py / wheel build scripts No (user constraint: do not edit build program without explicit order)
csrc/flash_attn/src/softmax.h Already contains A-1/A-2 refinements
csrc/flash_attn/src/utils.h Already contains A-2 fma_f32x2 + UNFUSE_FMA message
bench/check_sass_gates.py New validation script (A-2-R5)

2. Modified files

File Role in 2.9.1
flash_attn/__init__.py Authoritative package version read by setup.py via get_package_version()
csrc/flash_attn/src/softmax.h A-1 threshold constant, compile-time guards, scale_apply_exp2 packed path, comments
csrc/flash_attn/src/utils.h fma_f32x2 helper + UNFUSE_FMA compile-time warning
bench/check_sass_gates.py SASS gate: FFMA.X2 inside scale_apply_exp2 on sm_120 builds

Not modified for this release: setup.py, WindowsWhlBuilder_cuda.bat, CUDA arch lists, TORCH_CUDA_ARCH_LIST, or other build configuration.

3. Package version (flash_attn/__init__.py)

3.1 Meaning

setup.py does not hardcode the version. It reads __version__ from this file:

def get_package_version():
    with open(Path(this_dir) / "flash_attn" / "__init__.py", "r") as f:
        version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
    public_version = ast.literal_eval(version_match.group(1))
    ...
    return str(public_version)

So pip install and import flash_attn; flash_attn.__version__ report 2.9.1 after install. Optional env FLASH_ATTN_LOCAL_VERSION appends a local suffix (e.g. 2.9.1+dev).

Wheel filenames may still contain 2.9.0 until you run bdist_wheel again; the installed metadata uses __init__.py.

3.2 Full file (modified lines only)

__copyright__ = "Copyright (c) 2023, Tri Dao"
__version__ = "2.9.1"

from flash_attn.flash_attn_interface import (
    flash_attn_func,
    flash_attn_kvpacked_func,
    flash_attn_qkvpacked_func,
    flash_attn_varlen_func,
    flash_attn_varlen_kvpacked_func,
    flash_attn_varlen_qkvpacked_func,
    flash_attn_with_kvcache,
)

4. CUDA: csrc/flash_attn/src/softmax.h

This file is the single source of truth for FA2 forward softmax on this fork. Below: full text of every function/region touched by A-1/A-2 refinements (not the entire 470-line file).

4.1 scale_apply_exp2 (A-2 core + comment A-2-R6)

Purpose: After computing per-row max in log₂ space, apply exp2(scale * t - max_scaled) to score tiles before the final softmax normalization. On sm_100+, pairs of columns use fma.rn.f32x2 via fma_f32x2, then scalar exp2f.

// Apply the exp to all the elements.
// Note: this function is shared by:
//   - fwd via Softmax::softmax_rescale_o (both Is_first branches), and
//   - bwd via flash_bwd_kernel.h::compute_dq_dk_dv (Scale_max=false branch).
// Both paths benefit equally from the sm_100+ packed-FMA inner loop below.
template <bool Scale_max=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
__forceinline__ __device__ void scale_apply_exp2(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> const &max, const float scale) {
    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
    CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
    #pragma unroll
    for (int mi = 0; mi < size<0>(tensor); ++mi) {
        // If max is -inf, then all elements must have been -inf (possibly due to masking).
        // We don't want (-inf - (-inf)) since that would give NaN.
        // If we don't have float around M_LOG2E the multiplication is done in fp64.
        const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));
#if __CUDA_ARCH__ >= 1000 && !defined(UNFUSE_FMA)
        // Plan A-2: on sm_100+ (Blackwell), pair-process columns via fma.rn.f32x2.
        // The pre-exp2f term is fma(t, scale, -max_scaled); fma.rn.f32x2 does two
        // such FMAs in one instruction. exp2f itself stays scalar (MUFU.EX2 has
        // no f32x2 form). Rounding mode .rn matches scalar fmaf default.
        const float neg_max_scaled = -max_scaled;
        constexpr int N1 = decltype(size<1>(tensor))::value;
        static_assert(N1 % 2 == 0,
                      "scale_apply_exp2 packed-FMA path assumes N1 is even; "
                      "if a new MMA atom produces odd N1, restore the scalar tail loop.");
        #pragma unroll
        for (int ni = 0; ni < N1; ni += 2) {
            float t0 = tensor(mi, ni);
            float t1 = tensor(mi, ni + 1);
            float r0, r1;
            fma_f32x2(r0, r1, t0, t1, scale, scale, neg_max_scaled, neg_max_scaled);
            tensor(mi, ni)     = exp2f(r0);
            tensor(mi, ni + 1) = exp2f(r1);
        }
#else
        #pragma unroll
        for (int ni = 0; ni < size<1>(tensor); ++ni)  {
            // Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
            // max * log_2(e)) This allows the compiler to use the ffma
            // instruction instead of fadd and fmul separately.
            // The following macro will disable the use of fma.
            // See: https://github.com/pytorch/pytorch/issues/121558 for more details
            // This macro is set in PyTorch and not FlashAttention
            #ifdef UNFUSE_FMA
                tensor(mi, ni) = exp2f(__fmul_rn(tensor(mi, ni), scale) - max_scaled);
            #else
                tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
            #endif
        }
#endif
    }
}

Line-by-line (sm_100+ path):

Lines Meaning
67–69 Documents that forward (softmax_rescale_o) and backward (flash_bwd_kernel.h, Scale_max=false) both call this helper.
81–99 __CUDA_ARCH__ >= 1000 && !UNFUSE_FMA: use fma_f32x2 + paired exp2f; else scalar loop with optional UNFUSE_FMA PyTorch workaround.
88–90 static_assert(N1 % 2 == 0) — odd N1 is a compile error, not a silent scalar tail (A-2-R2).
91–98 Loop ni += 2: two FMAs + two exp2f per iteration.
100–114 Pre-sm_100 fallback: same math, scalar fmaf / exp2f; UNFUSE_FMA forces unfused path.

fma_f32x2 was removed from this file and lives in utils.h (A-2-R3).

4.2 Threshold constant and softmax_rescale_o (A-1)

// Threshold below which softmax_rescale_o skips the O / row_sum rescale.
// scores_scale = exp2(scaled_diff); scaled_diff is the negative excursion of
// row_max from one iteration to the next, in units of softmax_scale_log2.
// At -0.01f, scores_scale >= ~0.993 (worst-case relative error <= 0.7%).
inline constexpr float kSoftmaxRescaleSkipThreshold = -0.01f;
    template<bool Is_first, bool Check_inf=false, bool Use_rescale_threshold=false, typename Tensor0...
Read more
Loading

CUDA 13.0 → 13.2 Windows wheel build: errors, root causes, and code fixes

CUDA 13.0 → 13.2 Windows wheel build: errors, root causes, and code fixes

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 15 May 12:05
v1.3
b11cf5e

This document records how FlashAttention 2.9.0 (fork v1.2) was built on Windows when moving from a CUDA 13.0 PyTorch stack to CUDA 13.2 (torch 2.12.0+cu132). It is written for anyone reproducing the build or auditing why setup.py and docs changed.

Successful artifact (reference):

  • Wheel: dist/flash_attn-2.9.0+cu132torch2.12.0cxx11abiTRUE-cp313-cp313-win_amd64.whl
  • Built: 2026-05-15 19:39:32 (local time)
  • GPU tested: NVIDIA GeForce RTX 5060 Ti (sm_120)
  • Post-build tests: tests/test_a2_smoke.py (8/8 accuracy OK), tests/test_triton_smoke.py (JIT OK) — see md/2.9.0_COMPLETE_TEST_AND_VALIDATION_GUIDE.md §10

Git commits that implement the arch-policy fixes:

Commit Subject
33c55ee Drop sm_110 from default CUDA arch list and build gencodes
cf2f241 Document and enforce CUDA arch policy 80;90;100;120

1. Environment: what changed from cu130 to cu132

Item cu130 path (earlier) cu132 path (this port)
PyTorch e.g. 2.11.0+cu130 2.12.0+cu132
CUDA toolkit (CUDA_HOME) C:\...\CUDA\v13.0 C:\...\CUDA\v13.2
Wheel filename tag cu130torch2.11.0... cu132torch2.12.0...
Target GPU (this machine) — RTX 50-series → sm_120
Build logs _build_cu130*.log (if any) _build_cu132*.log, _build_cu132_zcp.log

Build command used (venv with cu132 torch):

set CUDA_HOME=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.2
set PATH=%CUDA_HOME%\bin;%PATH%
set MAX_JOBS=4
python setup.py build_ext --inplace

Canonical Windows wheel script (sets DISTUTILS_USE_SDK=1):

WindowsWhlBuilder_cuda.bat

Optional arch override (must match fork policy):

WindowsWhlBuilder_cuda.bat CUDA_ARCH 80;90;100;120

2. First errors that appeared (verbatim from build logs)

When CUDA_HOME pointed at 13.2 but setup.py still used the pre-fix default arch list 80;90;100;110;120 and Thor gencode logic, the first failing translation unit in log _build_cu132_zcp.log was flash_bwd_hdim128_bf16_causal_sm80.cu. Two independent failures showed up in the same ninja step.

2.1 Error A — CCCL / MSVC preprocessor (CUDA 13.2)

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.2\bin/../include/cccl\cuda/std/__cccl/preprocessor.h(23): fatal error C1189: #error:  MSVC/cl.exe with traditional preprocessor is used. This may lead to unexpected compilation errors. Please switch to the standard conforming preprocessor by passing `/Zc:preprocessor` to cl.exe. You can define CCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING to suppress this warning.

When it happens: nvcc compiles a .cu file and invokes MSVC cl.exe for host-side phases. CUDA 13.2 ships stricter CCCL headers under include/cccl/. Those headers reject the legacy (“traditional”) MSVC preprocessor.

What did not fix it: Ignoring the warning via CCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING only hides the diagnostic; the build still needs a conforming preprocessor for real compilation.

2.2 Error B — spurious sm_110 gencode (Thor) on this fork

The same nvcc command line included Thor family gencode (before commit 33c55ee):

-gencode arch=compute_80,code=sm_80
-gencode arch=compute_90,code=sm_90
-gencode arch=compute_100f,code=sm_100
-gencode arch=compute_120f,code=sm_120
-gencode arch=compute_110f,code=sm_110
-gencode arch=compute_120,code=compute_120

Follow-on failure:

fatal   : Could not open input file C:/Users/ussoe/AppData/Local/Temp/tmpxft_000070dc_00000000-25_flash_bwd_hdim128_bf16_causal_sm80.compute_110f.cpp1.ii
ninja: build stopped: subcommand failed.
RuntimeError: Error compiling objects for extension

Meaning: The fork’s FA2 CUDA sources are instantiated for sm_80 / sm_90 / sm_100 / sm_120 (Ampere, Hopper, Blackwell consumer/datacenter). There is no flash_*_sm110 (or Thor-specific) kernel variant in this tree. Asking nvcc to run the sm_80 template path while also passing -gencode ... sm_110 makes the toolchain generate intermediate files such as *.compute_110f.cpp1.ii that do not correspond to a supported kernel build → “Could not open input file”.

3. Root cause (essential)

# Symptom Root cause
A __cccl/preprocessor.h C1189 CUDA 13.2 + CCCL requires MSVC /Zc:preprocessor. PyTorch’s default nvcc/cl flags on Windows did not pass it. setup.py only added /Zc:__cplusplus when DISTUTILS_USE_SDK=1, which is insufficient for CCCL on 13.2.
B compute_110f.cpp1.ii missing Default FLASH_ATTN_CUDA_ARCHS included 110, and add_cuda_gencodes() emitted -gencode arch=compute_110f,code=sm_110 for toolkit ≥ 13.0. This fork does not build Thor (sm_101/sm_110) FA2 kernels; multi-arch wheels should be 80;90;100;120 only.

Both issues are orthogonal: fixing only A still leaves B; fixing only B still leaves A. A reproducible cu132 Windows build needs both.

4. Files modified for the 13.2 port

File Role
setup.py Arch defaults, Thor stripping, add_cuda_gencodes(), Windows MSVC flags
md/FA2_CHANGES_v1.2.md Arch policy table + Windows /Zc:preprocessor note
md/CHANGELOG.md v1.2 summary: default arch list
md/2.9.0_COMPLETE_TEST_AND_VALIDATION_GUIDE.md Env table: FLASH_ATTN_CUDA_ARCHS
WindowsWhlBuilder_cuda.bat Comment/default example 80;90;100;120
md/CUDA_13.0_TO_13.2_BUILD_FIX.md This document

Kernel / FA2 feature code (softmax.h, launch templates, etc.) is unchanged for the toolkit bump; only build policy and MSVC host flags changed.

5. Code changes (full text of modified sections)

5.1 setup.py — constants and cuda_archs() (after cf2f241)

Purpose: Single source of truth for which SASS targets this fork emits; strip Thor (101, 110) and unknown tokens with warnings; default env string 80;90;100;120.

# FA2 fork wheel build: SASS for Ampere/Hopper/Blackwell only (see md/FA2_CHANGES_v1.2.md).
FORK_SUPPORTED_CUDA_ARCHS = ("80", "90", "100", "120")
FORK_THOR_CUDA_ARCHS = frozenset({"101", "110"})


@functools.lru_cache(maxsize=None)
def cuda_archs() -> list[str]:
    raw = os.getenv("FLASH_ATTN_CUDA_ARCHS", "80;90;100;120")
    requested = [a.strip() for a in raw.split(";") if a.strip()]
    dropped_thor = [a for a in requested if a in FORK_THOR_CUDA_ARCHS]
    if dropped_thor:
        warnings.warn(
            "FLASH_ATTN_CUDA_ARCHS includes Thor GPU arch(es) "
            f"{dropped_thor}; this fork FA2 build emits SASS for "
            f"{list(FORK_SUPPORTED_CUDA_ARCHS)} only. Ignoring those entries.",
            stacklevel=2,
        )
    archs = [a for a in requested if a in FORK_SUPPORTED_CUDA_ARCHS]
    unknown = [
        a
        for a in requested
        if a not in FORK_SUPPORTED_CUDA_ARCHS and a not in FORK_THOR_CUDA_ARCHS
    ]
    if unknown:
        warnings.warn(
            f"FLASH_ATTN_CUDA_ARCHS entries ignored (not built by this fork): {unknown}",
            stacklevel=2,
        )
    if not archs:
        warnings.warn(
            "FLASH_ATTN_CUDA_ARCHS has no supported entries after filtering; "
            f"using default {list(FORK_SUPPORTED_CUDA_ARCHS)}.",
            stacklevel=2,
        )
        archs = list(FORK_SUPPORTED_CUDA_ARCHS)
    return archs

Before (33c55ee parent):

@functools.lru_cache(maxsize=None)
def cuda_archs() -> str:
    return os.getenv("FLASH_ATTN_CUDA_ARCHS", "80;90;100;110;120").split(";")

Meaning:

  • 110 in the default list caused every clean build to request Thor gencode on CUDA 13+ toolchains → Error B.
  • Returning a filtered list[str] lets add_cuda_gencodes() only see supported arches; user typos and Thor IDs fail loudly via warnings.warn instead of breaking ninja mid-compile.

5.2 setup.py — add_cuda_gencodes() (after 33c55ee + cf2f241)

Purpose: Map 80/90/100/120 to correct -gencode lines for CUDA 12.8+ / 12.9+ (100f, 120f on ≥12.9). No Thor branch.

def add_cuda_gencodes(cc_flag, archs, bare_metal_version):
    """
    Adds -gencode flags for this fork's supported CUDA arch list only.

    Requested arch tokens (after cuda_archs() filtering) map to nvcc targets as:
      - 80  -> compute_80, sm_80
      - 90  -> compute_90, sm_90  (CUDA >= 11.8)
      - 100 -> compute_100f, sm_100 on CUDA >= 12.9 else compute_100, sm_100 (CUDA >= 12.8)
      - 120 -> compute_120f, sm_120 on CUDA >= 12.9 else compute_120, sm_120 (CUDA >= 12.8)

    Thor / sm_101 / sm_110 are not built (see FORK_THOR_CUDA_ARCHS in cuda_archs()).
    PTX for the newest numeric arch is embedded for forward-compatible JIT.
    """
    # Always-regular 80
    if "80" in archs:
        cc_flag += ["-gencode", "arch=compute_80,code=sm_80"]

    # Hopper 9.0 needs >= 11.8
    if bare_metal_version >= Version("11.8") and "90" in archs:
        cc_flag += ["-gencode", "arch=compute_90,code=sm_90"]

    # Blackwell 10.x requires >= 12.8
    if bare_metal_version >= Version("12.8"):
        if "100" in archs:
            # CUDA 12.9 introduced "family-specific" for Blackwell (100f)
            if bare_metal_version >= Version("12.9"):
                cc_flag += ["-gencode", "arch=compute_100f,code=sm_100"]
            else:
                cc_flag += ["-gencode", "arch=compute_100,code=sm_100"]

        if "120" in archs:
            # sm_120 is supported in CUDA 12.8/12.9+ toolkits
            if bar...
Read more
Loading

FA2 2.9.0 Kernel Changes in v1.2

FA2 2.9.0 Kernel Changes in v1.2

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 15 May 11:50
v1.2
cf2f241

This document explains every code change applied to the FA2 (CUDA) path in this fork for fork release v1.2 (Python package flash_attn 2.9.0). The upstream repository has shifted focus to FA3 (Hopper) and FA4 (CuTeDSL); this fork continues FA2 development for sm_80+ on Windows with PyTorch 2.10+ / 2.12+ and CUDA 13+.

Files modified

File Purpose
csrc/flash_attn/src/softmax.h Core softmax math: rescale logic and exp2 scaling
csrc/flash_attn/src/flash_fwd_kernel.h Kernel dispatch: adds template flags to call sites
flash_attn/__init__.py Version marker for the fork feature line

Scope (this document): forward FA2 CUDA kernels only. Other v2.9.0 fork changes (Triton fix, split-KV launch countermeasure, build scripts) are documented in 2.9.0_COMPLETE_TEST_AND_VALIDATION_GUIDE.md.

1. A-1 — Rescale threshold skip (softmax_rescale_o)

What it does

When a new block’s row-max is virtually identical to the running row-max, the rescale factor exp2(scaled_diff) is approximately 1.0. Skipping the rescale saves one exp2 and N multiplies per row.

Why it is safe

The threshold (-0.01f) corresponds to scores_scale >= ~0.993, i.e. a worst-case relative error below 0.7%. The current block still uses the correct row_max for its own scale_apply_exp2; only the running O/row_sum rescale is approximated.

Forward-only

softmax_rescale_o is used on the forward path only. Backward kernels (flash_bwd_*.h) do not call this helper; A-1 does not affect backward numerics or build.

Changed code

csrc/flash_attn/src/softmax.h — Softmax::softmax_rescale_o signature

// Before:
template<bool Is_first, bool Check_inf=false, typename Tensor0, typename Tensor1>

// After (adds Use_rescale_threshold):
template<bool Is_first, bool Check_inf=false, bool Use_rescale_threshold=false, typename Tensor0, typename Tensor1>

csrc/flash_attn/src/softmax.h — rescale logic body

float scaled_diff = (scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2;
// Optionally skip the O / row_sum rescale when the new row_max is virtually
// the same as the previous one (scaled_diff is a very small negative number,
// so scores_scale = exp2(scaled_diff) ~= 1.0). The threshold -0.01 corresponds
// to scores_scale >= ~0.993 (worst-case relative error <= 0.7%).
if constexpr (Use_rescale_threshold) {
    constexpr float kRescaleSkipThreshold = -0.01f;
    if (scaled_diff >= kRescaleSkipThreshold) { continue; }
}
float scores_scale = exp2f(scaled_diff);
row_sum(mi) *= scores_scale;
// ... acc_o_rowcol(mi, ni) *= scores_scale;

2. A-2 — Packed FMA via fma.rn.f32x2 (scale_apply_exp2)

What it does

On sm_100 and sm_120 (and any arch with __CUDA_ARCH__ >= 1000, i.e. Blackwell-class), the kernel uses inline PTX fma.rn.f32x2 to compute two pre-exp2f FMA terms in one instruction. On older architectures (sm_80, sm_90, …) the helper falls back to two plain fmaf calls. exp2f remains scalar (no f32x2 MUFU form).

New helper — fma_f32x2

// Packed FMA helper (Plan A-2): computes (d0,d1) = (a0*b0+c0, a1*b1+c1).
// On sm_100 / sm_120+ (__CUDA_ARCH__ >= 1000), emits a single fma.rn.f32x2 instruction.
// Falls back to two plain FMAs on older arches and when UNFUSE_FMA is set.
__forceinline__ __device__ void fma_f32x2(
    float &d0, float &d1,
    float a0, float a1,
    float b0, float b1,
    float c0, float c1) {
#if __CUDA_ARCH__ >= 1000 && !defined(UNFUSE_FMA)
    asm volatile(
        "{\n\t"
        ".reg .b64 ra, rb, rc, rd;\n\t"
        "mov.b64 ra, {%2, %3};\n\t"
        "mov.b64 rb, {%4, %5};\n\t"
        "mov.b64 rc, {%6, %7};\n\t"
        "fma.rn.f32x2 rd, ra, rb, rc;\n\t"
        "mov.b64 {%0, %1}, rd;\n\t"
        "}\n"
        : "=f"(d0), "=f"(d1)
        : "f"(a0), "f"(a1), "f"(b0), "f"(b1), "f"(c0), "f"(c1));
#else
    d0 = fmaf(a0, b0, c0);
    d1 = fmaf(a1, b1, c1);
#endif
}

Changed code — scale_apply_exp2 (sm_100+ path)

const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));
#if __CUDA_ARCH__ >= 1000 && !defined(UNFUSE_FMA)
    const float neg_max_scaled = -max_scaled;
    constexpr int N1 = decltype(size<1>(tensor))::value;
    #pragma unroll
    for (int ni = 0; ni < N1 - 1; ni += 2) {
        float t0 = tensor(mi, ni);
        float t1 = tensor(mi, ni + 1);
        float r0, r1;
        fma_f32x2(r0, r1, t0, t1, scale, scale, neg_max_scaled, neg_max_scaled);
        tensor(mi, ni)     = exp2f(r0);
        tensor(mi, ni + 1) = exp2f(r1);
    }
    if constexpr (N1 % 2 != 0) {
        constexpr int last = N1 - 1;
        tensor(mi, last) = exp2f(tensor(mi, last) * scale - max_scaled);
    }
#else
    // Original single-element loop for sm_80 / sm_90 and when UNFUSE_FMA is set.
    ...
#endif

3. Kernel dispatch changes (flash_fwd_kernel.h)

Four call sites pass /*Use_rescale_threshold=*/true to softmax_rescale_o. All four are Is_first=false only — the masking_step == 0 branches keep Is_first=true and leave Use_rescale_threshold at its default false.

# Function Line (approx.) Loop Check_inf when threshold enabled
1 compute_attn_1rowblock 344 masking, masking_step > 0 Is_causal || Is_local
2 compute_attn_1rowblock 407 no masking on S Is_local
3 compute_attn_1rowblock_splitkv 918 masking, masking_step > 0 Is_causal || Is_local || !Is_even_MN
4 compute_attn_1rowblock_splitkv 985 no masking on S Is_local

Example (site 1 — standard forward):

// Before:
softmax.template softmax_rescale_o</*Is_first=*/false, /*Check_inf=*/Is_causal || Is_local>(...)

// After:
softmax.template softmax_rescale_o</*Is_first=*/false, /*Check_inf=*/Is_causal || Is_local, /*Use_rescale_threshold=*/true>(...)

Example (site 3 — split-KV masking branch; note extra !Is_even_MN on Check_inf):

masking_step == 0
    ? softmax.template softmax_rescale_o</*Is_first=*/true,  /*Check_inf=*/Is_causal || Is_local || !Is_even_MN>(...)
    : softmax.template softmax_rescale_o</*Is_first=*/false, /*Check_inf=*/Is_causal || Is_local || !Is_even_MN, /*Use_rescale_threshold=*/true>(...);

Related documentation

4. Version bump

flash_attn/__init__.py:

-__version__ = "2.8.4"
+__version__ = "2.9.0"

This marks the fork feature line for v1.2. The upstream package may continue its own numbering; this fork’s 2.9.0 only indicates that A-1 and A-2 are present in this tree.

Build / compatibility notes

  • PyTorch: >=2.10 required — enforced in setup.py install_requires as torch>=2.10 on the CUDA wheel path. Extension code uses <torch/extension.h> (PyTorch 2.10+ layout). Wheels and local builds are commonly tested with 2.12+cu132; the validation guide also documents runs on 2.11.0+cu130.
  • CUDA: Toolkit >=13.0 for native compilation; 13.2 is used for cu132 PyTorch builds. Default FLASH_ATTN_CUDA_ARCHS in setup.py: 80;90;100;110;120 (Ampere through Blackwell).
  • Windows: Supported. FA4 (CuTeDSL) remains unavailable on Windows due to missing win_amd64 native libraries. MSVC host compiles invoked by nvcc need /Zc:preprocessor (see setup.py when DISTUTILS_USE_SDK=1).
  • Fallback: A-1 is gated by the Use_rescale_threshold template flag at each call site. A-2 is gated by #if __CUDA_ARCH__ >= 1000 inside scale_apply_exp2 / fma_f32x2, so one binary runs on Ampere, Hopper, and Blackwell.
Loading

Flash Attention Fork v2.8.4 — Complete Changelog from v2.8.3

Flash Attention Fork v2.8.4 — Complete Changelog from v2.8.3

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 14 May 14:40
v1.1
5ea3401

1. Overview

Aspect Value
From v2.8.3 (tag v1.0, commit f680e35)
To v2.8.4 (commit 3038f78)
What happened Merged upstream main (9cee95f...484b981), then applied fork-specific reverts/fixes

2. Fork-Specific Change: Revert num_splits == 1 Branch in FA2 Launch Template

2.1 Background

During the upstream merge, commit b322ae2 ("expose num_splits for FA2 and add option for kernel blocksize alignment") introduced a branch in run_mha_fwd_splitkv_dispatch:

When num_splits == 1, use a smaller kBlockN_standard (128 for headdim<=64, 64 otherwise) to match the standard non-splitKV kernel's block size.

This was intended to give bitwise-identical numerics between the split-KV and standard paths.

2.2 Why It Was Reverted in This Fork

The num_splits == 1 branch instantiates a separate Flash_fwd_kernel_traits template with a different kBlockN. On this fork's target configuration (sm_80 Ampere + cutlass 4.3 + CUDA 13.x), this extra specialization caused ptxas to enter an infinite loop or crash during SASS generation for specific cases (hdim=32, fp16, causal).

2.3 Full Code Change

File: csrc/flash_attn/src/flash_fwd_launch_template.h

Before (v2.8.3):

template<typename T, int Headdim, bool Is_causal>
void run_mha_fwd_splitkv_dispatch(Flash_fwd_params &params, cudaStream_t stream) {
    constexpr static int kBlockM = 64;
    constexpr static int kBlockN = Headdim <= 64 ? 256 : (Headdim <= 128 ? 128 : 64);
    run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN, 4, false, false, T>, Is_causal>(params, stream);
}

After (v2.8.4):

template<typename T, int Headdim, bool Is_causal>
void run_mha_fwd_splitkv_dispatch(Flash_fwd_params &params, cudaStream_t stream) {
    constexpr static int kBlockM = 64;
    constexpr static int kBlockN = Headdim <= 64 ? 256 : (Headdim <= 128 ? 128 : 64);
    run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN, 4, false, false, T>, Is_causal>(params, stream);
}

What changed: The if (params.num_splits == 1) branch that upstream added between v2.8.3 and v2.8.4 was removed. The function now keeps the unified single-path behavior from v2.8.3.

Why this fixes the build: Removing the branch eliminates the extra kernel specialization that triggered the ptxas/register-layout conflict on sm_80.

2.4 Impact

  • num_splits=1 cases no longer get the direct-to-O optimization (they use O_partial like all other cases).
  • Build completes successfully on sm_80 + cutlass 4.3 + CUDA 13.x.

3. Submodule Update: CUTLASS v4.0 → v4.3.0

v2.8.3 v2.8.4
Commit dc481792 71275920 (v4.3.0-48-g71275920)
CUTLASS_MINOR 0 3
CUTLASS_PATCH — 4

Required by the upstream FA4/CuTeDSL code merged in this release. FA2 kernels in this fork do not directly use CuTeDSL but share the cutlass build environment.

4. Documentation

  • README.md — Restored fork-specific content after upstream merge overwrote it.
  • md/CHANGELOG.md — Added fork release notes.

5. Version Marker

# flash_attn/__init__.py
__version__ = "2.8.4"

Scope: Changes from this fork's v2.8.3 (f680e35) to v2.8.4 (3038f78).

Loading