[FA4] max_seqlen_q tensor causes an argument-annotation mismatch in varlen forward

Author: TorinLiCreated Aug 17, 2026Updated Aug 23, 2026

Summary

The FA4 varlen forward path can pass a scalar torch.Tensor for max_seqlen_q into cute.compile(), while the compiled function declares that parameter as Int32 | int | None. This makes compilation fail with an argument-annotation mismatch.

Reproduction

From a checkout of the current default branch, run:

bash
pytest tests/cute/test_flash_attn_fast.py -k test_flash_attn_varlen_tensor_max_seqlen_reuses_fwd_cache

Actual behavior

The test fails with an error equivalent to:

expects argument #26 (max_seqlen_q) to be one of (Int32, int, NoneType), but got torch.Tensor

Expected behavior

The cache-stabilization path should retain its intended behavior while passing a value consistent with the compiled function's type contract.

Possible direction

Normalize the scalar before invoking cute.compile() (for example, construct the appropriate CUTLASS DSL scalar type), or update the relevant API contract if a tensor is the intended representation.

Source: Dao-AILab/flash-attention