Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys by parambole · Pull Request #5758 · AI-Hypercomputer/maxtext

Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128

Rohan-Bierneni

authored and

parambole

committed

dede90a

[MaxTextEngine] Stop the gradient of frozen parameters inside fwd_bwd

parambole

committed

7f5dcc7

MoE: compute the ring-of-experts routing index maps without sorts (moe_routing_index_math)

parambole

committed

add4a1e

MoE: three exact glue rewrites (moe_unsplit_prefused_wi, moe_ragged_sort_unique_pair_argsort, moe_prefused_swiglu_vjp)

parambole

committed

048d617

MoE: one packed routing all-gather per layer in the pipelined ring of experts (moe_roe_packed_routing_gather)

parambole

committed

c88674b

MoE: SwiGLU in the wi GMM epilogue (moe_wi_gmm_fused_swiglu)

parambole

committed

4ab9f12

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

committed

63cba26

GDN: keep the Conv1D halo input in qkv's dtype (gdn_conv1d_input_in_qkv_dtype)

parambole

committed

cb709c9

Address review feedback on the MoE and GDN keys

parambole

committed

8cc954e

GDN: build the forward kernel's Conv1D output from dense per-token views (gdn_conv1d_dense_token_stack)

parambole

committed

20c877a

Navigation Menu

Repository navigation

Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys - #5758

Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys#5758
parambole wants to merge 10 commits into
parambole/rl-trainer-4600-basefrom
parambole/rl-trainer-4600

Commits on Oct 15, 2026

Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128

[MaxTextEngine] Stop the gradient of frozen parameters inside fwd_bwd

MoE: compute the ring-of-experts routing index maps without sorts (moe_routing_index_math)

MoE: three exact glue rewrites (moe_unsplit_prefused_wi, moe_ragged_sort_unique_pair_argsort, moe_prefused_swiglu_vjp)

MoE: one packed routing all-gather per layer in the pipelined ring of experts (moe_roe_packed_routing_gather)

MoE: SwiGLU in the wi GMM epilogue (moe_wi_gmm_fused_swiglu)

MoE: fuse the ring-of-experts down projection and unsort, with a backward that reads the GEMM input (moe_router_grad_from_wo_input)

GDN: keep the Conv1D halo input in qkv's dtype (gdn_conv1d_input_in_qkv_dtype)

Address review feedback on the MoE and GDN keys

GDN: build the forward kernel's Conv1D output from dense per-token views (gdn_conv1d_dense_token_stack)