From de35da78eba406d854b921f800b60f488052c897 Mon Sep 17 00:00:00 2001 From: Yiwen Xie Date: Tue, 28 Jul 2026 14:29:29 -0700 Subject: [PATCH] Add OSS regression test for GDN channelwise_gated_delta_rule chunked-prefill NaN (#21439) Summary: Adds an OSS unit test that guards the GDN2 chunked-prefill overflow NaN fixed in D113866825. The fused `torch.ops.llama.channelwise_gated_delta_rule` chunked-prefill kernel (`op_sdpa.cpp`) forms `eg_inv = exp(-cumsum(log decay))` per CHUNK_SIZE (32) tokens; for tiny positive decay the within-chunk cumulative `-log(decay)` exceeds the float `expf` overflow limit (~88.7), so `eg_inv` overflows to `+inf` and `eg * eg_inv` becomes `0 * inf = NaN`, poisoning the whole head's output. The `channelwise_gated_delta_rule_chunked_is_safe` guard routes such inputs to the exact (numerically stable) recurrence. `test_channelwise_gated_delta_rule_prefill_tiny_decay_is_finite` builds a prefill (T=64 > CHUNK_SIZE) input with tiny positive decay (1e-12) and asserts the op output and final_state are finite and match the token-by-token reference. Without the guard the `isfinite` assertions fail. Reviewed By: JakeStevens Differential Revision: D113943464 --- extension/llm/custom_ops/test_gated_delta.py | 42 ++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/extension/llm/custom_ops/test_gated_delta.py b/extension/llm/custom_ops/test_gated_delta.py index 9fc9af63d36..bac8909ee25 100644 --- a/extension/llm/custom_ops/test_gated_delta.py +++ b/extension/llm/custom_ops/test_gated_delta.py @@ -294,6 +294,48 @@ def test_channelwise_gated_delta_rule_multichunk_matches_reference(self): torch.allclose(actual_state, expected_state, atol=1e-3, rtol=1e-3) ) + def test_channelwise_gated_delta_rule_prefill_tiny_decay_is_finite(self): + # Regression test for the chunked-prefill overflow NaN in op_sdpa.cpp. + # The prefill kernel forms eg_inv = exp(-cumsum(log decay)) per + # CHUNK_SIZE (32) tokens; for tiny positive decay the within-chunk + # cumulative -log(decay) exceeds the float expf overflow limit (~88.7), + # so eg_inv overflows to +inf and eg * eg_inv becomes 0 * inf = NaN, + # poisoning the whole head's output. The + # channelwise_gated_delta_rule_chunked_is_safe guard must route such + # inputs to the exact recurrence, which stays finite (and correct) for + # tiny decay; without it the isfinite assertions below fail. + torch.manual_seed(0) + + # seq_len > CHUNK_SIZE (32) so the prefill route runs the internal chunk + # loop; a single chunk of tiny decay overflows the unguarded kernel. + query, key, value, decay, beta, initial_state = self._make_inputs(seq_len=64) + # decay.min ~6.5e-12 was observed in a real GDN2 checkpoint; + # -log(1e-12) ~27.6 per token, so a few tokens blow past the ~80 safe + # bound and exp(-gc) would overflow to +inf -> NaN. + decay = torch.full_like(decay, 1e-12) + + expected_output, expected_state = self._reference_channelwise_gated_delta_rule( + query, key, value, decay, beta, initial_state + ) + actual_output, actual_state = torch.ops.llama.channelwise_gated_delta_rule( + query, key, value, decay, beta, initial_state + ) + + self.assertTrue( + torch.isfinite(actual_output).all(), + "prefill output has NaN/inf for tiny decay", + ) + self.assertTrue( + torch.isfinite(actual_state).all(), + "prefill final_state has NaN/inf for tiny decay", + ) + self.assertTrue( + torch.allclose(actual_output, expected_output, atol=1e-3, rtol=1e-3) + ) + self.assertTrue( + torch.allclose(actual_state, expected_state, atol=1e-3, rtol=1e-3) + ) + def test_channelwise_gated_delta_rule_exports(self): class Module(torch.nn.Module): def forward(self, query, key, value, decay, beta, initial_state):