GitLab

* fix: update grad_output quant to avoid redundant work
Signed-off-by: kshitij12345 <kshitijkalambarkar@gmail.com>

* add test
Signed-off-by: kshitij12345 <kshitijkalambarkar@gmail.com>

* don't keep only columnwise quant if requires_dgrad=False
Signed-off-by: kshitij12345 <kshitijkalambarkar@gmail.com>

* fix stray merge
Signed-off-by: kshitij12345 <kshitijkalambarkar@gmail.com>

* fix for ctx.use_bias is True case
Signed-off-by: kshitij12345 <kshitijkalambarkar@gmail.com>

* Skip if FP8 not available
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* PDL for MXFP8 Quantize
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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



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

* Fixed integer overflow when computing offsets
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



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

* Fixed conflicts
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



* Minor code refactoring to avoid unnecessary checks
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 typo
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed dBias accumulation error due to initialization. Minor code refactoring
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



* Test case to reproduce the init 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



* Fixed rowwise dbias 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



* Changed ptx API
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 a struct for two packed FP8 values
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Rolled back to scalar code for columnwise scaling due to its better performance
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Minor corrections
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



* Rebased on main
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



* Fixes per code review
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 constexpr in C++ test suite to build faster
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Computed activations are now numerically truncated to InputType before scaling. Improved test suite.
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



* Minor refactoring
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Minor refactoring
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



* Modified mismatches checks of MXFP8 to address FP8 numerics
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Implemented Jeremy's fixes to JAX test suite with an intermediate downcast
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



* Reduced the dims of the test tensors to improve CI runtime
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed memory alignment issue. Compute dbias without downcast.
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 misaligned memory issue also in gated kernels. Reduced size of MXFP8 gated tests
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



---------
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.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>

