All pull requests

Pull requests

• Add quantized all-gather for the ring of experts combine backward

  pull ready
  #5429·shuningjin opened Sep 29, 2026Collaborator·

  ·
  4

• Sow router bias updates from the use_lineage DeepSeek-V3 path.

  #5428·copybara-service[bot] opened Sep 28, 2026Bot·

  ·
  1

• Add Pallas cost_estimate to forward fused_conv1d_gdn kernel

  #5427·suhithr opened Sep 28, 2026Collaborator·

  ·
  4

• Suppress new Pyrefly findings before upgrading to v1.3.0

  #5425·copybara-service[bot] opened Sep 28, 2026Bot·

  ·
  1

• [trainerperf] Merge forward and backward scans in vocab_tiling_nnx_loss

  #5424·suhithr opened Sep 28, 2026Collaborator·

  ·
  2

• test(qwen3-next): run the packing equivalence tests on a GatedDeltaNet stack

  #5422·A9isha opened Sep 28, 2026Collaborator·

  ·

• Scaffolding Diffusers in Maxtext-M3 (for Weaver) via WAN [5/5]

  #5421·lydhr opened Sep 28, 2026Collaborator·

  ·

• Scaffolding Diffusers in Maxtext-M3 (for Weaver) via WAN [4/5]

  #5420·lydhr opened Sep 28, 2026Collaborator·

  ·

• Scaffolding Diffusers in Maxtext-M3 (for Weaver) via WAN [3/5]

  #5419·lydhr opened Sep 28, 2026Collaborator·

  ·

• Scaffolding Diffusers in Maxtext-M3 (for Weaver) via WAN [2/5]

  #5418·lydhr opened Sep 28, 2026Collaborator·

  ·
  14

• Scaffolding Diffusers in Maxtext-M3 (for Weaver) via WAN [1/5]

  #5417·lydhr opened Sep 28, 2026Collaborator·

  ·
  7

• Migrate DiLoCo launch scripts and tutorial from XPK to Cluster Toolkit

  pull ready
  #5416·darisoy opened Sep 28, 2026Collaborator·

  ·
  16

• [NNX] Port the remaining Linen sharding helpers into MaxText

  #5415·ecnal-cienet opened Sep 28, 2026Collaborator·

  ·
  16

• [NNX] Swap remaining Linen aliases for jax.nn and fix stale Linen text

  #5414·ecnal-cienet opened Sep 28, 2026Collaborator·

  ·
  1

• [Draft] Parity: Numerical logprob alignment tools, MoE-only FP8, and router replay for Qwen3.5-35B-A3B

  #5413·wenxindongwork opened Sep 28, 2026Collaborator·

  ·
  6

• fix(tests): prevent CI test collisions on shared GCS output paths

  scheduled-only
  #5411·shralex opened Sep 28, 2026Collaborator·

  ·
  4

• Skip the GDN kernel's segment path for unpacked training rows

  #5410·parambole opened Sep 28, 2026Collaborator·

  ·
  2

• Speed up Qwen3.5-397B-A17B training on TPU v7x

  #5409·parambole opened Sep 28, 2026Collaborator·

  ·
  2

• Train: warmup_programs_in_init and log_step_diagnostics

  gemini-reviewpull ready
  #5406·ultrons opened Sep 28, 2026Collaborator·

  ·
  4

• MoE: split SparseCore all-gather pin (EP / FSDP / FSDP forward-only) and defer_small_all_reduces

  pull ready
  #5405·ultrons opened Sep 28, 2026Collaborator·

  ·
  4

• Fix/indexer sparse plumbing

  #5404·h1t35h opened Sep 28, 2026Contributor·

  ·
  2

• Introducing Absorbed-Query Latent MQA for MLA decode : 57x memory reduction for long context

  #5403·h1t35h opened Sep 28, 2026Contributor·

  ·
  11

• Use sliding window to track goodput metrics over time

  #5398·lydhr opened Sep 27, 2026Collaborator·

  ·
  4

• Exclude cuda-bindings and CUTLASS DSL from decoupled requirements

  #5396·ksapru opened Sep 27, 2026·

  ·
  1

• DiLoCo: fragment extraction/application and outer step for threaded streaming DiLoCo.

  #5394·copybara-service[bot] opened Sep 26, 2026Bot·

  ·
  1