Skip to content

Fix duplicated word in layer norm comment - #2744

Merged
drisspg merged 1 commit into
Dao-AILab:mainfrom
cupkk:docs/fix-layernorm-comment
Aug 6, 2026
Merged

Fix duplicated word in layer norm comment#2744
drisspg merged 1 commit into
Dao-AILab:mainfrom
cupkk:docs/fix-layernorm-comment

Conversation

@cupkk

@cupkk cupkk commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Fix a duplicated warp in the comment above the per-warp statistics store in ln_utils.cuh. This is a comment-only change and does not affect generated code.

Validation:

  • Confirmed the old phrase is absent and the corrected phrase is present
  • git diff --check

Copilot AI review requested due to automatic review settings July 29, 2026 00:30

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@cupkk
cupkk force-pushed the docs/fix-layernorm-comment branch from 9d597dd to 1543ce3 Compare July 31, 2026 01:36
@drisspg
drisspg merged commit d7e4dba into Dao-AILab:main Aug 6, 2026
StevenWang-CY pushed a commit to StevenWang-CY/flash-attention that referenced this pull request Aug 28, 2026
(cherry picked from commit d7e4dba)
Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com>
MatthewBonanni added a commit to vllm-project/flash-attention that referenced this pull request Sep 4, 2026
* 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>
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.

3 participants