All pull requests

Pull requests

• tpu: stack the query heads that share a KV head in the Tokamax forward

  #130·wconstab opened Oct 30, 2026Contributor·

  ·
  3

• tpu graph_split: choose which collectives fsdp_pipeline moves; pipeline Q235 4x4x4 reduce-scatters

  #128·wconstab opened Oct 30, 2026Contributor·

  ·

• tpu graph_split: let fsdp_pipeline own every backward FSDP collective

  #127·wconstab opened Oct 30, 2026Contributor·

  ·

• ci: pin the nightly image to the last good PyTorch nightly

  #126·jerryx99 opened Oct 30, 2026Contributor·

  ·

• Add opt-in traceable FlexShard Qwen3 TPU experiment

  #125·weifengpy opened Oct 30, 2026·

  ·

• [DO NOT MERGE] Flint 40B via the HF transformers backend on TorchTPU: 16.3 % MFU on v7x-8

  #123·vlasenkoalexey opened Oct 30, 2026Contributor·

  ·

• [DO NOT MERGE] Flint 40B on the TorchTPU native lane: 16.23 % MFU on v7x-8

  #122·vlasenkoalexey opened Oct 30, 2026Contributor·

  ·

• tpu: run the Tokamax backward per example for Q235

  #120·wconstab opened Oct 29, 2026Contributor·2/2·

  ·
  2

• tpu: per-example Tokamax forward (Tokamax 0.0.17) and vendored fused CE

  #119·wconstab opened Oct 29, 2026Contributor·1/2·

  ·
  3

• Keep SplashAttention MaskInfo on host NumPy to avoid JAX TPU races

  #117·copybara-service[bot] opened Oct 29, 2026Bot·

  ·

• tpu graph_split: finish expert all-gathers under Splash attention

  #114·wconstab opened Oct 28, 2026Contributor·

  ·

• [tpu] Support EP>1 with scan layers in the graph trainer

  #113·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Add a static-shape Ring-of-Experts EP transport

  #112·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Graph-split and scan remat fixes for EP collectives

  #111·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Add a TensorCore ragged row-gather Pallas kernel for MoE

  #110·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Support multi-mesh gradients in grad-norm clipping for EP>1

  #109·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Add a Qwen3-235B scan recipe with 4-way Ring-of-Experts EP

  #108·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Add a TensorCore ragged gather-reduce Pallas kernel for MoE combine

  #107·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Pipeline Ring-of-Experts EP over token chunks

  #106·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• [tpu] Use Ring-of-Experts as the TPU MoE dispatcher's EP>1 transport

  #105·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• No public description

  #104·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• Migrate torchtitan TPU experiments and profiler to public torch.tpu and torch.accelerator APIs

  #101·copybara-service[bot] opened Oct 28, 2026Bot·

  ·

• Bump libtpu to 0.0.51

  #95·jack8558 opened Oct 27, 2026Contributor·

  ·

• tpu qwen3 fp8: Add FP8 attention projections to qwen3_moe_235b_scan_fp8

  #93·alexanderb14 opened Oct 27, 2026Contributor·

  ·

• tpu qwen3 fp8: Add a TPU-native FP8 linear layer for torchtitan on TPU

  #92·alexanderb14 opened Oct 27, 2026Contributor·

  ·