[AMD] Allow aiter attention backend for Gemma-4 (#38758)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -548,12 +548,13 @@ def handle_model_specific_adjustments(server_args: Any):
|
||||
"ascend",
|
||||
"intel_xpu",
|
||||
"intel_amx",
|
||||
"aiter",
|
||||
)
|
||||
assert (
|
||||
prefill_backend in accepted_backends and decode_backend in accepted_backends
|
||||
), (
|
||||
"Gemma4 only supports trtllm_mha, triton, ascend, intel_xpu, or intel_amx "
|
||||
f"attention backend, got prefill={prefill_backend}, decode={decode_backend}"
|
||||
"Gemma4 only supports trtllm_mha, triton, ascend, intel_xpu, intel_amx, or "
|
||||
f"aiter attention backend, got prefill={prefill_backend}, decode={decode_backend}"
|
||||
)
|
||||
|
||||
# The quantization/moe_runner_backend resolution moved to the override
|
||||
|
||||
Reference in New Issue
Block a user