[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:
pytest tests/cute/test_flash_attn_fast.py -k test_flash_attn_varlen_tensor_max_seqlen_reuses_fwd_cacheActual 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.TensorExpected 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