vllm - 💡(How to fix) Fix [Bug]: TRITON_MLA grouped-decode kernel fails for Mistral Small 4 (kv_lora_rank=256): 'Cannot make_shape_compatible: 256 and 512'

Your current environment

• vLLM: 0.20.1
• GPU: NVIDIA GB10 (Grace Blackwell, SM121 / 12.1a) — DGX Spark
• CUDA: 13.0, PyTorch 2.x
• Model: Mistral Small 4 (119B MoE) — Mistral3/Pixtral wrapper over a DeepSeek-V2-style MLA + grouped-MoE text model, served via --load-format mistral
• Quantization is not relevant here: attention is BF16 in the checkpoint; the failure is in the MLA decode path.

🐛 Describe the bug

Serving Mistral Small 4 (MLA attention with kv_lora_rank=256) selects the TRITON_MLA attention backend, and the grouped MLA decode kernel fails to compile at the first decode step with:

ValueError: Cannot make_shape_compatible: incompatible dimensions at index 1: 256 and 512

The engine loads, warms up, and accepts requests; the error fires on the first token generation and kills the EngineCore (EngineDeadError). Prefill uses FlashAttention and is fine — it's specifically the TRITON_MLA decode kernel.

This appears to be independent of quantization (our attention tensors are BF16) and independent of the Pixtral vs deepseek_v2 model-class resolution — both routes land on TRITON_MLA and both crash identically. It looks like the grouped-decode kernel assumes the DeepSeek-V2 kv_lora_rank=512, and Mistral 4's kv_lora_rank=256 trips the shape check.

Model MLA config:

q_lora_rank=1024  kv_lora_rank=256
qk_nope_head_dim=64  qk_rope_head_dim=64  v_head_dim=128  head_dim=128
n_heads=32  n_kv_heads=32  dim=4096

(256 in the error = kv_lora_rank; 512 is DeepSeek-V2's kv_lora_rank.)

Backend selection log:

Using TRITON_MLA attention backend out of potential backends: ['TRITON_MLA']
Using FlashAttention prefill for MLA

Traceback (abridged):

.../model_executor/layers/attention/mla_attention.py", line 755, in forward_impl
    attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
.../v1/attention/backends/mla/triton_mla.py", line 200, in forward_mqa
.../v1/attention/ops/triton_decode_attention.py", line 500, in _decode_grouped_att_m_fwd
.../triton/compiler/compiler.py", line 80, in make_ir
triton.compiler.errors.CompilationError: at 152:12:
    v = tl.trans(k)
    ...
    acc += tl.dot(p.to(v.dtype), v)
    ^
ValueError('Cannot make_shape_compatible: incompatible dimensions at index 1: 256 and 512')

Repro

1. Serve Mistral Small 4 (NVFP4 or BF16) with --load-format mistral on a GB10/SM12x box (vLLM auto-selects TRITON_MLA).
2. Send any chat/completions request.
3. Crash on first generated token.

✅ Workaround

Setting VLLM_MLA_DISABLE=1 disables the MLA decode path; attention is then served via FLASH_ATTN and generation works correctly (validated end-to-end: coherent output, tool-calling, long generations). On SM12x we also use VLLM_NVFP4_GEMM_BACKEND=marlin + VLLM_USE_FLASHINFER_MOE_FP4=0 for the MoE GEMMs, but those are unrelated to this attention bug.

Ask

• Is TRITON_MLA's grouped-decode kernel expected to support kv_lora_rank != 512 (e.g. Mistral 4's 256)? If it's a real bug, the fix is presumably in triton_decode_attention.py _decode_grouped_att_m_fwd (the tl.trans/tl.dot shape handling around line 152).
• If it's a known limitation for this arch on SM12x, documenting "use VLLM_MLA_DISABLE=1 for Mistral 4" would save others the dig.

Happy to provide the full log or test a patch on GB10 hardware. Tested on 0.20.1 — can re-verify on main if helpful.