ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image by Johnsonms · Pull Request #2636 · Dao-AILab/flash-attention

Johnsonms · 2026-06-10T00:19:13Z

Summary

This PR fixes a Blackwell FA4 CI failure caused by stale deps in the CI SIF image.

CI installs FA4 with uv pip install --no-deps, so the dependency floors in flash_attn/cute/pyproject.toml are not enforced at test time:

nvidia-cutlass-dsl >= 4.5.2
quack-kernels >= 0.5.0

A stale SIF can still contain older deps, for example nvidia-cutlass-dsl 4.4.2 and quack-kernels 0.3.7, which fail on the new AuxData JIT arg with:

DSLRuntimeError: failed to generate argument #34 (aux_data)
for JIT function '...FlashAttentionForwardSm100...'

Changes

Commit	Change
`39d177f`	Add `tools/ci/assert_dsl_floor.py` to read floors from `pyproject.toml` and fail early if the SIF deps are too old.
`bdaeaed`	Rebake `fa4_image_cu130` to `flash-attn-cu13.0-26.06.10` with `cutlass-dsl 4.5.2` and `quack-kernels 0.5.0`.

Why

--no-deps is intentional to avoid clobbering the SIF-baked torch / cudnn, but it also skips dependency floor enforcement.

We add a guard instead of upgrading at runtime because run_step uses a RAM-backed --writable-tmpfs overlay. Reinstalling cutlass-dsl there hits ENOSPC and can corrupt the baked environment.

Validation

Verified on a real B200 system with 8× SM100 GPUs:

Scenario	Result
Stale SIF: `cutlass-dsl 4.4.2` / `quack-kernels 0.3.7`	Reproduces the `aux_data` `DSLRuntimeError`
Rebaked image: `flash-attn-cu13.0-26.06.10`	`DSL floor check OK`; benchmark and tests pass

Notes

Floors are read from pyproject.toml, so versions do not drift.
fa4_image_cu129 is unchanged because CUDA 13 runners use cu130; rebake separately if a sub-CUDA-13 runner is added.

