[rl] Flex-attention generator crashes with "Dynamo recompile limit exceeded (recompile_limit=8)"

[wwwjn]

Summary

The RL GRPO config rl_grpo_qwen3_0_6b_flex (Qwen3-0.6B, attn_backend="flex", compile.enable=True) crashes in the vLLM generator during rollout. vLLM's compiled flex-attention kernel recompiles on every new query length and exceeds the default torch._dynamo.config.recompile_limit=8. The region is fullgraph=True, so the limit is a hard failure instead of an eager fallback. The crash deadlocks the Monarch async loop, so it surfaces as a hang with GPUs pinned rather than a clean exit.

Reproduction

4 GPUs (trainer TP=2 + generator TP=2). Config: torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py:140.

python -m torchtitan.experiments.rl.train \
  --module alphabet_sort --config rl_grpo_qwen3_0_6b_flex

Error

torch._dynamo.exc.Unsupported: Dynamo recompile limit exceeded
  Explanation: ...exceeding the recompile_limit cache size limit (currently set to 8)...
  Developer debug context: Limit type: recompile_limit

The above exception was the direct cause of the following exception:

torch._dynamo.exc.FailOnRecompileLimitHit: Hard failure due to fullgraph=True

Crash site:

vllm/v1/attention/backends/flex_attention.py:1199, in forward
    out = flex_attention_compiled(...)

Root cause

Every recompile is on the generator (both TP ranks: [actor=...VLLMGenerator generator{'gpus': 0/2, 1/2}]); the trainer never reaches generation. Two fullgraph=True regions recompile past the limit:

1. flex_attention (vllm/v1/attention/backends/flex_attention.py:48: torch.compile(flex_attention, fullgraph=True), no dynamic= arg; torch._dynamo.try_mark_dynamic(query, 2) is commented out at line 1184). Guard failures:
   • query size mismatch at index 2 (query/seq length): 256 -> 128 -> 64 -> 16 -> 118 -> 1 ... (prefill lengths plus the 9 cudagraph decode-capture sizes [1, 2, 4, 8, 16, 32, 64, 128, 256]).
   • key dispatch-key-set mismatch: DispatchKeySet(CUDA, BackendSelect) vs (..., ADInplaceOrView). Non-shape; dynamic=True cannot fix this.
2. forward (torchtitan/models/qwen3/model.py:54, compiled via vLLM @support_torch_compile): x._local_tensor size mismatch at index 1 (2048 -> 256 -> 1) plus a type flip between Tensor and AsyncCollectiveTensor (torchtitan/protocols/module.py:551, _redistribute_inputs).

The crash first fires during generator cudagraph-capture warmup (capture reaches 8/9, dies on the 9th size) and then re-fires on every subsequent novel query length during rollout.

Notes / possible fixes

• We do not set torch.compile(dynamic=False) anywhere; all flex compiles use the default dynamic=None. Automatic-dynamic partially engages (symbolic guards like 128 <= query.size()[2] appear) but the query-length spread plus the ADInplaceOrView dispatch-key flip exceed 8 recompiles before it can generalize.
• torchtitan/models/gpt_oss/parallelize.py already raises torch._dynamo.config.recompile_limit to 10-12 for its sliding-window flex path; the RL flex path does not.
• Candidate fixes: mark the query seq dim dynamic (uncomment vLLM's try_mark_dynamic(query, 2) / pass dynamic=True); raise recompile_limit (does not fix the dispatch-key flips); or use the varlen backend for RL, which never routes through flex_attention_compiled.

Environment

• torch 2.13.0.dev20260608+cu130
• vllm 0.22.1rc1.dev256+g2ed0a9627.d20260610
• torch._dynamo.config.recompile_limit = 8 (default)

[wwwjn]

cc @aditvenk @liangel-02 found it might be related to #3737 ,

We should also add RL CI for flex attention too

[tianyu-l]

#3737 has _raise_dynamo_recompile_limit for gpt-oss, particularly because gpt-oss has both sliding window attention and full attention. I don't remember qwen3 having the same issue. Why would it suddenly break recompile limit? The core qwen3 CI tests (not sure if compile is enabled) didn't seem to need it?