What attention scheme for DFlash SWA? Non-Causal?

#7
by benchislett - opened

Hi, I'm the vllm maintainer implementing this in vLLM: https://github.com/vllm-project/vllm/pull/46104

Could you kindly disclose what attention mechanism should be used for the DFlash Sliding Window layers? Specifically: does it use non-causal attention (as is common in DFlash), or does it use causal-only attention? This changes how the attention pattern looks in the query block, and changes the kernel requirements (causal SWA has much better software support).

If non-causal SWA is used, is the window bi-directional? e.g. if a prefix token attends to some query token, does it always attend to the entire query? Or, could it attend to only a subset of the query tokens (those within the sliding window)?

If there is a specforge recipe that I could use as a reference, that would be helpful to analyze the pattern.

Yes, we use non-causal SWA attention for MiMo-v2.5 Pro DFlash, and the window is bi-directional. Prefix tokens are also limited in the sliding window. Currently we do not have a specforge-compatitble training codebase, but you may could refer to this: https://github.com/sgl-project/sglang/pull/27638

Sign up or log in to comment