[grug] Drive 90B-5.3BA MoE MFU on cw-us-east-02a from 3% to ~25%

[rjpower]

Description

The 90B-total/5.3B-active Grug MoE from the cw-us-east-02a bringup (#6292, PR #6293) trains at 218K tok/s on 256 H100 — 2.8% active-FLOPs MFU (6·5.3e9 FLOPs/token ÷ 256·989 TFLOP/s). This issue drives it toward ~25% (≈2.0M tok/s) and is the running log of every experiment along the way. Complements the H100 MoE MFU epic #4302.

Diagnosis (repo mining + external research)

Structurally communication-bound. At global batch 256×2048, each GPU computes 2,048 tokens/step (~66 ms ideal bf16) while FSDP/ZeRO-3 over the 32-node IB axis gathers ~45–130 GB/GPU/step of parameters (90B total, fp32 master weights) against 400 Gb/s/GPU — 1–2.6 s of comm vs 0.07 s of compute. The observed 2.4 s step matches this arithmetic. Secondary suspects: reference attention on GPU (attention_implementation unset → materialized fp32 logits ×48 layers ×2 with remat; FA4 blocked by #6226), full-block remat (+20%, #5328), XLA fused-CE fallback (#5510).

Batch/seq scaling — the single biggest lever — is gated on the replicated-reshard pathology: XLA GSPMD reshards the [batch,seq,hidden] activation through a fully-replicated intermediate at the FSDP boundary (58 GiB tile OOM at 512×4096; see launch_cw_scale.py). Nothing in-repo has tried Shardy yet.

Plan (ranked; profile first, change one thing at a time)

1. Baseline profile with the in-tree harness (ProfilerConfig + lib/marin/tools/profile_summary.py, per .agents/skills/profile-training/): split the step across comm / attention / dispatch / GMM / optimizer, per H100 x 8 MOE MFU perf #4302's method.
2. Kill the replicated reshard: try Shardy (jax_use_shardy_partitioner; default partitioner since JAX 0.7.1, targets exactly this), else restructure _batch_reshard placement (PR [Grug] Reshard MoE block before MLP #6069 pattern) or fold it into the EP backend.
3. Scale tokens/device 8–16× once unblocked (MaxText H100 configs run 8K tok/dev; OLMoE sustained ~19% MFU at 1.3B-active).
4. XLA flag baseline + PGLE (latency-hiding scheduler, pipelined/combined collectives, cuBLAS GEMM, double buffering — MaxText A3 + JAX GPU perf-tips set). Note TPU flag sweeps were a negative result (Experiment: MaxText XLA flags for Grug MoE on v5p-64 #3399); GPU is untested.
5. bf16 param gathers (fp32 master weights currently risk 2× FSDP bytes — verify in HLO whether convert sinks below the all-gather).
6. FA4 attention on the CoreWeave image ([gpu] Make FA4 CuTe and THD work on CoreWeave #6226: install cutlass/flash-attn-4, set attention_implementation).
7. EP topology + dispatch: hybrid ZeRO (shard only the 87B expert params cross-node), wider EP with resident experts (DeepSeek-V3 topology), assigned-token/DeepEP dispatch ([grug] Replace ring EP global buffers with assigned-token dispatch #6215, PR [Grug] Add assigned-token DeepEP MoE dispatch #6251). Caution: naive ragged_all_to_all was 1.71× slower than ring ([grug] Close B200 d5120/L8 MoE transport gap after #5815 #6139).
8. Remat policy using the existing _CHECKPOINT_* names instead of bare eqx.filter_checkpoint (~+21.8% measured, Epic: make local Grug MoE GMM/MLP fast via SonicMoE-style kernels #5328); fused-CE block-size fix (Pallas fused cross-entropy falls back to XLA on CoreWeave H100 Grug MoE canary #5510); fp8/TE last.

Definition of Done

≥25% active-FLOPs MFU on 256 H100 (≈2.0M tok/s) at the 90B-5.3BA shape, or a documented ceiling with the remaining gap attributed (kernel vs comm vs geometry), with every experiment logged below.

---

Running log

Baseline (from bringup, 2026-06-09): d3072 L48 E128 top-4, batch 256 × seq 2048, mesh (data=32, expert=8), fp32 params / bf16 compute, ring EP, reference attention, full-block remat, no XLA flags. 218K tok/s = 2.8% MFU, step ~2.4 s. Canary (9.5B replicated per node, no cross-node FSDP) hit 1.39M tok/s on identical hardware — consistent with the comm-bound diagnosis.

[claude]

🤖 Analysis — this is a multi-experiment performance epic / running log (drive 90B-5.3BA MoE from ~3% to ~25% MFU), not a self-contained bug. No PR per the size rule; below is the diagnosis verified against code plus the single most concrete first lever for a human to drive.

Verified against code

• Reference attention is the GPU default. attention_implementation=None falls through to reference_attention on non-TPU backends (attention/_core.py#L494-L498); the scale model never sets it (launch_cw_scale.py#L83-L96). Confirms the suspect in the diagnosis.
• Full-block remat, not the named policy. Each block is wrapped in a bare eqx.filter_checkpoint (model.py#L528). The named _CHECKPOINT_* policies (_moe/common.py#L39-L42) are only applied inside the EP backends, not at block level — matches the ~+21.8% lever in Epic: make local Grug MoE GMM/MLP fast via SonicMoE-style kernels #5328.
• Shardy is described but unwired. The OOM/replicated-reshard note says "pending Shardy" (launch_cw_scale.py#L50-L53), but jax_use_shardy_partitioner is set nowhere in the repo.

Most concrete first step (no source change)

TrainerConfig.jax_config (trainer.py#L872, applied via jax.config.update at #L1018-L1020) is the existing XLA/JAX flag pass-through. Both step 1 (baseline profile via ProfilerConfig, callbacks/profiler.py#L67-L92) and step 2 (Shardy) can run by populating jax_config from the launch path — jax_use_shardy_partitioner=True plus the latency-hiding/PGLE flags — with zero source change. Suggest landing that as the first logged experiment before any kernel work.

[rjpower]

🤖 Baseline profile captured — run cw-grugmoe-mfu-baseline-20260610-031312 (32 nodes / 256 H100, batch 256 × seq 2048, 30 steps, profiler on steps 8–18). W&B: https://wandb.ai/marin-community/marin_moe/runs/cw-grugmoe-mfu-baseline-20260610-031312 — artifact jax-profile-step-8-18:v0, analyzed with lib/marin/tools/profile_summary.py --breakdown-mode exclusive_global.

Steady state: 8 profiled steps in 18.97 s ⇒ 2.37 s/step → ~221K tok/s → ~2.8% active-FLOPs MFU (matches the bringup figure; the hypothesis in this issue is confirmed by the trace).

Exclusive-time breakdown (single-host xplane, global timeline):

bucket	share
communication (NCCL)	73.4%
stall (idle gaps)	18.9%
compute	7.7%

What the communication is:

• AllGather_RING_LL: 37,304 launches in 8 steps ≈ 583 all-gathers per GPU per step, avg 2.0 ms each. That count ≈ 48 layers × ~12 param tensors × fwd+bwd — i.e. per-tensor FSDP param gathers, uncombined, in the LL (small-message) protocol.
• ReduceScatter_Sum_bf16: 6,336 launches, avg 4.5 ms (grad reduce-scatter).
• SendRecv accumulates ~1.0 s average pre-op idle gap (×152) — host-side/data-dependency stalls.

Reading: pure FSDP at 2,048 tok/GPU/step cannot hide ~600 serialized small param gathers; we are latency-bound on collective launches, not just bandwidth-bound. This reorders the plan slightly — collective combining is the cheapest first win:

1. XLA flags: --xla_gpu_all_gather_combine_threshold_bytes / --xla_gpu_reduce_scatter_combine_threshold_bytes (collapse 583 gathers → dozens of large ones) + latency-hiding scheduler + the MaxText A3 flag set already recorded above, then PGLE.
2. Tokens/device ↑ (batch and/or seq; currently blocked ~2K tok/GPU by the replicated-reshard memory spike — attacking that unlocks this).
3. bf16 param all-gathers (halve gather bytes; params are f32 at rest).
4. Shardy / reshard restructure (PR [Grug] Reshard MoE block before MLP #6069 pattern) to kill the replicated intermediate.

Ops note: the run wedged after step 30 on the final S3 checkpoint barrier — 31/32 processes committed, one sat >35 min in "Starting commit" with no progress, holding all 256 GPUs at the barrier. Killed the job after capturing the profile (checkpoint disposable for throughput runs); teardown converged in <1 min. Will disable end-of-run checkpointing for subsequent MFU experiment runs and keep an eye on this S3-commit hang pattern.

[rjpower]

🤖 Experiment 1 — XLA collective-combining / latency-hiding flag set: NEGATIVE (~-7%)

Run cw-grugmoe-mfu-xlaflags-20260610-051204 (32 nodes / 256 H100, baseline shape: batch 256 × seq 2048, 30 steps, local checkpoints). One change vs baseline:

XLA_FLAGS=--xla_gpu_enable_latency_hiding_scheduler=true
  --xla_gpu_all_reduce_combine_threshold_bytes=1073741824
  --xla_gpu_all_gather_combine_threshold_bytes=1073741824
  --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824
  --xla_gpu_enable_pipelined_all_gather=true
  --xla_gpu_enable_pipelined_reduce_scatter=true
  --xla_gpu_enable_pipelined_all_reduce=true
  --xla_gpu_enable_while_loop_double_buffering=true
  --xla_gpu_enable_all_gather_combine_by_dim=false
  --xla_gpu_enable_reduce_scatter_combine_by_dim=false

Result: ~205K tok/s, p50 MFU 2.71% (levanter throughput/mfu) vs baseline 221K tok/s / ~2.8%. Step ~2.7 s/it. XLA compile also grew ~14 → ~25 min. Matches the TPU negative result on flag bundles (#3399): the combining/pipelining passes don't rescue the involuntary full-remat reshard (spmd_partitioner.cc:667 fires on bf16[1,2048,3072] going from batch-sharded(256) → (data=32, expert=8)-sharded), which dominates. PGLE alone (without this bundle) still untested but deprioritized.

Infra notes for subsequent runs (commit 7829cdd23f on weaver/moe-mfu):

• XLA_FLAGS/NCCL_*/JAX_* are now actually forwarded from the dispatcher to GPU pods (they previously silently weren't — any earlier flag experiments through this path were no-ops).
• SCALE_CHECKPOINTS=local writes checkpoints to node-local disk; the step-30 save committed in seconds at 32 nodes (vs the >35 min S3 barrier wedge on the baseline run).

Next (lever 2): tokens/device. At 2,048 tok/GPU the param gathers can't amortize; doubling global batch to 512 @ seq 2048 (4,096 tok/GPU) should raise MFU near-linearly if memory holds (bringup: 256×2048 fit; 512×4096 OOM'd; 512×2048 untested). Flags OFF to isolate the lever. Launching now.

[rjpower]

🤖 Experiment 2 — tokens/device ×2 (batch 512 × seq 2048): +85%, MFU 2.8 → 5.37%

Run cw-grugmoe-mfu-batch512-20260610-061526 (32 nodes / 256 H100, no XLA flags, local checkpoints). One change vs baseline: global batch 256 → 512 (4,096 tok/GPU).

Result: ~408K tok/s, p50 MFU ~5.37%, step ~2.7 s/it. Near-perfectly linear: 2× tokens cost only +14% step time (2.37 → 2.7 s), exactly what the latency-bound diagnosis predicts — the per-step FSDP gather cost is ~fixed, so tokens ride along free until compute or memory pushes back. No OOM at 4,096 tok/GPU.

Next: batch 1024 × seq 2048 (8,192 tok/GPU). Caveat: bringup saw 512 × 4096 OOM, and peak memory scales with batch×seq through the replicated-reshard intermediate, so 1024×2048 (same product) may hit the same wall — cheap to probe (~25 min); fallback is batch 768. The reshard fix itself (Shardy / PR #6069 restructure) remains the unlock for going past this ceiling.

experiment	config delta	tok/s	MFU
baseline	—	221K	2.8%
exp 1	XLA LHS+combine flags	205K	2.71% ❌
exp 2	batch 512 (4K tok/GPU)	408K	5.37% ✅

[rjpower]

🤖 Experiment 3 — batch 1024 × seq 2048 (8,192 tok/GPU): no OOM, MFU 9.32%

Run cw-grugmoe-mfu-batch1024-20260610-070245. ~709K tok/s, p50 MFU ~9.3%, step ~2.96 s. Notable: bringup's 512×4096 OOM was seq-driven, not a batch×seq product limit — the GSPMD full-remat replicates a per-example [1, seq, hidden] tile, so batch rides free while seq doubles both the replicated tile and attention memory. Batch is therefore the safe scaling axis under the current reshard.

Scaling is going mildly sub-linear (2× tokens → 1.74× throughput), as compute starts to matter. Launched exp 4: batch 2048 (16,384 tok/GPU) — linear-ish would put it near ~16% MFU.

experiment	config delta	tok/s	MFU
baseline	batch 256 (2K tok/GPU)	221K	2.8%
exp 1	+ XLA LHS/combine flags	205K	2.71% ❌
exp 2	batch 512 (4K tok/GPU)	408K	5.37%
exp 3	batch 1024 (8K tok/GPU)	709K	9.32%

Shardy finding (kills lever 1 as previously framed): JAX 0.10.0 defaults jax_use_shardy_partitioner=True and nothing in the repo disables it — the spmd_partitioner.cc:667 involuntary-full-remat fires under Shardy propagation (GSPMD still does the partitioning). "Enable Shardy" is a no-op; the reshard must be restructured in code (the [B,S,H] activation transitions batch-sharded(256) → hidden-sharded(32)×expert-replicated(8) at the shard_map boundary in lib/levanter/src/levanter/grug/grug_moe.py / experiments/grug/moe/model.py:469).

Ops note: between runs, CoreWeave's cw-hpc-verification NHC v3 pods grabbed 8 GPU on all 32 nodes for ~45 min; the queued gang showed Kueue QuotaReserved=False ... excluded: nvidia.com/gpu: 32 and looked wedged, but it admitted itself once NHC finished. Don't mistake this for the #6302 wedge.

Next zero-code probe after the batch ceiling: PGLE (JAX_ENABLE_PGLE=true, JAX_PGLE_PROFILING_RUNS=3) — now possible since JAX_* env forwarding landed.

[rjpower]

🤖 Experiment 4 — batch 2048 × seq 2048 (16K tok/GPU): OOM. Batch ceiling sits between 8K and 16K tok/GPU under the current reshard. Killed the job mid-retry (it would have OOM'd twice more) and launched the midpoint probe, batch 1536 (12K tok/GPU) as exp 5.

Two infra footnotes: (a) the first batch-2048 attempt (075954) died at import, my bug — the submit script passed -e JAX_USE_SHARDY_PARTITIONER "" and JAX's bool_env rejects empty strings; fixed by only passing runtime-tuning env vars that are actually set. (b) After the OOM ceiling is mapped, next zero-code probe is PGLE at the largest working batch (JAX_ENABLE_PGLE=true, JAX_PGLE_PROFILING_RUNS=3).

[rjpower]

🤖 Experiment 5 — batch 1536 × seq 2048 (12K tok/GPU): OOM. The tokens/device ceiling under the current reshard is batch 1024 (8,192 tok/GPU, 9.32% MFU) — 1536 and 2048 both exhaust HBM. Killed the retry loop and cleared the kill-wedge manually again (finalizer strip + pod delete; second occurrence of the #6302 sequence this session — the deployed controller predates PR #6305).

Launched exp 6: PGLE probe at batch 1024 (JAX_ENABLE_PGLE=true, JAX_PGLE_PROFILING_RUNS=3) — profile-guided latency estimates are the main reason the latency-hiding scheduler underwhelmed bare in exp 1; expect ~3 extra profiling steps + a recompile at startup.

Current state: 221K → 709K tok/s (2.8% → 9.32% MFU) purely from batch scaling. Getting past ~9% requires either PGLE/scheduling wins or attacking the replicated reshard (memory ceiling and comm volume share the same root cause).

[rjpower]

🤖 Experiment 6 — PGLE at batch 1024: OOM (negative). JAX_ENABLE_PGLE=true + JAX_PGLE_PROFILING_RUNS=3 ran 3 profiling-mode steps at ~74–104K tok/s then exhausted HBM — PGLE's profiling instrumentation + recompile double-buffering doesn't fit at the 8K tok/GPU ceiling, and it logged PGLE collected an empty trace (CUPTI contention) even while running, so there was likely no signal to gain. Not worth retrying at batch 512 (would need >1.7× from scheduling alone to beat exp-3's absolute throughput).

Zero-code levers are now exhausted. Scoreboard:

experiment	config	tok/s	MFU	verdict
baseline	batch 256	221K	2.8%	—
exp 1	XLA LHS/combine flags	205K	2.71%	❌
exp 2	batch 512	408K	5.37%	✅
exp 3	batch 1024	709K	9.32%	✅ best
exp 4	batch 2048	OOM	—	ceiling
exp 5	batch 1536	OOM	—	ceiling
exp 6	PGLE @ 1024	OOM	—	❌

Pivoting to code levers on weaver/moe-mfu, in order: (1) remat policy — replace the bare eqx.filter_checkpoint(block) (model.py:528) with a save-policy over the existing named checkpoints (grug_moe_dispatch_input/expert_hidden/dispatch_output/output in lib/levanter/src/levanter/grug/_moe/common.py), per the ~+21.8% result in #5328; lower recompute AND lower peak memory may also raise the batch ceiling. (2) Reshard restructure at the shard_map boundary (model.py:469, grug_moe.py) to kill the GSPMD full-remat. Each validated on a 1-node smoke, then 32-node at batch 1024.

Ops: third manual wedge-clear of the session after the PGLE kill (#6302 sequence; deployed controller predates PR #6305 — redeploy awaits approval). Cluster left clean (0 stray pods/workloads).

[rjpower]

🤖 Experiment 7 — save_moe remat at batch 1024: neutral. remat_mode=save_moe (commit e5e9cd6f2a: save_only_these_names over the MoE backends' tagged dispatch tensors, SCALE_REMAT knob) gave 712K tok/s / 9.37% MFU vs exp-3's 709K / 9.32% — within noise. The EP-dispatch recompute the policy avoids evidently isn't on the critical path at this shape (XLA likely already CSE'd most of it). Loss sane (7.33 @ step 29). Keeping the knob; it may matter once comm shrinks.

Experiment 8 launched — pure-bf16 params (SCALE_MP=params=bfloat16,..., commit e02986ae99): params at rest are f32 today, and the FSDP regather of ~45GB+/GPU/step is the dominant cost; bf16 storage halves gather bytes. Numerics are acceptable for a 30-step throughput probe; if it delivers, the production-quality version is bf16 gathers with f32 master weights in the optimizer. Predicted ~+30–50% if fully comm-bound.

experiment	config	tok/s	MFU
exp 3	batch 1024	709K	9.32%
exp 7	+ save_moe remat	712K	9.37% (~neutral)
exp 8	+ bf16 params	running	—

[rjpower]

🤖 Experiment 8 — pure-bf16 params (SCALE_MP=params=bfloat16,...): NEGATIVE — deterministic gang hang. Two attempts at batch 1024: first ran 9 steps at ~440K tok/s (already below the f32-param 709K, oddly) then all 32 ranks went silent mid-collective for 12+ min; the retry hung the same way at step 1. No NCCL warnings, no OOM, no crash — a clean deadlock, killed both. bf16-at-rest params are off the table as a cheap lever; if gather-bytes still matter after the reshard fix, the right form is bf16 gathers with f32 master weights (a levanter-level change), not bf16 storage.

Both hangs + kills required the manual #6302 wedge-clear again (5th time today). Cluster left clean.

Scoreboard (best: 712K tok/s / 9.37% — exp 7, batch 1024 + save_moe remat):

lever	result
XLA flags / PGLE	❌ negative / OOM
tokens/device (batch 256→1024)	✅ 2.8% → 9.32%
batch >1024	❌ OOM ceiling
save_moe remat	~neutral (+0.5%)
bf16 params at rest	❌ deterministic hang

Pivoting to the structural fix: the GSPMD involuntary full-remat at the expert-dispatch boundary (experiments/grug/moe/model.py _batch_reshard / lib/levanter/src/levanter/grug/grug_moe.py shard_map entry) — it is simultaneously the memory ceiling (replicated [1, seq, hidden] tile) and a comm tax. Implementation starting on weaver/moe-mfu.

[rjpower]

🤖 Root cause of the involuntary full-remat found and fixed (commit 451ba5108b on weaver/moe-mfu) — it was not the MoE dispatch: it's the token-embedding backward. The scatter-add of the [B,S,H] cotangent into the embed gradient ([V,H] sharded hidden-over-("replica_dcn","data")) needs a reshard from batch-sharded (256 tiles on dim0) to hidden-sharded (32 tiles on dim2, 8-way expert-replicated). GSPMD's all-to-all reshard requires matching tile counts (256 ≠ 32), so it replicated the full cotangent: 3 GiB/device at batch 1024, 18 GiB at 1536 — this was the OOM ceiling, and it shows in the profile as 41 s of AllReduce_Sum_u8 in only 72 calls.

Fresh profile at batch 1024 + save_moe (run cw-grugmoe-mfu-profile2-20260610-134610, artifact jax-profile-step-12-22:v0): communication 82.3% / compute 13.9% / stall 3.9%. Compute share rose 7.7→13.9% with batch scaling, but all-gather still dominates (126 s RING_LL across the window).

The fix: route the embed lookup through reshard hops GSPMD lowers natively (replicate int32 indices → gather along the table's hidden shard → all-to-all to batch shard → expert-axis slice); the transpose mirrors it (intra-node expert all-gather → all-to-all → fully local scatter-add). Validated with a CPU repro (--xla_force_host_platform_device_count=8, warning 2→0), bit-identical loss, 44/45 grad leaves bit-identical (embed grad differs only by summation order, 2.3e-10), all grug/moe tests green.

Cluster validation: 1-node smoke green — 5 steps, 0 SPMD warnings (previously always ≥1), loss bit-identical to the pre-fix smoke (11.1695 @ step 4). Exp-9 (32 nodes, batch 1024) launched; then a batch-1536 probe, since the replicated buffer that caused the OOM is gone.

[rjpower]

🤖 Experiment 9 — reshard fix (451ba5108b) at batch 1024: OOM. Unexpected: the fix removes the 3 GiB/device replicated cotangent, yet the same batch that fit pre-fix (exp-7, 712K/9.37%) now exhausts HBM. Hypothesis: the replacement path changes the allocation profile — the replicated int32 index tensor plus the all-to-all materializes its own intermediates, and removing the full-remat reshuffles XLA's buffer-assignment/scheduling so peak shifts elsewhere. The 41 s AllReduce_Sum_u8 replicate path should be gone, but we couldn't observe throughput before the OOM.

Killed the job at the first OOM rather than letting it burn 3 auto-retries (~20 min × 256 GPUs each), manually cleared the kill-wedge (6th time today — #6302 sequence, deployed controller predates PR #6305).

Exp 9b launched immediately: reshard fix at batch 512 (run cw-grugmoe-mfu-reshard512-…-151503) — the clean A/B against pre-fix exp-2 (batch 512 = 408K tok/s / 5.37%). Decision rule:

• 512-with-fix ≫ 408K → the comm win is real; chase the 1024 memory regression (mem-profile the new path, try batch 768).
• ~flat or worse → the fix trades the replicate for similar-cost all-to-alls and loses on memory; revert it for throughput purposes, keep exp-7 (712K / 9.37%) as best-known, and move to seq-4096 / hybrid-sharding / FA4 levers.

[rjpower]

🤖 exp-9b verdict: reshard fix is a throughput regression — reverted

A/B at batch 512, seq 2048, recompute_all, 30 steps (job iris-run-job-20260610-151503, completed cleanly):

	pre-fix (exp-5)	with fix (exp-9b)
tokens/s	408K	364K (−11%)
MFU	5.37%	4.79%
involuntary-full-remat warnings	2	0

The fix does eliminate the GSPMD full-remat (warnings 2→0, matching the CPU repro), but the explicit replicate-indices → hidden-sharded gather → all-to-all path is slower than GSPMD's replicate-and-scatter-add at this scale. Combined with exp-9 (fix @ batch 1024) OOMing where the unfixed path fit, the fix loses on both throughput and memory in practice. Reverted 451ba5108b (revert 2b1c412b70); the root-cause analysis and CPU repro remain valid documentation of why the AllReduce_Sum_u8 path exists.

Best known config stands: exp-7 — batch 1024 × seq 2048, save_moe remat, 712K tok/s / 9.37% MFU.

Next: exp-10 — seq 4096 × batch 512, save_moe, no fix. Same 2M tokens/step as exp-7 → identical per-step FSDP gather traffic, but ~2× attention FLOPs per token. In an 82%-comm-bound regime, raising arithmetic intensity per token is the direct MFU lever; if step time holds near exp-7's, MFU rises materially. Memory risk is moderate (activation footprint tracks token count, which is unchanged; the GSPMD replication buffer tracks tokens too, ~3 GiB/dev as at exp-6).