diff --git a/docs/docs/advanced_features/attention_backend.mdx b/docs/docs/advanced_features/attention_backend.mdx index 04e70b6de..f8a60ff1f 100644 --- a/docs/docs/advanced_features/attention_backend.mdx +++ b/docs/docs/advanced_features/attention_backend.mdx @@ -477,6 +477,25 @@ Internally, the DSA backend dispatches to different sub-backends for prefill and +#### 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. + + +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. + + 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) diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 14f8e910b..5cf5870a0 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -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