Skip to content

MOE Sequential Changes for me to see - #5

Draft
vthumbe1503 wants to merge 34 commits into
enable_extra_out_consumptionfrom
dispatch_combine_op
Draft

MOE Sequential Changes for me to see#5
vthumbe1503 wants to merge 34 commits into
enable_extra_out_consumptionfrom
dispatch_combine_op

Conversation

@vthumbe1503

Copy link
Copy Markdown
Owner

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

vthumbe1503 and others added 12 commits August 10, 2026 06:27
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Use compute_dtype for apples-to-apples TE comparisons, and set
output_to_caller=False so fused MOE does not materialize internal
dispatch extras for the Sequential caller.

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Enable runtime resolution of CUDA header path for NVRTC

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

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

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

* Update transformer_engine/common/util/cuda_runtime.cpp

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: fheinecke <23390735+fheinecke@users.noreply.github.com>

* Fix resolution not working properly when package is installed outside of site-packages directory

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* fix cuda include dir resolution when installing as an editable package

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* fix linter failure

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

---------

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: fheinecke <23390735+fheinecke@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…#3270)

* EP Dispatch with MXFP8

* Separate EP dispatch-forward and combine-backward quant recipes

* Simplify EP quantize path and guard zero-copy pool alloc under CUDA graph capture

*  Add eager-mode coverage for MXFP8 EP combine backward

* Supply persistent symm-mem recv buffers for zero-copy 1F1B CUDA-graph capture

* Enable MXFP8 EP combine backward under zero-copy

* Release symm-mem pool in ep_finalize

* Skip MXFP8 EP tests on pre-Blackwell devices

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

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…IA#3318)

FusedAdam.step() skipped a param group with no parameters before touching
its step counter, so a group that is empty on one data-parallel rank and
populated on another stopped counting on the empty ranks. Since step is
stored in param_groups it is checkpointed, and a rank that loads its shard
from a rank where the group was empty resumes with a stale step and a wrong
bias correction.

Move the counter update above the empty-group skip. Empty groups have no
parameter to read a device from, so the capturable tensor now falls back to
the device of the optimizer scratch buffer.

Fixes NVIDIA#1986

Signed-off-by: Aditya Singh <adisin650@gmail.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
NVIDIA#3281)

* Fix NVFP4 stochastic rounding on architectures without cvt.rs

The four e2m1 stochastic rounding helpers gate on
ARCH_HAS_STOCHASTIC_ROUNDING, which is sm_100 and sm_103, where the
cvt.rs instruction exists. The other branch is NVTE_DEVICE_ERROR, which
is printf plus assert(0), and release builds define NDEBUG, so on
sm_120 and sm_121 the quantize returns all-zero FP4 data while printing
once per thread. Replace that branch with a software e2m1 stochastic
rounder that takes 8 random bits per element, the same as the
instruction's rbits operand, and follows cvt.satfinite at the edges.
Architectures with cvt.rs keep the asm path.

Signed-off-by: David Kogan <davidkny22@gmail.com>

* Make the software e2m1 stochastic rounder branchless

The rounding ladder becomes a floor in units of the grid step at the
jittered value, and the isnan branch folds into the final sign select,
since fminf carries a NaN through to max_norm and the select keeps it
positive, matching cvt.satfinite. Output is bit-identical to the
previous form across all 2^32 float bit patterns at 16 rbits values and
across every float with magnitude in [1, 8) at all 256 rbits values. An
isolated microbenchmark on sm_121 runs 1.34x faster.

Signed-off-by: David Kogan <davidkny22@gmail.com>

---------

Signed-off-by: David Kogan <davidkny22@gmail.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

address review comment

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
YangFei1990 and others added 17 commits August 12, 2026 09:54
…IA#3344)

* fix swizzle scale output shape in variable shape case

Signed-off-by: YangFei1990 <feiw@nvidia.com>

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

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

---------

