Fix AMX GQA extend attention (#25180)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
2026-05-18 09:30:02 +08:00
parent 2a357071ec
commit 5147de26e4
2 changed files with 15 additions and 2 deletions
@@ -217,6 +217,18 @@ class TestExtendAttention(CustomTestCase):
            b_seq_len_extend=[5000],
        )

+    def test_extend_attention_gqa_partial_extend_with_prefix(self):
+        self._test_extend_attention_once(
+            B=1,
+            N_CTX=256,
+            H_Q=16,
+            H_KV=4,
+            D=128,
+            DV=96,
+            b_seq_len_prefix=[97],
+            b_seq_len_extend=[37],
+        )
+

 if __name__ == "__main__":
    unittest.main()