Skip to content

flash-attn-hopper: bump to 00756db for TE 2.17 window kwargs - #10

Merged
yueming-yuan merged 1 commit into
mainfrom
fa3-te-2.17-window-kwargs
Jul 26, 2026
Merged

flash-attn-hopper: bump to 00756db for TE 2.17 window kwargs#10
yueming-yuan merged 1 commit into
mainfrom
fa3-te-2.17-window-kwargs

Conversation

@yueming-yuan

Copy link
Copy Markdown
Owner

Problem

miles bumped transformer_engine to 2.17.0 (radixark/miles#1781). TE 2.17 moved
use_flash_attn_3 out of the window_size-tuple branch in
cp_p2p_fwd_flash_attn / cp_p2p_bwd_flash_attn:

# TE 2.12
if use_flash_attn_3 or (v2_3_plus and not v2_7_0_plus):
    fa_forward_kwargs["window_size"] = (-1, -1)        # FA3 took the tuple
elif v2_7_0_plus:
    fa_forward_kwargs["window_size_left"] = -1

# TE 2.17
if v2_3_plus and not v2_7_0_plus:
    fa_forward_kwargs["window_size"] = (-1, -1)
elif use_flash_attn_3 or v2_7_0_plus:                  # FA3 moved here
    fa_forward_kwargs["window_size_left"] = -1

The pinned revision fbf24f67 builds flash_attn_3 3.0.0b1, whose
_flash_attn_forward only accepts a single window_size tuple, so every
context-parallel path raises TypeError. There is no FA3-only kill switch in
TE — NVTE_FLASH_ATTN disables FA2 and FA3 together.

FA2 in the same image (2.7.4.post1) already takes the new kwargs, so only the
FA3 path is affected.

Change

  • Bump the hopper/ pin to 00756db, which builds 3.0.0.
  • Force a source build. Without it setup.py prints
    Guessing wheel URL: .../v3.0.0/flash_attn_3-3.0.0+cu122torch2.11...whl and
    downloads a prebuilt wheel built against a different CUDA/torch than the image.
  • Drop the separate flash_attn_interface.py fetch in test_wheels.py. It was a
    second pin that can drift from the one the extension module was built at. The
    wheel now ships the interface top-level and a flash_attn_3.flash_attn_interface
    re-export shim, so nothing extra is needed — and copying the full module over
    that shim would re-run @torch.library.custom_op("flash_attn_3::...") and
    double-register the ops.

Verification

Built and installed on an h200-sci devbox running the radixark/miles:pr-1795
image itself (torch 2.11.0+cu130, CUDA 13.0), so the wheel is compiled against
the toolchain it will ship with.

[INFO     | DotProductAttention]: Running with FlashAttention backend (version 3.0.0)
[PASS] all 5 symbols TE imports exist
[PASS] _flash_attn_forward schema takes window_size_left/right  ['SymInt window_size_left=-1', 'SymInt window_size_right=-1']
[PASS] _flash_attn_backward schema takes window_size_left/right  ['SymInt window_size_left=-1', 'SymInt window_size_right=-1']
[PASS] import transformer_engine.pytorch
[PASS] TE sees flash-attn-3 == 3.0.0  fa3_version=3.0.0 v3_0_0_beta=False
[PASS] TE bound a v3 forward
[PASS] TE's bound v3 fwd is the same op we checked
[PASS] fwd accepts window_size_left/right kwargs  out=(2, 512, 8, 128) lse=(2, 8, 512)
[PASS] returns (out, softmax_lse, ...) in the order TE indexes [0]/[1]
[PASS] TE DotProductAttention fwd+bwd on FlashAttention backend  out=(2, 512, 1024) grad_finite=True
[PASS] output matches SDPA reference  max_abs_err=0.001953
[PASS] FA3 registered as installed in TE
       fa2=2.7.4.post1 fa3=3.0.0 v3_installed=True

RESULT: ALL PASS

Notes on things checked so they are not rediscovered later:

  • The five symbols are load-bearing. TE guards only PackageNotFoundError
    around the version lookup, not ImportError around the imports themselves, so
    a wheel missing any of flash_attn_func, flash_attn_varlen_func,
    flash_attn_with_kvcache, _flash_attn_forward, _flash_attn_backward makes
    import transformer_engine.pytorch raise rather than degrade to cuDNN.
  • 3.0.0b1 -> 3.0.0 is benign. It flips
    FlashAttentionUtils.v3_0_0_beta to False; the only consumer
    (backends.py:1230) uses it to append an "update your flash-attn v3 beta"
    hint to an already-raised TypeError. No functional branch depends on it.
  • Arch coverage is unchanged: sm_80 + sm_90a, verified with cuobjdump
    against both the old and new _C.abi3.so. hopper/setup.py derives arch from
    source-file suffixes (_sm90.cu -> sm_90a) and ignores the
    TORCH_CUDA_ARCH_LIST that _setup_env sets, so that variable has never had
    any effect on FA3. FA3 has no Blackwell kernels here either way.
  • Return shape is compatible. 3.0.0 returns
    (out, softmax_lse, out_accum, softmax_lse_accum); TE's CP code reads indices
    [0] and [1] under use_flash_attn_3, which line up.

transformer_engine 2.17 moved use_flash_attn_3 out of the window_size-tuple
branch in cp_p2p_fwd/bwd_flash_attn, so it now passes window_size_left/right
to FA3. The pinned revision builds 3.0.0b1, whose _flash_attn_forward only
accepts a single window_size tuple, so every context-parallel path raises
TypeError. TE has no FA3-only kill switch -- NVTE_FLASH_ATTN disables FA2 too.

00756db builds 3.0.0, whose forward and backward are torch.library custom ops
with SymInt window_size_left/right in their schemas, and which returns
(out, softmax_lse, out_accum, softmax_lse_accum) -- indices 0 and 1, exactly
what TE reads. Verified on an H200 against the miles image: TE selects
'FlashAttention backend (version 3.0.0)' and matches an SDPA reference to
2e-3. Same sm_80 + sm_90a coverage as before; hopper/setup.py hardcodes those
from source-file suffixes and ignores TORCH_CUDA_ARCH_LIST.

Also force a source build, since setup.py otherwise falls back to downloading
a prebuilt release wheel built against a different CUDA/torch, and drop the
separate flash_attn_interface.py fetch: the wheel now ships it top-level plus
a flash_attn_3.flash_attn_interface re-export shim, and overwriting that shim
with the full module double-registers the flash_attn_3:: custom ops.
@yueming-yuan
yueming-yuan merged commit 7207a38 into main Jul 26, 2026
@yueming-yuan
yueming-yuan deleted the fa3-te-2.17-window-kwargs branch July 26, 2026 07:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant