add linearize scheduling to combine kernel for full cudagraph by liangel-02 · Pull Request #2692 · Dao-AILab/flash-attention

liangel-02 · 2026-07-01T20:18:42Z

referencing https://github.com/vllm-project/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h to adopt the improved vllm linearized combine kernel so that we can set max_seqlen_q = num_tokens for full cudagraph capture to prevent the case where max_seqlen_q at runtime > max_seqlen_q at capture but preserve performance

the 50% sparsity heuristic to choose standard vs linearize scales w/ B and does perform worse w/ larger B but shouldn't be a problem in real life cases

jayhshah · 2026-07-01T22:39:23Z

This change would cripple varlen perf with the combine kernel. We either can adopt the improved vllm linearized combine kernel or stick with max_seqlen_q as fixed for the cudagraph

need to bump for this change Dao-AILab/flash-attention#2692 Pull Request resolved: #191348 Approved by: https://github.com/drisspg, https://github.com/janeyx99

Brings the branch up to date with 35 upstream commits. 22 conflict hunks across 7 files. Redundant SM120 fixes dropped in favour of upstream's: - softmax.py: upstream Dao-AILab#2706 implements the all-masked-row (row_max == -inf) sink guard unconditionally and with a proper max-shift, so the is_sm120-gated version and its plumbing are removed (3 call sites in flash_fwd.py). - paged_kv.py: upstream landed the same page_entry_per_thread ceil-div. - interface.py: dropped the local maybe_contiguous, which was shadowing upstream's alignment-aware version everywhere in the file. Kept, with upstream's additions merged in: - flash_fwd.py / flash_bwd.py: TileSchedulerArguments gained cu_total_m_blocks alongside the existing is_split_kv. - flash_fwd_combine.py: upstream rewrote it for linearized scheduling (Dao-AILab#2692); the SM120 use_pdl gate (griddepcontrol.wait is illegal on sm_120) is re-applied on top. - bwd postprocess: compile key carries both upstream's cu_total_m_blocks / learnable_sink_dtype and the SM120 pack_gqa pair, in signature order. Forward config refactor: upstream extracted tile selection into _get_fwd_config(), now shared with the public get_scheduler_metadata(). The SM120 per-shape tuning moved into _get_fwd_config_sm120() behind an arch-12 branch, with an Sm120FwdContext for the extra inputs and FwdConfig carrying the SM120 outputs (num_threads, cp.async stages, Q-in-regs). Verified dispatch-identical to the pre-merge branch over 1080 shapes. Fixes required by the merge: - flash_fwd_sm120_tma.py: upstream added mCuTotalMBlocks/mCuTotalSplitsMBlocks to the arch 8/9/12 launch args, but the SM120 TMA kernel had no such parameters, so every d64 TMA forward failed to bind. Wired through to its SingleTileVarlenScheduler. - interface.py: current_stream moved into upstream's compile-miss block; the fp8 decode path compiles earlier and now makes its own. - interface.py: learnable-sink backward. Upstream Dao-AILab#2706 asserts SM90/SM100/110 whenever a sink is present, which the autograd backward always passes. That removed working SM120 behaviour: dSink is a pure side-output of the dQ postprocess and dq/dk/dv receive the sink only through LSE. The assert is now gated on compute_dsink, derived from needs_input_grad, so a frozen sink still backprops on SM120 and only a real dSink request errors out. Pre-existing bug fixed: _bwd_postprocess_dkv_sm120 unpacked 14 values from make_fake_bwd_tensors, which returns 15 (mScaleP). Tested on RTX PRO 6000 (sm_120): 2652 passed / 246 skipped across the SM120 suites, the combine suite and a 350-case random sample of test_flash_attn.py; forward and backward match SDPA for causal, GQA, D256, sliding-window, non-square and decode shapes.

…Lab#2692) * fix combine kernel bug for full cudagraph * linearzie kernel (cherry picked from commit 14c3779) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com>

