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.
fmha_interface_wrapperinsrc/tilegym/ops/attn_interface.pyacceptsattention_maskfrom the HF attention interface but never passes it on;fmhaitself 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.