GitLab

* Adding Amax Primitive and related args.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Enable local-amax for current-scaling and optionally run AR aross FSDP/TP/SP.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding doc for Amax Primitive.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the function name conflict.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Modification as feedback suggested.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix errors from lint.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the wrong amax-scope in the bwd.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Added more description for amax-scope
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the wrong attribute name.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Keep dim for AmaxCalcuation.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove keepDim and add shardy_rule
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix shardy_rule
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove extra-collective bytes from ref_coll_count due to local amax.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

* remove import jax.extend.ffi
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add swizzle in jax
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* added outer_impl
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* clean up FFI
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* Decouple recipe and scaling mode
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Expose global QuantizeConfig instance as a getter
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Format and lint
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Merge branch 'main' into dev/jberchtold/jax-scaling-mode-and-recipe-decoupling
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Rename UsageType to TensorSource
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Update test_layer.py
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

* enabled TE GEMM for all recipes
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add warnings
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

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



* fix lint
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Revert "[JAX] Disable TE Norm Custom Calls (#1993)"

This reverts commit 6c970612

.

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Disable Norm custom calls
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add manage_primitives() helper

* disable GEMM primitives for non-MXFP8 recipes

* implement the NVTE_JAX_CUSTOM_CALLS + deprecate NVTE_JAX_CUSTOM_CALLS_RE

* replace NVTE_JAX_CUSTOM_CALLS_RE with NVTE_JAX_CUSTOM_CALLS in TE tests and examples

* fix use_jax_gemm contextmanager
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Make primitive names more granular for better disabling granularity
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Add experimental Shardy support.

Production use is not yet recommended.

---------
Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

* refactor + mxfp8

* added grouped gemm

* rename linear to dense

* added cublas init phase for groupedGemm

* relax the tol of test encoder multiprocessing mxfp8 by 0.001
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Hua Huang <huah@nvidia.com>
Co-authored-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

* softmax custom calls with correct encapsulates

* rm jax deprecated features

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* Add enabled() to BasePrimitive

* Add layernorm/rmsnorm fallback

* Add cast_fp8 fallback

* Add transpose/cast_transpose XLA fall back

* Act_lu fallback

* Add transpose fallback

* Add softmax fallback

* Unify the use of _cast_fp8

* Add tests for NVTE_JAX_CUSTOM_CALLS_RE

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

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

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

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

* Splitted cpp_extensions.py, renamed mlp.py and fused_attn.py
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* fixed import in tests
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>