Sync upstream - #171
Open
MatthewBonanni wants to merge 8 commits into
Open
Conversation
MatthewBonanni
force-pushed
the
sync_upstream_20260722
branch
2 times, most recently
from
July 22, 2026 16:51
642cefe to
320f7c5
Compare
stack-info: PR: Dao-AILab#2696, branch: drisspg/stack/47 Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
…n 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>
… 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>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
) 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 .max().item() 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>
…void 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>
This reverts commit 1689202. Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
MatthewBonanni
force-pushed
the
sync_upstream_20260722
branch
from
July 22, 2026 16:56
320f7c5 to
aa6c57c
Compare
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
6 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
Dao-AILab/flash-attentionup to2409214a, on top ofvllm-project/main(ed4b7342).Notable upstream changes
tcgen.ld.redsupport (add tcgen.ld.red support to sm103a arch Dao-AILab/flash-attention#2696).max_seqlenis a tensor ([CuTe, Bwd] Fix backward compile key churn due to pickling, max_seqlen is a tensor Dao-AILab/flash-attention#2507).block_tablebounds checking (Add paged-KV block_table bounds check in mha_fwd_kvcache Dao-AILab/flash-attention#2711).max_offsetto avoid probability saturation ([CuTe, SM100] Fix FP8 e4m3 accuracy: make max_offset dtype-aware to avoid P saturation Dao-AILab/flash-attention#2717).Downstream follow-up
max_offset=4for E4M3 while retainingrescale_threshold=4.0; this supersedes the downstreamrescale_threshold=0.75workaround and preserves the rescale-skip optimization.Conflict resolution
flash_attn/cute/interface.py— preserved the fork fused-output-quantization path (output_quant_key) while threading upstreamkv_subtile_factorthrough both SM100 forward kernel variants.Validation
git diff --checkpasses.python -m compileall -q flash_attn/cute hopper/setup.py tests/cute tests/test_flash_attn.pypasses.pre-commit run --from-ref vllm-project/main --to-ref HEADpasses (ruff check,ruff format).