[XPU] fix correctness issue of GDN triton kernel for XPU (#26065)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
2026-05-25 13:18:11 +08:00
parent ec6fcb93cb
commit 2bd3ac0b5d
3 changed files with 102 additions and 42 deletions
@@ -210,6 +210,16 @@ class TestChunkGatedDeltaRule(unittest.TestCase):
    def test_large_pool(self):
        self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=512)

+    # ------------------------------------------------------------------
+    # Long prompts (many chunks; regression test for cross-chunk errors)
+    # ------------------------------------------------------------------
+    def test_long_prompt(self):
+        for B, T_per_seq in [(1, 1024), (1, 1536), (1, 2048), (2, 1024)]:
+            with self.subTest(T_per_seq=T_per_seq):
+                self._check_shape(
+                    B=B, T_per_seq=T_per_seq, H=16, K=128, V=128, pool_size=32
+                )
+
    # ------------------------------------------------------------------
    # Combined stress
    # ------------------------------------------------------------------