Repository: vllm-project/vllm

Issue:
Serving for Google Gemma model failing on AMD MI 300X GPUs
The vLLM serving for the Google Gemma models is failing on AMD MI 300X GPUs  with the latest version of the vLLM (0.3.2) and ROCm 6.0.2 version.

Reproducing steps:

1. Clone the vllm repo and switch to [tag v0.3.2](https://github.com/vllm-project/vllm/tree/v0.3.2)
2. Build the Dockerfile.rocm dockerfile with instructions from [Option 3: Build from source with docker -Installation with ROCm](https://docs.vllm.ai/en/latest/getting_started/amd-installation.html#build-from-source-docker-rocm)

    build command:
    ```sh
    docker build  -f Dockerfile.rocm -t vllm-rocm:latest .
    ```

3. The vLLM serving command used:
    ```sh
    python3 -m vllm.entrypoints.openai.api_server --model google/gemma-7b --dtype bfloat16
    ```

The error below:
```sh
 File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/vllm-0.3.2+rocm603-py3.9-linux-x86_64.egg/vllm/model_executor/models/gemma.py", line 160, in forward
    attn_output = self.attn(q, k, v, k_cache, v_cache, input_metadata)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/vllm-0.3.2+rocm603-py3.9-linux-x86_64.egg/vllm/model_executor/layers/attention.py", line 196, in forward
    out = xops.memory_efficient_attention_forward(
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/xformers/ops/fmha/__init__.py", line 244, in memory_efficient_attention_forward
    return _memory_efficient_attention_forward(
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/xformers/ops/fmha/__init__.py", line 341, in _memory_efficient_attention_forward
    out, *_ = op.apply(inp, needs_gradient=False)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/xformers/ops/fmha/flash.py", line 458, in apply
    out, softmax_lse, rng_state = cls.OPERATOR(
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/xformers/ops/fmha/flash.py", line 130, in _flash_fwd
    ) = _C_flashattention.varlen_fwd(
RuntimeError: FlashAttention forward only supports head dimension at most 128
```

**Issues:**
1. This issue is for serving both Gemma 2B and 7B variants on the AMD MI 300X GPUs


Discussion:
same on mi210
Failing for us on MI100 as well

Resolving patch:
diff --git a/vllm/attention/ops/triton_flash_attention.py b/vllm/attention/ops/triton_flash_attention.py
--- a/vllm/attention/ops/triton_flash_attention.py
+++ b/vllm/attention/ops/triton_flash_attention.py
@@ -677,8 +677,7 @@ def check_args(
     assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
     # TODO: Change assert if we support qkl f8 and v f16
     assert q.dtype == k.dtype and q.dtype == v.dtype
-    # TODO: Fix assert to check head size <=256 once supported
-    assert head_size <= 128
+    assert head_size <= 256
     assert o.shape == q.shape
     assert (nheads_q % nheads_k) == 0
 
@@ -729,7 +728,7 @@ def forward(
             o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
 
         # Get closest power of 2 over or equal to 32.
-        unpadded_head_dims = {32, 64, 128}
+        unpadded_head_dims = {32, 64, 128, 256}
         if head_size not in unpadded_head_dims:
             padded_d_model = None
             for i in unpadded_head_dims:


Test patch: