[C/Pytorch] Expand layout support for fused attention by cyanguwa · Pull Request #403 · NVIDIA/TransformerEngine

cyanguwa · 2023-08-25T21:54:11Z

This PR expands our layout support for fused attention and unifies the interface for all three attention backends. Specificially, it

Adds a third C/PyTorch API, _fused_attn_fwd_q_k_v and _fused_attn_bwd_q_k_v, which accept q, k, v as three separate tensors regardless of how they are laid out in memory.
Adds functionality to figure out the layout of q, k, v tensors and expand our support to not only bs3hd and bshd_bs2hd but also h3d, hd_h2d, hd_hd_hd based layouts. This saves some transpose kernels before and after fused attention.

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa · 2023-08-28T23:53:52Z

/te-ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa · 2023-09-15T15:11:01Z

Pipeline 9821365 has passed. tests/pytorch/test_onnx_export.py::test_export_core_attention is skipped due to ONNX not taking str arguments. The issue exists in our other ONNX tests as well and needs to be addressed separately.

timmoon10

Overall looks good. I see some artifacts of debugging and have some minor design suggestions. The biggest question is if we're breaking ONNX support.

…ze checks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa · 2023-09-22T20:12:03Z

Pipeline 9935088

timmoon10

LGTM, pending CI. We should also think carefully so we can settle on one mask type API.

timmoon10 · 2023-09-22T20:19:32Z

- if self_attn_mask_type is not None:
- warnings.warn(
- "Argument :attr:`self_attn_mask_type` has been moved to the `forward` method and"
- "is deprecated. It will be fully removed in future releases.",
- category=DeprecationWarning,
- )


We deprecated this option recently in #402, so I think it's worth a proper discussion if we want to bring it back. For the record, I think this API is the right one: allow the user to specify a mask type in the constructor and to override their choice in forward. This allows us to handle both the ONNX case with a static compute graph, as well as the case where the mask changes between training and inference.
Pinging @ksivaman.

Yeah, we reverted this because onnx doesn't take dynamic string arguments, so things like the mask type or qkv format would have to be specified in the init rather than forward.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman · 2023-09-23T18:01:45Z

/te-ci

* add flexible layout support Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add support for flexible qkv layout Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add more changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes for compiling Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove redudant file Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix options device error Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix typos Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more changes; WIP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more changes; WIP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes and tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes and wrong results Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * sb3hd/bs3hd working on top of 3xsbhd/bshd/thd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix dQ, dK, dV Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add nvtx Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove qkvso_strides on torch side; cover it in generateQKVStrides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * all 15 layouts pass Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add workspace optimization Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes and test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * removed most debug info/clean up Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add note to deprecate some qkv layouts Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix code for unit tests in test_fused_attn.py Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * further remove debug info Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove a couple more comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix numerics tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes for lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fp8 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix onnx for core attn; not fixed Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove nvtx and add env var for workspace opt Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove testing for env var Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace zeros/zeros_like with empty/empty_like Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix nvtx marker name for _q_k_v API Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove sm80 when compiling for h100 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add mapping from qkv layout to layout group and qkv format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up enums mapping and remove trailing spaces Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * simplify workspace opt control logic; only need env var Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fp8 test, and minor modifications for other tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * avoid overwriting model configs in unit test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * random fixes/improvements: get_qkv_format/etc, default values, docstrings, comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix minor issues: invalid syntax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * change workspace opt logic back to FORCE_WORKSPACE_OPT Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix FP8 tests and generateStrides function Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix get_backend logic for max512/arbitrary Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix unit tests; need cleanup Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up unit tests for layouts, and fix minor lint issue Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks for CI testing: onnx string issue and test fused attn first Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove one unsupported layout from max512 and add a check to qkvpacked API Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix te layer test; reduce test time Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert compiler option changes; add back sm80 for even h100 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove some unit tests or make them optional to reduce CI time Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove more unit tests temporarily Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove _q_k_v in naming and add NVTE_ERROR for FP8 Aux_CTX_Tensors size checks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add more deprecation notes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove temp tests from last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace with te::getenv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints from last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove redundant contiguous() Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove thd->bs3hd user warning to avoid GPU sync Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * adjust fused attn bs in tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * temporary fix for onnx issue; more fixes in PR 437 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unused variables Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Charlene Yang Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

