[release/2.12] - #3602
Draft
dnikolaev-amd wants to merge 7 commits into
Draft
Conversation
…mentation (pytorch#185573) This PR matches the activation `logsumexp`'s dtype within cpu flash SDPA's meta formula to the one used in the cpu eager implementation. Eager type mapping: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/OpMathType.h Meta type mapping: https://github.com/pytorch/pytorch/blob/7ed7ac4802b7a1e1cab54fcd1b2bbb10105f8285/torch/_prims_common/__init__.py#L1493-L1501 Fixes pytorch#185539 Pull Request resolved: pytorch#185573 Approved by: https://github.com/Skylion007 (cherry picked from commit 84287a9)
…184914) fixes pytorch#184330 today in attention.cpp, we have this check in `scaled_dot_product_attention`: https://github.com/pytorch/pytorch/blob/93c10b22b50317447f2ba046680a57c4c8681d4a/aten/src/ATen/native/transformers/attention.cpp#L739-L746 which guards against empty outputs by checking if any of q/k/v has numel()==0. however, this is too strict since even if head dim of q/k = 0 but v still has valid elements, then attention output is still defined (should just be the mean of all values). the guard should actually just check if any of v's dims are 0 (which checks for empty batch, head dims, 0 kv seqlen, and 0 v headdim) and if q seqlen = 0. note that in the head_dim=0 case, you will end up using math backend since flash/cudnn doesnt support this Pull Request resolved: pytorch#184914 Approved by: https://github.com/drisspg (cherry picked from commit cf37855)
…#186434) (pytorch#186434) Summary: The seqlenq_ngroups_swapped optimization in the ROCm CK flash-attn host wrapper (mha_fwd_ck.hip, from upstream PyTorch D65184638 / enabled for GQA in D70314996) is mis-ported: when it triggers (seqlen_q==1 && num_heads>num_heads_k, i.e. GQA single-query) it reassigns seqlen_q=ngroups / num_heads=num_heads_k for the kernel args, but then passes the original un-swapped q and allocates out=empty_like(q) with the original shape. The CK kernel strides over the swapped dims using the original tensor strides, reading/writing out of bounds across batch elements and producing finite garbage (up to ~FLT_MAX) that overflows to Inf/NaN downstream. This surfaced as ~100% 'Got NaN or Inf in the final output' on MI350 DISAGG_GPU traffic for the HSTU beam-decoder GQA cross-attention (8 q-heads / 2 kv-heads, head_dim 128). Fix the swap plumbing so it uses the swapped tensors: pass q_padded/k_padded/ v_padded to the kernel args, and allocate the output with the swapped shape. This keeps the optimization while producing correct results. Validated (micro-repro dper_lib/.../ck_fa_nan_repro.py): CK GQA seqlen_q==1 max-abs-diff vs fp32 3.21 -> 0.0016 with the swap active; seqlen_q>1 unchanged. Also confirmed end-to-end via offline vanguard replay (0 NaN). Test Plan: micro-repro + offline end-to-end test. Reviewed By: xw285cornell, royren622 Differential Revision: D107725962 Pull Request resolved: pytorch#186434 Approved by: https://github.com/alugorey, https://github.com/jeffdaily (cherry picked from commit de8d102)
Gate preferred_blas_library CK on ckGemmSupported (includes gfx90a) and keep SDPA/ROCm FA on ckSPDASupported. Update _is_ck_sdpa_available accordingly. Skip test_ck_blas_library_mm bfloat16 on MI200 (gfx90a). Pull Request resolved: pytorch#187267 Approved by: https://github.com/alugorey, https://github.com/jeffdaily Co-authored-by: Jeff Daily <jeff.daily@amd.com> (cherry picked from commit 793fe7f)
… OOB fault on tile-unaligned shapes (pytorch#187152) (pytorch#187152) Summary: PROBLEM & REPRODUCE GPU lowering of models on AMD MI350X (gfx950) crashes deterministically with "Memory access fault by GPU node-N ... Reason: Unknown" (SIGABRT) during the autotune forward passes. Reproduced exact crash as reported by an internal document. ROOT CAUSE on ROCm, F.scaled_dot_product_attention's flash path dispatches to the CK ck_tile::FmhaFwd kernel. For tile-unaligned attention shapes (here seqlen_q=4, not a multiple of kM0=16; seqlen_k=656, not a multiple of kN0=32) the CK forward dispatcher selects a no-seqlen-pad instance (kPadSeqLenQ= kPadSeqLenK=false). That kernel issues unpredicated full-tile loads/stores of round_up(seqlen, tile) rows, reading past the end of Q/K/V (and writing past O/LSE). The over-read is numerically harmless (the tail is masked out in softmax), but it faults when the tensor's allocation ends on an unmapped page. Decoded from the faulting kernel's kernarg: bf16, b=5, h=2, seqlen_q=4, seqlen_k=656, head_dim=128, bias_ptr=NULL -- i.e. NOT the attention-bias path. FIX 3 orthogonal candidates: app layer, host stopgap, device fix. This is the host stopgap. (app layer): A non-bug-exposer app could have its input data checked to be tile-aligned. (host-side stopgap): when seqlens are tile-unaligned, hand the kernel seqlen-padded allocations of Q/K/V (read) and O/LSE (write) -- the extra rows are mapped scratch -- while leaving seqlen_q/seqlen_k passed to the kernel unchanged, so the kernel's logical masking and numerics are identical. The real rows are copied back into the caller-visible output/LSE. Engages only on unaligned shapes; padding LCM (256) >= every fwd kM0/kN0 in the generated instance set. (Device side): CK should generate/select kPadSeqLen forward instances for unaligned seqlens upstream, after which this host guard can be removed. I will upstream this fix with a separate diff and PR. Test Plan: Repro: deterministic "Memory access fault by GPU node-N ... Reason: Unknown" (SIGABRT) during AOTT GPU lowering on MI350X (gfx950) for tile-unaligned attention shapes (e.g. seqlen_q=4 not a multiple of kM0=16; seqlen_k=656 not a multiple of kN0=32). Decoded faulting kernel: bf16 ck_tile::FmhaFwd, no bias (bias_ptr=NULL), kPadSeqLenQ=kPadSeqLenK=false, head_dim=128 -- an unpredicated tile tail over-read that faults only when the buffer is page-terminal. Fix verification (MI350X / gfx950): - Built the lowering binary with this fix and re-ran the exact failing lowering job: all autotune forward passes complete, 0 memory faults, scripted-model validation forward pass completed successfully, all lowered .so packed. CK fast path stays ENABLED (not disabled). - Standalone CK-backend SDPA at the faulting shape: outputs finite, max|abs-diff| vs fp32 math ~7e-4 (numerically unchanged before/after the fix). - arc lint: clean (fbcode + xplat copies). Reverse-repro (ablation, MI350X/gfx950): rebuilt the lowering pkg from the same base with ONLY the seqlen-pad guard removed -> the "Memory access fault ... Reason: Unknown" SIGABRT returns on autotune forward pass 5/5, while the with-guard build is clean. The guard is the sole variable that flips crash<->pass, confirming causality in both directions. Reviewed By: liangbeixu Differential Revision: D108370469 Pull Request resolved: pytorch#187152 Approved by: https://github.com/jeffdaily (cherry picked from commit a934255)
Fixes pytorch#191564 ## Summary `DeviceIndex` is `signed char`, while plain `char` is unsigned on ppc64le. Casting `q.get_device()` to plain `char` makes `CUDAGuard`'s list initialization convert a non-constant unsigned `char` value to `DeviceIndex`, which is a narrowing conversion. Construct `CUDAGuard` from `q.device()` at all 12 affected sites instead. This selects the existing `CUDAGuard(Device)` overload, preserves the tensor's device type and index, and fixes the root cause rather than suppressing the narrowing diagnostic with another cast. ## Testing - Compiled the affected ROCm translation units with AMD clang 23 on ppc64le. - Completed a full ROCm-based PyTorch build successfully. - Confirmed that no affected `char`-based `CUDAGuard` constructions remain under `aten/`, `torch/csrc/`, or `c10/`. - Ran the applicable local source checks on all six modified files, including the CUDA source. - CUDA compilation coverage is provided by CI. Pull Request resolved: pytorch#192005 Approved by: https://github.com/jeffdaily (cherry picked from commit 4892ba4)
trying to make things more symmetric between regular sdpa and varlen This commit was authored with an AI assistant. Variable-length attention previously selected eligible cuDNN internally regardless of the active `sdpa_kernel` context, and backward recomputed that decision. A forced Flash request could therefore run cuDNN, while a context change between forward and backward could select different backends. Select from the enabled cuDNN and Flash backends before entering the private custom op, report the constraints that prevent a forced cuDNN request, and save the selected backend for backward. Eligible cuDNN remains the default when both backends are enabled. Compiled graphs capture backend enablement at trace time, matching the existing behavior of the SDP enablement getters. ```bash gpu-run auto -- env PYTHONPATH=$PWD /home/dev/.venvs/dev/bin/pytest -q -rs test/test_varlen_attention.py -k "sdpa_kernel_backend_selection or sdpa_kernel_backend_errors or custom_op_compliance or custom_op_registration" gpu-run auto -- env PYTHONPATH=$PWD /home/dev/.venvs/dev/bin/pytest -q test/test_varlen_attention.py gpu-run auto -- env PYTHONPATH=$PWD /home/dev/.venvs/dev/bin/pytest -q -rs test/test_flop_counter.py -k varlen git diff --check spin quicklint lintrunner -a torch/nn/attention/varlen.py test/test_varlen_attention.py ``` Pull Request resolved: pytorch#193631 Approved by: https://github.com/liangel-02 ghstack dependencies: pytorch#193630 (cherry picked from commit 7476810)
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
7 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.
Summary
Cherry-pick upstream SDPA/ROCm attention fixes onto
release/2.12. These fixes landed onupstream/mainandrelease/2.13after the 2.12 branch cut and are needed for ROCm CI stability.Verified fixes (MI300X / gfx942, ROCm 10.1)
19 previously failing tests now pass:
test/test_transformers.pytest_flash_attention_ck_gqa_seqlen_q_1_cuda, 8×test_sdpa_zero_qk_head_dim_*test/test_meta.pytest_scaled_dot_product_flash_attention_for_cpu_logsumexp_dtype_*test/test_linalg.pytest_ck_blas_library_mm_{float32,bfloat16}_cuda_*Test plan
pytest test/test_transformers.py -k "flash_attention_ck_gqa_seqlen_q_1 or sdpa_zero_qk_head_dim"pytest test/test_meta.py -k "logsumexp_dtype"pytest test/test_linalg.py -k "test_ck_blas_library_mm"pytest test/test_varlen_attention.py -k "sdpa_kernel_backend_selection or sdpa_kernel_backend_errors or custom_op_compliance"