[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045

Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged

[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom
alan-hpc:feature/cutlass_group_gemm_support

Conversation

@alan-hpc

@alan-hpcalan-hpc commented Aug 8, 2025•
edited by yaox12
Loading

Copy link
Copy Markdown
Contributor

Description

Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.

Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.

run test script with

python tests/pytorch/test_group_gemm.py

Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%

(Update by @yaox12)
Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:

Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%

Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.

Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@alan-hpc
alan-hpcforce-pushed the feature/cutlass_group_gemm_support branch 2 times, most recently from d2a9a55 to b42385dCompareAugust 8, 2025 09:14
@phu0ngng

phu0ngng commented Aug 11, 2025•
edited
Loading

Copy link
Copy Markdown
Collaborator

Hi @cassiewilliam ,

Thank you for a great PR - it’s good to see such a clear performance improvement!

I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.

Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.

Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.

@alan-hpc

alan-hpc commented Aug 12, 2025

Copy link
Copy Markdown
ContributorAuthor

Hi @cassiewilliam ,

Thank you for a great PR - it’s good to see such a clear performance improvement!

I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.

Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.

Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.

I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?

@yaox12

yaox12 commented Aug 12, 2025

Copy link
Copy Markdown
Member

Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.

Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?

Please go ahead in this PR.

@alan-hpc

alan-hpc commented Aug 12, 2025

Copy link
Copy Markdown
ContributorAuthor

Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.

Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?

Please go ahead in this PR.

Got it — I’ll refactor the code to meet the requirements described above.

@alan-hpc
alan-hpcforce-pushed the feature/cutlass_group_gemm_support branch 12 times, most recently from 6f01bc8 to e832972CompareAugust 13, 2025 04:24
@alan-hpc

alan-hpc commented Aug 13, 2025

Copy link
Copy Markdown
ContributorAuthor

hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.

phu0ngng
phu0ngng reviewed Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py Outdated
Comment threadtransformer_engine/common/CMakeLists.txt Outdated
Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu Outdated
Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu Outdated
Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu Outdated
Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh Outdated
Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh Outdated
@alan-hpc
alan-hpcforce-pushed the feature/cutlass_group_gemm_support branch 7 times, most recently from a023c5f to a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others added 2 commits August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12

yaox12 commented Aug 27, 2025

Copy link
Copy Markdown
Member

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

@alan-hpc

alan-hpc commented Aug 27, 2025

Copy link
Copy Markdown
ContributorAuthor

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.
Could you please help take a look?

@yaox12

yaox12 commented Aug 28, 2025•
edited
Loading

Copy link
Copy Markdown
Member

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?

I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get

FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is

unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc

alan-hpc commented Sep 1, 2025

Copy link
Copy Markdown
ContributorAuthor

pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.

test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?

I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get

FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is

unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments

@phu0ngng

phu0ngng commented Sep 10, 2025

Copy link
Copy Markdown
Collaborator

@yaox12 could you help resolve the remaining mismatches?
Let's try to get this merge for TE 2.8. Thanks!

@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng

phu0ngng commented Sep 16, 2025

Copy link
Copy Markdown
Collaborator

/te-ci L0

@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

yaox12 commented Sep 17, 2025•
edited
Loading

Copy link
Copy Markdown
Member

Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.

Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.

@yaox12

yaox12 commented Sep 17, 2025

Copy link
Copy Markdown
Member

/te-ci L0

yaox12
yaox12 approved these changes Sep 17, 2025
View reviewed changes

@yaox12yaox12 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.

@yaox12

yaox12 commented Sep 18, 2025

Copy link
Copy Markdown
Member

CI passed except for some irrelevant failures.

@yaox12
yaox12 merged commit 8aee1bb into NVIDIA:mainSep 18, 2025
37 of 41 checks passed
vthumbe1503 pushed a commit to vthumbe1503/TransformerEngine that referenced this pull request Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman mentioned this pull request Sep 30, 2025
Merged
13 tasks
@timmoon10timmoon10 mentioned this pull request Oct 3, 2025
Merged
13 tasks
@alan-hpcalan-hpc mentioned this pull request Jun 1, 2026
Open
8 tasks
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10
+2 more reviewers
@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

2.8.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model - #2045
Merged
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Sep 18, 2025
Merged
[Pytorch] Add Cutlass GroupGEMM Support for fine-grained MoE Model#2045
yaox12 merged 27 commits into
NVIDIA:mainfrom 
alan-hpc:feature/cutlass_group_gemm_support
Conversation
@alan-hpc
@alan-hpcalan-hpc
commented
Aug 8, 2025•

edited by yaox12

Loading

Copy link


Copy Markdown

Contributor
Description
Add Cutlass Group GEMM Support for H100(SM90), which provides greater performance advantages compared to the current Multi-Stream implementation in Fine-Grained MoE models. Currently, this PR only supports FP16 and BF16 scenarios, and FP8 support is not yet available. The implementation is limited to the standard MoE Module (Bias and other related features have not been validated yet). Please take note.
Initial performance test results are as follows, and the testing method can be found in file test_group_gemm.py.
run test script with
python tests/pytorch/test_group_gemm.py
Shape(g,m,n,k)TE V2.2 (TFLOPs)Cutlass-Opt-V1(TFLOPs)Speed-Up
(8, 4096, 768, 2048)508.77568.6311.77%
(16, 2048, 768, 2048)398.81534.7534.08%
(Update by @yaox12)

Benchmark the performance of general_grouped_gemm with TE 2.7 and CUDA 12.9:
Shape(g,m,n,k)TE V2.7 (TFLOPs)Cutlass (TFLOPs)Speed-Up
(8, 4096, 768, 2048)627.903650.5743.6%
(16, 2048, 768, 2048)553.945636.92714.98%
Add the system environment variable NVTE_USE_CUTLASS_GROUPED_GEMM to toggle between the two GEMM implementations. Setting export NVTE_USE_CUTLASS_GROUPED_GEMM=0 selects the original Multi-Stream cuBLAS GEMM, while setting export NVTE_USE_CUTLASS_GROUPED_GEMM=1 enables the newly added CUTLASS Group GEMM. The default value is 0.
Add a environment variable NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, when it's set, it will give warnings when falling back to the cuBLAS implementaion.
Type of change
  •  Documentation change (change only to the documentation, either a fix or a new content)
  •  Bug fix (non-breaking change which fixes an issue)
  •  New feature (non-breaking change which adds functionality)
  •  Breaking change (fix or feature that would cause existing functionality to not work as expected)
  •  Infra/Build change
  •  Code refactoring
Changes
Please list the changes introduced in this PR:
  • Change A
  • Change B
Checklist:
  •  I have read and followed the contributing guidelines
  •  The functionality is complete
  •  I have commented my code, particularly in hard-to-understand areas
  •  I have made corresponding changes to the documentation
  •  My changes generate no new warnings
  •  I have added tests that prove my fix is effective or that my feature works
  •  New and existing unit tests pass locally with my changes
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
2 times, most recently
from
d2a9a55 to
b42385dCompareAugust 8, 2025 09:14
@phu0ngng
phu0ngng
commented
Aug 11, 2025•

edited

Loading

Copy link


Copy Markdown

Collaborator
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Hi @cassiewilliam ,
Thank you for a great PR - it’s good to see such a clear performance improvement!
I have one suggestion - I think we should refactor the change slightly to minimize modifications in the TE framework extensions.
Currently, we have two separate C APIs: nvte_multi_stream_cublas_gemm and nvte_cutlass_grouped_gemm. The PyTorch extensions call these individually, and we would need to do the same on the JAX side. Since they share the same function signature, we could unify them into a single API - nvte_multi_tensor_gemm - and deprecate nvte_multi_stream_cublas_gemm.
Within nvte_multi_tensor_gemm, we can determine the GPU architecture and enable CUTLASS GroupedGEMM for FP16/BF16 on Hopper. This way, future changes to the GroupedGEMM implementation or backend would not require modifications to the PyTorch/JAX extensions.
I fully agree with your suggestion — keeping the code architecture clean is very important. Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
@yaox12
yaox12
commented
Aug 12, 2025

Copy link


Copy Markdown

Member
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
@alan-hpc
alan-hpc
commented
Aug 12, 2025

Copy link


Copy Markdown

ContributorAuthor
Agree with @phu0ngng. We could unify the API and do the dispatch (based on GPU arch/data type/env variable) on the TE/common side.
Will you be handling the refactor on your side, or should I go ahead and make the changes directly in the current PR?
Please go ahead in this PR.
Got it — I’ll refactor the code to meet the requirements described above.
@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
12 times, most recently
from
6f01bc8 to
e832972CompareAugust 13, 2025 04:24
@alan-hpc
alan-hpc
commented
Aug 13, 2025

Copy link


Copy Markdown

ContributorAuthor
hello @phu0ngng@yaox12 The nvte_multi_tensor_gemm interface has been fully refactored. Please review the implementation for correctness and compliance with the updated design.
phu0ngng
phu0ngng
reviewed
Aug 13, 2025
View reviewed changes
Comment threadtests/pytorch/test_group_gemm.py
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cublaslt_gemm.cu
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

Comment threadtransformer_engine/common/gemm/cutlass_groupgemm.cuh
Outdated

@alan-hpc
alan-hpcforce-pushed
the
feature/cutlass_group_gemm_support
branch
7 times, most recently
from
a023c5f to
a76e1cdCompareAugust 18, 2025 03:58
Min Yangand others
added 2 commits
August 26, 2025 22:04
Min Yang
feat: refactor host workspace malloc
b3ef3c5
Signed-off-by: Min Yang <min.yang@shopee.com>
@alan-hpc
Merge branch 'main' into feature/cutlass_group_gemm_support
579c539
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
@yaox12
yaox12
commented
Aug 27, 2025

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@alan-hpc
alan-hpc
commented
Aug 27, 2025

Copy link


Copy Markdown

ContributorAuthor
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.log
test_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd.

Could you please help take a look?
@yaox12
yaox12
commented
Aug 28, 2025•

edited

Loading

Copy link


Copy Markdown

Member
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log
It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

@alan-hpc
alan-hpc
commented
Sep 1, 2025

Copy link


Copy Markdown

ContributorAuthor
pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
@cassiewilliam I'm still seeing some failures in test_grouped_gemm. See https://github.com/NVIDIA/TransformerEngine/actions/runs/17169254634/job/48715768402 (expand the Log session). Other failures are irrelevant.
test_grouped_linear_accuracy_cutlass.logtest_grouped_gemm.log

It runs fine on my local H100 GPU machine, but the unit tests still fail, which feels a bit odd. Could you please help take a look?
I can reproduce the failure. I use the docker image nvcr.io/nvidia/pytorch:25.06-py3, install your branch inside it, run pytest -v -s tests/pytorch/test_numerics.py::test_grouped_gemm and get
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-TN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.float16-shape3] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape0] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape1] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape2] - AssertionError: Tensor-likes are not close!
FAILED tests/pytorch/test_numerics.py::test_grouped_gemm[True-True-NN-torch.bfloat16-shape3] - AssertionError: Tensor-likes are not close!

