Repository navigation
attention_mask is discarded in fmha_interface_wrapper: is a varlen or padding-mask prefill kernel wanted? #191
Description
fmha_interface_wrapper in src/tilegym/ops/attn_interface.py accepts attention_mask from the HF attention interface but never passes it on; fmha itself only supports dense causal input. So padded batches at batch_size > 1 attend over padding tokens on the patched path.
Is a varlen-style (cu_seqlens packed, FlashAttention-varlen-like) or padding-mask-aware prefill kernel something you would want as an external contribution to ops/cutile/experimental/, in the style of the SWA prefill from #107? Or is this territory already covered by the roadmap's "Autoregressive Flash Attention (WIP Internal)" or the planned Flex Attention work?
Asking for scope before writing anything; happy to write up a concrete design in this thread if it is wanted.
Activity
hannahli-nv commentedon Sep 16, 2026
Hi @Aryan95614, apologies for the slow reply, and thank you for catching this!
You're right,
fmha_interface_wrappertakesattention_maskand never uses it, so padded batches with batch_size > 1 do end up attending over padding tokens. That's a real gap in our code.On the kernel side, there's an update: fmha_varlen (cuTile) is now available in TileGym. It takes padded (B, H, S, D) tensors plus per-batch q_lens / kv_lens, and is forward-only for now. The HF attention interface doesn't route to it yet. We'll follow up to fix it shortly.
Thanks again for the clear and well-scoped report.