Sync upstream#141
Merged
Merged
Conversation
* [ROCm Windows] fix triton requirement * pin triton-windows>=3.6.0
* varlen block-sparsity for forward Squashed forward-path varlen support: extends BlockSparseTensors usage to [num_heads, total_m_blocks] / [num_heads, total_n_blocks] layouts, threads cu_seqlens / cu_total_m_blocks / cu_total_n_blocks through the kernel and compute_block_sparsity, and routes through get_curr_blocksparse_tensors and get_total_block_count for shape-aware indexing. * rename cu_total_n_blocks to cu_block_idx_offsets; move cu_total_m_blocks/cu_block_idx_offsets into BlockSparseTensors instead of threading them as standalone parameters; drop the two <tensor>[-1].item() syncs in normalize_block_sparse_config
…ao-AILab#2515) num_splits_heuristic divides num_SMs by total_mblocks, which collapses to 0 when seqlen_q == 0 or batch_size == 0 (e.g. CUDA graph padding or empty microbatches). The existing seqlen_k == 0 early-exit in _flash_attn_fwd does not cover these cases. - Extend the early-exit to also cover total_q == 0, using the same zero-output / -inf-LSE contract. total_q is batch_size * seqlen_q (dense) or q.shape[0] (varlen), so a single predicate handles both code paths. - Add a defensive total_mblocks == 0 guard inside num_splits_heuristic itself so the function is safe in isolation. - Add regression tests covering dense (batch=0, seqlen_q=0) and varlen (total_q=0) paths under both causal and non-causal masks. Fixes Dao-AILab#2503.
* split out varlen batch search into utils * more descriptive name
stack-info: PR: Dao-AILab#2536, branch: drisspg/stack/38
) Summary: Extract the inline `AttentionMask` construction in `FlashAttentionForwardSm100` and `FlashAttentionBackwardSm100` into an overridable `_generate_attention_mask_cls` method. This allows subclasses to inject a custom `AttentionMask` without modifying the base kernel code. For example, a custom attention kernel can override the mask to add a `causal_q_divisor` field for scaling the `row_idx` value. ``` class CustomAttentionMask(AttentionMask): causal_q_divisor: cutlass.Constexpr[int] = 1 @cute.jit def apply_mask_sm100(self, acc_S, m_block, n_block, ...): # Custom causal logic using causal_q_divisor row_idx = (tScS_t2r[0][0] + m_block * self.tile_m) // self.causal_q_divisor ... class CustomFlashAttentionForwardSm100(FlashAttentionForwardSm100): def __init__(self, *args, causal_q_divisor=1, **kwargs): super().__init__(*args, **kwargs) self.causal_q_divisor = causal_q_divisor def _generate_attention_mask_cls(self, window_size_left, window_size_right): return partial( CustomAttentionMask, self.m_block_size, self.n_block_size, window_size_left=window_size_left, window_size_right=window_size_right, bottom_right=self.is_bottom_right, causal_q_divisor=self.causal_q_divisor, ) ``` Test Plan: ``` $ pytest tests/cute/test_flash_attn_fast.py -v ================ 240 passed, 4139 warnings in 984.24s (0:16:24) ================ ``` Reviewers: Subscribers: Tasks: Tags:
…10 (Dao-AILab#2532) * Fix: Remove misleading py_limited_api=cp39 wheel tag for PyTorch extension * Implement dynamic ABI tagging for PyTorch versions Add dynamic ABI tag based on PyTorch version for correct and improved naming of the wheel. * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Update Python version requirements based on torch metadata * Refactor setup.py for dynamic ABI and CUDA settings Refactor dynamic ABI tag and Python version requirements based on installed PyTorch version and streamline CUDA extension arguments. * Update CUDAExtension compile arguments Restored some accidentally removed content * Update setup.py * Updated setup.py minor fix: cleaned up the comments * Brought back Py_LIMITED_API flag to CUDA extension compilation * Minor fix * Update setup.py for Python version requirements Updated the wheel tag to cp310 and python_requires=">=3.10". --------- Co-authored-by: aw920h <alien@alien.alien> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
* vectorized mask mod application for existing mask mod signatures * add vectorized mask mod examples, get vectorized evaluation and application working * guard sm80/90/120 against mask_vec_size > 2 * thread mask_vec_size thru sm80/90/120 kernel * Small tweaks coverign sm90 * Small tweaks coverign sm90 --------- Co-authored-by: drisspg <drisspguessous@gmail.com>
…b#2572) (Dao-AILab#2590) * Fix bwd postprocess 2CTA gating to include sm_11x The 2CTA gating in flash_bwd_postprocess.py used `arch // 10 == 10`, which only matches SM 10.x (B100/B200/B300) and misses SM 11.x (Thor). The rest of the codebase (e.g. interface.py:549, 563, 834) consistently gates Blackwell-family 2CTA features as `arch // 10 in [10, 11]`. Bring the two postprocess sites in line with that convention. Flagged by @jayhshah in Dao-AILab#2572 follow-up discussion. * Include sm_110 in interface.py Blackwell-family heuristics Three sites in interface.py gate Blackwell-family behavior using `arch // 10 == 10`, which appears inconsistent with the rest of the file's `arch // 10 in [10, 11]` convention (used at lines 549, 563, 834, 974, 1035, etc.): - L533: `q_stage` heuristic for Blackwell forward - L579: `use_dedicated_hd256_kernel` (forward) - L1335: `use_dedicated_hd256_kernel` (backward) The dispatch in `_flash_attn_fwd` already routes both sm_10x and sm_11x through the same `FlashAttentionForwardSm100` / MLA classes, so these gates likely should treat them the same. NOTE FOR REVIEWERS: I'm not certain these are all oversight vs. intentional SM100-only paths. If any of them is intentional, please flag so I can revert just that hunk. The FP8 assert at L480 is left untouched on purpose — its error message reads as deliberate. * Apply ruff format to flash_bwd_sm100.py Pre-existing format drift surfaced by pre-commit. Not in the cute_exclude pattern, so it gets auto-fixed when other files in flash_attn/cute/ are touched in the same commit chain.
* Use is_family_of for sm_90 and sm_103 arch checks Follow-up to Dao-AILab#2572 — apply the same is_family_of pattern to the two remaining range-style arch checks for consistency: - flash_fwd_sm90.py:69 (SM 9.x assert) - flash_fwd_sm100.py:195 (is_sm103 flag) Same semantic narrowing as Dao-AILab#2572: bare-base SMs (sm_90, sm_103) are excluded. These kernels rely on wgmma / UMMA / 2CTA paths that require the a/f PTX variant anyway, so bare-base targets could not compile. * Clarify is_sm103 forward-inclusive semantics is_family_of(sm_103f) also matches any future sm_10x with x > 3, not just sm_103a/f. This was raised in PR review (@ocss884) — adding an inline comment clarifying that this forward-inclusive behavior is intentional: the flag gates ex2 emulation, sm_103 (B300) has fast hardware ex2, and later Blackwell variants in the same family are assumed to inherit it. No code-behavior change.
* Bump aiter submodule commit Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com> * Bump aiter submodule to 3b2e6f48ce97e1d494e8b3f1af5c65f74e304b28 (vllm-project#2) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com>
…ao-AILab#2594) * Clamp kv_stage to avoid SMEM overflow for small head_dims on SM100 Fixes Dao-AILab#2591. The unbounded formula at flash_fwd_sm100.py:335 ignores per-stage state (mbarriers, sScale, pipeline counters) and yields kv_stage values that overflow the sm_100a 227 KB SMEM cap when head_dim_padded=16 (head_dim in {8, ..., 16}). Repro: hd=8/16 + seqlen >= 256 + bf16 fails with cudaErrorInvalidValue ("launch shared memory exceeds current GPU arch sm_100a allowed. Allocated: 233472 bytes. Max: 232448 bytes."). Clamp kv_stage at 32. Surgical to the broken case: the unbounded formula maxes at 26 stages for head_dim_padded >= 32, and the 2CTA gate at interface.py:572 restricts 2CTA to hd_padded in {128, 192} (both no-op), so the clamp only fires at hd_padded in {8, 16}. Verified across 24 configs (hd in {8,16,32,64,96,128} x causal in {T,F} x seqlen in {128,2048}) on B200 with max_err vs torch SDPA <= 0.0078. * Add test_flash_attn_small_head_dim regression test The main test_flash_attn_output parametrizes d over {64, 96, 128, 192, 256} and never exercises head_dim < 64, even though _validate_head_dims accepts head_dim >= 8 for sm_100/110. That coverage gap let the SMEM-overflow bug in Dao-AILab#2591 slip through. This focused test covers d in {8, 16, 32} x causal x seqlen in {128, 2048}. The seqlen=2048 cases push q_stage 1->2 (the actual bug trigger); the seqlen=128 cases also exercise the q_stage=1 boundary that fits on main today but is structurally adjacent. d=32 serves as a canary against any future tighter kv_stage clamp regressing it.
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
LucasWilkinson
approved these changes
May 29, 2026
This was referenced May 30, 2026
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
15 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Sync with upstream