All pull requests

Pull requests

• Document shard_mode performance and the explicit-sharding tuning flags

  #5650·NuojCheng opened Oct 11, 2026Collaborator·3 of 4 tasks·

  ·

• Contract the tied LM head over the embedding table's feature axis

  #5649·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·
  2

• Build YaRN cis values from cos and sin instead of complex exp

  #5648·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·

• Turn lm_head_vocab_parallel on by default wherever it applies

  #5647·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·
  2

• Add an opt-in vocab-parallel layout for the untied LM head

  #5646·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·
  3

• Add save_fsdp_gathered_weights to keep gathered kernels for backward

  #5645·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·
  2

• Add dense_weight_grad_in_kernel_order for attention and MLP layers

  #5644·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·

• Add lm_head_weight_grad_in_kernel_order for the untied LM head

  #5643·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·
  3

• Add a kernel-order weight-gradient option to DenseGeneral

  #5642·NuojCheng opened Oct 11, 2026Collaborator·4 tasks done·

  ·
  2

• [Draft][Qwen3.5][MoE] Support TPU_ATTN_MOE_OVERLAP on flax_nnx path

  #5641·wenxindongwork opened Oct 10, 2026Collaborator·

  ·
  3

• [MoE] Forward num_valid_tokens in fused_moe_matmul under TPU_MOE_SKIP_PADDED_TOKENS

  #5640·wenxindongwork opened Oct 10, 2026Collaborator·

  ·
  2

• [Fix] Fix per-token activation quantization in native_fp8_dot_general and respect unquantized_modules for serve_fp8_weight (atwigg/mlperf)

  #5639·sierraisland opened Oct 10, 2026Collaborator·4 tasks done·

  ·
  2

• feat(checkpoint): support loading unquantized BF16 checkpoints into fp8_moe models

  #5638·sierraisland opened Oct 10, 2026Collaborator·4 tasks done·

  ·
  3

• Sync Lineage DeepSeek-V3 modules into experimental/lineage, including the branched splash attention kernel, and fix an import-time failure under jaxtyping >= 0.3.

  #5637·copybara-service[bot] opened Oct 10, 2026Bot·

  ·
  1

  

• feat(vllm): support Decode Context Parallelism (DCP) for Qwen3.5 in MaxTextForCausalLM (atwigg/mlperf)

  #5635·wenxindongwork opened Oct 10, 2026Collaborator·4 tasks done·

  ·
  3

• [MLPerf][Serving] Enable DCP attention routing and is_decode propagation for MaxText vLLM serving (#5621)

  #5634·wenxindongwork opened Oct 10, 2026Collaborator·

  ·
  3

• Migrate DeepSeek e2e tests to to_maxtext and fix scanned MoE gate stacking & Grain iterator cleanup

  #5633·SurbhiJainUSC opened Oct 9, 2026Collaborator·4 tasks done·

  ·
  4

• Support checkpoint loading for fp8_moe inference model

  #5629·wenxindongwork opened Oct 9, 2026Collaborator·4 tasks done·

  ·
  7

• [GDN] Layout fixes for the Pallas kernel path: bf16 qkv input, natural-layout backward kernel, head-major backward glue (opt-in)

  #5628·tissue3 opened Oct 9, 2026Collaborator·4 tasks done·

  ·
  7

• [GDN] Sequence-CP two-pass forward: pass-1 states-only, in-kernel M_local, pass-2 t_inv reuse (opt-in)

  #5627·tissue3 opened Oct 9, 2026Collaborator·4 tasks done·

  ·
  4

• MoE: invert sort permutations with an O(N) scatter

  gemini-review
  #5617·systalyze-ai opened Oct 8, 2026·4 tasks done·

  ·
  5

• [DO NOT REVIEW] Release pipeline v2

  #5616·SurbhiJainUSC opened Oct 8, 2026Collaborator·4 tasks·

  ·
  2

• [DO NOT REVIEW] Add a Github workflow that builds release candidate Docker images weekly

  #5615·SurbhiJainUSC opened Oct 8, 2026Collaborator·4 tasks·

  ·
  2

• Optimize OLMo 3.5 Small (olmoe3-3p5b) training on TPU.

  #5614·copybara-service[bot] opened Oct 8, 2026Bot·

  ·
  

• fix(load_dynamic): Correct weight dtype and -Instruct lookup

  #5610·ecnal-cienet opened Oct 8, 2026Collaborator·4 tasks done·

  ·
  1