[CuTe, Fwd] Fix forward compile key churn when max_seqlen is a tensor

[eamonn-zh]

Problem

When max_seqlen_q / max_seqlen_k are CUDA tensors, as in Hugging Face varlen and padding-free attention, _flash_attn_fwd can place a fresh 0-D tensor in its compile key. Identical forward calls then miss the JIT cache and recompile on every invocation, significantly slowing down training.

This is the forward counterpart of #2507.

Root Cause

max_m_blocks_leq_one inherits the tensor type from max_seqlen_q:

max_m_blocks_leq_one = seqlen_q_packgqa <= q_stage * tile_m

When this predicate is false, the Python boolean expression used to compute is_static_persistent can return the tensor itself instead of a Python bool. Since is_static_persistent is included in the forward compile key, each newly created tensor produces a different cache key.

Fix

Only use the max_m_blocks_leq_one static-persistent shortcut when the predicate is not a tensor.

This keeps the compile key host-scalar-only and avoids calling .item(), so the fix does not introduce a CPU-GPU synchronization. Tensor-valued callers conservatively use the non-static-persistent path.

Verification

Added a regression test that:

• computes a fresh CUDA tensor max_seqlen for each call;
• invokes varlen forward twice with identical inputs;
• verifies that both outputs are identical;
• verifies that the forward compile cache contains only one entry;
• verifies that no tensor is present in the compile key.

Tested on NVIDIA RTX PRO 6000 Blackwell (SM120):

CUDA_VISIBLE_DEVICES=0 \
FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 \
pytest -q -s \
tests/cute/test_flash_attn_fast.py::test_flash_attn_varlen_tensor_max_seqlen_reuses_fwd_cache

Result:
1 passed

The full FakeTensor fast suite also passes:
192 passed, 49 skipped

[drisspg]

btw this breaks because we end up passing a pytorch tensor into the kernel for the max_seqlen_q slot even though we are running the persistent kernel