My installing command is
unset PIP_CONSTRAINT && NVTE_CUDA_ARCHS="90" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/cassiewilliam/TransformerEngine.git@feature/cutlass_group_gemm_support

Very strange, the accuracy seems to be inconsistent in different docker environments
@phu0ngng
phu0ngng
commented
Sep 10, 2025

Copy link


Copy Markdown

Collaborator
@yaox12 could you help resolve the remaining mismatches?

Let's try to get this merge for TE 2.8. Thanks!
@yaox12
Merge branch 'main' into feature/cutlass_group_gemm_support
99d95a5
@yaox12
update cutlass
35d916c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
update cutlass
d7f0dc0
Signed-off-by: Xin Yao <xiny@nvidia.com>
@phu0ngng
phu0ngng
commented
Sep 16, 2025

Copy link


Copy Markdown

Collaborator
/te-ci L0
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12

This comment was marked as outdated.

Sign in to view
@yaox12
yaox12
commented
Sep 17, 2025•

edited

Loading

Copy link


Copy Markdown

Member
Further relaxing the threshold solved the mismatch. We think it's because we're comparing two different BF16 results. Maybe it's better to compare with the FP32 baseline. But the unit tests were originally designed for cuBLAS grouped gemm to bit-wise match the non-grouped version. I think the current version is fine as we mentioned that this accumulating pattern is rarely used.
Renamed the env var to use CUTLASS implementation to NVTE_USE_CUTLASS_GROUPED_GEMM. And add another env var NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK, we only warn falling back when it's set, otherwise there will be too many warnings.
@yaox12
yaox12
commented
Sep 17, 2025

