Substantial numerical differences with Minimax MSA attention vs Minimax's Reference
Author: CoffeeVampir3
Pull request: no
Labels: name: bug

### System Info

-> Minimax-M3/references/transformers/modular_minimax_m3_vl.py:550

  The problematic operation is:
```
  block_scores = scores.amax(dim=-1).amax(dim=1)
```

That first max pools within the 128-token block, which is expected. The second max, amax(dim=1), collapses H_idx, producing one shared selected block list per query. This is operationally different than minimax's reference which keeps things per-head. -> [minimax reference asserts](https://github.com/MiniMax-AI/MSA/blob/main/python/fmha_sm100/api.py#L207-L216)

I ran some tests to quantify the divergence between methods:
```
Selection divergence: 11759/11808 (99.5%) 
-- 99.5% of the time, the HF code produces a different result than the minimax reference

Mean overlap: 10.4 / 16 ~ 35% of each head selected blocks are being excluded from the HF method.

distinct blocks reached / token: 27.1 VS 16 ~ huggingface is touching nearly double the number of blocks minimax reference does.
```

Some useful references here: 
https://github.com/MiniMax-AI/MSA/blob/main/python/fmha_sm100/api.py#L1060-L1063

https://github.com/MiniMax-AI/MSA/blob/main/python/fmha_sm100/api.py#L1179-L1192

https://github.com/MiniMax-AI/MSA/blob/main/python/fmha_sm100/api.py#L1219-L1223

@ArthurZucker @vasqu

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Run minimax m3 model.

### Expected behavior

The attention produced by the reference. (VLLM also matches minimax M3's method)

Comments:
Thanks for the analysis, I dug into the code and can confirm this. The bug isn't isolated to that one line, the head-collapse is assumed by every consumer of the selection.

## Root cause
[modular_minimax_m3_vl.py:550](https://github.com/huggingface/transformers/blob/main/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py#L550):

```
block_scores = scores.amax(dim=-1).amax(dim=1)  # -> [B, S_q, num_key_blocks]
```

scores is `[B, H_idx, S_q, num_key_blocks, block_size]`:
 - `.amax(dim=-1)` pools within the 128-key block
 - .`amax(dim=1)` collapses `H_idx` producing one shared block list per query

The MSA reference's `sparse_topk_select` emits `(total_qo_len, num_qo_heads, topk)` (per-head), where in the MQA-proxy path `max_score.shape[0] == num_kv_heads_real`. With this config `index_n_heads == num_key_value_heads == 4`, so each index head owns one KV group (`qheads_per_kv = 64/4 = 16`). Collapsing the head axis is exactly what drives the ~99.5% selection divergence and the inflated distinct-block count.


It's a shared-selection assumption in 3 places, not one line:
1. Indexer ([modular:550-564](https://github.com/huggingface/transformers/blob/main/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py#L550-L564)): collapses heads, returns `[B, S_q, topk]`.
2. eager/sdpa mask ([modular:588](https://github.com/huggingface/transformers/blob/main/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py#L588)): build_block_mask does `.unsqueeze(1)`, broadcasting one mask to all query heads.
3. MSA kernel path ([msa_attention.py:204](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/msa_attention.py#L204)): `q2k ... .unsqueeze(0).expand(num_kv_heads, -1, -1)` literally broadcasts the same list to every KV head instead of indexing per head.

So even the kernel path feeds the deployment kernel a head-broadcast selection.

## Proposed fix
 Keep the head axis end-to-end (index head <-> KV head, 1:1):
- `Indexer.forward`: drop `.amax(dim=1)`, keep `block_scores = [B, H_idx, S_q, num_key_blocks]`, and apply the local-block boost and topk per head -> return `[B, H_idx, S_q, topk]`.
- `build_block_mask`: scatter per index head and `repeat_interleave(num_q_heads // H_idx, dim=1)` so it returns `[B, num_q_heads, S_q, S_k]`.
- `_sparse_attention`: replace the `expand(num_kv_heads, …)` with a real per-head reshape, e.g. `q2k = block_indices.permute(1, 0, 2).reshape(num_kv_heads, bsz * q_len, topk)`.

Plus an equivalence test against the reference selection (the overlap metric above is a good basis).

I'm happy to open a PR for this.

#46719 fixes this