Skip to content

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
· 155 commits to main since this release
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).

Assets 2
Loading