From 76ef679cb1f720913b122c75f2da994641a3f38d Mon Sep 17 00:00:00 2001
From: Vignesh Sethuraman <vignesh.premseth@outlook.com>
Date: Thu, 10 Sep 2026 23:37:03 -0700
Subject: [PATCH] [AMD] aiter: fail loudly on cross-layer KV sharing in
 target_verify (#38755)

---
 .../srt/layers/attention/aiter_backend.py     | 19 ++++++++
 test/srt/test_aiter_cross_layer_kv_guard.py   | 45 +++++++++++++++++++
 2 files changed, 64 insertions(+)
 create mode 100644 test/srt/test_aiter_cross_layer_kv_guard.py

diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py
index 04cdd2487..09b720e2f 100755
--- a/python/sglang/srt/layers/attention/aiter_backend.py
+++ b/python/sglang/srt/layers/attention/aiter_backend.py
@@ -2285,6 +2285,23 @@ class AiterAttnBackend(AttentionBackend):
         )[:2]
         return output, lse.transpose(0, 1).contiguous()
 
+    @staticmethod
+    def _reject_target_verify_cross_layer_kv(k, v):
+        """Reject cross-layer KV sharing on the legacy ragged target-verify path.
+
+        Cross-layer KV sharing (e.g. Gemma4) passes ``k=v=None`` so the kernel
+        reads K/V from the pool. The legacy ``extend_attention_fwd`` path takes
+        ragged K/V and has no pool-reading fallback, so it would raise an opaque
+        ``AttributeError`` on ``.contiguous``. Fail loudly instead. Inert when
+        real K/V is passed.
+        """
+        if k is None or v is None:
+            raise ValueError(
+                "aiter target_verify does not support cross-layer KV "
+                "sharing (k/v are None). Use the unified verify path "
+                "(speculative_eagle_topk=1 and SGLANG_AITER_UNIFIED_VERIFY=1)."
+            )
+
     def forward_extend(
         self,
         q: torch.Tensor,
@@ -2797,6 +2814,8 @@ class AiterAttnBackend(AttentionBackend):
                     )
                     return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
+                self._reject_target_verify_cross_layer_kv(k, v)
+
                 self.extend_attention_fwd(
                     q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
                     k.contiguous(),
diff --git a/test/srt/test_aiter_cross_layer_kv_guard.py b/test/srt/test_aiter_cross_layer_kv_guard.py
new file mode 100644
index 000000000..f8a156971
--- /dev/null
+++ b/test/srt/test_aiter_cross_layer_kv_guard.py
@@ -0,0 +1,45 @@
+"""PR2 guard: aiter target-verify rejects cross-layer KV sharing (k/v is None).
+
+Cross-layer KV sharing (Gemma4) passes ``k=v=None``. The legacy ragged
+``extend_attention_fwd`` target-verify path has no pool-reading fallback, so it
+must fail loudly with a ``ValueError`` naming the unified verify path rather
+than an opaque ``AttributeError`` on ``.contiguous``. The guard is inert when
+real K/V is passed.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+
+class TestAiterCrossLayerKVGuard(CustomTestCase):
+    def test_raises_when_k_is_none(self):
+        with self.assertRaises(ValueError) as ctx:
+            AiterAttnBackend._reject_target_verify_cross_layer_kv(None, torch.empty(0))
+        self.assertIn("cross-layer KV", str(ctx.exception))
+        self.assertIn("unified verify path", str(ctx.exception))
+
+    def test_raises_when_v_is_none(self):
+        with self.assertRaises(ValueError):
+            AiterAttnBackend._reject_target_verify_cross_layer_kv(torch.empty(0), None)
+
+    def test_raises_when_both_none(self):
+        with self.assertRaises(ValueError):
+            AiterAttnBackend._reject_target_verify_cross_layer_kv(None, None)
+
+    def test_inert_when_kv_present(self):
+        # Real K/V present -> no raise (the guard must not fire for standard
+        # models that pass ragged K/V into the legacy target-verify path).
+        AiterAttnBackend._reject_target_verify_cross_layer_kv(
+            torch.empty(0), torch.empty(0)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()