MOE Sequential Changes for me to see - #5
Draft
vthumbe1503 wants to merge 83 commits into
Draft
Conversation
Signed-off-by: Varun Thumbe <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]>
…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 <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
* Enable runtime resolution of CUDA header path for NVRTC Signed-off-by: Fred Heinecke <[email protected]> * [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 <[email protected]> * Fix resolution not working properly when package is installed outside of site-packages directory Signed-off-by: Fred Heinecke <[email protected]> * fix cuda include dir resolution when installing as an editable package Signed-off-by: Fred Heinecke <[email protected]> * fix linter failure Signed-off-by: Fred Heinecke <[email protected]> --------- Signed-off-by: Fred Heinecke <[email protected]> Signed-off-by: fheinecke <[email protected]> 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 <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]>
…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 <[email protected]> Co-authored-by: Przemyslaw Tredak <[email protected]>
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 <[email protected]> * 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 <[email protected]> --------- Signed-off-by: David Kogan <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]> address review comment Signed-off-by: Varun Thumbe <[email protected]> address review comments Signed-off-by: Varun Thumbe <[email protected]>
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 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: YangFei1990 <[email protected]> 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 <[email protected]>
…support (NVIDIA#3224) * improve device grouped linear Signed-off-by: Zhongbo Zhu <[email protected]> * proper support for grouped bias Signed-off-by: Zhongbo Zhu <[email protected]> * [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 <[email protected]> # Conflicts: # tests/pytorch/test_grouped_mlp.py * fix unit test failure Signed-off-by: Zhongbo Zhu <[email protected]> * fix group linear UT Signed-off-by: Zhongbo Zhu <[email protected]> * fix group mlp UT by disabling wrong fallback Signed-off-by: Zhongbo Zhu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint Signed-off-by: Zhongbo Zhu <[email protected]> * fix CI errors Signed-off-by: Zhongbo Zhu <[email protected]> * rename toggle resolve comments Signed-off-by: Zhongbo Zhu <[email protected]> * [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 <[email protected]> * test other fp8 recipes on hopper Signed-off-by: Zhongbo Zhu <[email protected]> * fix test failure Signed-off-by: zhongboz <[email protected]> * relax atol for bf16 sum errors Signed-off-by: zhongboz <[email protected]> * fix CI Signed-off-by: Zhongbo Zhu <[email protected]> * continue to fix edge cases Signed-off-by: Zhongbo Zhu <[email protected]> * unify grouped tensor check Signed-off-by: zhongboz <[email protected]> * special case handling for all empty moe inputs for fp8 CS Signed-off-by: Zhongbo Zhu <[email protected]> * add another specific case guard Signed-off-by: zhongboz <[email protected]> * fix cublas setup meta data issue for empty inputs Signed-off-by: zhongboz <[email protected]> * mcore integration fixes Signed-off-by: Zhongbo Zhu <[email protected]> * chore fix after rebase Signed-off-by: Zhongbo Zhu <[email protected]> * chore Signed-off-by: Zhongbo Zhu <[email protected]> * chore fix after rebase Signed-off-by: Zhongbo Zhu <[email protected]> * resolve comments, enable save original input Signed-off-by: Zhongbo Zhu <[email protected]> * fix cudnn version gurad Signed-off-by: Varun Thumbe <[email protected]> * comments Signed-off-by: Zhongbo Zhu <[email protected]> * update benchmark script Signed-off-by: Zhongbo Zhu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Zhongbo Zhu <[email protected]> Signed-off-by: zhongboz <[email protected]> Signed-off-by: Varun Thumbe <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Varun Thumbe <[email protected]>
Signed-off-by: andrewwhitecdw <[email protected]> Co-authored-by: andrewwhitecdw <[email protected]>
Signed-off-by: andrewwhitecdw <[email protected]> Co-authored-by: andrewwhitecdw <[email protected]>
Squashed to single commit for review. Original PR: andrewwhitecdw#14 Signed-off-by: andrewwhitecdw <[email protected]> Co-authored-by: andrewwhitecdw <[email protected]>
* 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 <[email protected]> * 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 <[email protected]> * [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 <[email protected]> * 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 <[email protected]> * 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 <[email protected]> * 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 <[email protected]> * 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 <[email protected]> --------- Signed-off-by: Sudhakar Singh <[email protected]> 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 <[email protected]> Co-authored-by: andrewwhitecdw <[email protected]>
… Op in TE Sequential (NVIDIA#3320) * produce/consume extra output Signed-off-by: Varun Thumbe <[email protected]> * allow for fusions with producer/consumer being part of same fuser with error handling tests Signed-off-by: Varun Thumbe <[email protected]> * cleanup Signed-off-by: Varun Thumbe <[email protected]> * minor cleanup Signed-off-by: Varun Thumbe <[email protected]> * dispatch combine impl Signed-off-by: Varun Thumbe <[email protected]> * fusible ops test Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * cleanup with residual tests Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * update to cleaner documentation Signed-off-by: Varun Thumbe <[email protected]> * address review comments Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * update docs Signed-off-by: Varun Thumbe <[email protected]> * pin channels through channel version Signed-off-by: Varun Thumbe <[email protected]> * unecessary handling removal Signed-off-by: Varun Thumbe <[email protected]> * simplify Signed-off-by: Varun Thumbe <[email protected]> * doc update + extra_grad = None case Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * no need to check staleness in every forward call Signed-off-by: Varun Thumbe <[email protected]> * remove redundant tests Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * keep simple Signed-off-by: Varun Thumbe <[email protected]> * unecessary checks Signed-off-by: Varun Thumbe <[email protected]> * minor doc Signed-off-by: Varun Thumbe <[email protected]> * fix lint Signed-off-by: Varun Thumbe <[email protected]> * Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> * Update docs/examples/op_fuser/op_fuser.rst Co-authored-by: Tim Moon <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> * Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> * address review comments + extra output being configurable to be outputted Signed-off-by: Varun Thumbe <[email protected]> * cleanup Signed-off-by: Varun Thumbe <[email protected]> * address review comments Signed-off-by: Varun Thumbe <[email protected]> * not picklable Signed-off-by: Varun Thumbe <[email protected]> * simplify.. lock it permanently Signed-off-by: Varun Thumbe <[email protected]> * a bit of doc Signed-off-by: Varun Thumbe <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <[email protected]>
* Support MXFP8 2D quantization Signed-off-by: kunlunl <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix comments Signed-off-by: kunlunl <[email protected]> * Fix comments Signed-off-by: kunlunl <[email protected]> * [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 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: kunlunl <[email protected]> Signed-off-by: Przemek Tredak <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemek Tredak <[email protected]>
…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 <[email protected]> * [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 <[email protected]> Signed-off-by: Tim Moon <[email protected]> * Remove unnecessary comment Signed-off-by: Tim Moon <[email protected]> --------- Signed-off-by: Cael Ling <[email protected]> Signed-off-by: Tim Moon <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <[email protected]>
…3382) Move ffi type definition before ffi target definitions Signed-off-by: Jeremy Berchtold <[email protected]>
Signed-off-by: Vladimir Cherepanov <[email protected]>
use pinned memory; use explicit cuda sync Signed-off-by: YangFei1990 <[email protected]> Co-authored-by: Phuong Nguyen <[email protected]>
* [Common] Ensure quantization kernels handle noop properly Signed-off-by: Kaining Zhong <[email protected]> * nit Signed-off-by: Kaining Zhong <[email protected]> --------- Signed-off-by: Kaining Zhong <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]>
…nd JAX (NVIDIA#3396) * Add CI workflow building TE against torch 2.1 Signed-off-by: Pawel Gadzinski <[email protected]> * Remove MAX_JOBS=1 from pytorch 2.1 build workflow Signed-off-by: Pawel Gadzinski <[email protected]> * Use MAX_JOBS=2 in pytorch 2.1 workflow, unbounded build OOMs the runner Signed-off-by: Pawel Gadzinski <[email protected]> * Point TE build at apt cudnn9, torch 2.1 pip deps shadow it with cudnn 8.9 Signed-off-by: Pawel Gadzinski <[email protected]> * Avoid passing std::optional to at::get_generator_or_default, torch 2.1 needs c10::optional Signed-off-by: Pawel Gadzinski <[email protected]> * Gate NCCL EP in torch extension on nccl_ep_enabled, matching common CMake Signed-off-by: Pawel Gadzinski <[email protected]> * Pass comm streams as raw cudaStream_t handles, torch 2.1 pybind lacks c10::Stream caster Signed-off-by: Pawel Gadzinski <[email protected]> * Revert "Pass comm streams as raw cudaStream_t handles, torch 2.1 pybind lacks c10::Stream caster" This reverts commit 04b4b0f. Signed-off-by: Pawel Gadzinski <[email protected]> * Revert "Avoid passing std::optional to at::get_generator_or_default, torch 2.1 needs c10::optional" This reverts commit 061834f. Signed-off-by: Pawel Gadzinski <[email protected]> * Retarget CI workflow to torch 2.8 Signed-off-by: Pawel Gadzinski <[email protected]> * Rename workflow to Minimum supported PyTorch, parametrize torch version Signed-off-by: Pawel Gadzinski <[email protected]> * Add Minimum supported JAX job, rename workflow to minimum_versions Signed-off-by: Pawel Gadzinski <[email protected]> * Add packaging to min-jax job deps Signed-off-by: Pawel Gadzinski <[email protected]> * Skip NCCL EP in torch extension when torch lacks symm-mem headers; build min-torch job on sm90 Signed-off-by: Pawel Gadzinski <[email protected]> * Fix min-versions jobs: pydantic for jax import, libcuda stub for sm90 import Signed-off-by: Pawel Gadzinski <[email protected]> * Extract torch symm-mem EP gate into nccl_ep_supported_by_torch helper Signed-off-by: Pawel Gadzinski <[email protected]> * Tidy nccl_ep_supported_by_torch Signed-off-by: Pawel Gadzinski <[email protected]> * Compile out EP zero-copy path when torch lacks symm-mem headers Per review: only the zero-copy path needs torch's c10d symm-mem (2.11); guard the includes with __has_include so the rest of EP builds on older torch, expose ep_zero_copy_supported() and raise in ep_bootstrap when zero_copy=True without support. Drop the build-time torch gate. Signed-off-by: Pawel Gadzinski <[email protected]> --------- Signed-off-by: Pawel Gadzinski <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
…in HLO (NVIDIA#3412) * [JAX] Fix counting of synced and wrapped collectives in HLO Signed-off-by: Alex Y. Chan <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alex Y. Chan <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…A#3392) * [Common] Make ptx.cuh arch agnoistic by moving arch specific helpers to a new header Signed-off-by: Kaining Zhong <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * whoops Signed-off-by: Kaining Zhong <[email protected]> * Update transformer_engine/common/util/ptx.cuh Co-authored-by: Tim Moon <[email protected]> Signed-off-by: Kaining Zhong <[email protected]> --------- Signed-off-by: Kaining Zhong <[email protected]> Signed-off-by: Kaining Zhong <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <[email protected]>
* [Common/PyTorch] Fused grouped MXFP8 requantization Replace the group_dequantize -> group_quantize(columnwise) -> grouped_swizzle(rowwise scales) chain in group_requantize_inplace with a single kernel (NVTE_FUSED_GROUP_REQUANTIZE=0 restores the unfused path). Co-authored-by: Oleg Goncharov <[email protected]> Signed-off-by: YangFei1990 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix zero tensor dim Signed-off-by: YangFei1990 <[email protected]> * rename func; read less attrs Signed-off-by: YangFei1990 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: YangFei1990 <[email protected]> Co-authored-by: Oleg Goncharov <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Michał Marcinkiewicz <[email protected]>
…uncher pass count (NVIDIA#3415) [PyTorch] Scale test_multi_process_ep outer timeout with launcher pass count Signed-off-by: Phuong Nguyen <[email protected]> Co-authored-by: fheinecke <[email protected]>
…assing epbuffer from dipsatch to combine op via extra input Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
…Fused Adam. (NVIDIA#3414) * Add the support of variable dtypes of params and grads in capturable Fused Adam. Signed-off-by: Przemek Tredak <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <[email protected]>
…merEngine into dispatch_combine_op
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <[email protected]>
* Add support in lower level JAX API for returning max logit and softmax aux to the user from TE JAX fused attn output Signed-off-by: Kshitij Lakhani <[email protected]> * Add support for returning reduced per head max logit. Plumb max logit and softmax through the JAX fused attn primitives Signed-off-by: Kshitij Lakhani <[email protected]> * Add max logit to JAX fused attn FFI and set it in the workspace Signed-off-by: Kshitij Lakhani <[email protected]> * Add first pass tests for max logit and softmax aux tensor outputs in JAX fused attn tests Signed-off-by: Kshitij Lakhani <[email protected]> * Reject aux returns with score_mod Signed-off-by: Kshitij Lakhani <[email protected]> * Handle SM120 max-logit layout Signed-off-by: Kshitij Lakhani <[email protected]> * Drop softmax aux return Signed-off-by: Kshitij Lakhani <[email protected]> * Modify static args in fused attn tests for jax Signed-off-by: Kshitij Lakhani <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Inline the choice of what is to be returned and remove redundant function for it Signed-off-by: Kshitij Lakhani <[email protected]> * Expose JAX max logit Signed-off-by: Kshitij Lakhani <[email protected]> * Support CP max logit Signed-off-by: Kshitij Lakhani <[email protected]> * Expand JAX max logit tests Signed-off-by: Kshitij Lakhani <[email protected]> * Remove JAX max logit integration tests Signed-off-by: Kshitij Lakhani <[email protected]> * Broaden JAX CP max logit tests Signed-off-by: Kshitij Lakhani <[email protected]> * Reduce JAX max logit across DP Signed-off-by: Kshitij Lakhani <[email protected]> * Simplify JAX max logit return Signed-off-by: Kshitij Lakhani <[email protected]> * Document JAX max logit reductions Signed-off-by: Kshitij Lakhani <[email protected]> * Refine JAX max logit tests Signed-off-by: Kshitij Lakhani <[email protected]> * Rename JAX max logit buffer Signed-off-by: Kshitij Lakhani <[email protected]> * Check JAX attention tensor pack capacity Signed-off-by: Kshitij Lakhani <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add distributed Muon optimizer Signed-off-by: Vladimir Cherepanov <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Muon closure and reference test Signed-off-by: Vladimir Cherepanov <[email protected]> * Fix Muon optimizer distributed API handling Signed-off-by: Vladimir Cherepanov <[email protected]> * Fix Muon optimizer docs and params typing Signed-off-by: Vladimir Cherepanov <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tensor-parallel Newton-Schulz wrapper Signed-off-by: Vladimir Cherepanov <[email protected]> * Move Newton-Schulz wrapper into optimizers Signed-off-by: Vladimir Cherepanov <[email protected]> * Use tensor-parallel Newton-Schulz in Muon Signed-off-by: Vladimir Cherepanov <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Limit Muon branch to Newton-Schulz TP Signed-off-by: Vladimir Cherepanov <[email protected]> * Run Newton-Schulz distributed cases in one launch Signed-off-by: Vladimir Cherepanov <[email protected]> * Add single-GPU Newton-Schulz coverage Signed-off-by: Vladimir Cherepanov <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Newton-Schulz compatibility and replicated coverage Signed-off-by: Vladimir Cherepanov <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Vladimir Cherepanov <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…IA#3424) fix clamped swiglu test Signed-off-by: Varun Thumbe <[email protected]>
…A#3380) * Unify eager and non-eager EP dispatch into one fused prepare+dispatch cpp op * Trim eager EP dispatch and combine autograd operands to cut per-call Python overhead * Call the backend directly for eager EP dispatch, combine to cut per-call CPU overhead --------- Signed-off-by: Phuong Nguyen <[email protected]>
Signed-off-by: Varun Thumbe <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <[email protected]>
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
20 participants
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: