GitLab

Signed-off-by: wenjh <wenjh@sugon.com>

Signed-off-by: wenjh <wenjh@sugon.com>

Signed-off-by: wenjh <wenjh@sugon.com>

Signed-off-by: wenjh <wenjh@sugon.com>

Signed-off-by: wenjh <wenjh@sugon.com>

Signed-off-by: wenjh <wenjh@sugon.com>

Signed-off-by: wenjh <wenjh@sugon.com>

Add env to chose blocklen of blockwise quantize.
Signed-off-by: wenjh <wenjh@sugon.com>

Fix pytest of blockwise error
Signed-off-by: wenjh <wenjh@sugon.com>

Resolve new api in  int8 gemm test
Signed-off-by: wenjh <wenjh@sugon.com>

Fix incorrect launch parm
Signed-off-by: wenjh <wenjh@sugon.com>

Fix 1D blockwise(64) acc error
Signed-off-by: wenjh <wenjh@sugon.com>

* add support for head dim > 128
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* remove debugging
Signed-off-by: Charlene Yang <charleney@nvidia.com>

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

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



* raise tols slightly to tolerate 1/2048 mismatches
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix is_training for test_te_layer
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add bprop support for blackwell
Signed-off-by: Charlene Yang <charleney@nvidia.com>

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

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



* minor tweak for format
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix backend selection results
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* bump sm100 to sm100+
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add sq=1 test for MLA
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* enable sq=1 for bprop
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor tweak in comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix head_dim logic and remove pytest skip
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add FE fix for d>128
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* update FE again to take in small fixes
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add cuDNN version info in L0 tests
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* increase tols for Unfused + large dim
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* Revert "add cuDNN version info in L0 tests"

This reverts commit 3e1b426ca5319a2c0540b9e73bba7047d0e583e5.
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix tols for Unfused
Signed-off-by: Charlene Yang <charleney@nvidia.com>

---------
Signed-off-by: Charlene Yang <charleney@nvidia.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>

* Added support of FP4 data type
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring to BitsNum in progress
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed compilation errors. All C++ tests passed
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed a typo
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Added FP4 guard to TMA tensor descriptor data type
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Fixed errors in JAX C++ extensions
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Removed dummy NVFP4 C++ test file
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

* Refactored the code per the review notes. Fixed JAX build error.
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Removed unnecessary static casts
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Typo fix
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>

* Pass correct num bits to create_2D_tensor_map; fixes CI
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

---------
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add FP8 current scaling to te.Sequential tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Helper function for test/ref tensors does not produce quantized tensor by default
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add FP8 current scaling to distributed te.Sequential tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add FP8 current scaling to Userbuffers te.Sequential tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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



* Debug MXFP8 tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Do not initialize quantized weights with column-wise usage in inference mode
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix bug in test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use no-grad mode instead of inference mode in tests
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

* Flatten basic op params during fuser init
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 949abe97070721b1da5117903067608250f5fb61)

* Add caching for is_non_tn_fp8_gemm_supported
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit fd830ae24ffbd2d0727010b1a8a119ca72f61ce5)

* Pass fuser to _OperationFuserAutogradFunction.forward and moving computation to __init__
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit fd808991993958b670726896254b82fcb967fa07)

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

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



* Pass basic_op_kwargs and is_grad_enabled as parameters rather than in fuser
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Add support for overlapping wgrad NCCL AG with dgrad GEMM
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* Remove unused wait on memcpy API from UB
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* Add better commenting to MXFP8 overlap
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

---------
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
Co-authored-by: dastokes <dastokes@dastokes-dvt-01.nvidia.com>

* Support L2Norm basic op
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Add L2Norm module wrapper
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Expose qk_norm to MHA nd transformer laayer
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Move tests into separate file
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* fix pass
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Add license
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Remove  module
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Resollve comments
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



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

* Added double buffering support initial commit
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>

* Fixed bugs
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>

* Make only one double buffer creation
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>

* Fixed bug
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>

* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>

* Fixed flag setting
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>

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

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



* Merge conflict
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>

* fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* lint fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <anandaraj@wisc.edu>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche02.ptyche.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

typo fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Fix for loading old ckpt formats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* rm not needed funcs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Fix JAX build in baremetal env
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update install inst in readme
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* docstring fix
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>

* Support MLA (qk_dim != v_dim) for AttnFuncWithCPAndKVP2P
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>

* add UT for MLA CP
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>

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

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



* refine the code
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>

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

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



* refine the code
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>

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

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



---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

* Manage deps and add einops
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Lora spike
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos02.a51.clusters.nvidia.com>

* Added FP8 param support
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos02.a51.clusters.nvidia.com>

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

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



* lint fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos02.a51.clusters.nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos02.a51.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>

only get skip_fp8_weight_update in fp8_graph_capturing
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* refactor the multi_stream utils + implement nvte_multi_tensor_quantize in TE/Common

* implement GroupedQuantizer and grouped_quantize in jaxx

* fix logical_axes_names for transpose tensor in ScaledTensor
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: Ming Huang <mingh@nvidia.com>

[PyTorch] FP8 Subchannel Recipe With FP8 Gather And Configurable Scaling Factor Tensor Swizzling (#1707)

* functional kernel for columnwise + no-transpose option, still hacky
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* pass all quantizer unit tests
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* refactor, add gemm ready api
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* make format options private members, simplify api
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* swizzle scales right before gemm
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* bug fix of single layer test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* attempt to fix lint issue
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* fp8 gather pass, need minor refine
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* fix return_layernorm_output_gathered case
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* remove special cases, add sanity check before gemm
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix lint
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* lint ungrouped imports
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* Implement dequantize for compact 1D blocks.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* add more unit test with dequantize compact supported
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* lint again
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* make ag for subchannel respect async
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* zero tolerance in distributed test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix zero tolerance test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* resolve rebase issues
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* lint & format
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix lint
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* clean up
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* bug fix
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* relax rtol for fp32 distributed test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix some ci issue
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix ci test failure in debug mode
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Force row-wise and column-wise data to have same data format

Prototype "all-gather usage" in quantizer.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove dead logic for high-precision AGs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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



* Debug FP8 block-wise tensor tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug distributed test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Handle case where LayerNormLinear returns gathered norm output
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix debug mode
Signed-off-by: zhongboz <zhongboz@nvidia.com>

---------
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Keith Wyss <kwyss@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Keith Wyss <kwyss@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>