* ad tcgen.ld.red support to sm103a arch (Dao-AILab#2696) stack-info: PR: Dao-AILab#2696, branch: drisspg/stack/47 Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit 89b44b5) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, Bwd] Fix backward compile key churn due to pickling, max_seqlen is a tensor (Dao-AILab#2507) * Fix backward compile key instability when max_seqlen is a tensor When max_seqlen_q/max_seqlen_k are passed as torch.Tensor (e.g. by HuggingFace Transformers _prepare_from_posids), the arithmetic in _flash_attn_bwd produces tensor results that leak into the compile key tuple. Since pickle.dumps(torch.Tensor) produces a unique hash per object, every backward call generates a new compile key, causing infinite kernel recompilation and filling the persistent JIT cache with identical .o files. Cast max_seqlen_q/k to int() before they enter the seqlen_q/k computation path, ensuring the compile key contains only Python scalars. * Replace int() with host-scalar guard to avoid CPU-GPU sync Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit d145d0c) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * hopper/setup.py: harden tarfile extraction against path traversal and symlink escape (Dao-AILab#2702) * hopper/setup.py: harden tarfile extraction against path traversal and symlink escape download_and_copy() extracted NVIDIA toolchain archives with a bare tarfile.extractall() into the predictable ~/.flashattn/nvidia/<name> cache, allowing arbitrary file write at build time via a pre-planted symlink or a malicious archive member (issue Dao-AILab#2637). - Add safe_extractall(): use the PEP 706 data filter when available (3.12, backported to 3.10.12/3.11.4), else fall back to per-member path containment and link rejection (stream-safe, single pass). - Refuse extraction into a symlinked cache path, closing the primary pre-planted-symlink vector on all Python versions. Signed-off-by: Aryan Putta <aryansputta@gmail.com> * hopper/setup.py: allow in-destination links in extractall fallback Address review on Dao-AILab#2702: 1. The no-data-filter fallback rejected every link member, which regressed real builds: the cuda_nvcc archives ship intra-package symlinks (e.g. libnvvm.so -> libnvvm.so.4) that the data filter permits. Allow links whose resolved target stays inside the extract dir instead, matching the data-filter behavior, and keep rejecting escaping and absolute-target links. 2. Harden the cache-path check: os.path.islink only inspects the leaf, so also require the fully resolved tmp_path to stay under the cache root, catching a symlinked parent directory. Signed-off-by: Aryan <aryansputta@gmail.com> --------- Signed-off-by: Aryan Putta <aryansputta@gmail.com> Signed-off-by: Aryan <aryansputta@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit d51aea5) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Enable 2CTA for SM100 block-sparse backward (Dao-AILab#2661) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit 7bedd8e) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Add paged-KV block_table bounds check in mha_fwd_kvcache (Dao-AILab#2711) The split-KV kernel (compute_attn_1rowblock_splitkv) indexes block_table[n_block * kBlockN / page_block_size], bounded only by actual_seqlen_k. In the kvcache path actual_seqlen_k is seqlens_k[b] + seqlen_knew, but block_table only has max_num_blocks_per_seq columns per sequence. If a caller passes a cache_seqlens (or appends new keys) exceeding max_num_blocks_per_seq * page_block_size, the kernel reads block_table out of bounds with no in-kernel check (see issue Dao-AILab#2709). Validate the caller contract host-side and raise a clear error instead. The maximum-length readback sync is only paid on the paged-KV path. Add test_flash_attn_kvcache_paged_block_table_bounds covering both the cache-length overflow and the appended-new-keys overflow, plus a positive control exactly at capacity. Co-authored-by: yunweili3 <yunweili3@users.noreply.github.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit 374dfa1) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, SM100] Fix FP8 e4m3 accuracy: make max_offset dtype-aware to avoid P saturation (Dao-AILab#2717) * [CuTe, SM100] Make FP8 max_offset dtype-aware to avoid e4m3 P saturation With rescale_threshold=4 the online-softmax row max can be stale by up to 4 (in log2 units), so P reaches 2^(max_offset + 4). max_offset=8 puts that at 4096, past e4m3fn's 448 ceiling: the largest probabilities saturate on the f32->fp8 satfinite convert and e4m3 accuracy degrades below e5m2 (up to 1.6x worse rel_l2, growing with seqlen). Cap max_offset at 4 for e4m3 so the worst case is 2^8 = 256 <= 448; e5m2 keeps 8 (57344 ceiling absorbs the overshoot). B200: restores e4m3 to ~2x lower error than e5m2 across seqlen 256-4096, uniform and peaked softmax, matching quantization-only emulation; LSE consistent; fwd timing unchanged (0.387 vs 0.390 ms, hd128 s4096). Related: Dao-AILab#2716 * [CuTe, Tests] Unrot the FP8 dtype path in test_flash_attn_output Running the suite with dtype=float8_e4m3fn has bit-rotted: - the test sets requires_grad on fp8 tensors, which the interface now rejects (FP8 is forward-only); gate it on non-fp8 dtypes. - it generates random descales and applies them in attention_ref, but the flash_attn_func call site has no descale kwargs (only _flash_attn_fwd takes them), so kernel and reference disagreed by construction; stop generating them. With these, the fp8 sweep runs cleanly (378 cases on SM100 with the e4m3 max_offset fix; 190 of them fail without it). fp8 stays out of the default dtype parametrize. Related: Dao-AILab#2716 Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit be8c0e0) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Revert "Fix threshold (#166)" This reverts commit 1689202. Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit aa6c57c) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, Flex] Allow score mod use in varlen backward (Dao-AILab#2547) * allow varlen score mod in backward; add tests and examples * add recompute fastdiv_mods to sm90 bwd * remove softcap != 0 limitation in test * fix linter error * guard use 2cta against softcap in bwd * undo formatting in test_flash_attn.py * reset test_flash_attn * update tests for score mod varlen bwd, guard blocksparse varlen bwd * aux_tensors -> aux_data; unpack args in test * aux_tensors -> aux_data for sm90 backward * make_fragment -> make_rmem_tensor in score_mod_definitions * predicate on aux_data.tensors, not aux_data * relax test tolerance in vectorized score mod tests - bitwise equality failing on sm103 though within tolerance * revert erroneous test reformatting to main (cherry picked from commit b54df16) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Expand FLASHATTENTION_DISABLE_DROPOUT to not bring in unneeded headers (Dao-AILab#2669) * Expand FLASHATTENTION_DISABLE_DROPOUT to not bring in unneeded headers Summary: Previously, using the FLASHATTENTION_DISABLE_DROPOUT flag still pulled in unneed dependencies from ATen for at::Generator and Philox related headers. This change sets up the codebase so that using the flag will not pull in these unnecessary headers. There are two major changes of note: 1. We remove needing an RNG gen in the schema--the Python frontend always passed in None so this should not be BC breaking to most users. 2. Instead of referencing the PhiloxState directly, in order to detach dependencies when dropout is not needed, we introduce an opaque buffer that will hold the philox state when dropout is desired. Test Plan: pytest tests/test_flash_attn.py::test_flash_attn_output -k "113-203-64 and dtype0 and mha" pytest tests/test_flash_attn.py::test_flash_attn_varlen_output -k "113-203-64 and dtype0 and mha" g++ -c -O1 -std=c++17 -D_GLIBCXX_USE_CXX11_ABI=1 <torch+cutlass+cuda -I flags> \ csrc/flash_attn/flash_api.cpp -o /tmp/fa.o nm -C /tmp/fa.o | grep -E 'mha_(fwd|bwd|varlen)\(' | grep -c Generator returns 0 g++ -E -DFLASHATTENTION_DISABLE_DROPOUT <same -I flags> csrc/flash_attn/flash_api.cpp \ | grep -c 'CUDAGeneratorImpl.h\|philox_unpack.cuh' also returns 0 Reviewers: Subscribers: Tasks: Tags: Add trivially copyable assert Add back gen * use mark.skipIf (cherry picked from commit 00756db) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * add linearize scheduling to combine kernel for full cudagraph (Dao-AILab#2692) * fix combine kernel bug for full cudagraph * linearzie kernel (cherry picked from commit 14c3779) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Numeric tweaks to fp8 (Dao-AILab#2731) stack-info: PR: Dao-AILab#2731, branch: drisspg/stack/49 (cherry picked from commit 849f660) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Remove SM100 Functions from Hopper (Dao-AILab#2746) (cherry picked from commit c75d019) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe,Sm100] Varlen Dynamic Persistent scheduler and metadata (Dao-AILab#2559) * add dynamicpersistentvarlenscheduler to flash_fwd_sm100 and prepare kernel * mild refactor to tile scheduler protocol, guard num_m_blocks_ptr for sm100, update tests to use scheduler metadata * rename varlen_batch_idx -> virtual_batch_idx, because it is relevant for non-varlen blocksparse batch sorting * split out VarlenSchedulerBase to share code between SingleTile and DynamicPersistent schedulers * add benchmark script for varlen dynamic persistent scheduler * minor clean up * updates to has_work logic, tile scheduler selection, and varlen test suite * fix tile scheduler dispatch logic * integrate binary batch search for single tile varlen * refactor tile scheduler for compositionality * work PR 2520 into interface and kernels * fix linter errors * wip: modify scheduler metadata public api * clean up scheduler metadata API; add docstrings; split out _get_fwd_config method; remove cluster_size==1 restriction; guard architectures against unused scheduler metadata args * address driss' comments * fix compute_tile_cumsum guards in interface * simplfiy benchmark, guard against _compute_tile_cumsum with small batch size * fix compute_tile_cumsum guard * update to 4.6.0 * fix linter error * fix rebase bug * add cu_blocks_kernel to replace _compute_tile_cumsum * fix linter error * address comments on PR * add seqlen_k_per_split, add single tile varlen scheduler to combine kernel * add blocks to batch idx O(1) lookup path to varlen scheduler * fix lint errors * various legibility improvements and bug fixes (cherry picked from commit c46b814) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [AI] Add doc on debug methodology (Dao-AILab#2753) * add methodology doc * revise method * concision pass (cherry picked from commit df61ab6) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [ROCm] Fix CK varlen_fwd binding argument mismatch (Dao-AILab#2742) (cherry picked from commit 4a948e9) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, SM100] Sparse MLA bwd: don't scatter dK/dV at -1 sentinel indices (Dao-AILab#2755) * [CuTe,Sm100] Sparse MLA bwd: skip dK/dV scatter at -1 sentinel indices The sparse-MLA (gather_kv_indices) backward scatter epilogues atomically accumulated dV/dK at row indices read straight from gather_kv_indices with no validity guard, while every load path treats -1 (the documented sentinel for invalid top-k slots, which any causal top-k index tensor contains as padding) as invalid and predicates the gather. The masked math is correct -- p/dS are exactly 0.0 at sentinel slots -- but the atomic itself is destructive: index -1 addresses one row before the (batch-sliced) buffer base, and red.add.f32 flushes subnormal destinations to +0.0 and canonicalizes NaN payloads even when the addend is 0.0. For batch 0 this lands out of bounds in whatever tensor the caching allocator placed before dv/dk (silently corrupting e.g. int32 tensors, whose small values are all subnormal fp32 bit patterns); for later batches it lands in the previous batch's last row. Symptoms depend purely on allocation layout: bitwise-correct results, silently wrong grads, or IMA. Fix: skip the atomic when the index is negative, mirroring the load-side guard. The skipped contribution is mathematically 0.0, so numerics for valid slots are unchanged. Also fix _flash_attn_bwd_sparse_mla discarding caller-supplied dq=/dk= buffers (dq = dk = None right after recording prealloc_dq/dk, after which the reallocation is skipped because prealloc is set, so passing dq=/dk= crashed). * [CuTe,Sm100] Test sparse MLA bwd with -1-padded gather_kv_indices Adds test_flash_attn_mla_sparse_bwd_sentinel and a varlen counterpart (causal x shared_kv, seqlen 512/1024 non-varlen, packed docs [512, 4, 1024] varlen): builds causal top-k indices with -1 tail padding, checks out/lse/grads against attention_ref through the public autograd path, then reruns the backward with preallocated dk/dv buffers surrounded by int32 canaries (values 1..N, all subnormal fp32 bit patterns, so one misdirected red.add.f32 -- even of +0.0 -- flushes them to zero) and asserts the canaries are untouched. The varlen kernels are separate compile-time specializations, and the dK epilogue guard must apply to the doc-relative index before seqlen_k_offset is added; the varlen test pins that down (doc 0's row -1 is the canary-visible case) and includes a doc shorter than topk_len whose index rows are almost entirely sentinels. Fails deterministically without the sentinel-scatter guard; existing sparse-MLA tests never hit the bug because they generate gather_kv_indices as full argsort permutations with no -1 slots. Also makes attention_ref's top-k mask sentinel-aware: its scatter_ used to route -1 indices into key 0 (unmasking it) and trip scatter's bounds check; out-of-range indices now also map to the dummy column, and the mask applies regardless of topk_len vs seqlen_k (equivalent for permutation indices). --------- (cherry picked from commit c68c592) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe,Bwd,Sm90] Fix: wait for bwd_preprocess on the block-sparse path, matching the dense path (Dao-AILab#2756) Co-authored-by: qaf <qaf@example.org> (cherry picked from commit 4460ebb) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [Cute, bwd, sm90/100/110] Support learnable sink in backward (Dao-AILab#2706) * Support learnable sink backward * Simplify learnable sink backward plumbing * Format learnable sink postprocess * Support learnable sink with hd256 and frozen QKV * Format standalone sink reduction * Tighten learnable sink backward checks * Keep learnable sink scope lean * Remove standalone learnable sink varlen test * Always test learnable sink backward * Simplify learnable sink test setup * Detach learnable sink test tensors * Use Tuple return type for backward * Remove sink backward tensor wrapper * Clarify sink reduction CTA selection * Refine learnable sink backward handling * Handle sink-only rows in backward * Address learnable sink review feedback * Separate learnable sink dtype coverage * Use dtype-aware learnable sink gradient tolerance * Handle empty learnable sink backward * Address learnable sink follow-up feedback (cherry picked from commit 5579b12) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, FA4] Preserve first-tile flag during scheduler reconstruction (Dao-AILab#2705) * [CuTe, FA4] Preserve first-tile flag during scheduler reconstruction * Pin nvidia-cutlass-dsl to 4.7.0 * Keep scheduler fix separate from DSL upgrade (cherry picked from commit 7a08d7a) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Fix duplicated word in layer norm comment (Dao-AILab#2744) (cherry picked from commit d7e4dba) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, SM100] Fix deadlock in varlen + block-sparse + SplitKV forward (Dao-AILab#2761) When scheduler metadata provides per-batch dynamic num_splits, the varlen tile schedulers pack num_splits into the top 16 bits of split_idx. The dense path unpacks it inside BlockInfo.get_n_block_min_max, but the SM100 block-sparse paths pass split_idx to the block-sparse helpers as is: the load and MMA warps passed the packed value, for which split_block_range yields an empty block range, while the softmax and correction warps passed the unpacked value, yielding a non-empty range. The warps then disagree on whether a tile has work, the softmax/correction/MMA mbarrier handshake never completes, and the kernel spins forever. Every varlen + block-sparse + SplitKV forward hangs this way. Also pass the dynamic per-batch num_splits (instead of the static maximum) to the block-sparse helpers so the split ranges cover the whole block list, matching the dense path. (cherry picked from commit 3fa8105) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe, Fwd] Stabilize tensor max_seqlen compile key (Dao-AILab#2762) (cherry picked from commit 1cc7ff6) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe] Fix forward dynamic-shape correctness (Dao-AILab#2745) Fix forward issues exposed by dynamic-shape and layout fuzzing. Canonicalize unaligned inputs, distinguish static broadcast and auxiliary tensor ABIs in the compile cache, and compile SplitKV combine optional operands exactly as called. Use target-SKU SM metadata during fake selection. Size aliased SM100 K/V shared memory for the larger staged layout and ceil-divide non-TMA paged-loader entries so partial row waves receive page pointers. Add one focused regression for each underlying bug. (cherry picked from commit 69e1bcb) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Fix CLC fuzz scheduler expectations (Dao-AILab#2766) PR Dao-AILab#2559 added static- and dynamic-persistent dispatch for varlen SM100 kernels, but test_clc_fuzz still required the pre-change single-tile scheduler, so the scheduler assertions failed before numerical validation could run. Validate each scheduler class against its scheduling mode, account for dynamic SplitKV dispatch, and update the two static fallback expectations. STATIC mode accepts both StaticPersistentTileScheduler and SingleTileVarlenScheduler for varlen because the dispatch depends on whether every batch fits in a single m-block. Verification: the full tests/cute/test_clc_fuzz.py suite passes on SM100 (B200, CC 10.0) — 189 passed — both on this branch's base (1cc7ff6) and cherry-picked onto c68c592. The GQA + SplitKV varlen cases were confirmed to still select the CLC SingleTileVarlenScheduler, matching the assertion precedence. (cherry picked from commit a369df7) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Fix removed Quack packed subtraction API (Dao-AILab#2787) Quack no longer exports sub_packed_f32x2 from quack.activation. Use the CuTe DSL primitive directly at the FlashAttention call sites and remove the obsolete imports. This prevents AttributeError when compiling with current Quack. Fixes Dao-AILab#2782 Co-authored-by: Jiaxuan Bai <JiaxuanBai@users.noreply.github.com> (cherry picked from commit 145b101) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe] Relax CUTLASS DSL requirement (Dao-AILab#2798) * [CuTe] Relax CUTLASS DSL requirement ## Human Note ## Agent note FA4 still required the 4.6.0 development build after support moved to the released 4.6.2 package and CUTLASS DSL 4.7 compatibility landed in Dao-AILab#2787. Replace the exact base and cu13 requirements with a 4.6.2 floor so the resolver can select the version supported by the installed QuACK release. Keep CI prerelease resolution enabled for future development-build floors and update the documented dependency versions. The same SM100 forward, backward, and ex2-emulation checks passed with CUTLASS DSL 4.6.2 and 4.7.0; the 4.7.0 run used QuACK main because its matching release is not on PyPI yet. ## Test Plan ```bash uv build flash_attn/cute --out-dir agent_space/dist-fa4-pr ``` * [CuTe] Stabilize Ruff CI ## Human Note ## Agent note Ruff 0.16 enables a much broader rule set when this project leaves the selection implicit, causing the main-branch lint job to report hundreds of pre-existing violations. Select the previous E4/E7/E9/F baseline explicitly and format the two packed-subtraction call sites changed by the CUTLASS DSL 4.7 compatibility fix. ## Test Plan ```bash ruff check flash_attn/cute/ --extend-exclude "flash_attn/cute/flash_bwd.py,flash_attn/cute/flash_fwd.py,flash_attn/cute/flash_fwd_sm100.py,flash_attn/cute/interface.py" ruff format --check flash_attn/cute/ --exclude "flash_attn/cute/flash_bwd.py,flash_attn/cute/flash_fwd.py,flash_attn/cute/flash_fwd_sm100.py,flash_attn/cute/interface.py" ``` (cherry picked from commit 0251105) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * [CuTe] Speed up scalar SM100 mask compilation (Dao-AILab#2819) ## Human Note ## Agent note Use a fully unrolled dynamic CuTeDSL range for the 128-element scalar SM100 mask-mod loop. The previous static Python expansion emitted a DSLOptimizationWarning and spent substantially more time in cold compilation. The loop index supports dynamic fragment access, so this retains the same fully unrolled execution while avoiding the slow frontend expansion. ## Performance On a B200 with PyTorch 2.15.0.dev20260819+cu132 and CUTLASS DSL 4.6.2, the VSA Flex/Flash scalar path improved from 11.02 s to 7.58 s mean cold first-call time across three runs, a 31.2% reduction. CUDA-graph steady-state latency was unchanged within noise (15.23 us baseline versus 15.17 us candidate). ## Test Plan ```bash pytest test/test_vsa.py::test_vsa_flex_flash_backend_sm100_matches_reference test/test_vsa.py::test_vsa_flex_flash_backend_sm100_masks_partial_kv_subblocks ``` The two tests passed both with the nightly default vectorized mask path and with a local pytest plugin forcing `FlexFlashConfig(mask_mod_vec_size=1)` to exercise the changed loop. (cherry picked from commit ce088ab) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Align ROCm sources with synchronized submodules Update the CK and AITER gitlinks together with the CK wrapper signatures expected by those revisions. Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Reconcile FA4 scheduler changes with vLLM integrations Preserve the fork's SM90 FP8 KV path, SM100 head-dimension 256 tuning, fused output quantization, and compile-only warmup support while adopting the upstream dynamic varlen scheduler. Keep SM90 on its supported fixed SplitKV scheduler and add focused compile-only and runtime regression coverage. Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Fix FA4 backward return arity The fused-output integration adds out and output_scale to FlashAttnFunc.apply. Return two matching trailing gradients so ordinary backward calls preserve the autograd input contract. Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Preserve flash_attn_combine keyword compatibility Keep virtual_batch_idx as the canonical parameter while accepting the existing varlen_batch_idx keyword as a compatibility alias. Reject ambiguous calls that provide both names, and cover canonical, compatibility, and positional use. Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> * Keep FlashAttention 3 frozen Signed-off-by: Chuyue Wang <stevenwang0805@outlook.com> --------- Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com> Signed-off-by: Aryan Putta <aryansputta@gmail.com> Signed-off-by: Aryan <aryansputta@gmail.com> Signed-off-by: Chuyue Wang <stevenwang0805@outlook.com> Co-authored-by: Driss Guessous <32754868+drisspg@users.noreply.github.com> Co-authored-by: michaelxu-msft <115037246+michaelxu-msft@users.noreply.github.com> Co-authored-by: aryan <aryansputta@gmail.com> Co-authored-by: Yunwei Li <yunweili372423@gmail.com> Co-authored-by: yunweili3 <yunweili3@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Reuben Stern <107093092+reubenconducts@users.noreply.github.com> Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com> Co-authored-by: liangel-02 <liangel@meta.com> Co-authored-by: ankutalev <31923880+ankutalev@users.noreply.github.com> Co-authored-by: jayhshah <jayhshah@gmail.com> Co-authored-by: Hosang Yoon <156028780+hyoon1@users.noreply.github.com> Co-authored-by: Lequn Chen <abcdabcd987@users.noreply.github.com> Co-authored-by: Fugoes <fugoes.qa@gmail.com> Co-authored-by: qaf <qaf@example.org> Co-authored-by: Henry Tsang <henrylhtsang@meta.com> Co-authored-by: dongxiao <xiaod@nvidia.com> Co-authored-by: Haijie Zhi <133995660+cupkk@users.noreply.github.com> Co-authored-by: Jiaxuan Bai <jbai@nvidia.com> Co-authored-by: Yiming Zhang <49868620+eamonn-zh@users.noreply.github.com> Co-authored-by: Jiaxuan Bai <JiaxuanBai@users.noreply.github.com>

liangel-02 force-pushed the combine_kernel branch from 433abbc to d9cd00c Compare July 1, 2026 20:50

liangel-02 marked this pull request as ready for review July 1, 2026 20:50

drisspg requested a review from jayhshah July 1, 2026 21:00

fix combine kernel bug for full cudagraph

565aa20

liangel-02 force-pushed the combine_kernel branch from d9cd00c to 565aa20 Compare July 1, 2026 21:34

liangel-02 marked this pull request as draft July 9, 2026 20:33

liangel-02 force-pushed the combine_kernel branch 2 times, most recently from 1dcf54d to be00c0b Compare July 9, 2026 20:53

liangel-02 changed the title ~~fix combine kernel bug for full cudagraph~~ add linearize scheduling to combine kernel for full cudagraph Jul 9, 2026

liangel-02 mentioned this pull request Jul 10, 2026

enable varlen full cudagraph pytorch/torchtitan#3893

Merged

liangel-02 marked this pull request as ready for review July 13, 2026 17:24

drisspg reviewed Jul 14, 2026

View reviewed changes

Comment thread hopper/flash_fwd_combine_kernel.h

Comment thread hopper/flash_fwd_combine_kernel.h

Comment thread hopper/flash_fwd_combine_kernel.h

Comment thread hopper/flash_fwd_combine_kernel.h Outdated

Comment thread hopper/flash_fwd_combine_kernel.h Outdated

liangel-02 force-pushed the combine_kernel branch from be00c0b to 20c7f45 Compare July 17, 2026 02:02

liangel-02 requested a review from drisspg July 21, 2026 20:24

liangel-02 force-pushed the combine_kernel branch from 20c7f45 to e106c8b Compare July 22, 2026 20:24

linearzie kernel

04eb8d1

liangel-02 force-pushed the combine_kernel branch from e106c8b to 04eb8d1 Compare July 22, 2026 20:43

jayhshah approved these changes Jul 27, 2026

View reviewed changes

drisspg merged commit 14c3779 into Dao-AILab:main Jul 27, 2026

liangel-02 mentioned this pull request Jul 28, 2026

bump fa submodule hash pytorch/pytorch#191348

Closed

StevenWang-CY mentioned this pull request Aug 28, 2026

Sync Dao-AILab FlashAttention through ce088ab9 vllm-project/flash-attention#188

Merged

Navigation Menu

add linearize scheduling to combine kernel for full cudagraph - #2692

add linearize scheduling to combine kernel for full cudagraph#2692
drisspg merged 2 commits into
Dao-AILab:mainfrom
liangel-02:combine_kernel

liangel-02 commented Jul 1, 2026 •

edited

Loading

Uh oh!

jayhshah commented Jul 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

liangel-02 commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jayhshah commented Jul 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

liangel-02 commented Jul 1, 2026 •

edited

Loading