[Fix] Skip padded state slots in the chunked GDN kernel (#33431)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

test/registered/attention/test_chunk_gated_delta_rule.py

@@ -294,6 +294,53 @@ class TestChunkGatedDeltaRule(unittest.TestCase):
             sequential_indices=True,
         )
 
+    # ------------------------------------------------------------------
+    # Padded rows (idle DP ranks under breakable CUDA graph prefill)
+    # ------------------------------------------------------------------
+
+    def test_padded_state_index_is_skipped(self):
+        """Rows carrying state index -1 must be skipped, not addressed.
+
+        Without the guard the sentinel reaches pointer arithmetic and
+        addresses before the state pool.
+        """
+        device = get_device()
+        dtype = torch.bfloat16
+        B, T_per_seq, H, K, V, pool_size = 4, 64, 4, 128, 128, 8
+        T = B * T_per_seq
+        torch.manual_seed(0)
+
+        pool_init = (
+            torch.randn(pool_size, H, V, K, dtype=torch.float32, device=device) * 0.1
+        )
+        cu_seqlens = torch.zeros(B + 1, dtype=torch.long, device=device)
+        cu_seqlens[1:] = (
+            torch.arange(1, B + 1, dtype=torch.long, device=device) * T_per_seq
+        )
+        q = torch.randn(1, T, H, K, dtype=dtype, device=device)
+        k = torch.randn(1, T, H, K, dtype=dtype, device=device)
+        v = torch.randn(1, T, H, V, dtype=dtype, device=device)
+        g = torch.nn.functional.logsigmoid(
+            torch.randn(1, T, H, dtype=dtype, device=device)
+        )
+        beta = torch.sigmoid(torch.randn(1, T, H, dtype=dtype, device=device))
+
+        # "all padded" is the idle-DP-rank case that faulted.
+        for label, indices in (("mixed", [0, -1, 2, -1]), ("all_padded", [-1] * B)):
+            with self.subTest(label):
+                cache_indices = torch.tensor(indices, dtype=torch.int32, device=device)
+                o, pool = self._run_chunk(
+                    pool_init, cache_indices, q, k, v, g, beta, cu_seqlens
+                )
+                # Forces a sync: an out-of-bounds access surfaces here.
+                self.assertTrue(torch.isfinite(o.float()).all())
+
+                untouched = [s for s in range(pool_size) if s not in indices]
+                self.assertTrue(
+                    torch.equal(pool[untouched], pool_init[untouched]),
+                    f"{label}: padded rows wrote into the state pool",
+                )
+
 
 if __name__ == "__main__":
     unittest.main()