[RL] Changes to enable compilation for trainer

[Lucaskabela]

Summary

In this PR, we enable naive, JIT style torch.compile for the RL policy trainer. This is the first step towards speeding up the trainer model. Changes are:

1. Wiring through compilation config:

• Added TrainerCompileConfig dataclass with enable (bool) and backend (str, default "eager") fields
• Added compile field to PolicyTrainer.Config with compile and aot_eager backend
• Added _compile_model() method that calls .compile(backend=..., fullgraph=True) on each transformer
  layer -> This is crticial, as torch.compile() results in logit changes

2. config_registry.py — Enable compile by default in configs

• Both rl_grpo_qwen3_0_6b and rl_grpo_qwen3_debug configs now set
  compile=TrainerCompileConfig(enable=True). Default backend is 'eager'

4. vllm_compat/models/attention.py — Make flash-attention compile-compatible

• Moved the FlashAttnWithBackward autograd function out of the forward() method (nested classes
  can't be traced by the compiler) into a module-level FlashAttnVarlenFunction
• Registered the flash-attention forward as a torch.library.custom_op (rl::flash_attn_varlen_fwd)
  with a fake implementation, so AOT Autograd can trace through it with FakeTensors
• Simplified the call site in VLLMCompatibleFlashAttention.forward() to use the new function

Test Plan

python torchtitan/experiments/rl/unified/simple_grpo_sum_digits.py --module rl.unified --config rl_grpo_qwen3_0_6b --hf_assets_path=torchtitan/experiments/rl/example_checkpoint/Qwen3-0.6B

Results in the same losses as on main - the timing is now like:

Main

[2026-03-13 12:13:11] INFO simple_grpo_sum_digits.py:401: [actor=<root>] Cumulative Timing | Generator: 22.6s | Optimizer: 0.1s | Trainer: 148.6s | WeightSync: 119.4s | Total: 290.7s

Changes

[2026-03-13 12:03:26] INFO simple_grpo_sum_digits.py:401: [actor=<root>] Cumulative Timing | Generator: 22.1s | Optimizer: 0.1s | Trainer: 103.2s | WeightSync: 119.2s | Total: 244.6s

So we save ~50s of runtime

[tianyu-l]

if we go with pytorch varlen, do we still need to worry about this file? cc @wwwjn