[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>

* Changed the Tensor allocation strategy
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Disable debug flag
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



* Fix the double free error
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

* Fixed pyTorch recipe extension
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

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

* Hide TensorAllocator and fix the usage in LayerNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

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

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



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

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

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

* Add GEMM logic for blockwise quantized tensors.

GEMM test cases included in pytorch integration.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update NVTE_BLOCK_SCALING for GEMM.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Gate feature on CUDA 12.9
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Gemm typo.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Remove unecessary type converter change.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Reflect epilogue availability and test supported epilogues.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* GEMM simplifications from recipe branch.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Format py code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update GEMM DGelu tests to match support depending on output dtype.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Force pow2Scales in GEMM
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add GEMM test to pytorch test suite.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add copyright to GEMM test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update import for GEMM test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add license.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update test gemm supported predicate.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use sgemm like interfaces and naming.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Rewrite GEMM comment.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR Feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Recipe setup for Linear modules.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use 12.9 feature test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Run against tensor dumps from internal library.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update FIXME to TODO with linked issue.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update full recompute feature to save recipe.

The recompute context uses the same recipe
and fp8 settings as the original fwd pass.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR Feedback. Avoid reusing quantizer objects.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update logic in module.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Format py.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update for PP bug.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update test numerics.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update force_power_of_2 scales in the recipe.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update usage method to satisfy upstream changes.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* fix subchannel recipe in distributed test with bf16 gather
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Edit and cleanup BF16 gather code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update test import.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* support columnwise only mode to 1D quantize kernel
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Format and move enum
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Skip alloc.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* try async bf16 gather
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Format python code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Document and type code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update pytorch lint errors.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Dont set high precision dtype.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add test for sanity and CG; fix CG for sequential?
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Keep make_quantizers API stable

Update num_quantizers instead to pass cuda_graph tests.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Fix import name.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Rename recipe method.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Skip grouped linear sanity test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Set usage before BF16 gather.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* refactor for nvte_quantize_v2
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Format code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Cleanup nvte_quantize_v2
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Test fp32 scales.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Disable CUDA graph.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Simplify layernorm linear
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Cleanup layernorm linear.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* LayerNorm linear bwd gather logic.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Communication updates.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update transformer_engine/pytorch/ops/op.py

Apply MR comment change.
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: kwyss-nvidia <kwyss@nvidia.com>

* Lint fix.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Enable cuda graph tests.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Reduce chance of spurious failure and reword.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Review suggestions from @timmoon10
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update CPP tests.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update common.h
Signed-off-by: Xin Yao <yaox12@outlook.com>

* Update test_float8blockwisetensor.py
Signed-off-by: Xin Yao <yaox12@outlook.com>

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

* fix cpp warning
Signed-off-by: Xin Yao <xiny@nvidia.com>

* more fix
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

* Blockwise float8 quantizer and quantized tensor class.

The classes are configurable for 128x128 blocksize
and 1x128 blocksize via setting block_scaling_dim == 2,1 respectively.

Scale tensors are stored in a format emenable for matrix multiplication,
however the integration of matmul is deferred as a separate story.

Fusions of quantization and DBIAS or activation functions are not yet
implemented, and the dequantization is currently implemented in torch.

Tests for quantization are included in C++ and pytorch layers, with
exact comparison to reference quantizer behavior as well as an attempt
to hit interesting branches through the API such as tensor creation
in pytorch and CPP and dequantization of row and columnwise usage.

Two CUDA kernels for quantization are included, and are direct ports
of equivalents in the kitchen repository, where a subchannel recipe
has been used for end to end training.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Apply linting changes.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Alignment for 1D scaling for GEMM edge case.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Change API name.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Fix merge conflict with name change.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use common tensor map API.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Change API to use two scaling mode enums.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Fix typo.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update some call sites.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Tests for torch tensor API surface.

Since the quantized tensor is a tensor
subclass, these tests exercise torch hooks.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Reuse scale calculation between quantizer refs.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Save memory by dropping reference to saved tensors.

Issues previously observed are solved.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Remove constexpr parameters from kernel.

Code size is reduced with fewer constexpr params.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Merge conflict from rebase.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add shape implementations for block scaling.

nvte_shape was added upstream. Logic added
for block scaled fp8.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Move benchmark to te_playground
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Remove amax_epsilon and pow_2_scales from tensor.

Hardcodes the default values.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Lint changes.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Fixup MR changes that broke.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Safer ifdef in kernel.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Documentation prose.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Reuse compute_scale function from Current Scaling.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Bugfix on inf_value scale refactor.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Remove qopt calls from test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update pytest list.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add copyright to reference scale calc.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use ptx.cuh functions instead of cde.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update shape logic with allocation and reuse shape.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Usage defaults MR feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Copyright and header guard.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Updating torch dispatch code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Fix exception type.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use TypeInfo
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update CS scale update test to use updated ref impl
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update JAX scaling mode enum
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Skip tests on Lovelace
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Keith Wyss <kwyss@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* check in per-tensor current scaling full recipe
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

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

setup basics of current scaling quantizer in python level
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

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

add test case for current scaling dequantize
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

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

finish linear layer fwd bwd test, determined error with bf16
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

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

achieved zero tolerance for Linear by specify gemm use_split_accumulator config
Signed-off-by: zhongboz <zhongboz@nvidia.com>

enable layernormlinear with current scaling, pass bitwise test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

refactor test case code
Signed-off-by: zhongboz <zhongboz@nvidia.com>

make current scaling quantizers distrbuted, pass distributed linear&layernormlinear tests
Signed-off-by: zhongboz <zhongboz@nvidia.com>

bug fix: use cached fp8 recipe in backward
Signed-off-by: zhongboz <zhongboz@nvidia.com>

fix layernorm_mlp with current scaling, fix activation_helper with current scaling
Signed-off-by: zhongboz <zhongboz@nvidia.com>

support detailed numerical settings from recipe to quantization kernel
Signed-off-by: zhongboz <zhongboz@nvidia.com>

resolving MR comments
Signed-off-by: zhongboz <zhongboz@nvidia.com>

recipe naming
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



* resolve mr comments, remove IS_CURRENT_SCALING template from kernels
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



* resolve mr comments, make current scaling c++ test cases
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* add current scaling to test_numerics.py, skip act recomp and grouped linear
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



* add benchmark for quantizer
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



* add benchmarks for linear layer
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



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

* resolve more mr comments
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* avoid potential race condition by not using from_blob to construct amax tensor in C++
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* resolve more comments
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Debug linter warnings and license check
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 import error in FP8 tensor test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug compilation error with CUDA 12.1 for Turing
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



* resolve mr comments, fix activation cast fusion
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* resolve comments, add NVTEQuantizationParams for compute scale
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



* remove is_current_scaling check totally from common folder
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* remove benchmarks, will contribute in another repo
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* adjust cs default recipe config
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



* adjust comments in test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Remove current scaling mode from core lib
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Refactor current-scaling-specific logic in core C++ lib

Move amax and scale update functions out of casting functions, and put into dedicated current-scaling source file. Add general API for accessing quantization config object.
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



* Add missing header in C++ tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable test config with FP8 transpose on Blackwell
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix compilation error in C++ test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: zhongboz <zhongboz@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>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

Added constexpr checks of tensor boundaries
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added TMA alignment check to cast_fp8_1D
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



* Use tensor const-ref instead of tensor const-ptr
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Oleg Goncharov <ogoncharov@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: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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