GitLab

* rowwise colwise RHT group quant v1
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* remove local array RW
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* change wait_barrier
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fast math options
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* use mult to replace div
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* format
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* bulk move random states
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* greptile
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

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

* revert to use divides
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* avoid fp32 bf16 round-trip in RHT cast fusion
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* trigger fastmath by toggle NVTE_RHT_CAST_FUSION_USE_FAST_MATH
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* integrate row col rht fusion, functional
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* numerics aligned
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* style
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* remove device sync
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* 128 padding
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* revert colwise rng state creation because of row-col fused kernel
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix CI, linter
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* refactor RS for generating two random values
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* Avoid invalid configs with templated kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix acc pipeline init with 0 arrival count
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* restore rowwise-only mode
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* switch to dynamic atomic scheduler
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* Avoid instantiating group RHT+cast kernel without row-wise or col-wise output
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Include fast math option in quantization config
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix linter warnings and review nits
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



* Use TE license
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix bug where kernel is always launched on stream
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



* Restore BF16 intermediate downcast in fused RHT-cast kernels
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix numerical test of grouped kernel
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* Make sure row-wise and col-wise quantization use different RNG seeds
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Restore autoformatter
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



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

* minor fix of torch view dtype
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* multi-tensor RHT amax, compiles
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* setup multi_tensor_quantize_nvfp4_impl
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* wire things up and run without crash
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* numerical test
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* unit test passing
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* finish unit test of split quantize api
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* bump up padding to 64 for nvfp4 grouped quantize
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix stochastic rounding
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

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

* change error message
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

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

* enable multi-amax without RHT
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix col-only quantize mode
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* improve benchmark script
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add NCU example script
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add larger test case
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add contiguous_data_and_scale check to bulk allocator
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* unified naming and differentiate between group_ and multi_
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* move regular amax into multi_tensor.h
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* Disentangle logic for split-quantize and general multi-tensor quantize
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use size_t for split sections
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Suggestions from @greptile-apps
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



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

* pipeclean, fix nvfp4 padding of 32 alignment
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

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

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



* numerical test passed
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix CI failure with test_cast_master_weights_to_fp8 (in a hacky way)
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* found CUDA mis-aligned address error in training in multi-swizzle, hack the vec_load_size to 1 to unblock
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* leave comments about alignment issue
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fused bulk alloc nvfp4
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix RHT sign mask CPU overhead
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

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

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

* Remove incorrect logic that treats 0-D tensor as uninitialized

Tensor shape logic still requires treating 0-D tensor as uninitialized.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix invalid conversion from tensor to int
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Zhongbo Zhu <zhongboz@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>

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

* Change all imports and api
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

* fix docs, tests, and make Jax change as well
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change internal uses of fp8_autocast
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* CG function, and small test fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change instances of make_graphed_callables internally
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Fix test and add more docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Cleanup test imports and minimize internal file imports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Better docs and better api
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Apply suggestions from code review
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

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

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

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

* multi-tensor swizzle
Signed-off-by: Xin Yao <xiny@nvidia.com>

* pad zeros in swizzle kernels
Signed-off-by: Xin Yao <xiny@nvidia.com>

* unify single- and multi-tensor swizzle
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



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

* fix bug for col swizzle
Signed-off-by: Xin Yao <xiny@nvidia.com>

* check context & fix signifiers
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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

* before zero padding in mxfp8 swizzle, use torch zeros to malloc for now
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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

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

* finish python ref impl for bulk alloc
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* c++ bulk alloc worked, still draft version
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

* resolve rebase conflict
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 license
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* use shared_ptr to auto manage reference count
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



* attempt to fix misc training error
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



* attempt to handle case where experts get zero token
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* updated with fused C++ function calls
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

* experiment with reducing py object construction time
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix seg fault bug in inference mode
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

* fuse torch split into bulk alloc
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

* rebase to latest main
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix unit test failure
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

* refactor create_tensor to use get_scale_shape
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* refactor quantize to call quantize_cpp
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Implement separate functions for multi-tensor quantize and split + multi-tensor quantize
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update grouped linear module with fused split+quantize func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move multi-tensor quantize func to cast.cpp
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Do not expose quantizer helper function externally
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix linter warnings
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



* Revert cuDNN frontend commit
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix corner cases with zero tokens
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

---------
Signed-off-by: zhongboz <zhongboz@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>