[AMD] Allow aiter attention backend for Gemma-4 (#38758)

This commit is contained in:
Vignesh Sethuraman
2026-09-11 18:48:26 -07:00
committed by GitHub
parent e1d364dd1f
commit 0d1bea77da
2 changed files with 22 additions and 2 deletions
@@ -477,6 +477,25 @@ Internally, the DSA backend dispatches to different sub-backends for prefill and
</tbody>
</table>
#### Gemma 4 on `aiter` (ROCm)
Gemma 4's hybrid attention mixes sliding-window layers (`head_dim` 256) with
full-attention layers (`head_dim` 512) and shares KV across layers. The `aiter`
backend supports this and is allow-listed for Gemma 4, so you can serve it with
`--attention-backend aiter` on MI300X/MI350X. Under the hood `aiter` routes the
`head_dim`-512 prefill through the Triton `unified_attention` kernel (the CK
batch-prefill path caps `head_dim` at 256), honors each layer's softmax scale,
and resolves the sliding-window KV pool for speculative-decode draft workers.
Non–sliding-window models (`head_dim` ≤ 256) are unaffected and keep using the
CK path.
<Warning>
For Gemma 4 on `aiter`, use bf16 KV cache (`--kv-cache-dtype auto`). Do **not**
use `--kv-cache-dtype fp8_e4m3`: batched decode currently returns incorrect
output under fp8 KV on this path (a pre-existing aiter fp8-decode issue, tracked
separately). bf16 KV is correct at every batch size.
</Warning>
For deployment examples, see the [DeepSeek V3.2 deployment guide](/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2).
### Hybrid attention (different backends for prefill vs decode) (Experimental)