Skip to content

GitLab

Switch branch/tag
  1. 28 Jan, 2026 1 commit
  3. 02 Jan, 2026 1 commit
  5. 29 Oct, 2025 1 commit
  7. 14 Oct, 2025 1 commit
  9. 28 Aug, 2025 1 commit
  11. 20 Aug, 2025 1 commit
  13. 29 Jul, 2025 1 commit
  15. 16 Jul, 2025 1 commit
  17. 07 Feb, 2025 1 commit
  19. 02 Jan, 2025 1 commit
  21. 16 Oct, 2024 1 commit
  23. 13 Aug, 2024 1 commit
  25. 06 Aug, 2024 1 commit
  27. 14 Jun, 2024 1 commit
  29. 12 Jun, 2024 1 commit
  31. 07 Jun, 2024 1 commit
  33. 06 Jun, 2024 1 commit
  35. 12 Apr, 2024 2 commits
  37. 03 Feb, 2024 1 commit
  39. 03 Jan, 2024 1 commit
  41. 31 Oct, 2023 1 commit
  43. 05 Oct, 2023 1 commit
  45. 04 Oct, 2023 1 commit
  47. 25 Sep, 2023 1 commit
  49. 23 Sep, 2023 1 commit
  51. 26 Aug, 2023 1 commit
  53. 19 Aug, 2023 1 commit
  55. 16 Aug, 2023 1 commit
  57. 11 Aug, 2023 1 commit
    • Neta Zmora's avatar
      Improve softmax ONNX export tests (#370) · a0f44354
      Neta Zmora authored 
      

* Add dynamically shaped input mask in test_export_softmax
* Fix test_softmax_mask_fn - use env. var `NVTE_ONNX_KVCACHE_MAX_SEQ_LEN` to control whether the test uses the default mask generation function or dynamic TRILU mask slicing.
* Change core_attention ONNX export test: use "no_mask" as attn mask type when testing `te.attention.DotProductAttention` w/o masking.
* Use ORT CUDA backend by default.
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>
      a0f44354
  59. 27 Jul, 2023 1 commit
  61. 13 Jul, 2023 1 commit
  63. 26 Jun, 2023 1 commit
  65. 20 Jun, 2023 1 commit
  67. 16 Jun, 2023 2 commits
  69. 13 Jun, 2023 1 commit
  71. 06 Jun, 2023 1 commit
  73. 13 May, 2023 1 commit
    • Neta Zmora's avatar
      Add env. var. for efficient text-generation in inference (#214) · 68f60b89
      Neta Zmora authored 
      

* Dynamically-generated causal attention mask (for ONNX export)

TE's default causal mask is square (seq_len, seq_len) and is
dynamically allocated for different sequence sizes. Dynamic
allocation and dictionary lookups are not supported by ONNX.
GPT generative phase uses rectangular masks.

This commit forces softmax to use `forward_torch_softmax` and
to dynamically generate an attention mask when exporting to ONNX.
The mask is generated w/o using conditional control-flow by generating
a  (k_seq_len, k_seq_len) mask and slicing it to (q_seq_len, k_seq_len)

An alternate implementation is to pre-allocate a mask of shape
(max_seq, max_seq) and to slice that. This solution is more performant
at the expense of space, but the problem is the TE doesn't have a concept
of max_seq.

* Add to test_export_softmax a test for te.softmax.FusedScaleMaskSoftmax.
* Add test_softmax_mask_fn to test that TE's default attention mask and
the new ONNX-compatible mask produce the same behavior.
* Add test_export_gpt_generation to test that the ONNX model can correctly
handle inputs with different shapes and that the attention mask it adjusted
on-the-fly to different sequence lengths.

Misc:
* Add a PRNG seeding fixture for more stability in tests.
* Add dynamic shapes for ONNX input/output tests.
* Allow validate_result to compare ORT output to pre-computed TE outputs.
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>

* Add NVTE_ONNX_KVCACHE_MAX_SEQ_LEN for efficient text-generation in inference

* Introduce an environment variable (NVTE_ONNX_KVCACHE_MAX_SEQ_LEN) to set the maximum sequence length.
In ONNX inference with KV-Cache optimizations for GPT text generation, the attention mask shape can be square (context-phase) or rectangular (generation-phase).
When exporting to ONNX and this variable is set, TE preallocates an upper triangular (k=1) matrix with a size as prescribed by the variable, and dynamically slices the mask for the required shape.
TE models can be exported to ONNX when NVTE_ONNX_KVCACHE_MAX_SEQ_LEN is not configured, but the attention masking is always square and not fit for efficient text generation.

* Work-around torch.onnx.export bug that incorrectly folds
layer_norm(data, scale=add(gamma,1)) to layer_norm(data, scale=gamma)
when we use LN with zero-centered gamma.

* ONNX export tests
  * Add a fixture (seed_default_rng) to seed the PRNG
  * Add a fixture (set_max_seq_len) to set the max sequence length when exporting to ONNX for GPT text generation
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>

* Fix linting errors
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>

* Remove immutable default values from a couple of function signatures
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>

* Add @skip_FP8 to test_export_gpt_generation
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>

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

* Fix CI error for softmax export
Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>

* Lint
Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: default avatarNeta Zmora <nzmora@nvidia.com>
Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
      68f60b89
  75. 09 May, 2023 1 commit