cyanguwa added 24 commits August 10, 2023 17:17

add flexible layout support
4b5c08f
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add support for flexible qkv layout
b9b2eec
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add more changes
09cfaba
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fixes for compiling
b2cab71
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove redudant file
bcef882
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix options device error
cad4cc0
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix typos
fe7a45e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

merge main and fix conflicts
d2cffbd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

more changes; WIP
aeae983
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

more changes; WIP
e51c45c
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fixes and tests
b41292e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fixes and wrong results
d48915b
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

sb3hd/bs3hd working on top of 3xsbhd/bshd/thd
fb879a3
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix dQ, dK, dV
fe6a9a8
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add nvtx
2d9473b
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove qkvso_strides on torch side; cover it in generateQKVStrides
9796df4
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

all 15 layouts pass
c1bd49a
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add workspace optimization
25edaca
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into fused_attn/dpa_layout
3c7ad01
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

minor fixes and test
2e2bb08
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

removed most debug info/clean up
b21fab2
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add note to deprecate some qkv layouts
f3b6da1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix code for unit tests in test_fused_attn.py
c1ebbed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

further remove debug info
00a7355
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa added 5 commits August 28, 2023 17:05

remove a couple more comments
93028f8
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix numerics tests
1324ced
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fixes for lint
3589d0a
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix fp8 tests
dd9956e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix onnx for core attn; not fixed
accb77e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove more unit tests temporarily
0b40cd1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

timmoon10 reviewed Sep 18, 2023
View reviewed changes

Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h

timmoon10 self-requested a review September 18, 2023 23:18

zlsh80826 reviewed Sep 19, 2023
View reviewed changes

cyanguwa added the 0.13.0 label Sep 20, 2023

timmoon10 reviewed Sep 21, 2023
View reviewed changes

Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/pytorch/attention.py Outdated
Comment threadtransformer_engine/pytorch/attention.py Outdated

timmoon10 self-requested a review September 21, 2023 00:36

timmoon10 requested changes Sep 21, 2023
View reviewed changes

Comment threadtests/pytorch/test_onnx_export.py Outdated
Comment threadtests/pytorch/test_fused_attn.py Outdated
Comment threadtests/pytorch/test_fused_attn.py

remove _q_k_v in naming and add NVTE_ERROR for FP8 Aux_CTX_Tensors si…
727651d
…ze checks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

timmoon10 mentioned this pull request Sep 22, 2023
[Paddle] Some refactor and optimization on fused attention #411
Merged

cyanguwa added 10 commits September 22, 2023 09:52

add more deprecation notes
aa62706
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove temp tests from last commit
aa4e2bc
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

replace with te::getenv
838cd1e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove prints from last commit
b72c3ec
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove redundant contiguous()
4aa72a6
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove thd->bs3hd user warning to avoid GPU sync
ca1be26
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

adjust fused attn bs in tests
d165a62
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

temporary fix for onnx issue; more fixes in PR 437
4a11d56
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into fused_attn/dpa_layout
e2128cc
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

remove unused variables
6fd9262
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

timmoon10 approved these changes Sep 22, 2023
View reviewed changes

timmoon10 mentioned this pull request Sep 23, 2023
Debug CI tests on Ada #397
Merged

fix and resolve
82cd3e8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman merged commit 76669cd into NVIDIA:mainSep 25, 2023

LeonardoEmili mentioned this pull request Nov 27, 2023
Nemotron-3-8b finetuning NVIDIA-NeMo/Speech#7939
Closed

hungryGeek16 mentioned this pull request May 31, 2026
fix unfused padding causal sdpa #3063
Open

Uh oh!

Conversation

cyanguwa commented Aug 25, 2023• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cyanguwa commented Aug 28, 2023

Uh oh!

cyanguwa commented Sep 15, 2023• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cyanguwa commented Sep 22, 2023

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

timmoon10Sep 22, 2023• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwaSep 22, 2023

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Sep 23, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

cyanguwa commented Aug 25, 2023•
edited
Loading

cyanguwa commented Sep 15, 2023•
edited
Loading

timmoon10 left a comment •
edited
Loading

timmoon10Sep 22, 2023•
edited
Loading