Commits · eb64ec2a26b1091a78af424deb7d5204bee17cd2 · OpenDAS / TransformerEngine

18 Sep, 2023 1 commit
- Add an option for TP-only AMAX reduction (#431) · eb64ec2a
  Sangkug Lym authored Sep 18, 2023
```
Signed-off-by: Sangkug Lym <slym@nvidia.com>
```
  eb64ec2a
12 Sep, 2023 1 commit

[C/PyTorch] Add workspace optimization for fused attention arbitrary_seqlen backend (#396) · a150d286

cyanguwa authored Sep 12, 2023



* add workspace optimization for arbitrary_seqlen fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix whitespace for lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add use_workspace_opt to cudnn plan cache and fix workspace estimate
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* modify workspace opt logic; move zero fill to FP8 API only; other minor fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix try/catch
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix std string error when input is nullptr
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Add = for required vs allowed workspace comparison
Signed-off-by: cyanguwa <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>

a150d286

26 Aug, 2023 1 commit

[PyTorch] move mask types to fprop (#402) · 6aa1fcc8

Kirthi Shankar Sivamani authored Aug 25, 2023



* API change and some test fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* more test fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fixed fused attention tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rm duplicate test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

6aa1fcc8

25 Aug, 2023 1 commit

Fix rng_state issue and minor compiler warning (#395) · b90a8d3a

cyanguwa authored Aug 25, 2023



fix rng_state issue and minor compiler warning
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

b90a8d3a

19 Aug, 2023 2 commits

[PyTorch] Catch misaligned address errors in softmax (#390) · 3b7b7c68
Kirthi Shankar Sivamani authored Aug 19, 2023
```
Catch misaligned address errors
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
3b7b7c68

PyTorch MultiheadAttention API (#387) · 8aa2da17

Kirthi Shankar Sivamani authored Aug 19, 2023



* PyTorch MultiheadAttention API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix ONNX export tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Expose MultiheadAttention for import
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Expand mask type and add no mask numerical test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

8aa2da17

18 Aug, 2023 1 commit
- fix for amax_and_scale_update when reduce_amax=False (#386) · f29efb77
  Sudhakar Singh authored Aug 18, 2023
```
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
```
  f29efb77
16 Aug, 2023 2 commits

Add CPU initialization method (#368) · 2451c8d1

Kirthi Shankar Sivamani authored Aug 16, 2023



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

* Fix default value
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change API and add to RMSNorm
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

2451c8d1

[PyTorch] FP8 fixes (#380) · 2e0bfbd9

Kirthi Shankar Sivamani authored Aug 16, 2023



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

* Reorder methods by purpose
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Save full global state
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* More fixes to test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

2e0bfbd9

11 Aug, 2023 1 commit

Miscellaneous fixes for core attention (#344) · cbfb8c6b

cyanguwa authored Aug 10, 2023



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

* add back pytorch csrc extensions.h
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add unit tests for dpa checkpointing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove seqlen%32/64 checks for now
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix tests for core attn bias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add tests for changes regarding rng_state in aux_ctx_tensor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* reuse rng tracker from numerics in fused attn; skip checkpointing if FAv2 in numerics
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* uncomment comments used for testing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix pre/post scale bias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Update transformer_engine/pytorch/attention.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

* remove skipifs for FAv2 check after PR366
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove checkpointing tests for transformer layer; dpa tests still provide coverage
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust random number range for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Add upper bound to FA version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Check backend only when using FusedAttention
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove imports/variables related to FAv2 checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further fix random number ranges for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix variable referenced before assignment error
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: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

cbfb8c6b

10 Aug, 2023 4 commits

Add mixed use of cuDNN fprop and flash-attn v2 bprop (#349) · ecd4f808

cyanguwa authored Aug 10, 2023



* Add support for cuDNN fprop and FAv2 bprop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* skip activation recompute tests if FAv2
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* restrict the use of FAv2 bprop to H100 only
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move use_FAv2_bwd check to init
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove skipifs for FAv2 in test numerics
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix typos and wording for deterministic checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Remove variables related to FAv2 skipifs
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: cyanguwa <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>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ecd4f808

AMP support for LN and RMSNorm (#371) · 3cc2c1d2
Kirthi Shankar Sivamani authored Aug 10, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
3cc2c1d2
Fix pybind strings for RMSNorm (#372) · 88c88654
Kirthi Shankar Sivamani authored Aug 10, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
88c88654

fix linear sequence parallel when weight is frozen (#373) · e056664f

Jason Wang authored Aug 10, 2023



* fix linear sequence parallel when weight is frozen
Signed-off-by: jasonwan <jasonwan@nvidia.com>

* Extend fix to all APIs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

e056664f

09 Aug, 2023 1 commit

Disable FAv2 for deterministic use (#366) · 7804d116

Kirthi Shankar Sivamani authored Aug 08, 2023



* Disable FAv2 for deterministic use
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Also disable FusedAttention backend with deterministic
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

7804d116

08 Aug, 2023 2 commits

Optimize calls to .cpu() during checkpointing (#363) · 86d148f9

Kirthi Shankar Sivamani authored Aug 08, 2023



* Optimize calls to .cpu() during checkpointing
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes for ONNX
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

86d148f9

Fix for the RMSNorm tests/doc/ONNX export to match the actual implementation (#364) · a0e1cf99
Przemyslaw Tredak authored Aug 09, 2023
```
Fix for the RMSNorm tests/doc/ONNX export to match the actual
implementation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
a0e1cf99

02 Aug, 2023 1 commit
- Store FP8 checkpointing data in CPU (#351) · 85928d08
  Kirthi Shankar Sivamani authored Aug 02, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  85928d08
29 Jul, 2023 1 commit

Add support for multi-query and grouped-query attention (#338) · 1cb4b25a

cyanguwa authored Jul 28, 2023



* add support for multi-query/grouped-query attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* revert to flash-attn 1.0.6 and build 2.0.0.post1 manually in CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add keyword name for DPA input
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix skipif for pytest
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Update transformer_engine/pytorch/attention.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update tests/pytorch/test_fused_attn.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix TP and SP case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add skipifs for pytest
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove higher limit for flash-attn version
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: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

1cb4b25a

28 Jul, 2023 1 commit
- [PyTorch] Minor bug fix (#343) · 59c0f096
  Kirthi Shankar Sivamani authored Jul 28, 2023
```
Minor bug fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  59c0f096
27 Jul, 2023 1 commit

Exposing RMSNorm in pyTorch (#306) · fdc09f42

Przemyslaw Tredak authored Jul 27, 2023



* Exposing RMSNorm in pyTorch extensions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* First pass at the Python API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Small fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added numerics tests and fixed issues
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Lint fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added RMSNorm to LayerNormMLP
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added ONNX export and tests for RMSNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix python lint
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix BERT case
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added normalization option to the TransformerLayer
Added tests
Fixed test failures
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix documentation
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix kwarg bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix IMA and invalid type error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Increase RMSNorm threshold for bf16 case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix ONNX tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fdc09f42

26 Jul, 2023 1 commit
- Make QK layer scaling opt-in (#339) · 06d5fa97
  Kirthi Shankar Sivamani authored Jul 25, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  06d5fa97
20 Jul, 2023 1 commit
- Relax FA 2.0 checks for Ada (#331) · f0ddab82
  Kirthi Shankar Sivamani authored Jul 19, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  f0ddab82
19 Jul, 2023 3 commits

FlashAttention 2.0 support (#329) · faefacb8

Kirthi Shankar Sivamani authored Jul 18, 2023



* FA v2.0 support
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

faefacb8

Fix numerics for activation recompute (#327) · ded8b9bd
Kirthi Shankar Sivamani authored Jul 18, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
ded8b9bd

FA does not support head_dim > 64 on Ada (#328) · 2a81e939

Kirthi Shankar Sivamani authored Jul 18, 2023



* FA does not support head_dim > 64 on Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add cc8.7 to no FA list
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

2a81e939

18 Jul, 2023 1 commit

layer number bug fixes (#326) · 21162d06

Kirthi Shankar Sivamani authored Jul 18, 2023



* layer number bug fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

21162d06

15 Jul, 2023 1 commit

Disable TorchDynamo optimizations in PyTorch modules (#312) · c6538d6e

Tim Moon authored Jul 15, 2023



* Disable TorchDynamo optimizations in PyTorch modules
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add test for Torch Dynamo
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add torch.dynamo test to qa
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Skip torch.compile test for <v2.0
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

c6538d6e

14 Jul, 2023 2 commits

[PyTorch] Fix FP8 checkpointing for non forward execution cases (#323) · f351191b
Kirthi Shankar Sivamani authored Jul 14, 2023
```
Bug fix for checkpointing
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
f351191b

Deprecate unused APIs (#321) · 58d2ebab

Kirthi Shankar Sivamani authored Jul 14, 2023



* Deprecate unused APIs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

58d2ebab

13 Jul, 2023 4 commits

Fix dtype for KV inference cache (#319) · b172bad8
Kirthi Shankar Sivamani authored Jul 13, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
b172bad8

Fix FP32 LayerNorm ONNX export (#313) · 6bd35bf9

Neta Zmora authored Jul 13, 2023



* Fix FP32 LayerNorm ONNX export

When running inference use a fwd method that is registered with torchscript.
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

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

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

6bd35bf9

Remove buffer registration for FSDP like cases (#318) · 11c5d588
Kirthi Shankar Sivamani authored Jul 13, 2023
```
Remove extra buffers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
11c5d588

Catch cublas FP8 errors (#317) · 8c3110d1

Kirthi Shankar Sivamani authored Jul 12, 2023



* Better dimension assert for FP8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

8c3110d1

26 Jun, 2023 1 commit

Get default dtype from pytorch (#300) · 79a78cae

Kirthi Shankar Sivamani authored Jun 26, 2023



* Get default dtype from pytorch
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

79a78cae

23 Jun, 2023 1 commit

Fix layer_norm ONNX export (#293) · 574f1b41

Neta Zmora authored Jun 23, 2023



* Fix ONNX export of layer_norm

ONNX has a spec bug: ConstantOfShape supports all dtypes except for BF16.
To WAR we use dtype FP32 and then cast to BF16.

Will also issue a PR to the ONNX sig committee to change the spec in opset 20.
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

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

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

574f1b41

22 Jun, 2023 1 commit

Add long sequence support for fused attention (#237) · 5c58beaa

cyanguwa authored Jun 22, 2023



* add long sequence support and unify three backends for fused attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update cudnn-frontend to v0.9.1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace cpu_float2half_rn with __float2half_rn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix backend selection and NVTEDType
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* make cudnn plan caches thread_local
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* replace cuDNN throw with NVTE_CHECK
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix replacement of cuDNN throw with NVTE_CHECK
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force dropout probablity to 0 in inference mode
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change negInfinity to be consistent with m512 fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove float2half conversion for scale_dropout
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add back runtime api for sm detection
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add gemm3 to enums FP8Fwd/BwdTensors
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change dropout from no to yes for fmha_v1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove output_rng_state in m512 kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix elts_per_thread calculation in kvpacked fwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove dropout=0.0 restriction for m512 fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove output_rng_state completely from m512 kernels
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: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

5c58beaa

21 Jun, 2023 1 commit
- Fix BF16 ONNX export for successful ONNX Runtime Verification (#290) · 804f1203
  asfiyab-nvidia authored Jun 20, 2023
```
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
```
  804f1203
20 Jun, 2023 1 commit
- Consistent docs for fuse_wgrad_accumulation (#289) · 0426feb6
  Kirthi Shankar Sivamani authored Jun 20, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  0426feb6
16 Jun, 2023 1 commit

Fix softmax ONNX export (#282) · 3fbded65

Neta Zmora authored Jun 17, 2023



* Fix softmax ONNX export

* BF16 is validated using "fake i/o": ie. instead of using BF16 as input/output, use FP32 input/output and convert to/from BF16 in the forward method.

* Wrap softmax symbolic functions with conversion to/from FP32 to produce the same semantics as TE's softmax (compute is performed at FP32 precision regardless of input/output data type).
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* ONNX export Code refactoring

Share function compute_in_fp32 between softmax.py (softmax symbolic functions) and te_onnx_extensions.py (the rest of the symbolic functions).
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

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

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

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

3fbded65