feat(model): add boundary-aware FA3 cross-attention
#4
by z2z23n0 - opened
This adds an opt-in path for the specialized MOSS-VL FA3 cross-attention kernel.
With cross_attention_implementation="flash_attention_3", the model converts each frame-level mask row to an int32 visible-KV boundary and passes it to FA3. This path no longer materializes a dense text-by-vision mask. Existing eager, SDPA, and FA2 behavior is unchanged. Setting global attn_implementation="flash_attention_3" alone also keeps the previous cross-attention dispatch.
The separate override lets Ampere runs keep FA2 or SDPA for self-attention while using the MOSS FA3 kernel for cross-attention.
Tested:
- 11 CPU unit tests covering mask parity, invalid masks, config serialization, dispatch, and default compatibility
- A800 BF16 forward/backward parity at the model head shape: 32 query heads, 8 KV heads, head dimension 128
- MOSS-VL 11B two-image prefill and four cached decode steps: 60 FA3 calls, no dense mask expansion, stable 48-token vision KV cache, and matching greedy top-1 against SDPA at all five positions
- Standard 16-token
generate(): 192 FA3 calls and no dense mask expansion - Synthetic 4096-text-token, 2048-frame mask: 352 MB dense mask versus a 16 KB boundary
The full checkpoint test used SDPA for self-attention and FA3 only for cross-attention. The A800 kernel build used OpenMOSS/MOSS-VL#16.