Repository navigation

All pull requests

Add fused fp8 QKV/out projection kernels with QK-RMSNorm and RoPE for Qwen3 attention.
#6305copybara-service[bot] opened Nov 13, 2026Bot
chore(deps): upgrade qwix to 0.1.10
#6304darisoy opened Nov 13, 2026Collaborator8 tasks done
4
Add config knobs and fp8 helpers for the Qwen3-235B v7x training stack.
#6303copybara-service[bot] opened Nov 13, 2026Bot
[MoE replication 4/4] Wire online hot-expert replication into the ragged all-to-all MoE path
#6302Shuwen-Fang opened Nov 12, 2026Collaborator4 tasks done
1
[MoE replication 3/4] Add replica weight-move ops for online hot-expert replication
#6301Shuwen-Fang opened Nov 12, 2026Collaborator4 tasks done
1
[MoE replication 2/4] Add replica planners for online hot-expert replication
#6300Shuwen-Fang opened Nov 12, 2026Collaborator4 tasks done
1
Add log_ra2a_required_rbf debug flag for ragged all-to-all MoE
#6299Shuwen-Fang opened Nov 12, 2026Collaborator4 tasks done
6
Add synthetic MoE routing benchmark harness under benchmarks/moe_synthetic_routing
#6298Shuwen-Fang opened Nov 12, 2026Collaborator4 tasks done
1
Remove legacy XPK wrapper and XPK instructions for OLMo-3 7B
#6297cj-klein opened Nov 12, 2026
1
Support absmax calibration for moe_quantize_token_all_gather.
#6296copybara-service[bot] opened Nov 12, 2026Bot4 tasks done
1
Address post-submission review follow-ups for MaxTextTrainingEngine.with_fwd_only_fn / fwd_only:
#6293copybara-service[bot] opened Nov 12, 2026Bot
1
Use the configured PRNG impl in standalone_checkpointer
#6292NuojCheng opened Nov 12, 2026Collaborator4 tasks done
4
Raise instead of assuming 1 device when config validation can't get the device count
#6291NuojCheng opened Nov 12, 2026Collaborator4 tasks done
3
Internal change
#6290copybara-service[bot] opened Nov 12, 2026Bot
1
Support full attention on TPU Ulysses CP in MaxText.
#6289copybara-service[bot] opened Nov 12, 2026Bot
1
Support dynamic (absmax) weight scaling for MoE QWAG
#6288copybara-service[bot] opened Nov 12, 2026Bot4 tasks done
MoE: add moe_wi_dlhs_merged_k to compute the routed wi_0/wi_1 input gradient with one gmm_v2 DLHS. Default off.
#6287copybara-service[bot] opened Nov 12, 2026Bot
1
Adapt the experimental Lineage DeepSeek-V3 671B (+ 1 MTP layer) trainer to run on a single 64-chip TPU v7x slice (tpu7x:4x4x4, 128 TensorCores) and add tuned 8K and 16K long-context recipes under src/maxtext/experimental/lineage/recipes/.
#6286copybara-service[bot] opened Nov 12, 2026Bot
1
Expand the MaxText m3 conformance suite with four additional checks (CheckId 9 through 12), wire qwen3-0.6b and glm5.3-flash conformance specs to pass all 12 checks, and add fault-injection negative controls in conformance_test.py:
#6285copybara-service[bot] opened Nov 12, 2026Bot
Add tiled block-sparse FlashAttention with 32-bit packed bitmap mask for m3 Qwen3.8 QSA
#6283AusarYao28 opened Nov 12, 2026Collaborator4 tasks done
13
feat(m3): adopt shared MoEExecutor for GraniteMoE and multi-rank EP GLM-5.3
#6282AusarYao28 opened Nov 12, 2026Collaborator4 tasks done
8
feat(m3): add capacity-bounded dropless MoE EP overflow buffer
#6281AusarYao28 opened Nov 12, 2026Collaborator4 tasks done
5
Store skip_step_on_spikes flag as int32 for Pathways persistence.
#6280copybara-service[bot] opened Nov 12, 2026Bot
1
MoE: add moe_fold_expert_scale_into_wo to apply per_expert_scale to the sharded wo instead of gathering it per token. Default off.
#6279copybara-service[bot] opened Nov 12, 2026Bot
1
fix(testing): restore relative-tolerance loss check for checkpoint_save_restore
#6277ryanliuCienet opened Nov 12, 2026Collaborator4 tasks done
5

Navigation Menu

Repository navigation

All pull requests

Add fused fp8 QKV/out projection kernels with QK-RMSNorm and RoPE for Qwen3 attention.

chore(deps): upgrade qwix to 0.1.10

Add config knobs and fp8 helpers for the Qwen3-235B v7x training stack.

[MoE replication 4/4] Wire online hot-expert replication into the ragged all-to-all MoE path

[MoE replication 3/4] Add replica weight-move ops for online hot-expert replication

[MoE replication 2/4] Add replica planners for online hot-expert replication

Add log_ra2a_required_rbf debug flag for ragged all-to-all MoE

Add synthetic MoE routing benchmark harness under benchmarks/moe_synthetic_routing

Remove legacy XPK wrapper and XPK instructions for OLMo-3 7B

Support absmax calibration for moe_quantize_token_all_gather.

Address post-submission review follow-ups for `MaxTextTrainingEngine.with_fwd_only_fn` / `fwd_only`:

Use the configured PRNG impl in standalone_checkpointer

Raise instead of assuming 1 device when config validation can't get the device count

Internal change

Support full attention on TPU Ulysses CP in MaxText.

Support dynamic (absmax) weight scaling for MoE QWAG

MoE: add moe_wi_dlhs_merged_k to compute the routed wi_0/wi_1 input gradient with one gmm_v2 DLHS. Default off.

Adapt the experimental Lineage DeepSeek-V3 671B (+ 1 MTP layer) trainer to run on a single 64-chip TPU v7x slice (`tpu7x:4x4x4`, 128 TensorCores) and add tuned 8K and 16K long-context recipes under `src/maxtext/experimental/lineage/recipes/`.

Expand the MaxText m3 conformance suite with four additional checks (`CheckId` 9 through 12), wire `qwen3-0.6b` and `glm5.3-flash` conformance specs to pass all 12 checks, and add fault-injection negative controls in `conformance_test.py`:

Add tiled block-sparse FlashAttention with 32-bit packed bitmap mask for m3 Qwen3.8 QSA

feat(m3): adopt shared MoEExecutor for GraniteMoE and multi-rank EP GLM-5.3

feat(m3): add capacity-bounded dropless MoE EP overflow buffer

Store skip_step_on_spikes flag as int32 for Pathways persistence.

MoE: add moe_fold_expert_scale_into_wo to apply per_expert_scale to the sharded wo instead of gathering it per token. Default off.

fix(testing): restore relative-tolerance loss check for checkpoint_save_restore

Pull requests