Fix flash/sage varlen prep under torch.compile with dynamic shapes

[ShivamShrirao]

What does this PR do?

Fixes #11957.

That report is the .item() graph break on the no-mask path (Flux passes no attention mask), which
change 2 below removes entirely — the no-mask helper no longer calls .item() at all. #11970 was an
earlier attempt at the same issue and was closed unmerged. Masked varlen still keeps one necessary
.item() for the data-dependent key lengths, as noted after the two changes.

Makes the flash-attn / sage varlen attention backends usable under torch.compile with dynamic shapes. Two changes in the _prepare_for_flash_attn_or_sage_varlen_* helpers, with no numerical change:

1. cu_seqlens is now built with torch.arange instead of cumsum(full(...)). Inductor rewrites the cumsum-of-a-constant pattern internally, and that rewrite fails when the fill value is a symbolic sequence length. The lengths are uniform here (every sequence in the batch has the same length), so torch.arange(0, (batch_size + 1) * seq_len, seq_len) produces identical offsets directly — and is cheaper.

   Minimal repro (fails on torch 2.6.0 and 2.12.1):

   import torch
   
   @torch.compile(dynamic=True, fullgraph=True)
   def f(x):
       b, s = x.shape[0], x.shape[1]
       seqlens = torch.full((b,), s, dtype=torch.int32, device=x.device)
       cu = torch.zeros(b + 1, dtype=torch.int32, device=x.device)
       cu[1:] = torch.cumsum(seqlens, dim=0)
       return cu
   
   f(torch.empty(2, 77, 8))
   # BackendCompilerFailed: inductor raised:
   # TypeError: unsupported operand type(s) for *: 'FakeTensor' and 'Node'

2. max_seqlen_q/max_seqlen_k are returned as the already-known Python int instead of seqlens.max().item(). The .item() forces a GPU→CPU sync and a graph break; the max of uniform lengths is just seq_len itself, which the caller already has as an int. Same type as before (.item() also returned an int), so all call sites are unaffected.

In the masked helper only the query side changes: key lengths are data-dependent (attn_mask.sum(dim=1)), so their cumsum does not hit the inductor rewrite and their .item() is inherent to varlen with padding.

Found while enabling flash-attn varlen + torch.compile for the Bria Fibo pipelines, but the helpers are shared by every model using these backends.

Verification

• The new construction is bit-identical to the old one across batch/seq-length combinations, for both the masked and unmasked helpers.
• The repro above fails identically on torch 2.6.0 (cu124) and torch 2.12.1 (cu130).
• With the fix, the real flash_attn_varlen_func (flash-attn 2.8.3, H200) compiles with torch.compile(dynamic=True, fullgraph=True) across changing batch/sequence shapes and matches eager output exactly (max diff 0.0).

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc?

Self-review notes

Two-function diff reviewed against .ai/review-rules.md. Verdict: READY. No API change: same return tuple structure, same values, and the max-seqlen entries keep the type .item() produced (Python int). Checked all 10 call sites in attention_dispatch.py, including the hub-kernel forward/backward ops added recently — the masked-path callers already discard the helper's max_seqlen_q and use seq_len_q locally, consistent with this change; unmasked-path callers consume the returned ints unchanged (ctx scalar attributes, never save_for_backward). Equivalence verified numerically against the previous cumsum construction, masked and unmasked. Left for the reviewer: no CPU unit test covers these helpers directly; the behavior is exercised through the GPU flash-attn backend tests.

Who can review?

@sayakpaul @DN6

[github-actions]

Hi @ShivamShrirao, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[sayakpaul]

Can you provide a minimal reproducer with the diffusers functions included?

[ShivamShrirao]

Reproducer with diffusers varlen backends call:

import torch
from diffusers.models.attention_dispatch import _prepare_for_flash_attn_or_sage_varlen

@torch.compile(dynamic=True, fullgraph=True)
def prep(hidden_states):
    batch_size, seq_len, _ = hidden_states.shape
    return _prepare_for_flash_attn_or_sage_varlen(
        batch_size, seq_len, seq_len, attn_mask=None, device=hidden_states.device
    )

