[Bugfix][MiMo] Apply vision attention sinks in the window attention path

[almogtavor]

Fixes #47864.

MiMoVisionAttention allocates self.sinks only when the block is not in fullatt_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.5 ships visual.blocks.N.attn.sinks for exactly those blocks, 24 of its 28, 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. Biasing one key only rescales that key's softmax term, so the output is corrected in closed form from the weight the kernel already gave to key 0, recovered from the softmax LSE flash attention already returns:

p0   = exp(logit0 - lse)
out' = (out + p0 * (exp(s) - 1) * v0) / (1 + p0 * (exp(s) - 1))

Only the first window_size + 1 queries of a sequence can attend to its key 0, so only those rows are gathered and rewritten. Indices are built from cu_seqlens on device, so there is no host sync and no Python loop over sequences.

Real weights. One window block (block 1) of XiaomiMiMo/MiMo-V2.5, its own attn.sinks, compared against MiMoVisionAttention from the repo's modeling_mimo_v2.py, H100, bf16, 1024 tokens in 2 sequences, rotary made identity on both sides so only the attention math is compared:

	rel. error vs reference
main (sinks dropped)	8.7e-03
this PR	1.9e-03

1.9e-03 is the bf16 noise floor. The checkpoint's sink values are small (-0.10 to 0.59), so the per block effect is under 1%. It applies to 24 of 28 blocks; I did not measure the compounded error at the tower output.

Shape coverage. Same comparison with synthetic weights and larger sinks, over sequence lengths 128 to 4096, single and multi sequence batches, GQA, and TP=2 on both ranks: 1.8e-02 to 1.0e-01 before, 2.2e-03 after.

Test. tests/models/multimodal/test_mimo_v2_omni.py compares _forward_window_attn against a dense fp32 windowed softmax with the sink on key 0, for MHA and GQA, with one sequence shorter than the window and one longer. On this branch both cases land at 2.3e-03; on main they fail at 1.2e-01 and 2.1e-01.

$ pytest tests/models/multimodal/test_mimo_v2_omni.py
2 passed, 14 warnings in 17.66s

Cost. Window attention on 4096 tokens, 16 heads, head dim 80, window 64: 0.038 ms before, 0.248 ms after. On 16384 tokens: 0.134 ms before, 0.289 ms after. The overhead is roughly constant, so it is launch bound rather than data bound, and it applies once per window block.

Two alternatives, in case a reviewer prefers one:

• Passing the sinks as s_aux, as suggested in the issue, is a different operation. s_aux adds a null logit to the softmax denominator while the reference biases key 0. Measured on the synthetic setup it moves the error to 5.2e-02, slightly worse than dropping the sinks entirely.
• Recomputing the affected rows with SDPA also matches the reference, but mixes two attention implementations in one call.

The tradeoff here is that the key 0 logit is recomputed outside the kernel, so a little work the kernel already did is repeated, and the correction depends on the LSE convention of the varlen kernel. Adding real additive bias support to the flash attention wrapper would be cleaner but is a much larger change.

On the LSE convention: this call site takes the DEFAULT_FA_VERSION = 2 path, and cascade attention already feeds softmax_lse from the same function into merge_attn_states under any version, so the convention is shared in tree. Checked directly on H100 anyway, against a fp32 logsumexp of the masked scaled logits, max abs difference (nheads, total_q):

	FA2	FA3	FA4
lse vs fp32 reference	4.8e-07	4.8e-07	7.2e-07
corrected output vs reference	2.1e-03	2.1e-03	2.1e-03

The omni checkpoint is not public, so this is validated against the MiMo-V2.5 vision tower, whose config surface and weight names match what this file implements.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.