[PyTorch] Minor fixes for TE 2.2 by cyanguwa · Pull Request #1589 · NVIDIA/TransformerEngine

cyanguwa · 2025-03-19T06:20:12Z

Description

This PR fixes a few things for TE 2.2:

truncate cu_seqlens to the actual_batch_size + 1 length for bshd/sbhd formats for cuDNN 9.8+
after the above 9.8+ fix, re-enable padding_causal_bottom_right and multi-layer tests for FusedAttention
revert the name change in InferenceParams's max_sequence_length parameter from max_seqlen_kv in Add KV cache for paged/non-paged attention #1355
rename test_paged_attn tests to test_kv_cache since it include both paged and non-paged tests
add :members: to InferenceParams in pytorch.rst docs
include CUDNN::cudnn_all in link libraries in CMakeLists.txt so that custom CUDNN_PATH would work
remove two redundant None returns in FusedAttnFunc.backward() which has been around for a few TE versions but didn't do any harm
add NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 to the error message when no backend is found

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

See Description

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

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

cyanguwa · 2025-03-19T06:20:38Z

/te-ci pytorch L0

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

for more information, see https://pre-commit.ci

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

cyanguwa · 2025-03-20T02:46:22Z

/te-ci pytorch L1

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

cyanguwa · 2025-03-21T22:18:01Z

/te-ci pytorch L1

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

cyanguwa · 2025-03-21T23:36:53Z

/te-ci pytorch L1

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

cyanguwa · 2025-03-22T18:31:32Z

CI: 25860235 for 9.9, and 25863870 for 9.8. Both passed, with expected failures (unrelated to this PR).

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

cyanguwa · 2025-03-24T22:18:45Z

/te-ci pytorch L0

cyanguwa · 2025-03-25T18:25:41Z

/te-ci jax L0

cyanguwa · 2025-03-25T22:40:22Z

CI looks good, expect some expected failures (unrelated to this PR).

* skip cuDNN 9.8 for KV caching Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert from max_seqlen_kv to max_sequence_length for InferenceParams Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * rename test_paged_attn to test_kv_cache Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove redundant None returns in bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add debug flags when no backend is found Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * skip kv_cache_accuracy tests for cuDNN 9.8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * truncate length of cu_seqlens for consistency with q/k/v shape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add back padding_brcm for fused attn tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * re-enable kv_cache_accuracy test for 9.8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix cuDNN search dir Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes based on review Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove extra empty line Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

skip cuDNN 9.8 for KV caching
6674ded
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa requested a review from ptrendx March 19, 2025 20:03

cyanguwa added the 2.2.0 label Mar 19, 2025

cyanguwaand others added 7 commits March 19, 2025 16:50

revert from max_seqlen_kv to max_sequence_length for InferenceParams
fdc5ff3
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

rename test_paged_attn to test_kv_cache
9bcfa4e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove redundant None returns in bwd
ff5e26c
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add debug flags when no backend is found
678039f
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into skip_9.8
e390b4e

[pre-commit.ci] auto fixes from pre-commit.com hooks
144967b
for more information, see https://pre-commit.ci

skip kv_cache_accuracy tests for cuDNN 9.8
c9cc1cd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa mentioned this pull request Mar 21, 2025
Add KV cache for paged/non-paged attention #1355
Merged
13 tasks

cyanguwa added 4 commits March 21, 2025 13:52

truncate length of cu_seqlens for consistency with q/k/v shape
c8c4f11
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add back padding_brcm for fused attn tests
a051ed5
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into skip_9.8
e5b1bcc

re-enable kv_cache_accuracy test for 9.8
b810bdc
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix cuDNN search dir
edb4e8c
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa changed the title ~~[PyTorch] Skip cuDNN 9.8 for KV caching~~[PyTorch] Minor fixes for TE 2.2Mar 22, 2025

Merge branch 'main' into skip_9.8
a0da31f
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa mentioned this pull request Mar 22, 2025
Fix issues in fused_attn_bwd #1574
Merged
13 tasks

cyanguwa requested review from ptrendx and removed request for ptrendx March 24, 2025 18:07

ptrendx reviewed Mar 24, 2025
View reviewed changes

Comment threadtests/pytorch/test_numerics.py Outdated
Comment threadtransformer_engine/common/CMakeLists.txt Outdated
Comment threadtransformer_engine/pytorch/attention.py

cyanguwa added 3 commits March 24, 2025 14:59

fixes based on review
169ffeb
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

remove extra empty line
b4501e9
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into skip_9.8
5efbdb6

cyanguwa requested review from KshitijLakhani and ptrendx March 24, 2025 22:18

ptrendx approved these changes Mar 25, 2025
View reviewed changes

cyanguwa mentioned this pull request Mar 25, 2025
[CMake] Add CUDNN::cudnn_all target to include CUDNN_INCLUDE_DIR #1605
Closed
8 tasks

ksivaman reviewed Mar 25, 2025
View reviewed changes

Comment threadtransformer_engine/pytorch/dot_product_attention/inference.py

cyanguwa merged commit 65c2798 into NVIDIA:mainMar 25, 2025

cyanguwa mentioned this pull request Mar 26, 2025
[PyTorch] Make breaking change in InferenceParams.init more explicit #1619
Merged
13 tasks

pramodk mentioned this pull request Jun 7, 2025
Compiling on Slurmcluster fatal error: cudnn.h: No such file or directory #918
Closed

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

Navigation Menu

Uh oh!

[PyTorch] Minor fixes for TE 2.2 - #1589

[PyTorch] Minor fixes for TE 2.2#1589
cyanguwa merged 17 commits into
NVIDIA:mainfrom
cyanguwa:skip_9.8

cyanguwa commented Mar 19, 2025•
edited
Loading

Uh oh!

cyanguwa commented Mar 19, 2025•
edited
Loading

Uh oh!

cyanguwa commented Mar 20, 2025

Uh oh!

cyanguwa commented Mar 21, 2025

Uh oh!

cyanguwa commented Mar 21, 2025•
edited
Loading

Uh oh!

cyanguwa commented Mar 22, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cyanguwa commented Mar 24, 2025

Uh oh!

cyanguwa commented Mar 25, 2025

Uh oh!

Uh oh!

cyanguwa commented Mar 25, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

cyanguwa commented Mar 19, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

cyanguwa commented Mar 19, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cyanguwa commented Mar 20, 2025

Uh oh!

cyanguwa commented Mar 21, 2025

Uh oh!

cyanguwa commented Mar 21, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cyanguwa commented Mar 22, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cyanguwa commented Mar 24, 2025

Uh oh!

cyanguwa commented Mar 25, 2025

Uh oh!

Uh oh!

cyanguwa commented Mar 25, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cyanguwa commented Mar 19, 2025•
edited
Loading

cyanguwa commented Mar 19, 2025•
edited
Loading

cyanguwa commented Mar 21, 2025•
edited
Loading