Repository navigation

[Qwen3.5-397B MLPerf] Consolidated top-ROI trainer & GDN/MoE optimizations on atwigg/mlperf - #5721

Merged

Rohan-Bierneni merged 16 commits into

atwigg/mlperffrom

rbierneni/mlperf-397b-trainer-perf

Oct 16, 2026

Merged

[Qwen3.5-397B MLPerf] Consolidated top-ROI trainer & GDN/MoE optimizations on atwigg/mlperf#5721
Rohan-Bierneni merged 16 commits into
atwigg/mlperffrom
rbierneni/mlperf-397b-trainer-perf

Commits on Oct 15, 2026
- Port MoE, megablox, and ragged kernel prerequisites from PR #5698 onto atwigg/mlperf
  Rohan-Bierneni
  committed
  56d4363
- Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128
  Rohan-Bierneni
  committed
  e2db560
- [MaxTextEngine] Stop the gradient of frozen parameters inside fwd_bwd
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  92dc9f7
- MoE: compute the ring-of-experts routing index maps without sorts (moe_routing_index_math)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  c9f0298
- MoE: three exact glue rewrites (moe_unsplit_prefused_wi, moe_ragged_sort_unique_pair_argsort, moe_prefused_swiglu_vjp)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  f049c2a
- MoE: one packed routing all-gather per layer in the pipelined ring of experts (moe_roe_packed_routing_gather)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  18f50cd
- MoE: SwiGLU in the wi GMM epilogue (moe_wi_gmm_fused_swiglu)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  8b8939a
- MoE: fuse the ring-of-experts down projection and unsort, with a backward that reads the GEMM input (moe_router_grad_from_wo_input)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  d629844
- GDN: keep the Conv1D halo input in qkv's dtype (gdn_conv1d_input_in_qkv_dtype)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  d5f99f8
- Address review feedback on the MoE and GDN keys
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  49fd113
- GDN: build the forward kernel's Conv1D output from dense per-token views (gdn_conv1d_dense_token_stack)
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  cfdaf7d
- GDN forward: skip the zero fill of the prefill kernel's output buffer
  
  khatwanimohit
  authored and
  Rohan-Bierneni
  committed
  6b00e91
- Skip the GDN kernel's segment path for unpacked training rows
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  1f639cc
- Add flag-gated gather_scanned_kernels_once for scanned MoE router and attention query kernels
  Rohan-Bierneni
  committed
  3b126c0
Commits on Oct 16, 2026
- Add cp-as-ep-dense-no-cp and cp-as-ep-dense-fsdp mesh rules
  
  parambole
  authored and
  Rohan-Bierneni
  committed
  d66a820
- Complete the #5698 port that the trainer code in this branch depends on
  parambole
  committed
  22fa0c5

Commits

Commits on Oct 15, 2026

Commits on Oct 16, 2026