[Bug][ROCm]: BF16 MLA with ROCM_AITER_FA broken (gfx950) (fmha opus)

[simondanielsson]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

🐛 Describe the bug

Issue: Running DSv3 with BF16 KV throws an error during JIT due to incompatible argument types.

For context, this stems from two interlinked problems:

1. Issue 1: pure causal prefills are currently always run with a FP8 PS ASM kernel, no matter the KV dtype.
   1. Fix: [Bugfix][ROCm] Only run FP8 AITER MLA prefill when using FP8 KV #48712
2. Issue 2 (this issue): BF16 causal prefills in ROCM_AITER_FA (with Opus) fails during JIT.

Since issue 1 ensures that causal prefills are never routed to the BF16 Opus kernel, we don't see this issue by default. However, Issue 2 is surface after fixing Issue 1 in #48712 and/or when running a non-default MLA backend like TRITON_MLA

Reproducer

Surface by running a non-default MLA backend. (Equivalently, you can run in #48712)

docker run --rm --name repro-mla \
    --init --network host --ipc host --privileged \
    --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
    --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
    --group-add video --group-add render \
    --device /dev/kfd --device /dev/dri --device /dev/infiniband \
    -v /sys:/sys \
    -v "/root/.cache/huggingface/:/root/.cache/huggingface" \
    -e HF_HOME=/root/.cache/huggingface -e HF_TOKEN=$HF_TOKEN \
    -e HF_HUB_OFFLINE=1 -e HF_HUB_ENABLE_HF_TRANSFER=0 \
    -e VLLM_ROCM_USE_AITER=1 \
    vllm/vllm-openai-rocm:nightly-ac7509e2b1db40fec2f03dde1ed4e9dfdc2338c9 \
    deepseek-ai/DeepSeek-V3-0324 \
      --tensor-parallel-size 8 --enable-expert-parallel \
      --gpu-memory-utilization 0.8 --max-model-len 16K \
      --attention-backend TRITON_MLA

We can see how ROCM_AITER_FA prefill backend is used:

INFO:inference_testing.servers.docker_utils: (Worker_TP0_EP0 pid=1065) INFO 08-14 09:24:03 [rocm.py:638] Using TRITON_MLA backend (selected via --attention-backend).
INFO:inference_testing.servers.docker_utils: (Worker_TP0_EP0 pid=1065) INFO 08-14 09:24:03 [selector.py:190] Using ROCM_AITER_FA MLA prefill backend.

and we get this error after sending a basic request to /v1/chat/completions:

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "deepseek-ai/DeepSeek-V3-0324",
    "messages": [
      {"role": "system", "content": "You are a helpful assistant."},
      {"role": "user", "content": "Tell me a short joke."}
    ],
    "temperature": 0.7,
    "max_tokens": 10
  }'

INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033] Traceback (most recent call last):
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/executor/multiproc_executor.py", line 1025, in worker_busy_loop
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     output = func(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]              ^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/worker_base.py", line 351, in execute_model
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self.worker.execute_model(scheduler_output)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 124, in decorate_context
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return func(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_worker.py", line 1086, in execute_model
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     output = self.model_runner.execute_model(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 124, in decorate_context
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return func(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_model_runner.py", line 4554, in execute_model
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     model_output = self._model_forward(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]                    ^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_model_runner.py", line 3970, in _model_forward
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self.model(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/compilation/cuda_graph.py", line 254, in __call__
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self.runnable(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self._call_impl(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return forward_call(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v2.py", line 1894, in forward
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     hidden_states = self.model(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]                     ^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/compilation/decorators.py", line 520, in __call__
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self.aot_compiled_fn(self, *args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/aot_compile.py", line 240, in __call__
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self.fn(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v2.py", line 1427, in forward
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     def forward(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/compilation/caching.py", line 225, in __call__
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self.optimized_call(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "<string>", line 505, in execution_fn
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "<string>", line 9, in __vllm_inlined_submods__2
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 1275, in __call__
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self._op(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/attention/kv_transfer_utils.py", line 40, in wrapper
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return func(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/attention/mla_attention.py", line 1266, in unified_mla_attention_with_output
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     layer.forward_impl(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/attention/mla_attention.py", line 851, in forward_impl
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     self.impl.forward_mha(  # type: ignore[attr-defined]
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/attention/mla_attention.py", line 2842, in forward_mha
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     output_prefill = prefill_metadata.prefill_backend.run_prefill_new_tokens(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py", line 83, in run_prefill_new_tokens
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     result = self.flash_attn_varlen_func(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/ops/mha.py", line 3649, in flash_attn_varlen_func
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return FlashAttnVarlenFunc.apply(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 596, in apply
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return super().apply(*args, **kwargs)  # type: ignore[misc]
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/ops/mha.py", line 3282, in forward
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]                                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/ops/mha.py", line 2918, in _flash_attn_varlen_forward
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     out = fmha_fwd_bf16_opus_varlen_fwd(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/ops/mha.py", line 429, in fmha_fwd_bf16_opus_varlen_fwd
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     _fmha_fwd_bf16_opus_fwd(
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/jit/utils/torch_guard.py", line 289, in wrapper_custom
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     getattr(torch.ops.aiter, f"{loadName}")(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 1275, in __call__
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return self._op(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/jit/utils/torch_guard.py", line 312, in outer_wrapper
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     wrapper(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/jit/utils/torch_guard.py", line 207, in wrapper
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return func(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/jit/core.py", line 1784, in custom_wrapper
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return wrapper(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]   File "/usr/local/lib/python3.12/dist-packages/aiter/jit/core.py", line 1780, in wrapper
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     return op(*args, **kwargs)
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]            ^^^^^^^^^^^^^^^^^^^
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033] TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments. The following argument types are supported:
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]     1. (q: aiter_tensor_t, k: aiter_tensor_t, v: aiter_tensor_t, out: aiter_tensor_t, causal: bool, softmax_scale: typing.SupportsFloat | typing.SupportsIndex, seqstart_q: aiter_tensor_t | None = None, seqstart_k: aiter_tensor_t | None = None, seqstart_q_pad: aiter_tensor_t | None = None, seqstart_k_pad: aiter_tensor_t | None = None, max_seqlen_q: typing.SupportsInt | typing.SupportsIndex = 0, max_seqlen_k: typing.SupportsInt | typing.SupportsIndex = 0) -> None
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033] Invoked with: <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291d8b0>, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291f6f0>, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291fe70>, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291ff30>, True, 0.1352337788608801, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291c630>, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291e770>, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291da30>, <aiter.jit.module_aiter_core.aiter_tensor_t object at 0x7f023291fc30>, 4096, 4096
INFO:inference_testing.servers.docker_utils: (Worker_TP5_EP5 pid=1070) ERROR 08-14 09:27:31 [multiproc_executor.py:1033]

Short term fix

AITER_DISABLE_FMHA_OPUS=1

Potential solution

Pre-compile the opus kernel in AITER: i.e. remove it from the exlcude list here: https://github.com/ROCm/aiter/blob/1e8c33fd6262a678960d575465b100edda566759/setup.py#L256.

Important note

The error is not thrown when using the default AITER MLA backend (i.e. omitting --attention-backend TRITON_MLA). However, this is due to a separate issue; the AITER_MLA currently incorrectly defaults to using an FP8 ASM prefill kernel during causal prefill, even with BF16 KV. This is fixed in #48712. This means that during causal prefills, BF16 KV is truncated to FP8 and then run through an FP8 kernel. This is not correct.

Profiling proof of this issue

Running the default MLA backend with BF16 KV:

docker run --rm --name repro-mla \
    --init --network host --ipc host --privileged \
    --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
    --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
    --group-add video --group-add render \
    --device /dev/kfd --device /dev/dri --device /dev/infiniband \
    -v /sys:/sys \
    -v "/root/.cache/huggingface/:/root/.cache/huggingface" \
    -e HF_HOME=/root/.cache/huggingface -e HF_TOKEN=$HF_TOKEN \
    -e HF_HUB_OFFLINE=1 -e HF_HUB_ENABLE_HF_TRANSFER=0 \
    -e VLLM_ROCM_USE_AITER=1 \
    vllm/vllm-openai-rocm:nightly-ac7509e2b1db40fec2f03dde1ed4e9dfdc2338c9 \
    deepseek-ai/DeepSeek-V3-0324 \
      --tensor-parallel-size 8 --enable-expert-parallel \
      --gpu-memory-utilization 0.8 --max-model-len 16K

Causal prefills: FP8 asm kernel is running! This is not what we would expect with BF16 KV

with parent:

In contrast, chunked prefills run both the causal and non-causal part with the (correct) BF16 AITER FA impl:

So pure causal prefills have their BF16 KV are truncated to FP8 and ran in lower precision!

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.

[github-actions]

CC @hongxiayang @tjtanaa @vllmellm @giuseppegrossi for ROCm-related issue

[github-actions]

Hi @simondanielsson, thanks for reporting this ROCm issue!

To help us investigate, please make sure the following information is included:

• A minimal reproducer (code snippet or script that triggers the issue)
• The full error message or traceback
• How you installed vLLM (Docker image name, pip install, or build from source steps)
• The command you used to launch vLLM (e.g., vllm serve ... or the Python script)
• Your GPU model and GFX architecture (e.g., MI300X / gfx942) — run rocminfo | grep gfx

Please provide any unchecked items above. This will help us reproduce and resolve the issue faster. Thank you!