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.