[jialei777]

What does this PR do?

Ports the four TPU TorchTitan training performance optimizations (TPUSplashAttention, SimpleFSDP, per-block torch.compile(backend="tpu"), and DEFER_AND_FUSE eager execution mode, plus vectorized TPU log-prob computation) into verl-hardware-plugin (TorchTitanTPUEngineWithLMHead and verl_hardware_plugin.engines.tpu_utils), bringing compiled TPU trainer throughput (~6.4s/step vs ~61.4s/step eager on Qwen3-0.6B GSM8K GRPO) to upstream verl + verl-hardware-plugin without requiring changes to upstream verl.

• Why this is not duplicating an existing PR: Searched open PRs in verl-project/verl-hardware-plugin and jialei777/verl-hardware-plugin; no open PR ports TPUSplashAttention, SimpleFSDP, configure_torch_compile_for_tpu, or DEFER_AND_FUSE into verl-hardware-plugin.
• AI assistance disclosure: This PR was developed with AI assistance and verified with unit tests and live 2-slice TPU v6e-8 GRPO training runs.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: https://github.com/verl-project/verl-hardware-plugin/pulls?q=is%3Apr+is%3Aopen+splash+OR+simple_fsdp
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)

Test

1. Pre-commit & static checks:

pre-commit run --files verl_hardware_plugin/engines/tpu_utils.py verl_hardware_plugin/engines/torchtitan_tpu.py verl_hardware_plugin/platforms/platform_tpu.py tests/test_tpu_engine.py

All 7 hooks (ruff, ruff-format, mypy, check-license, check-no-bytedance-copyright, check-naming-conventions, compileall) pass.

2. Unit tests (tests/test_tpu_engine.py):

Ran all 16 unit tests in tests/test_tpu_engine.py on ray-tpu-v6e-cluster-head-9v898, covering splash_block_size_for, apply_splash_attention_tpu (including CPU/unaligned fallback), pad_packed_inputs_for_tpu(..., build_attention_mask=False) + cross-rank aligned bucket lengths, configure_torch_compile_for_tpu, extend_torchtitan_engine_config, and resolve_tpu_torchtitan_options.

3. 2-Slice TPU v6e-8 E2E GRPO Benchmark (Qwen3-0.6B on GSM8K):

Validated on jialeic-tpu-v6e8-2s-spot (8 trainer v6e chips + 8 rollout v6e chips) with use_torch_compile=True, use_splash_attention=True, use_simple_fsdp=True, and tpu_eager_mode=DEFER_AND_FUSE.

API and Usage Example

When actor_rollout_ref.actor.torchtitan.use_torch_compile=True is set on TPU, TorchTitanTPUEngineWithLMHead automatically enables use_splash_attention=True, use_simple_fsdp=True (for pure FSDP/HSDP where TP/PP/CP/EP == 1), and tpu_eager_mode=DEFER_AND_FUSE. Each option can also be explicitly configured via actor_rollout_ref.actor.torchtitan.* or VERL_TPU_* environment variables:

bash examples/tpu/grpo/run_qwen3_0_6b_torchtitan.sh \
  actor_rollout_ref.actor.use_torch_compile=True \
  actor_rollout_ref.actor.torchtitan.use_torch_compile=True \
  +actor_rollout_ref.actor.torchtitan.use_splash_attention=True \
  +actor_rollout_ref.actor.torchtitan.use_simple_fsdp=True \
  +actor_rollout_ref.actor.torchtitan.tpu_eager_mode=DEFER_AND_FUSE

Design & Code Changes

• verl_hardware_plugin/engines/tpu_utils.py:
  • extend_torchtitan_engine_config() & resolve_tpu_torchtitan_options(): registers optional use_splash_attention, use_simple_fsdp, and tpu_eager_mode fields on TorchtitanEngineConfig and resolves TPU defaults when use_torch_compile=True.
  • configure_torch_compile_for_tpu(): configures torch._dynamo.config (and _config[name].default so settings persist across Ray actor RPCs) with automatic_dynamic_shapes=False, assume_static_by_default=True, capture_scalar_outputs=True, and recompile_limit=64.
  • splash_block_size_for(), TPUSplashAttention, apply_splash_attention_tpu(): integrates the Pallas splash_sdpa kernel with dynamic power-of-two block sizes (512, 256, 128) and packed-document segment_ids masking, skipping dense [1, 1, S, S] causal mask construction in pad_packed_inputs_for_tpu(..., build_attention_mask=False).
  • make_simple_fsdp_parallelize_fn() & align_micro_batch_shapes_across_ranks(): wraps model_spec.parallelize_fn via tpu_torchtitan_config_overrides(use_simple_fsdp=...) to shard parameters with TorchTitan SimpleFSDP (Shard(0) DTensors traced into each compiled TransformerBlock), marks blocks spmd_safe, preserves tied embedding weights, and aligns micro-batch sequence/response bucket lengths across ranks with a single CPU all_reduce(MAX).
  • tpu_eager_mode_context(): wraps forward/backward and optimizer steps in torch_tpu._internal.execution_mode.set_eager_mode(EagerMode.DEFER_AND_FUSE).
