Skip to content

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
· 145 commits to main since this release
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.
Assets 2
Loading