From 955704544c60e920672aa434cefa2ce78c0ceb4c Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Tue, 18 Aug 2026 13:07:19 -0700 Subject: [PATCH] [Fix] Skip padded state slots in the chunked GDN kernel (#33431) Co-authored-by: Claude Opus 5 (1M context) --- .../attention/test_chunk_gated_delta_rule.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/test/registered/attention/test_chunk_gated_delta_rule.py b/test/registered/attention/test_chunk_gated_delta_rule.py index c521fac2e..3ef2ab9b1 100644 --- a/test/registered/attention/test_chunk_gated_delta_rule.py +++ b/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()