[Bugfix][MiMo] Apply vision attention sinks in the window attention path - #49815
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
0babb74 to
bbd5d51
Compare
| @@ -229,8 +229,68 @@ def _forward_window_attn( | |||
| softmax_scale=self.scale, | |||
| causal=False, | |||
| window_size=[w, w], | |||
| return_softmax_lse=self.sinks is not None, | |||
| ) | |||
There was a problem hiding this comment.
Why not pass s_aux=self.sinks?
There was a problem hiding this comment.
@Isotr0py since FA2 rejects it with NotImplementedError: FA2 does not support s_aux so it'll only work on FA3 / FA4 (flash_attn_interface.py passes it as learnable_sink=s_aux). It's also a different operation since s_aux appends a null logit to the softmax denominator, so every weight shrinks by the same factor and nothing is added to the numerator. MiMoVisionAttention instead adds sinks[h] to the logit of key 0, which rescales key 0's softmax weight by exp(sinks[h]) (head h gets sinks[h] at key 0 for all queries and zero at every other key), so the value vector of key 0 (v0) contributes more or less to the output depending on the sign.
Forcing fa_version=3 so it runs at all, H100 in bf16 against that reference:
| error | |
|---|---|
| main, sinks dropped | 0.048 |
| s_aux=self.sinks | 0.052 (FA3 only) |
| this PR | 0.0022 |
0.0022 is the bf16 noise floor, the same call with no sinks measures 0.0020 against fp32 SDPA. A real additive bias on key 0 would need new kernel code in vllm-project/flash-attention for FA2, FA3 and FA4, so this PR reconstructs the result that a kernel-native bias would produce, using the softmax LSE that flash attention already returns. It is a little slower than doing it in the kernel of course.
There was a problem hiding this comment.
I think you can fall back to Triton kernels for pre-Hopper devices:
vllm/vllm/v1/attention/ops/triton_prefill_attention.py
Lines 198 to 211 in 439f336
There was a problem hiding this comment.
@Isotr0py Thanks, I looked at that kernel. It computes the same thing as s_aux, so it changes where the code can run rather than what it computes.
At lines 99-103 it initialises m_i to the sink and l_i to 1.0 with acc at 0, so the sink is one extra term in the denominator that never reaches the numerator, because no value vector is paired with it and acc starts empty. That is the GPT-OSS null logit, but MiMoVisionAttention adds sinks[h] to the logit of key 0, which does reach the numerator through v0.
H100, bf16, same input, measured against both formulations:
| err vs MiMo (key 0 bias), we want it lower | err vs null logit (lower / higher doenst matter) | |
|---|---|---|
| main, sinks dropped | 0.048 | 0.022 |
s_aux=self.sinks on FA3 (issue suggestion) |
0.052 | 0.0020 |
sinks=self.sinks on triton (your suggestion) |
0.052 | 0.0020 |
| LSE correction (this PR) | 0.0021 | 0.053 |
Besides the baseline, each row reproduces either the key 0 bias or the null logit to within bf16 rounding (never both. bf16 rounding is about 0.002). The two sink paths land on the null logit and agree with each other in four digits and both are 0.052 from MiMo, a little further than simply dropping the sinks. As a wiring check i saw the triton kernel with no sinks matches flash attention with no sinks to 0.
We could implement in the triton kernel support for having the key 0 bias directly, but it would mean routing this path through triton on every device including Hopper.
There was a problem hiding this comment.
I think using/correcting the triton kernel should be fine, LSE correction will increase the maintainence effort.
|
@Isotr0py I tried your suggestion properly and it comes out ahead, so I would like to switch this PR to it. Two things were needed. The kernel treats sinks as a null logit, so I added a mode that biases key 0 instead, matching H100, bf16, 8 heads, head dim 64, window 64. Error is against a dense fp32 windowed softmax with the sink on key 0, latency is the mean of 20 iterations:
Same accuracy, roughly 2x to 4x faster, and the model file loses the ~50 line correction. The 20 existing cases in Only caveat is that it moves this path off flash attention and onto the Triton kernel for every device. Happy to rework the PR that way if you prefer it, once #50776 is settled. |
`MiMoVisionAttention` allocates `self.sinks` only for blocks that are not in `fullatt_block_indexes`, which is exactly the set of blocks running `_forward_window_attn`. That path never read the parameter, so the sink weights were loaded from the checkpoint and dropped. `XiaomiMiMo/MiMo-V2.5` ships `visual.blocks.N.attn.sinks` for 24 of its 28 blocks, and none for the full attention blocks [0, 9, 18, 27]. The reference adds `sinks[h]` to the logit of each sequence's first key. The Triton prefill kernel already carries sinks, but as a null logit in the denominator, which is a different operation. Add a mode that biases key 0 instead and use it here, so the softmax normalizes over the biased scores in one pass with no post-hoc correction. Measured on H100 in bf16 against a dense fp32 windowed softmax with the sink on key 0, window 64, 8 heads, head dim 64: | tokens | error | |--------|-------| | 512 | 2.14e-03 | | 1024 | 2.14e-03 | | 4096 | 2.14e-03 | 2.14e-03 is the bf16 noise floor. Dropping the sinks gives 1.2e-01 to 2.1e-01 on the regression test's shapes. Signed-off-by: almogtavor <almogtavor@gmail.com>
bbd5d51 to
1412668
Compare
github-actions
Bot
added
the
ready
|
/ci run |
|
✅ Triggered Buildkite CI #82840 for commit |
|
/ci retry |
|
✅ Queued 3 failed job(s) for retry in Buildkite CI #82840. |
|
/ci retry |
|
✅ Triggered Buildkite CI #83215 for commit |
2 participants
Fixes #47864.
MiMoVisionAttentionallocatesself.sinksonly when the block is not infullatt_block_indexes, which is exactly the set of blocks that run_forward_window_attn. That path never read the parameter, so the sink weights were loaded from the checkpoint and dropped.XiaomiMiMo/MiMo-V2.5shipsvisual.blocks.N.attn.sinksfor exactly those blocks, 24 of its 28, and none for the full attention blocks[0, 9, 18, 27].Approach
The reference adds
sinks[h]to the logit of each sequence's first key.triton_prefill_attentionalready accepts sinks, but treats them as a null logit that only inflates the denominator, which is a different operation (see the measurements below). This adds asinks_bias_key0mode that biases key 0 instead, and routes the window path through that kernel, so the softmax normalizes over the biased scores in one pass.An earlier revision of this PR corrected flash attention's output after the fact using the softmax LSE. That worked, but as @Isotr0py noted it carries ongoing maintenance cost, and the kernel change is both simpler and faster.
Numbers
H100, bf16, window 64, 8 heads, head dim 64. Error is against a dense fp32 windowed softmax with the sink on key 0. Latency is the mean of 20 iterations.
2.14e-03 is the bf16 noise floor. The kernel path is equal on accuracy, 2x to 4x faster, and drops ~50 lines of correction from the model file.
For reference, the two sink formulations are not interchangeable. Measured earlier on the same setup, against MiMo's key 0 reference and against a null logit reference:
s_auxon FA3Test
tests/models/multimodal/test_mimo_v2_omni.pycompares_forward_window_attnagainst the dense fp32 reference for MHA and GQA, with one sequence shorter than the window and one longer. Both pass here at 2.3e-03 and fail on main at 1.2e-01 and 2.1e-01.tests/kernels/attention/test_triton_prefill_attention.py, 20 passed, so the existing null-logit behaviour is unchanged.AI assistance was used for this change.