Afmoe flash attention

[samsja]

Implement Flash Attention (FA2/FA3) for the AFMoE model with sliding window support.

This enables the use of Flash Attention for AFMoE, which was previously unsupported due to perceived conflicts with sliding window and RoPE constraints, allowing for potential performance improvements.

---

 

---

Note

Introduces Flash Attention for AFMoE and wires it into the model, including sliding window support via the varlen API.

• New AfmoeFlashAttention implementing FA2/FA3 using flash_attn_varlen_func/flash_attn_3_varlen_func; added to AFMOE_ATTN_IMPL2CLASS
• Computes seqlens, cu_seqlens, and max_seqlen from position_ids when using flash attention; disables explicit causal masks for FA paths
• Plumbs cu_seqlens/max_seqlen through AfmoeModel → AfmoeDecoderLayer → attention forward
• Supports local/sliding attention by setting window_size in FA kwargs; preserves RoPE application
• Marks _supports_flash_attn = True and adjusts unsupported-impl error message

^{Written by Cursor Bugbot for commit d37a047. This will update automatically on new commits. Configure here.}

[cursor]

Cursor Agent can help with this pull request. Just @cursor in comments and I'll start working on changes in this branch.
_{Learn more about Cursor Agents}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

[cursor]

Flash Attention missing dropout_p parameter

Medium Severity

The AfmoeFlashAttention class does not pass dropout_p to the flash attention function, while AfmoeSDPAAttention properly handles attention_dropout with dropout_p = self.attention_dropout if self.training else 0.0. If the model config has non-zero attention_dropout, SDPA will apply dropout during training but Flash Attention will not, causing inconsistent training behavior between the two attention implementations.

 

[cursor]

cu_seqlens computation ignores actual batch size

High Severity

When position_ids is None and batch_size > 1, the cu_seqlens computation produces incorrect results. The default position_ids has shape [1, seq_len], so flat_position_ids.view(-1) yields only seq_len elements. The resulting cu_seqlens (e.g., [0, seq_len]) indicates a single sequence, but hidden_states contains batch_size * seq_len tokens. Flash attention then processes only the first seq_len tokens, silently ignoring the rest of the batch. The fallback in AfmoeFlashAttention.forward() would correctly handle this, but it's bypassed because cu_seqlens is already provided.

Additional Locations (1)

• src/prime_rl/trainer/models/afmoe/modeling_afmoe.py#L439-L441