run_fa4_ci.py installs FA4 with --no-deps (to keep the SIF's baked torch/cudnn), so the nvidia-cutlass-dsl>=4.5.2 / quack-kernels>=0.5.0 floors in flash_attn/cute/pyproject.toml are not enforced at install time. A SIF baked before a floor bump keeps a stale dep — e.g. cutlass-dsl 4.4.2, which can't convert the AuxData JIT arg and dies with a cryptic DSLRuntimeError deep in SM100 kernel launch (reproduced on B200). Upgrading the dep in-place is not viable: the --writable-tmpfs overlay is RAM-backed and too small for a cutlass-dsl reinstall (ENOSPC, and a partial removal corrupts the baked torch). So instead of installing, add assert_dsl_floor.py — it reads the floors from pyproject (no hardcoded version to drift) and fails with an actionable "rebake the image" message when the installed cutlass-dsl/quack are below them. Wired into run_step right after the editable install. The durable fix is to rebake the image at the current floors and bump the digest in .github/workflows/ci.yml; this guard makes future drift fail fast instead of silently.

Copilot

Pull request overview

This PR addresses FA4 Blackwell (SM100) CI failures caused by stale nvidia-cutlass-dsl / quack-kernels versions baked into the CI Apptainer (SIF) image when FA4 is installed with uv pip install --no-deps.

Changes:

Add a CI-time dependency floor guard (tools/ci/assert_dsl_floor.py) that reads required minimum versions from flash_attn/cute/pyproject.toml and fails early with a rebake instruction if the image is stale.
Invoke the floor guard in tools/ci/run_fa4_ci.py for each CI step (fresh --writable-tmpfs overlay each time).
Update the CUDA 13.0 FA4 CI image digest/tag in .github/workflows/ci.yml to a rebaked image containing the newer deps.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
`tools/ci/run_fa4_ci.py`	Adds an early per-step dependency floor check after editable install to catch stale SIF deps before tests run.
`tools/ci/assert_dsl_floor.py`	New script that compares installed `nvidia-cutlass-dsl` / `quack-kernels` against floors from `flash_attn/cute/pyproject.toml`.
`.github/workflows/ci.yml`	Points FA4 cu130 CI to the rebaked image tag+digest.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Johnsonms · 2026-06-10T02:02:55Z

+import sys
+import tomllib
+from importlib.metadata import PackageNotFoundError, version
+
+from packaging.requirements import Requirement
+from packaging.version import Version


Good catch — addressed in 00e8933: tomllib now falls back to tomli on Python 3.10, and if neither is present the script exits with a clear "needs a TOML parser" message instead of a raw ModuleNotFoundError. (In CI the guard runs inside the cu130 SIF on Python 3.12, so it uses stdlib tomllib; the fallback only matters for local 3.10 runs — verified on a 3.10 host that it now fails predictably.)

drisspg · 2026-06-10T01:59:16Z

is there a way to have landed pyproject.toml changes trigger a SIF build?

drisspg · 2026-06-10T01:59:37Z

ehh thats a lil sketch actually this is fine I think

Johnsonms · 2026-06-10T02:07:59Z

@drisspg assert_dsl_floor.py is only the passive/detection half: it runs once at test time and fails loudly if the SIF is below the floor, but it doesn't build anything. What you're describing is the active half — auto-building a new SIF when a floor lands — which is a nice follow-up on top of this guard.

* [Fwd,Sm100] fix: decode↔prefill exp2 emulation consistency (Dao-AILab#2595) apply_exp2_convert selected the exp2 implementation based on mask_fn presence: hardware ex2.approx.ftz for causal-masked tiles, polynomial emulation for unmasked tiles. Different q_stage values (1 for decode, 2 for prefill) compute different m_block for the same logical Q row, shifting which tiles are processed with vs without mask_fn. The same K tile could receive different exp2 methods across variants. Fix: always pass self.ex2_emu_freq regardless of mask_fn presence. Add regression test for decode↔prefill bitwise consistency on MLA (192,128) shapes. * replace deprecated apis (Dao-AILab#2602) * Bump nvidia-cutlass-dsl to >=4.5.2 and quack-kernels to >=0.5.0 (Dao-AILab#2605) cutlass 4.5.2 is safe to update, and quack 0.5.0 has been published, so bump the FA4 (flash_attn/cute) requirement floors to match. Updates the dependencies and the cu13 extra in pyproject.toml, and the documented versions in CLAUDE.md. Verified on NVIDIA GB300 (SM100, CUDA 13.2): deps resolve cleanly (nvidia-cutlass-dsl 4.5.2 base+cu13, quack-kernels 0.5.0), imports OK, and a representative GPU sample of tests/cute/test_flash_attn.py passes (6 passed / 6 skipped / 0 failed across hd 64/96/128/192, causal, mha/gqa/mqa, fwd+bwd). * [CuTe,Fwd,Sm100] refactor mla sm100 forward and add page table (Dao-AILab#2558) * refactor mla sm100 forward * add benchmark; address deprecation warnings; tweak ptx gemm dispatch * update interface and tests * ci: bump Jimver/cuda-toolkit to v0.2.35 for CUDA 13.2 support (Dao-AILab#2617) v0.2.30 only ships URLs up to CUDA 13.1.0; bumping to v0.2.35 adds 13.1.1, 13.2.0, and the matching aarch64 SBSA installers. Signed-off-by: oliver könig <okoenig@nvidia.com> * [ROCm] Bump Triton to >=3.6.0 and aiter submodule (Dao-AILab#2614) * [Triton] Fix graph capture issues and env var (Dao-AILab#2620) * graph capture fix * rm env flag * [CuTe,Bwd,Sm100] allow 2cta with score mod and mask mod in bwd (Dao-AILab#2557) * [CuTe] Fix lint failures (Dao-AILab#2625) stack-info: PR: Dao-AILab#2625, branch: drisspg/stack/42 * [CuTe] Fix lint failure in flash_bwd_sm100.py (Dao-AILab#2627) ruff format flagged flash_attn/cute/flash_bwd_sm100.py (trailing whitespace in a comment and an over-split call). It was missed by the lint sweep in Dao-AILab#2625. * fix: add weights_only=True to all torch.load call sites (Dao-AILab#2622) Passing weights_only=False (the pre-2.4 default) to torch.load allows arbitrary Python object deserialization from the checkpoint file. A malicious .pt/.pth file can execute arbitrary code on the machine loading it — a well-known PyTorch deserialization vector (CWE-502). Four call sites updated: training/src/utils/checkpoint.py load_checkpoint() training/src/eval.py eval checkpoint loader flash_attn/utils/pretrained.py partial(torch.load, ...) loader flash_attn/models/llama.py state_dicts_from_checkpoint() weights_only=True restricts deserialization to tensors, dicts, lists, tuples, and other primitive types — no arbitrary Python objects. Requires PyTorch >= 1.13; FA4's CuTeDSL dependency already requires a modern PyTorch 2.x build, so no compatibility regression. Fixes Dao-AILab#2583 * use correction warps if not tma store; remove outdated packgqa guard (Dao-AILab#2629) * Add aux-scalars to interface to enable dynamic ints and floats in expressions (Dao-AILab#2616) stack-info: PR: Dao-AILab#2616, branch: drisspg/stack/41 * fix: build and select cu13.2 prebuilt wheels (Dao-AILab#2618) * ci: use 1 ninja job for cu13.2 Signed-off-by: oliver könig <okoenig@nvidia.com> * fix(setup): request cu13 prebuilt wheels for CUDA 13 torch get_wheel_url() binned every CUDA >= 12 to major '12', so under a CUDA 13 torch it requested cu12 wheels and never matched the published cu13 artifacts, falling back to a multi-hour source build. Add a CUDA 13 branch so the guessed wheel name uses cu13, matching WHEEL_CUDA_VERSION in _build.yml. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: oliver könig <okoenig@nvidia.com> --------- Signed-off-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image (Dao-AILab#2636) * ci(fa4): assert cute dep floors in CI; fail loudly on a stale SIF run_fa4_ci.py installs FA4 with --no-deps (to keep the SIF's baked torch/cudnn), so the nvidia-cutlass-dsl>=4.5.2 / quack-kernels>=0.5.0 floors in flash_attn/cute/pyproject.toml are not enforced at install time. A SIF baked before a floor bump keeps a stale dep — e.g. cutlass-dsl 4.4.2, which can't convert the AuxData JIT arg and dies with a cryptic DSLRuntimeError deep in SM100 kernel launch (reproduced on B200). Upgrading the dep in-place is not viable: the --writable-tmpfs overlay is RAM-backed and too small for a cutlass-dsl reinstall (ENOSPC, and a partial removal corrupts the baked torch). So instead of installing, add assert_dsl_floor.py — it reads the floors from pyproject (no hardcoded version to drift) and fails with an actionable "rebake the image" message when the installed cutlass-dsl/quack are below them. Wired into run_step right after the editable install. The durable fix is to rebake the image at the current floors and bump the digest in .github/workflows/ci.yml; this guard makes future drift fail fast instead of silently. * ci(fa4): bump cu130 image to 26.06.10 (cutlass-dsl 4.5.2 / quack 0.5.0) * ci(fa4): fall back to tomli when tomllib is unavailable (Python 3.10) * Fix SM100 FP8 fwd with cutlass-dsl >=4.5.2 (MmaF8F6F4Op) (Dao-AILab#2640) cutlass-dsl >=4.5.2 changed make_trivial_tiled_mma to build plain FP8 MMAs as MmaF8F6F4Op (its _F8F6F4_TYPES branch) instead of the now-legacy MmaFP8Op. The two are siblings under MmaOp, so _tcgen05_mma_kind's isinstance(op, MmaFP8Op) check missed the new type and raised "Unsupported tcgen05 MMA op kind: MmaF8F6F4Op", breaking the FP8 forward path on Blackwell. Worked on 4.4.2. Accept both ops in the f8f6f4 branch (both map to kind::f8f6f4). mma_op_to_idesc only reads generic op attrs and is unaffected. Validated on B200: FP8 fwd passes for all configs in the issue (incl. hd=64) plus hd=128, causal and non-causal; mean abs err vs bf16 ~0.002-0.01. Fixes Dao-AILab#2639 * [cute] Fix int32 overflow in SM100 LPT tile scheduler for long context (Dao-AILab#2662) The LPT tile scheduler sizes its L2 swizzle from seqlen_k * (headdim + headdim_v) * element_size in int32. For long context this overflows once it exceeds 2**31 (seqlen_k > ~4M for hdim-128 bf16), making size_one_head negative. That corrupts the swizzle and the L2 divmods, so get_current_work decodes an out-of-bounds batch_idx and the kernel performs an illegal memory access (cudaErrorIllegalAddress) on SM100. Compute the byte size in int64. swizzle stays small and is cast back to int32 for the device-side divmods, so there is no behavior or perf change for non-overflowing shapes. Fixes both SingleTileLPTScheduler (forward; selected for causal/local) and SingleTileLPTBwdScheduler (backward; its extra seqlen_k * headdim * 4 term overflows even sooner). Repro on SM100 (e.g. GB200), causal forward at seqlen_k = 2**22: import torch from flash_attn.cute.interface import flash_attn_func sq, sk = 2048, 4_194_304 # seqlen_k = 2**22 -> int32 overflow q = torch.randn(1, sq, 8, 128, dtype=torch.bfloat16, device="cuda") k = torch.randn(1, sk, 1, 128, dtype=torch.bfloat16, device="cuda") v = torch.randn(1, sk, 1, 128, dtype=torch.bfloat16, device="cuda") out = flash_attn_func(q, k, v, causal=True) torch.cuda.synchronize() # cudaErrorIllegalAddress here before the fix Crashes before this change, runs clean after; seqlen_k = 2**22 - 128 is clean both ways (the int32 boundary). Verified clean under compute-sanitizer memcheck. * [Fwd,Sm100] Tune FP8 causal hd128 ex2_emu_freq (8 vs inherited 16) (Dao-AILab#2642) FP8 fwd is MUFU/ex2-bound on Blackwell, so the optimal exp2-emulation frequency differs from bf16. The causal hd128 key (False,True,128,False) had no FP8 entry and inherited bf16's freq=16; freq=8 offloads more exp from the MUFU unit. Thermally-matched back-to-back A/B on B200 (locked-ish clock, hot GPU, median of 300 iters, nheads=16 = benchmark default) across the official benchmark's causal hd128 shapes: b s f16 TFLOP f8 TFLOP delta 32 512 500.6 516.3 +3.1% 16 1024 796.5 832.5 +4.5% 8 2048 1124.6 1175.1 +4.5% 4 4096 1407.9 1481.3 +5.2% 2 8192 1604.1 1661.7 +3.6% 1 16384 1683.7 1726.0 +2.5% Accuracy-neutral (FP8-vs-bf16 mean-abs-err unchanged; benchmark --check passes 24/24). Keyed on is_causal=True only: freq=8 would regress non-causal hd128 (0.94x), which keeps its existing freq=10. * Make q_subtile_factor default to identity (Dao-AILab#2660) * fix(hd256/sm100): make q/k/v contiguous before dedicated hd256 kernel (Dao-AILab#2666) The BlackwellFusedMultiHeadAttentionForward kernel builds tensor layouts with hardcoded contiguous strides computed from shape dimensions, so non-contiguous inputs (e.g. from .transpose()) cause wrong memory accesses and silently corrupt outputs on B200 (SM100) with head_dim=256. maybe_contiguous() only guarantees stride(-1)==1; add explicit full contiguity checks in both the forward and backward paths when the hd256 dedicated kernel is selected. Fixes: Dao-AILab#2665 * [Cute,Bwd,Sm100] add sparse MLA (Deepseek v4) backward kernels (Dao-AILab#2621) * add backward sparse mla kernels * add dk gemm * fix errors * fix dq errors * rename bwd kernels * refactor interface * fix predicate error in dq kernel * update tests * mla fwd fixes * improve varlen fwd perf * use cluster idx scheduling in fwd * use packed scheduler for mqa 128 * fix int32 overflow in swizzle * simplify bwd preprocess * refactor bwd * simplify preprocess * update benchmark script * add safety check * remove test code * ruff format * ensure scale is 0 for masked out rows * fix: sync callers with new _flash_attn_fwd 4-tuple return signature (Dao-AILab#2674) * Fix compatibility issues with CuTe DSL 4.6.0+ (Dao-AILab#2648) * Prepare for 4.6 release * Bump version * Update pyproject.toml * Update nvidia-cutlass-dsl version in pyproject.toml * Pass tmem scalar fields as .ptr to TmemAllocator on SM100 (Dao-AILab#2679) The DSL now warns when a struct scalar is used directly as a pointer ("Use explicit struct.scalar.ptr for pointer instead"), so these fire on every tmem_holding_buf / dealloc mbar access. Just pass .ptr like the other SM100 kernels already do. * Add FLASHATTENTION_DISABLE_SPLIT_ALIGNMENT (Dao-AILab#2680) * ci: rebake cu130 image for cutlass-dsl 4.6.0.dev0 floor (Dao-AILab#2684) PR Dao-AILab#2648 bumped the flash_attn/cute/pyproject.toml floor to nvidia-cutlass-dsl==4.6.0.dev0, but the CI image (26.06.10) still ships 4.5.2. assert_dsl_floor.py correctly fails every push to main with "installed 4.5.2 does not satisfy floor ==4.6.0.dev0", so FA4 CI has been red since Dao-AILab#2648 landed. - Dockerfile: add --prerelease=allow to the FA4 install. The dev-build floor pulls transitive pre-releases (nvidia-cutlass-dsl-libs-base== 4.6.0.dev0 ...) that uv refuses without it; the old stable 4.5.2 floor didn't need it. - ci.yml: bump fa4_image_cu130 to the rebaked 26.06.27 image (cutlass-dsl 4.6.0.dev0, quack-kernels 0.5.3, torch 2.12.1). E2e verified on B200: assert_dsl_floor passes, compile + run + benchmark all green (run_fa4_ci.py, exit 0). * Update FA4 cute quack compatibility (Dao-AILab#2676) * Update FA4 cute quack compatibility * Use quack 0.5.3 make_smem_layout instead of vendored copy Tri re-added the major_mode_size arg to quack.sm90_utils.make_smem_layout in quack 0.5.3 (commit 68888e2), so FA4 no longer needs the local sm90_layout helper. Revert the 4 backward call sites to quack's helper and bump the floor to >=0.5.3 (0.5.2 lacks the arg). --------- Co-authored-by: Johnsonms <lizhaofu@gmail.com> * ci: install cutlass-dsl/quack at runtime to decouple from the baked image (Dao-AILab#2685) * [Cute,Bwd,Sm100] Assume 16B stride divisibility for LSE/dPsum bulk-copy inputs (Dao-AILab#2686) The SM100 backward stats (LSE, dPsum) are loaded via cp.async.bulk (CopyBulkG2SOp), which - unlike cp.async.bulk.tensor - needs the source pointer alignment provable at compile time. After slicing, the newer cute-dsl can't deduce 16B alignment unless the input strides carry the divisibility assumption, so the bulk copy fails to compile on real tensors (the FakeTensor path masks it). - flash_bwd_mla_sm100.py: add mdPsum to the new_stride divisibility list (it already covered ScaleP and the other stats; mdPsum was omitted). - flash_bwd_sm100.py: the ordinary backward had no divisibility assumption at all; add it for both mLSE and mdPsum. Only these two SM100 kernels use CopyBulkG2SOp; the SM90/SM80/SM120 and MLA dK/dQ backward kernels use other copy paths and are unaffected. Addresses the dPsum stride-divisibility finding (Finding 1) in Dao-AILab#2677. * fix(hd256/sm100): forward reads actual input strides, drop .contiguous() patch (Dao-AILab#2670) * follow up to Dao-AILab#2666: fixing the layouts in the sm100 hd256 kernels and removing the temporary fix of calling .contiguous everywhere * respond to PR comments * respond to PR comments-2: move to utils file * Add tests --------- Co-authored-by: drisspg <drisspguessous@gmail.com> * ci: run MLA backward cases so CI exercises flash_bwd_mla_sm100.py (Dao-AILab#2690) FA4_TEST_FILTER selected no MLA test, so the MLA backward kernels (flash_bwd_mla_sm100.py + dq_dqv + dk) had zero CI coverage. Add four small test_flash_attn_mla_absorbed cases covering the distinct backward paths: sparse (kv_sparsity=True) non-causal and causal, dense (kv_sparsity=False), and shared_kv=True. The ordinary SM100 backward is already covered by the existing test_flash_attn_output cases. Cold-cache cost on B200 (full 8-case filter): pass-1 compile ~4:54, GPU run ~1:03 — well under the 60-min job timeout. Stacked on Dao-AILab#2685 (runtime cutlass-dsl/quack install). * Parallelize splitkv alignment templated kernels, remove flag (Dao-AILab#2683) * [FA3] uv installation support (Dao-AILab#2458) * Expose flash_attn_3 as package so imports work correctly. * Add flash_attn_config package shim and fix uv packaging details Builds on the flash_attn_3 package exposure so both import styles work for downstream frameworks and uv/pyproject.toml installs: - Add flash_attn_3/flash_attn_config.py re-export so `from flash_attn_3 import flash_attn_config` works (previously only the top-level module was importable), matching the interface shim. - Un-ignore the committed shim in .gitignore; the bare `flash_attn_config.py` pattern (for the build-time generated top-level file) also matched the package shim and would have silently dropped it from the commit. - Read flash_attn_3.__version__ from installed package metadata with a fallback, avoiding drift from setup.py's version source. - README: move `dependencies` under `[project]` so the uv snippet is valid PEP 621. Verified on H100 (SM90): editable `uv pip install -e .` now succeeds (fails on main), both `import flash_attn_interface` and `from flash_attn_3 import flash_attn_interface` resolve, `flash_attn_config` imports both ways, and fp16 hdim128 forward matches a torch reference (max_abs_err <= 2e-3). ruff check passes. --------- Co-authored-by: Johnsonms <lizhaofu@gmail.com> * [AMD ROCm] Enable RDNA backward and adopt CK unified workspace (Dao-AILab#2675) * Add sink_ptr/d_sink_ptr to fmha_bwd_args to match updated CK submodule Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * update submodule * [CK_TILE] Use Unified Workspace for FMHA BWD (#182) * [CK_TILE] Use Unified Workspace for FMHA BWD Bump composable_kernel submodule to mono-split/users/yiding12/fmha-bwd-workspace HEAD and adapt the FMHA BWD host wrappers to the new unified workspace API: - Replace dq_acc tensor argument with workspace_ptr in get_ck_fmha_bwd_args / get_ck_fmha_varlen_bwd_args - Drop dq_acc strides that have been removed from fmha_bwd_args - In mha_bwd / mha_varlen_bwd, allocate the device workspace based on fmha_bwd_launcher::workspace_size and call launcher.prepare_workspace() - Invoke launcher.run(args, stream_config) instead of fmha_bwd(...) * Update CK pin as ROCm/rocm-libraries#6152 merged * [CK_TILE] FMHA BWD: stream-async workspace prepare (#183) * [CK_TILE] FMHA BWD: stream-async workspace prepare Bump composable_kernel submodule to mono-split/users/yiding12/fmha-bwd- async-prepare HEAD and adapt the FMHA BWD host wrappers to the new async workspace prepare API (CK PR #7331): - Replace launcher.prepare_workspace() with prepare_workspace_async(), which enqueues the full workspace setup (dq_acc zero, group-mode D2H of seqstart, host-side metadata pack via hipLaunchHostFunc, H2D back to device) on the caller's stream. No host-blocking sync remains in the BWD launch path. - Pass a pinned_host_alloc lambda backed by PyTorch's CachingHostAllocator (torch::empty(..., pin_memory=true)). The launcher keeps the returned shared_ptr alive via a stream-tail hipLaunchHostFunc keepalive so the pinned buffer is not recycled while async copies are still in flight. - mha_varlen_bwd: drop the cu_seqlens_q.cpu() / cu_seqlens_k.cpu() host copies; the launcher now reads device seqstart directly via async D2H. get_ck_fmha_varlen_bwd_traits no longer takes seqstart_qs/ks. * [CK_TILE] FMHA BWD: bump CK submodule to develop tip (#7331 merged) ROCm/rocm-libraries#7331 (async workspace prepare for FMHA BWD launcher) landed on develop. Move csrc/composable_kernel from the pre-merge fork tip ce838e19e5 to ROCm/composable_kernel develop tip 83566edb0f, which is the split commit for #7331 (rocm-libraries 5692db0). * [CK_TILE] FMHA BWD: explicit at::kCPU on pinned host TensorOptions * Update CK and enable RDNA backward --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Yi DING <yi.ding@amd.com> Co-authored-by: Hosang Yoon <hosang.yoon@amd.com> * Fix CuTe SM120 compile-time argument handling (Dao-AILab#2671) * Fix CuTe SM120 compile-time argument handling * clean up * guard empty SM120 local backward tiles --------- Co-authored-by: Kevin-Li-2025 <2242139@qq.com> Co-authored-by: drisspg <drisspguessous@gmail.com> * [NVIDIA][CuTe,Fwd,sm120] Implement Pack-GQA on SM120 (+ graceful SplitKV fallback) (Dao-AILab#2656) * [CuTe,Fwd,sm120] Fix use_tma_O crash on SM120 (issue Dao-AILab#2649) On SM120 (Blackwell GeForce / RTX PRO 6000 / DGX Spark) the forward kernel set `use_tma_O = self.arch >= Arch.sm_90`, enabling the TMA-based O-store epilogue. But SM120 does not build the TMA store atom (tma_atom_O is None), so any forward call crashes in cpasync.tma_partition with: AttributeError: 'NoneType' object has no attribute '_trait' This makes the CuTe-DSL forward unusable on every SM120 GPU. Restrict the TMA O-store to sm_90..sm_119, which is where the WGMMA-era epilogue path is actually available: self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120 SM120 falls back to the non-TMA register->gmem O store (already used for the SM80 path), which is correct and what the CpAsync SM120 kernel expects. Verified on RTX PRO 6000 Blackwell (sm_120, cc 12.0), torch 2.12.0+cu130, nvidia-cutlass-dsl 4.5.2: forward now runs and matches PyTorch SDPA reference for hdim 64/96/128, causal and non-causal (max abs err <= 8e-3 in bf16). Before this fix every SM120 forward call raised the AttributeError above. * [CuTe,Fwd,sm120] Implement Pack-GQA on SM120; graceful SplitKV fallback Pack-GQA was only half-wired in the SM80/SM120 CpAsync forward: the epilogue referenced PackGQA.store_O/store_LSE, but the Q-load and head-indexing used the plain (unpacked) path. So pack_gqa=True crashed in pack_gqa.store_O (crd2idx on a packed (h_idx, m_idx) coordinate against an unpacked mO layout). This implements Pack-GQA end to end on SM120 (and SM80), mirroring the SM90 path: - Reshape mQ/mO (head_idx=2) and mLSE (head_idx=1) via pack_gqa_layout so qhead_per_kvhead folds into the seqlen mode ((qhead, seqlen)). - Scheduler args use cute.size(mQ.shape[0]) (packed total rows) and seqlen_q_static = mQ.shape[0][1] (logical seqlen), so causal/mask q_idx stay correct. - Kernel head-indexing: when pack_gqa, num_head from the scheduler already indexes the KV head (mQ/mK share nheads_kv); no division. - Q-load: gather rows via PackGQA.load_Q (per-row (h_idx, m_idx) gmem pointers) instead of the contiguous local_tile path. SplitKV (num_splits>1) is an SM100-only feature (SM80/SM90 also assert it unsupported); SM120 has no forward+combine path. Fall back to num_splits=1, which is numerically correct, instead of crashing in _check_type on the fp32 partials. Verified on RTX PRO 6000 Blackwell (sm_120): pack_gqa=True matches PyTorch SDPA GQA/MQA reference (err <= 8.4e-3 bf16) AND is bit-identical to the unpacked path (max |packed - unpacked| = 0.0) across MHA/GQA/MQA, causal/non-causal, hd 64/128, seqlen 512-2048. num_splits=3 falls back and matches reference (err 6.8e-4). Stacked on the SM120 use_tma_O fix (Dao-AILab#2649). * re-enable SM120 pack-gqa after rebase * clean up SM120 pack-gqa split handling * fix SM120 varlen pack-gqa offset --------- Co-authored-by: drisspg <drisspguessous@gmail.com> * Fix pre-commit Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> --------- Signed-off-by: oliver könig <okoenig@nvidia.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: 鐘天楽 <tianle.zhong@bytedance.com> Co-authored-by: brandonsun <brandons@nvidia.com> Co-authored-by: Johnsonms <lizhaofu@gmail.com> Co-authored-by: jayhshah <jayhshah@gmail.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Michael Melesse <micmelesse@gmail.com> Co-authored-by: Reuben Stern <107093092+reubenconducts@users.noreply.github.com> Co-authored-by: Driss Guessous <32754868+drisspg@users.noreply.github.com> Co-authored-by: aryan <aryansputta@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: sryap <17482891+sryap@users.noreply.github.com> Co-authored-by: Yunwei Li <yunweili372423@gmail.com> Co-authored-by: Zihao Wang <rekind133@outlook.com> Co-authored-by: Anakin(Yancheng) Zheng <103552181+anakinxc@users.noreply.github.com> Co-authored-by: Prashant Kumar <prashant.kumar@cohere.com> Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com> Co-authored-by: Omar Attia <oy.attia@gmail.com> Co-authored-by: drisspg <drisspguessous@gmail.com> Co-authored-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com> Co-authored-by: rocking <ChunYu.Lai@amd.com> Co-authored-by: Yi DING <yi.ding@amd.com> Co-authored-by: Hosang Yoon <hosang.yoon@amd.com> Co-authored-by: Yin Li <kxl474@student.bham.ac.uk> Co-authored-by: Kevin-Li-2025 <2242139@qq.com> Co-authored-by: Johnny <johnnynuca14@gmail.com>

Johnsonms marked this pull request as draft June 10, 2026 00:22

Johnsonms force-pushed the johnson/ci-pin-cutlass-dsl-452 branch 2 times, most recently from b4f0c9a to 1a8647f Compare June 10, 2026 00:23

Johnsonms changed the title ~~ci(fa4): pin nvidia-cutlass-dsl>=4.5.2 in the SIF test overlay~~ ci(fa4): resolve cute deps from pyproject in the SIF test overlay Jun 10, 2026

Johnsonms force-pushed the johnson/ci-pin-cutlass-dsl-452 branch from 1a8647f to 06b276c Compare June 10, 2026 00:31

Johnsonms force-pushed the johnson/ci-pin-cutlass-dsl-452 branch from 06b276c to 39d177f Compare June 10, 2026 00:43

Johnsonms changed the title ~~ci(fa4): resolve cute deps from pyproject in the SIF test overlay~~ ci(fa4): assert cute dep floors in CI; fail loudly on a stale SIF Jun 10, 2026

ci(fa4): bump cu130 image to 26.06.10 (cutlass-dsl 4.5.2 / quack 0.5.0)

bdaeaed

Johnsonms changed the title ~~ci(fa4): assert cute dep floors in CI; fail loudly on a stale SIF~~ ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image Jun 10, 2026

Johnsonms marked this pull request as ready for review June 10, 2026 01:54

Johnsonms requested a review from Copilot June 10, 2026 01:54

Copilot started reviewing on behalf of Johnsonms June 10, 2026 01:54 View session

Johnsonms requested review from drisspg, jayhshah and tridao June 10, 2026 01:54

Copilot AI reviewed Jun 10, 2026

View reviewed changes

ci(fa4): fall back to tomli when tomllib is unavailable (Python 3.10)

00e8933

drisspg approved these changes Jun 10, 2026

View reviewed changes

Johnsonms merged commit fb02fc8 into Dao-AILab:main Jun 10, 2026

MatthewBonanni mentioned this pull request Jun 16, 2026

Sync upstream vllm-project/flash-attention#149

Merged

Navigation Menu

ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image - #2636

ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image#2636
Johnsonms merged 3 commits into
Dao-AILab:mainfrom
Johnsonms:johnson/ci-pin-cutlass-dsl-452

Johnsonms commented Jun 10, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Johnsonms Jun 10, 2026

Uh oh!

drisspg commented Jun 10, 2026

Uh oh!

drisspg commented Jun 10, 2026

Uh oh!

Johnsonms commented Jun 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Johnsonms commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Why

Validation

Notes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Johnsonms Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

drisspg commented Jun 10, 2026

Uh oh!

drisspg commented Jun 10, 2026

Uh oh!

Johnsonms commented Jun 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Johnsonms commented Jun 10, 2026 •

edited

Loading