MOE Sequential Changes for me to see - #5
Draft
vthumbe1503 wants to merge 34 commits into
Draft
Conversation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
…to dispatch_combine_op
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>
for more information, see https://pre-commit.ci
* 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>
vthumbe1503
force-pushed
the
dispatch_combine_op
branch
from
August 12, 2026 05:20
517d0ff to
3c37643
Compare
…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>
for more information, see https://pre-commit.ci
…merEngine into dispatch_combine_op
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: