Skip to content

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
· 122 commits to main since this release
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, typename Tensor1>
    __forceinline__ __device__ void softmax_rescale_o(Tensor0 &acc_s, Tensor1 &acc_o, float softmax_scale_log2) {
        static_assert(!(Is_first && Use_rescale_threshold),
                      "Use_rescale_threshold has no effect when Is_first=true; "
                      "remove the flag from the Is_first=true call site.");
        // ... reshape scores, acc_o_rowcol ...
        } else {
            // ...
            #pragma unroll
            for (int mi = 0; mi < size(row_max); ++mi) {
                float scores_max_cur = !Check_inf
                    ? row_max(mi)
                    : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));
                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).
                // This skip does not compound across iterations: normalize_softmax_lse
                // divides acc_o by row_sum, so the skipped factor cancels between
                // numerator and denominator, bounding the relative error by the threshold.
                if constexpr (Use_rescale_threshold) {
                    if (scaled_diff >= kSoftmaxRescaleSkipThreshold) { continue; }
                }
                float scores_scale = exp2f(scaled_diff);
                row_sum(mi) *= scores_scale;
                #pragma unroll
                for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scores_scale; }
            }
            FLASH_NAMESPACE::scale_apply_exp2(scores, row_max, softmax_scale_log2);
            // ...
        }

A-1 behavior:

  • Use_rescale_threshold=true only on Is_first=false call sites (see §7).
  • When scaled_diff >= -0.01f, skip updating row_sum and acc_o for that row — saves work when max barely moves.
  • normalize_softmax_lse still divides by row_sum, so skipped mass cancels in numerator/denominator (no drift across blocks).

5. CUDA: csrc/flash_attn/src/utils.h

5.1 fma_f32x2 (full)

// Packed FMA helper (Plan A-2): computes (d0,d1) = (a0*b0+c0, a1*b1+c1).
// On sm_100+ (Blackwell), 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
}

PTX: single fma.rn.f32x2 with .rn (round-to-nearest-even), matching default fmaf rounding.

5.2 UNFUSE_FMA compile-time message (A-2-R4)

#if defined(UNFUSE_FMA) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
    #pragma message("UNFUSE_FMA is defined while compiling for sm_100+; " \
                    "the Blackwell-only fma.rn.f32x2 path in scale_apply_exp2 is disabled. " \
                    "This is the documented ablation path (see AI/FA2_BACKPORT_FROM_FA4_PLAN.md " \
                    "section 8.2 ablation row 'A-2 off' and section 10.2 rollback). " \
                    "Remove UNFUSE_FMA to re-enable the packed FMA path.")
#endif

Meaning: If you define UNFUSE_FMA (PyTorch workaround for issue #121558), you silently lose A-2 on sm_100+ unless you remove the macro. The message makes that visible at compile time.

6. Tooling: bench/check_sass_gates.py

Purpose: Enforce the Phase-1 exit criterion from AI/FA2_BACKPORT_FROM_FA4_PLAN.md §6: on a sm_120 build, FFMA.X2 must appear inside demangled scale_apply_exp2* symbols.

Full script:

#!/usr/bin/env python3
"""SASS hard-gate checker for Flash-Attention A-1/A-2 optimizations.

Parses ``cuobjdump --dump-sass`` output to verify the presence or absence
of specific SASS instructions within ``scale_apply_exp2`` symbols.

Exit codes:
    0  gate passes
    1  gate fails
    2  usage / environment error
"""
# ... (see repository file for complete implementation: argparse, wheel/pyd extraction,
# cuobjdump auto-detect, streaming parser, gates a2 and b1)

Gate a2: symbol matches scale_apply_exp2, instruction FFMA.X2, must be present.

Usage example:

python bench/check_sass_gates.py \
  --whl dist/flash_attn-2.9.0+cu132torch2.12.0-cp312-cp312-win_amd64.whl \
  --arch sm_120 \
  --gate a2

Manual procedure is also documented in md/2.9.0_COMPLETE_TEST_AND_VALIDATION_GUIDE.md §11.

7. Call sites (unchanged, for context)

A-1’s Use_rescale_threshold flag is wired only at these four template instantiations in flash_fwd_kernel.h:

Line (approx.) Is_first Use_rescale_threshold
~344 false true
~407 false true
~918 false true
~985 false true

Is_first=true paths use default Use_rescale_threshold=false (correct: no prior max to compare on first tile).

Backward uses scale_apply_exp2<Scale_max=false> in flash_bwd_kernel.h — benefits from A-2 on sm_100+ without a separate template flag.

8. Build, install, and verification

8.1 Version bump only (no rebuild)

If you only change __init__.py to 2.9.1 without rebuilding the extension:

  • import flash_attn → __version__ is 2.9.1
  • The .pyd / .so on disk is still the old binary (still 2.9.0 feature code if that is what was built)
  • Wheel filename may still say flash_attn-2.9.0+... until the next bdist_wheel

8.2 Rebuild wheel (recommended for consistent 2.9.1 artifacts)

Use your existing Windows workflow (WindowsWhlBuilder_cuda.bat or equivalent). Do not ask the assistant to edit setup.py unless you explicitly want arch or dependency changes.

After build:

  1. pip install --force-reinstall dist/flash_attn-2.9.1+...whl (or your naming convention)
  2. python -c "import flash_attn; print(flash_attn.__version__)" → 2.9.1
  3. python tests/test_a2_smoke.py (if GPU available)
  4. python bench/check_sass_gates.py --whl <wheel> --arch sm_120 --gate a2

8.3 What this release does not include

  • Thor sm_101 / sm_110 in default TORCH_CUDA_ARCH_LIST (separate user/setup decision; see commit 65cc9db history on this fork)
  • Changes to md/CHANGELOG.md or md/FA2_CHANGES_v1.2.md (reverted when user requested no unsolicited doc edits)
  • FA3 (hopper/) or FA4 (flash_attn/cute/) trees — out of scope for this FA2 fork doc

9. References

Document Path
A-1/A-2 refinement plan AI/A1_A2_REFINEMENTS_PLAN.md
FA2 backport plan (A-1/A-2 source) AI/FA2_BACKPORT_FROM_FA4_PLAN.md
Fork change log md/FA2_CHANGES_v1.2.md
Test & validation guide md/2.9.0_COMPLETE_TEST_AND_VALIDATION_GUIDE.md
SASS gate procedure md/2.9.0_COMPLETE_TEST_AND_VALIDATION_GUIDE.md §11

End of document. Package version for this feature line: 2.9.1 (flash_attn/__init__.py).

Assets 4
Loading