Copy link


Copy Markdown

Member
/te-ci L0
yaox12
yaox12
approved these changes
Sep 17, 2025
View reviewed changes
@yaox12yaox12
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Pending CI.
@yaox12
yaox12
commented
Sep 18, 2025

Copy link


Copy Markdown

Member
CI passed except for some irrelevant failures.
@yaox12
yaox12
merged commit 8aee1bb
into
NVIDIA:mainSep 18, 2025

37 of 41 checks passed

vthumbe1503 pushed a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
@pre-commit-ci@vthumbe1503
[pre-commit.ci] auto fixes from pre-commit.com hooks
de3080e
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Lower precision gated-act to accelerate FP8 current-scaling. (#2153)
* Applying the original precision as Norm outputs' and activation compuations.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding knob to control norm output precision.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing the knob and applying lower-precision norm with current-scaling only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix the error when quantizer==None
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
[PyTorch] Support activation CPU offloading in fusible ops (#2158)
* Add CPU offloading logic to ops. Fix test to compute dgrad.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure grads are contiguous in op backwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add op-based MLP to CPU offloading 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
* Handle different weight cache behavior on Hopper/Blackwell
Add MXFP8 to CPU offload 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
* Remove MXFP8 test
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Do not use normalization forward + amax fusion if cuDNN backend is requested (#2174)
* Do not use norm fwd + amax fusion if cudnn backend is requested
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
* Read envirornment vairable directly to avoid include error
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Fix unjoined comm stream in UB communicator (#2160)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
FP8 Output Quantization for GEMM (#2123)
* Test working as I think it should work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* revert accidental change
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Restrict the number of cases for unfused quantization, some fp8->fp8 cases are handled by cublas
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
bug: missed a } in the code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
mxfp8 unfused quant support, refined unit test, remove unecessary quantization code
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
missed a quant code removal
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
minor bug fix
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Add cuBLASMp-backed GEMM-like API to TE common (#1824)
* Pick up cuBLASMp during build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Change lib order to fix link error
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Context creation, incomplete...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Test fixure
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A sanity AgGemm test, failing...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix axes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Take care of uneven distribution
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use MPI to get position of local matrices
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Refactor & fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-RS
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Gemm-AR, not working...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fixes
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Setting all-reduce epilogue for gemm-ar
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use supported shapes for GEMM-AR
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tolerance
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* First shot at fp8
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use TensorHolder in tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Support comm_sm_count
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Parametrize dtypes for A, B and D separately
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak scaling
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Amax ptr
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Flags parity with cublas_gemm, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Cleanup
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Bias tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix bias test
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Aux, saving...
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* aux_ld
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* A fix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Use test::Tensor
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Set scale inv
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove unsupported test configs
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Replace libcal with NCCL
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add NVTX markers to API functions
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Tweak GemmAr tests
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More test config
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix merge fallout
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove MPI dependency, comment API, add algo parameter
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem dependency
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix nvshmem build
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Excluse CommGemm tests from L0_cppunittest
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add cpp_distributed sh file for CI
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Adapt tp TensorAllocator
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip GemmAr test on unsupported HW
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Oversibscribe is needed on some clusters
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Fix incomplete libcal removal
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Move CI tests to L1
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Rename context to include NVTE prefix
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove leftover code
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* NVTE_WITH_CUBLASMP off by default
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed NVTE_CHECK diag
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Comment API
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Include stdbool header for legacy C compilers
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Remove now unused argument
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Abstract away cuBLASMp algo behind our own enum
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* More detailed shape diag messages
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/common/include/transformer_engine/comm_gemm.h
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
* Add license
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
---------
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <56651474+mk-61@users.noreply.github.com>
Co-authored-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue. (#2086)
* FP8 AllGather in FP8 GroupedGEMM
1. Support current scaling FP8 quantation with a given amax.
2. Support FP8 AG in fwd and BF16 RS in bwd.
3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Slightly refactor
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding documents of new args.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding unit-tests.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding license.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move unit-tests to L1.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Move quantizaer store/reset into FP8 only.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding all layout support for Blackwell+
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adopt the feedback from code-review.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed the wrong stream used by d2d in groupedGEMM FFI.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Delay MeshResource validation until first usage (#2124)
Delay MeshResource validation until first usage
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Decouple Recipe and ScalingMode (#1728)
* 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>
[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128)
* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
[JAX] Add amax input to DBiasQuantizePrimitive and FFI (#2118)
* add amax input to DBiasQuantizePrimitive and FFI
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
* make sure amax is init with zero
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* fix sharding rule
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>
Further relax constraints to cuDNN 9.13 for disabling fused attn for kv caching (#2121)
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Temporarily remove comm_gemm tests (#2133)
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
[PyTorch] Disable determinism for sm100 (#2130)
* disable determinism for sm100+ and cudnn<9.14
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix remaining CI failures
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* revert some changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert more changes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove sm100 from determinism table
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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>
[PyTorch] ONNX export of FP8 Current Scaling (#2068)
* Compute amax in normalization forward in current scaling in untuned kernels
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
* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* 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>
* code drop
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>
* apply tims suggestions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch][MOE] Tentative Fix For Replacing from_blob with empty for experts receiving zero tokens (#2134)
use torch empty for empty shape instead of from_blob
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
build: pull cached wheels (#2127)
* build: pull cached wheels
Signed-off-by: oliver könig <okoenig@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update setup.py
Signed-off-by: oliver könig <okoenig@nvidia.com>
---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
feat: Add support for multiple quantization modes in the UB communicators (#2043)
[Common] Add checks to CUDA kernel launch and CUDA API calls (#2074)
* add checks to cuda kernel launch and cuda API calls
Signed-off-by: Xin Yao <xiny@nvidia.com>
* Remove exceptions from destructors
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix weired dispatch in ln/rmsnorm
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@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>
[PyTorch] Support bf16+fp8 cudagraph (#2098)
* support bf16+fp8 model
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>
---------
Signed-off-by: Robin Zhang <robinz@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>
Dropout with 8-bit RNG (#2014)
* Add dropout kernel with 8-bit RNG
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
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
* Fix license
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid ambiguous types
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Do not enforce dropout prob is representable in 8 bits
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Expand error message
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
* Fix small statistical bug from using less-equal instead of less-than
Refactor kernel implementations and add comments. Interpret masks as bytes rather than 16-bit uints.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix linter warning
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
* Remove unnecessary helper function in PyTorch extensions
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: 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>
Create GPU reload buffers on main stream (#2131)
* Create GPU relaod buffers on main stream
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.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: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
* Fixed typo
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-preos01.a51.clusters.nvidia.com>
Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
minor code cleanup
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor cosmetics
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Address review comment
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
minor comment update
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Fix CI failures for UB overlap changes (#2149)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
minor bug: quantizer should not be none for unfused quantization
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[JAX] Fix failing fused attn tests for dropout=0.1 and bias for sm100 (#2135)
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix the skip message
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Assert in fused attn bwd pass for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Add check for sm100
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add support to get all devs in the process for jax
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Code clean up
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Make get_all_device_compute_capability more pythonic, thereby avoiding unnecessary type conversion
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* Represent attn bias using enum instead of string
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix linting error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[PyTorch][CUDA Graph] Fix FP8 Weight Quantization Cache under CUDA Graph (#2119)
* add noop to comp amax
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* fix for fp8 blockwise recipe
Signed-off-by: zhongboz <zhongboz@nvidia.com>
* resolve comments
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>
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>
address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Update test_multi_process_distributed_grouped_gemm.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* Update dense.py
change accidentally added while merging
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address revie comments
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bug solved: delayed scaling quantization with mxfp8 inputs didnt work
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the unit test error
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* just to trigger ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* fix merge conflict
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
address review comments: quantization inside gemm and outside both should exactly match for fp32 accumulation
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
TE Gemma tutorial attempt#2 (#1839)
* add tutorial files and other local changes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove extraneous code for easy debu
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* make cuda graphs work with non-paged and paged attention
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* perf imp for kv cache ops
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add code for calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* changes to make quantizers work with fp8_calibration
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* use quantizer only if needed
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* functionality of the tutorial tested and perf checked
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove files and update headers/licenses
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* update header/license
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* update tutorial for review
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make weights downloadable on the fly; remove extra print statements
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix lint and update comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add comma back, typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* sequence_start_positions should be None for training
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add paged attention numberes and update requirements.txt file
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* make tutorial work on blackwell
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* remove gemma FT tutorial for now
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixing the headings placement and rewording attention -> kv caching
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fixes from comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix the images
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* misc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add more comments to te_gemma.py and cleanup utils.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add more information about the hierarchy of the classes used in the tutorial
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add better cuda graphs picture
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* addd updated cuda graphs pictures
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add illustrated cuda graphs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* small fixes in documentation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add torch.no_grad() to force reduced memory usage
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* some fixes from recent comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* more fixes from remaining comments
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* add te_rope_emb to class desc
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* fix tutorial wording; add calibration fix to grouped_linear.py
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---------
Signed-off-by: Sudhakar Singh <sudhakars@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>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix memory overhead of linear layer when all gather from sequence parallel (#2125)
* fix memory overhead of all gather from sequence parallel
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
* Update transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* quick fix the errors when for UB buffers
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
* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Avoid deallocating FP8 scale-invs since they are reused
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Yuzhong Wang <yuzhongw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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 <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Fix incorrect TP rank calculation when using data parallel (#2179)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
[Pytorch] Add Cutlass Grouped GEMM Support for fine-grained MoE Model (#2045)
* feat: add cutlass group gemm support
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor multi tensor gemm interface
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refactor: refactor nvte_multi_stream_cublas_gemm func and add license info
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add unit test for cutlass group gemm
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: add cutlass support type protect
Signed-off-by: Min Yang <min.yang@shopee.com>
* add tests and fix lint
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
* feat: fix unit tests error
Signed-off-by: Min Yang <min.yang@shopee.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* feat: refactor host workspace malloc
Signed-off-by: Min Yang <min.yang@shopee.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update cutlass
Signed-off-by: Xin Yao <xiny@nvidia.com>
* further relex threshold and add a env var to warn fall back
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: Min Yang <min.yang@shopee.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: alan yang <89962857+cassiewilliam@users.noreply.github.com>
Co-authored-by: Min Yang <min.yang@shopee.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: Phuong Nguyen <phuonguyen@nvidia.com>
[PyTorch] Support FA3 for MLA and with CP (#1907)
feature(FA3,MLA,CP):
1. Update FA3 to commit-id 3ba6f82 (tag 2.8.0.post2 with compile error fixed), PR-1604 support hdimQK != hdimV backward
2. Update get_attention_backend method because FA3 support MLA now
3. Add CP MLA support for FA3
4. Add unit tests for FA3 MLA CP
5. Update attention doc
Signed-off-by: zhujian <zhujian.whu.cs@gmail.com>
Fix cuDNN version checks when getting backend and for sm89 kv cache (#2185)
* Fix cudnn version checks for kv cache for sm89. Add cudnn version check in preparation for 9.14 when getting backend
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Minor fix for cuDNN version condition check
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivamanksivaman
mentioned this pull request
Sep 30, 2025
 Merged

13 tasks
@timmoon10timmoon10
mentioned this pull request
Oct 3, 2025
 Merged

13 tasks
@alan-hpcalan-hpc
mentioned this pull request
Jun 1, 2026
 Open

8 tasks
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@phu0ngngphu0ngngphu0ngng left review comments
@timmoon10timmoon10Awaiting requested review from timmoon10

+2 more reviewers

@leefigeleefigeleefige left review comments
@yaox12yaox12yaox12 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels

2.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@alan-hpc@phu0ngng@yaox12@zhongbozhu@leefige