Signed-off-by: YangFei1990 <feiw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
bump nccl-extensions commit

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…support (NVIDIA#3224)

* improve device grouped linear

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

* proper support for grouped bias

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

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

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

* fixes and improvements

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

# Conflicts:
#	tests/pytorch/test_grouped_mlp.py

* fix unit test failure

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

* fix group linear UT

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

* fix group mlp UT by disabling wrong fallback

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

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

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

* lint

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

* fix CI errors

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

* rename toggle resolve comments

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

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

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

* don't always set optimize_for_gemm=True

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

* test other fp8 recipes on hopper

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

* fix test failure

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

* relax atol for bf16 sum errors

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

* fix CI

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

* continue to fix edge cases

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

* unify grouped tensor check

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

* special case handling for all empty moe inputs for fp8 CS

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

* add another specific case guard

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

* fix cublas setup meta data issue for empty inputs

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

* mcore integration fixes

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

* chore fix after rebase

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

* chore

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

* chore fix after rebase

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

* resolve comments, enable save original input

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

* fix cudnn version gurad

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* comments

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

* update benchmark script

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

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

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

---------

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Squashed to single commit for review.

Original PR: andrewwhitecdw#14

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
* Gate FA4 and stabilize attention test imports

FA4 can be installed on SM8x even though its current implementation rejects those GPUs. Disable selection and skip dedicated FA4 tests there so A100 and L40 use supported attention backends.

FA4 and CUTLASS can also expose a generic utils package on sys.path. Prepend the Transformer Engine test helper directory in the context-parallel test so collection resolves the intended utilities.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Isolate FlashAttention CI backends

Moving images can install FA4 alongside older FlashAttention generations, which mixes a shared Python namespace and can make context-parallel reference runs compile an unsupported backend. Isolate the L3 version matrix, keep current CP comparisons on FA2/FA3, and temporarily reject symmetric D512 FA4 on Blackwell until upstream kernel support is complete.

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

* Prepend PyTorch test utility imports

FA4 and its CUTLASS dependency expose a top-level utils module after Transformer Engine imports. Appending the test root can therefore bind these late imports to the installed module and fail collection. Give the repository helper precedence in the four test files that exhibited this ordering.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Recognize FA3 sliding-window CP support

The all-gather and a2a guards use the FA2 package version check to recognize FlashAttention support, so an isolated FA3 run is rejected even though FA3 implements sliding-window attention. Accept the explicit FA3 backend in both guards.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Keep CP backend selection at CI boundaries

The CP runner must honor an explicit generation selected by its caller, particularly the existing B200 L3 FA4 lane. Remove its internal V4 override, restore the L3 SM100 selection changed in 0f6c71e, and disable V4 only for the L1 suite that still targets FA2/FA3. This keeps per-generation L3 isolation intact without making the shared runner silently override directed coverage.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Align FlashAttention CI coverage by architecture

Keep L0 on the mature FA2 path while L3 owns newer-generation coverage. Restrict H100 L3 to FA3 and B200 L3 to non-CP FA4 so unsupported H100 FA4 kernels and mislabeled Blackwell CP results do not obscure the intended signal. Make FA4-specific tests honor backend enablement to prevent silent fallback under an FA4 label.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Guard FA4 against incompatible CUTLASS installs

Package metadata can report FA4 present even when a later dependency install leaves its transitive CUTLASS stack unusable. Reject the known b24/CUTLASS combination below the stable 4.6.2 release and treat a nested interface ImportError as an unavailable optional backend so unrelated Transformer Engine imports can continue.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Squashed to single commit for review.

Original PR: andrewwhitecdw#8

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
… Op in TE Sequential (NVIDIA#3320)

* produce/consume extra output

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* allow for fusions with producer/consumer being part of same fuser with error handling tests

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* minor cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* dispatch combine impl

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* fusible ops test

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

* keep just ops infra changes

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* cleanup with residual tests

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>

* update to cleaner documentation

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* 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

* some cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* update docs

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* pin channels through channel version

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* unecessary handling removal

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* simplify

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* doc update + extra_grad = None case

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

* test cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* no need to check staleness in every forward call

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* remove redundant tests

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

* revert from bad names

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* keep simple

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* unecessary checks

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* minor doc

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* fix lint

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* Update transformer_engine/pytorch/ops/fuser.py

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Update docs/examples/op_fuser/op_fuser.rst

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Update transformer_engine/pytorch/ops/fuser.py

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* address review comments + extra output being configurable to be outputted

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* not picklable

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* simplify.. lock it permanently

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* a bit of doc

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

---------

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>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Support MXFP8 2D quantization

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

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

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

* Fix comments

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

* Fix comments

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

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

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

* Support explicit 2D MXFP8 grouped quantization

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

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

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

---------

Signed-off-by: kunlunl <kunlunl@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
…ear (NVIDIA#3324)

* [PyTorch] Enable NVFP4 row-scaled (per-token) backward for GroupedLinear

Extend the row-scaled NVFP4 support added for dense Linear to the
MoE GroupedLinear module, so the wgrad is computed in NVFP4 instead of
falling back to high precision.

Signed-off-by: Cael Ling <caell@nvidia.com>

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

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

* Remove unhelpful comments

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

* Remove unnecessary comment

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------

Signed-off-by: Cael Ling <caell@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>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
…3382)

Move ffi type definition before ffi target definitions

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
use pinned memory; use explicit cuda sync

Signed-off-by: YangFei1990 <feiw@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
* [Common] Ensure quantization kernels handle noop properly

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* nit

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

---------

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
pre-commit-ci Bot and others added 5 commits August 17, 2026 16:43
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.