• verl_hardware_plugin/engines/torchtitan_tpu.py:
  • Wires configure_torch_compile_for_tpu, TPUSplashAttention, SimpleFSDP, align_micro_batch_shapes_across_ranks, and tpu_eager_mode_context into TorchTitanTPUEngineWithLMHead.
  • Adds TPUEngineEvalModeCtx to guard submodule.reshard() with isinstance(submodule, FSDPModule) when SimpleFSDP is active.
  • Scales loss by 1.0 / dp_size before synchronize_tpu_loss(loss) and loss.backward() under SimpleFSDP.
  • Uses vectorized logprobs_from_logits_naive on TPU in prepare_model_outputs instead of row-by-row logprobs_from_logits_v2.
• verl_hardware_plugin/platforms/platform_tpu.py:
  • Calls extend_torchtitan_engine_config() at platform load time so driver processes have the extended TorchtitanEngineConfig dataclass fields.
• tests/test_tpu_engine.py:
  • Adds unit tests for splash_block_size_for, apply_splash_attention_tpu, pad_packed_inputs_for_tpu(..., build_attention_mask=False), configure_torch_compile_for_tpu, and resolve_tpu_torchtitan_options.

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code.

---

End-to-End 10-Step GRPO Benchmark on TPU v6e-8 (2 Slices: 8 Trainer + 8 Rollout Chips)

Tested with Qwen/Qwen3-0.6B on GSM8K (examples/tpu/grpo/run_qwen3_0_6b_torchtitan.sh, train_batch_size=32, rollout.n=8, max_prompt_length=512, max_response_length=1024) on gke_tpu-pytorch_us-central2_jialeic-tpu-v6e8-2s-spot:

Steady-State Comparison (Steps 5–10 Mean)

Metric	verl tpu-main (d435b6ad)	verl PR #34 + verl-hardware-plugin main (Eager)	verl PR #34 + This PR (#2) (Compiled + Splash + SimpleFSDP)	Speedup vs. Eager Plugin
timing_s/old_log_prob	1.32s	13.76s	1.86s	7.39x faster
timing_s/ref	1.08s	12.76s	1.05s	12.17x faster
timing_s/update_actor	3.97s	38.52s	3.87s	9.94x faster
Trainer Total (old_log_prob + ref + update_actor)	6.38s	65.04s	6.78s	9.59x faster
timing_s/step (Full Step Time)	72.81s	75.51s	64.06s	1.18x faster
perf/throughput (tokens/s)	178.9 tok/s	169.0 tok/s	196.2 tok/s	+16.1%
actor/grad_norm	0.469	0.501	0.491	Matches
rollout_actor_probs_pearson_corr	0.997	0.997	0.997	Identical
critic/rewards/mean	0.272	0.303	0.309	Matches
Step 10 GSM8K Val Acc (val-core/openai/gsm8k/acc/mean@1)	—	—	47.66% (0.4766)	Verified

Per-Step Breakdown (PR #34 + This PR #2)

Step	gen (s)	old_log_prob (s)	ref (s)	update_actor (s)	Trainer Total (s)	update_weights (s)	step (s)	grad_norm	pearson_corr	reward_mean
1 (compile)	56.10	3.82	2.10	10.52	16.44	41.83	114.77	0.5294	0.99701	0.3164
2	12.04	2.68	1.58	8.60	12.85	40.96	66.05	0.5461	0.99693	0.2266
3	8.03	2.07	1.04	3.98	7.09	55.14	70.45	0.4664	0.99668	0.2500
4	8.04	1.85	1.04	3.78	6.68	48.32	63.24	0.4734	0.99679	0.2422
5	10.04	1.87	1.05	3.78	6.70	46.90	63.85	0.4952	0.99686	0.2578
6	10.04	1.85	1.05	4.41	7.31	48.40	65.96	0.5649	0.99659	0.2500
7	10.04	1.89	1.06	3.79	6.73	44.99	61.96	0.4623	0.99695	0.3047
8	8.04	1.85	1.03	3.73	6.61	43.83	58.68	0.4893	0.99670	0.3633
9	10.04	1.86	1.05	3.76	6.67	49.12	66.03	0.4611	0.99682	0.3125
10	10.04	1.87	1.05	3.77	6.68	50.99	67.90	0.4723	0.99690	0.3672