From f4158719d176fa6caf7086f4e4e187e538a09906 Mon Sep 17 00:00:00 2001
From: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
Date: Wed, 19 Aug 2026 00:17:40 -0700
Subject: [PATCH] [AMD] Let the diffusion AITer backend take grouped-query K/V
 (fix Cosmos3-Nano startup) (#34485)

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Michael <michaelzhang-ai@users.noreply.github.com>
---
 .../layers/attention/backends/aiter.py        | 11 +++--
 .../test/unit/test_aiter_attention_impl.py    | 41 +++++++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 python/sglang/multimodal_gen/test/unit/test_aiter_attention_impl.py

diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
index 52210c07f..f897214b7 100755
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
@@ -128,9 +128,14 @@ class AITerImpl(AttentionImpl):
         dropout_p: float = 0.0,
         **extra_impl_args,
     ) -> None:
-        if num_kv_heads is not None and num_kv_heads != num_heads:
-            raise NotImplementedError(
-                "AITer backend does not support Grouped Query Attention yet."
+        # aiter's mha entry points take GQA/MQA K/V directly (they broadcast
+        # each KV head across its group of query heads), so the only
+        # requirement is an even split. The FP8 ASM path is MHA-only and
+        # already routes grouped shapes back to BF16 below.
+        if num_kv_heads is not None and num_heads % num_kv_heads != 0:
+            raise ValueError(
+                f"AITer backend requires num_heads ({num_heads}) to be a "
+                f"multiple of num_kv_heads ({num_kv_heads})."
             )
         self.causal = causal
         self.dropout_p = dropout_p
diff --git a/python/sglang/multimodal_gen/test/unit/test_aiter_attention_impl.py b/python/sglang/multimodal_gen/test/unit/test_aiter_attention_impl.py
new file mode 100644
index 000000000..074560965
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_aiter_attention_impl.py
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AITer attention impl construction guards (ROCm-only; skipped elsewhere)."""
+
+import pytest
+
+HEAD_SIZE = 128
+
+
+def _impl_cls():
+    # `aiter` ships with ROCm only, and importing the backend module needs it.
+    pytest.importorskip("aiter", reason="AITer is a ROCm-only dependency")
+    from sglang.multimodal_gen.runtime.layers.attention.backends.aiter import AITerImpl
+
+    return AITerImpl
+
+
+def _build(num_heads: int, num_kv_heads: int | None):
+    return _impl_cls()(
+        num_heads=num_heads,
+        head_size=HEAD_SIZE,
+        softmax_scale=HEAD_SIZE**-0.5,
+        num_kv_heads=num_kv_heads,
+    )
+
+
+@pytest.mark.parametrize("num_kv_heads", [32, 8, 1, None])
+def test_accepts_grouped_and_multi_query_kv_heads(num_kv_heads):
+    # aiter's mha entry points broadcast each KV head across its group of
+    # query heads, so Cosmos3-style GQA cross-attention is supported.
+    assert _build(32, num_kv_heads).softmax_scale == pytest.approx(HEAD_SIZE**-0.5)
+
+
+def test_rejects_kv_heads_that_do_not_divide_the_query_heads():
+    with pytest.raises(ValueError, match="multiple of num_kv_heads"):
+        _build(32, 5)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))