All pull requests

Pull requests

• Add a TPU-friendly chunked delta rule for Gated DeltaNet (gdn_chunk_impl=fast)

  gemini-review
  #5525·junjieqian opened Oct 3, 2026Collaborator·4 tasks done·

  ·
  6

• Remove obsolete force_on_policy_ratio and log_sampler_trainer_agreement RL configuration flags from RLSettings (src/maxtext/configs/types.py) and AgenticGrpoConfig instantiation (src/maxtext/trainers/post_train/rl/train_rl.py), following their removal and in-loss agreement fusion unification in Tunix (cl/992453250).

  #5521·copybara-service[bot] opened Oct 3, 2026Bot·

  ·
  1

• Add a lineage_expert_permutation config option for per-step expert

  #5520·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

• [MaxText] Fix the MLPerf parallelism, precision and config_filename mllog disclosure for Lineage runs

  #5519·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• PR #5456: [NNX] Delete unused ToLinen wrappers and move fp8 and HF conversion to NNX

  #5518·copybara-service[bot] opened Oct 2, 2026Bot·4 tasks done·

  ·
  1

• [MaxTextEngine] Optionally skip a checkpoint save while the previous one is still being written

  #5516·A9isha opened Oct 2, 2026Collaborator·4 tasks done·

  ·
  2

• Support joint EP and FSDP expert-axis sharding for TE MoE

  #5515·abhinavgoel95 opened Oct 2, 2026Contributor·4 tasks done·

  ·
  2

• Add experimental Kimi-K3 text-only four-layer support

  #5514·hengtaoguo opened Oct 2, 2026Collaborator·2 of 3 tasks·

  ·
  6

• [DRAFT] Pipelined SparseCore ragged_gather + plan to speed up SC ragged kernels

  #5513·NuojCheng opened Oct 2, 2026Collaborator·2 of 4 tasks·

  ·
  3

• Fix remaining fixable critical and high docker image vulnerabilities

  #5512·KevinYenky opened Oct 2, 2026Collaborator·4 tasks done·

  ·
  3

• [TC sort 4/4] All-gather / combine MoE tokens in the TC kernels' 3D layout

  #5507·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  4

• Train loop and input pipeline hooks for training a data-parallel replica on a submesh.

  #5506·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Threaded (non-SPMD) streaming DiLoCo trainer.

  #5505·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• DiLoCo: threaded streaming DiLoCo recipe and docs.

  #5504·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Profiling and logging support for threaded DiLoCo.

  #5503·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Quantize the GMM bwd gradient once under fixed calibration

  #5502·RissyRan opened Oct 2, 2026Collaborator·2 of 4 tasks·

  ·
  4

• [TC sort 3/4] Keep the TC-sorted MoE token buffer in 3D layout through the gmm_v2 k…

  #5498·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  6

• [TC Sort 2/4] Apply MoE routing weights to the activation with the TC ragged sort

  #5497·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  5

• [TC Sort 1/4] Add TensorCore ragged sort for ring-of-experts MoE

  pull ready
  #5496·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  9

• Forward-pass velocity numerical checker for WeaverOmniTransformer

  #5495·parsley9877 opened Oct 1, 2026Collaborator·3/3·4 tasks done·

  ·
  5

• Add module tests to isolate quantization divergence

  #5490·pv97 opened Oct 1, 2026·

  ·
  6

• Add checkpoint conversion and parameter mapping for WeaverOmniTransformer

  #5488·parsley9877 opened Oct 1, 2026Collaborator·2/3·4 tasks·

  ·
  5

• [Draft] [GPU] DSv4 Support for TE MoE Block

  #5487·dinodeep opened Oct 1, 2026Collaborator·4 tasks done·

  ·
  9

• Migrate MaxText benchmark runner and configs from XPK to Cluster Toolkit

  #5483·KevinYenky opened Oct 1, 2026Collaborator·4 tasks done·

  ·
  31

• Restore Orbax checkpoint Cloud Logging for ML Goodput after the Orbax V1 migration

  #5478·lydhr opened Sep 30, 2026Collaborator·4 tasks done·

  ·
  9