Your current environment

• GPU: 8x AMD Instinct MI325X (256GB HBM3e each)
• CPU: AMD EPYC 9575F
• vLLM image: vllm/vllm-openai-rocm@sha256:cdea9cf61b3415bfdb4214ff56253f68645b0820d713a1745a21f0addeef4bd9
• mistral-common: 1.11.0
• ROCm: bundled in image

🐛 Describe the bug

Summary

Mistral Small 4 (mistralai/Mistral-Small-4-119B-2603) can't load on AMD MI325X GPUs with the current vLLM ROCm image. Two separate issues are currently preventing startup. I have listed them below with greater detail. Keen to try and work on the solution, unless I am doing something incorrectly? Please let me know if that is the case!

Environment

• GPU: 8x AMD Instinct MI325X (256GB VRAM each)
• vLLM image: vllm/vllm-openai-rocm@sha256:cdea9cf61b3415bfdb4214ff56253f68645b0820d713a1745a21f0addeef4bd9
• ROCm: (bundled in image)
• Model: mistralai/Mistral-Small-4-119B-2603 (119B MoE, 128 experts, 4 active, 32 attention heads)
• mistral-common: upgraded to 1.11.0 (required for v15 tokenizer)
• Tensor parallelism: TP=2 across 2x MI325X

For context, I used:

python3 -m vllm.entrypoints.openai.api_server \
  --model mistralai/Mistral-Small-4-119B-2603 \
  --tensor-parallel-size 2 \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.90 \
  --port 8000 \
  --trust-remote-code \
  --enable-prefix-caching \
  --enable-chunked-prefill \
  --max-num-seqs 128 \
  --max-num-batched-tokens 8192 \
  --swap-space 8 \
  --enable-auto-tool-choice \
  --tool-call-parser mistral

Issues

TLDR:

1. Hard assertion blocks 32 attention heads: rocm_aiter_mla.py:214 only allows 16 or 128 heads. Mistral Small 4 has 32. TP=2 works around it but shouldn't be required.
2. MoE kernel compilation exceeds startup timeout: The 128-expert FP8 MoE kernels take 15-20+ minutes to JIT compile on MI325X. vLLM's engine startup times out before compilation finishes, even with VLLM_ENGINE_READY_TIMEOUT_S=1800.

Issue 1: AITER MLA attention head assertion (fixed with TP=2)

With TP=1 (single GPU, 32 heads), the ROCm aiter MLA backend fails:

File "vllm/v1/attention/backends/mla/rocm_aiter_mla.py", line 214, in __init__
    assert num_heads == 16 or num_heads == 128
AssertionError: Aiter MLA only supports 16 or 128 number of heads.
Provided 32 number of heads.

Workaround: Setting --tensor-parallel-size 2 splits 32 heads into 16 per GPU, passing the assertion. However, I assume this shouldn't be required - the backend should either support 32 heads natively or fall back to a compatible backend.

Suggested fix: Expand the assertion to include all valid head counts (16, 32, 64, 128) or remove the hard assertion and let the kernel handle it. Excuse my naivety here on whether there is an existing fix.

Issue 2: MoE JIT kernel compilation timeout (no workaround)

With TP=2, the model gets past the attention head check but fails during ROCm JIT kernel compilation for the 128-expert MoE layer:

[aiter] waiting for baton release at .../lock_module_moe_ck2stages_f8_f8_preshuffle_off_b16_silu_per_tensor_mulWeightStage2_

The compilation takes longer than VLLM_ENGINE_READY_TIMEOUT_S (default 600s). The engine core process times out:

RuntimeError: Engine core initialization failed. See root cause above. Failed core proc(s): {}

Reproduce:

from vllm import LLM, SamplingParams

llm = LLM(
    model="mistralai/Mistral-Small-4-119B-2603",
    tensor_parallel_size=2,  # required to work around Issue 1
    max_model_len=32768,
    trust_remote_code=True,
    enable_auto_tool_choice=True,
    tool_call_parser="mistral",
)

# Never reaches here - crashes during engine init^
outputs = llm.generate(["Hello, my name is"], SamplingParams(temperature=0.8))

Setting VLLM_ENGINE_READY_TIMEOUT_S=1800 did not appear to take effect (the container still crashed at the same point). The MoE kernel compilation for 128 experts with FP8 weights on MI325X appears to take 15-20+ minutes, exceeding the engine startup timeout.

Suggested fix: not sure, ensure VLLM_ENGINE_READY_TIMEOUT_S actually propagates to engine core subprocesses? I am limited with ideas here however.

Related Issues

#37553 - Mistral Small 4 fails on 8x RTX 3090 (SM 8.6) with MLA backend PR #37081 - Mistral tool calling/reasoning parser for post-v15 tokenizers

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.