Repository navigation

Add a PyTorch EP=1 backend for MoE dispatch and combine - #7

Draft

0z5a wants to merge 119 commits into

vthumbe1503:nccl_ep_opsfrom

0z5a:ep1-pytorch-dispatch-combine

Draft

Add a PyTorch EP=1 backend for MoE dispatch and combine#7
0z5a wants to merge 119 commits into
vthumbe1503:nccl_ep_opsfrom
0z5a:ep1-pytorch-dispatch-combine

Commits on Sep 28, 2026
- Add a PyTorch EP=1 backend for MoE dispatch and combine
  0z5a
  committed
  8d8f8dd
- [pre-commit.ci] auto fixes from pre-commit.com hooks
  pre-commit-ci[bot]
  committed
  6ce761b
Commits on Sep 30, 2026
- [Common] row-scaled nvfp4 path: add single-launch group fused amax (#3467)
  
  authored
  b57646c
- Isolate FlashAttention 2/3 imports so a broken install cannot break TE (#3553)
  
  0z5a
  and
  ptrendx
  authored
  f6d1ecb
- [common] Reject mismatched bias/out dtype for non-FP8/FP4 GEMM (#3565)
  
  authored
  08ba459
- [Docs] Fix broken links (#3597)
  sbhavani
  authored
  ef2bf75
Commits on Oct 1, 2026
Commits on Oct 2, 2026
Commits on Oct 4, 2026
- Prototype NVFP4 with FP8 UE5M3 block scales (#3325)
  
  authored
  d0b4b32
Commits on Oct 5, 2026
- Bump NCCL EP submodule commit (#3619)
  phu0ngng
  authored
  8cd5e20
- Fix AArch64 NCCL EP linking with BOLT enabled (#3627)
  aybchan
  authored
  f68a64a
- [JAX] Add Triton dependency to JAX extension (#3610)
  jberchtold-nvidia
  authored
  2cd6893
Commits on Oct 6, 2026
- [Common] Add launch bounds to multi-tensor swizzle kernels (#3622)
  
  ravimajeti
  and
  claude
  authored
  b47a356
- Fix MXFP8 ONNX block quantization attributes (#3592)
  pggPL
  authored
  9d4bd38
- Skip FA4 for deterministic backward on SM12x (#3595)
  janbernloehr
  authored
  6d6a500
- [Common] Remove CuTeDSL first launch protection (#3620)
  kainzhong
  authored
  ce29d24
Commits on Oct 7, 2026
- Support building docs for publishing via GHA (#3515)
  fheinecke
  authored
  bba2d2b
- [PyTorch] Enable LayerNormLinear UB AG overlap under no_grad (#3546)
  
  ravimajeti
  and
  ptrendx
  authored
  6eb41bd
- Preserve tiny sigmoid scores in biased fused routing (#3637)
  
  philipcmonk
  and
  ksivaman
  authored
  5e99464
Commits on Oct 8, 2026
Commits on Oct 9, 2026
Commits on Oct 11, 2026
- [Common] Swizzle MXFP8 scaling factors with one CTA per scale tile (#3658)
  
  sraman-rgb
  and
  claude
  authored
  33e7ff3
Commits on Oct 12, 2026
Commits on Oct 13, 2026
- Bump nccl ep (#3677)
  phu0ngng
  authored
  20e97da
- [PyTorch] Distributed tests set device in NCCL init and manually shut down NCCL (#3655)
  timmoon10
  authored
  e394b12
- [PyTorch] Allow ZeroTensor dummy weight gradients (#3670)
  
  sraman-rgb
  and
  claude
  authored
  fc48751
- [PyTorch] Return the argmax from parallel_cross_entropy (#3671)
  
  sraman-rgb
  and
  claude
  authored
  39a8b87
- Change to enable CuteDSL DenseGEMM for eval (#3666)
  vasunvidia
  authored
  0c4d961
- [PyTorch] Refactor GroupedLinear forward and backward (#3606)
  
  pggPL
  and
  pre-commit-ci[bot]
  authored
  e6afa9d
- [PyTorch] Fix deferred initialization in fusible ops (#3327)
  denera
  authored
  53e241c
Commits on Oct 14, 2026
- [Common][JAX] Reconstruct EpHandle via ncclEpImportHandle (#3639)
  phu0ngng
  authored
  6ca7b5a
- [PyTorch] Set device ID in distributed cross entropy + argmax test (#3684)
  timmoon10
  authored
  cfdd228
- Bump flash attention 2 (#3403)
  
  francesco-bertolotti
  and
  cyanguwa
  authored
  c72db38
- Update VERSION.txt to 2.22.0.dev0
  fheinecke
  committed
  66ba897
- KDA attention (#3548)
  
  ksivaman
  and
  pre-commit-ci[bot]
  authored
  b9715e5
- [All] Guard THD learnable dSink on older cuDNN (#3470)
  KshitijLakhani
  authored
  c3081a5
Commits on Oct 15, 2026
- [PyTorch] Deprecate grouped linear and MLP environment gates (#3584)
  
  authored
  9d9bd9e
- Fix ONNX custom-op metadata and LayerNorm execution (#3659)
  
  authored
  33f3b35
Commits on Oct 16, 2026
Commits on Oct 19, 2026
Commits on Oct 20, 2026
Commits on Oct 21, 2026
- cuBLASMp NVFP4 support (#3634)
  
  authored
  6b5cd7e
- Merge main and preserve local EP alignment defaults
  0z5a
  committed
  95b69f7
- Integrating TE Triton mHC kernels into JAX (#3680)
  dinodeep
  authored
  53d1a3c
Commits on Oct 22, 2026
- [PyTorch] Support FP8 weight caching in compiled Linear (#3505)
  
  pggPL
  and
  ksivaman
  authored
  9ea05cf
- [PyTorch] Make the Triton mask-map permute respect num_out_tokens (#3347)
  
  truong-v
  and
  denera
  authored
  c9c9225
Commits on Oct 23, 2026
Commits on Oct 26, 2026
Commits on Oct 27, 2026
- Support row-only MXFP8 distributed master-weight casts (#3488)
  
  authored
  3850986
- [PyTorch] Set device during NCCL init in symmetric memory test (#3744)
  timmoon10
  authored
  0f8d92f
- [Common][PyTorch] Add NVFP4 row-scaled grouped cast with graph-safe path (#3649)
  
  authored
  89668d9
- TE MoE Return Expert Counts and Custom Routing for DSv4 Support (#3596)
  dinodeep
  authored
  83d8c4e
- [Common] Don't call CUDA from static destructors at exit (#3733)
  
  ConnorBaker
  and
  claude
  authored
  dd628d3
- Fix DistributedWeight handling under saved-tensor hooks (#3730)
  
  authored
  fa4b30f
- Remove device side validation of QB bounds (#3734)
  
  ksivaman
  and
  pre-commit-ci[bot]
  authored
  96c0fb0
- [PyTorch] FP8BS grouped GEMM test reference fix (#3722)
  denera
  authored
  d7cd302
- Abstract CUDA hardcodes into configurable te_device_type / te_platform (#3513)
  
  authored
  f8f34c6
- Guard fused MoE auxiliary loss integer offsets (#3569)
  
  arcusbuilds
  and
  denera
  authored
  268fdd5
- Rely on Recipe object for fusion (#3717)
  
  vthumbe1503
  and
  timmoon10
  authored
  9b7e959
Commits on Oct 28, 2026
Commits on Oct 29, 2026
- [Docs] Add Mixture of Experts guide (#3494)
  
  authored
  a0a7e92
- [PyTorch]Validate grouped split alignment on device (#3745)
  
  vthumbe1503
  and
  timmoon10
  authored
  6f6b8f7
- [PyTorch] Skip FA 2.8.4 for unsupported arches in L3 (#3772)
  cyanguwa
  authored
  5f3abe1
- [JAX] Disable TE C API CuTeDSL backend at import time (#3784)
  jberchtold-nvidia
  authored
  c83e856
Commits on Oct 30, 2026
- Add code quality guidelines to AGENTS.md (#3773)
  
  timmoon10
  and
  claude
  authored
  43f61f7
- [PyT] Disable TE torch compilation in KV-cache CI tests (#3781)
  KshitijLakhani
  authored
  45a6d7c
- Grouped Linear CPU Overhead Reductions By Removing Unnecessary Python Logic (#3613)
  
  authored
  aac5cb4
- Deterministic MoE Router Kernel By Removing Global Atomic Ops (#3587)
  
  authored
  fd2ab0c
- [PyTorch] Remove EP handle cache and reconstructing handle on the fly (#3760)
  phu0ngng
  authored
  6320b5f
- [Common] Make it easier to debug C++ UT failures (#3783)
  
  kainzhong
  and
  ksivaman
  authored
  d3a16e5
- Accumulate delayed wgrads into shared (tied) parameters (#3552)
  
  0z5a
  and
  denera
  authored
  35a8e2f
Commits on Oct 31, 2026
Commits on Nov 1, 2026
- Merge remote-tracking branch 'upstream/main' into ep1-pytorch-dispatch-combine
  
  0z5a
  and
  claude
  committed
  3df8821

Commits

Commits on Sep 28, 2026

Commits on Sep 30, 2026

Commits on Oct 1, 2026

Commits on Oct 2, 2026

Commits on Oct 4, 2026

Commits on Oct 5, 2026

Commits on Oct 6, 2026

Commits on Oct 7, 2026

Commits on Oct 8, 2026

Commits on Oct 9, 2026

Commits on Oct 11, 2026

Commits on Oct 12, 2026

Commits on Oct 13, 2026

Commits on Oct 14, 2026

Commits on Oct 15, 2026

Commits on Oct 16, 2026

Commits on Oct 19, 2026

Commits on Oct 20, 2026

Commits on Oct 21, 2026

Commits on Oct 22, 2026

Commits on Oct 23, 2026

Commits on Oct 26, 2026

Commits on Oct 27, 2026

Commits on Oct 28, 2026

Commits on Oct 29, 2026

Commits on Oct 30, 2026

Commits on Oct 31, 2026