prep(torch.randn(2, 77, 64, device="cuda"))

On main:

torch._inductor.exc.InductorError: TypeError: unsupported operand type(s) for *: 'FakeTensor' and 'Node'

With the fix:

((tensor([77, 77], device='cuda:0', dtype=torch.int32),
  tensor([77, 77], device='cuda:0', dtype=torch.int32)),
 (tensor([  0,  77, 154], device='cuda:0', dtype=torch.int32),
  tensor([  0,  77, 154], device='cuda:0', dtype=torch.int32)),
 (77, 77))

Tested with torch 2.12.1 and 2.13.0+cu129.

[sayakpaul]

Sorry, but this is bit vague to determine which attention backend is causing this, though.

[sayakpaul]

Why not return max_seqlen_q like other?

Cc: @zhtmike do you have any comments here?

[zhtmike]

I think using seq_len_q is fine , since query is always a full sequence.

But according to the PR description, the mask path is still broken under torch compile?

Plus, I think we need a test to guard this

[ShivamShrirao]

Yes. queries are always full length, so max_seqlen_q is seq_len_q, the caller already has it as an int.

The masked path compiles too and is not broken. The remaining .item() is on the key side only, and it doesn't block fullgraph=True. On torch 2.13.0+cu129 with this PR:

no-mask: PASS cu_q=[0, 77, 154] cu_k=[0, 77, 154] max=(77, 77)
masked:  PASS cu_q=[0, 77, 154] cu_k=[0, 77, 137] max=(77, 77)

The description only meant that the key-side .item() stays, since padded lengths are data-dependent.

[sergereview]

🤗 Serge says:

Solid, well-scoped fix — the arange construction is numerically identical to cumsum(full(...)) for uniform lengths (0, s, 2s, …, batch_size*s, batch_size+1 elements, same int32 dtype), and returning seq_len_q / seq_len_kv matches what seqlens.max().item() produced while removing the sync/graph break.

Correctness

• All five call sites were checked: the masked callers (_flash_attn_varlen_*, _flash_varlen_attention_hub) already discard the helper's max_seqlen_q or recompute it as seq_len_q, so the returned values are unchanged in every path. lru_cache_unless_export on the no-mask helper is unaffected since the new returns derive purely from the cache-key arguments.
• Only nit: torch.arange(..., step=seq_len_*) raises on a zero step, whereas the old cumsum path degenerated gracefully to all-zeros for a zero-length sequence. Almost certainly unreachable in practice, but it is a behavioral difference worth being aware of.

Tests

• The verification described (bit-identical outputs, compile with dynamic=True, fullgraph=True) is not reflected in the repo — no test exercises these helpers directly. A small unit test comparing the new cu_seqlens against the cumsum reference for a few (batch_size, seq_len) combos would lock this in cheaply, since the helpers are shared by every flash/sage varlen backend.

Minor

• In the no-mask helper, seqlens_q/seqlens_k are still materialized via torch.full but every caller discards them (the unmasked branches immediately set seqlens_k = None). Out of scope for this PR, but they are dead allocations on the hot path if you ever want to trim the return tuple.

serge v0.1.0 · model: claude-opus-5 · 6 LLM turns · 9 tool calls · 68.4s · 93703 in / 5078 out tokens

[sergereview]

Behavioral edge case worth noting: torch.arange with step=0 raises RuntimeError: step must be nonzero, so a zero-length query/key sequence now errors where the previous cumsum(zeros) path returned an all-zero cu_seqlens. I don't think any dispatcher path can reach here with seq_len == 0, so this is informational rather than blocking — just flagging it since the PR claims strict numerical/behavioral equivalence.

[ShivamShrirao]

The helper is shared by every varlen backend, so it's all of them rather than one:

• _flash_varlen_attention
• _flash_varlen_attention_3
• _sage_varlen_attention
• _flash_varlen_attention_hub and the hub forward ops _flash_varlen_attention_hub_forward_op / _flash_attention_3_varlen_hub_forward_op

All of them go through _prepare_for_flash_attn_or_sage_varlen to build cu_seqlens, so any of them fails the moment you wrap it in torch.compile(dynamic=True, fullgraph=True).