Fix NaN in _chunk_scaled_dot_kkt_fwd_eager across all 3 gated deltanet kernels

Zero out g_diff outside the strict lower triangle before calling exp(),
preventing inf * 0 = NaN when upper-triangle g differences overflow.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: yf225

problems/helion/gated_deltanet_chunk_fwd_h_py/reference.py

@@ -19,8 +19,9 @@ def _chunk_scaled_dot_kkt_fwd_eager(k, g_cumsum, beta, chunk_size):
     g_c = g_cumsum.float().reshape(B, NT, C, H).permute(0, 1, 3, 2)
     beta_c = beta.float().reshape(B, NT, C, H).permute(0, 1, 3, 2)
     kkt = k_c @ k_c.transpose(-1, -2)
-    g_diff = g_c.unsqueeze(-1) - g_c.unsqueeze(-2)
     strict_lower = torch.tril(torch.ones(C, C, device=k.device), diagonal=-1)
+    g_diff = g_c.unsqueeze(-1) - g_c.unsqueeze(-2)
+    g_diff = g_diff * strict_lower
     A = kkt * beta_c.unsqueeze(-1) * torch.exp(g_diff) * strict_lower
     return A.permute(0, 1, 3, 2, 4).reshape(B, T, H, C).to(torch.float32)
 

problems/helion/gated_deltanet_chunk_fwd_o_py/reference.py

@@ -19,8 +19,9 @@ def _chunk_scaled_dot_kkt_fwd_eager(k, g_cumsum, beta, chunk_size):
     g_c = g_cumsum.float().reshape(B, NT, C, H).permute(0, 1, 3, 2)
     beta_c = beta.float().reshape(B, NT, C, H).permute(0, 1, 3, 2)
     kkt = k_c @ k_c.transpose(-1, -2)
-    g_diff = g_c.unsqueeze(-1) - g_c.unsqueeze(-2)
     strict_lower = torch.tril(torch.ones(C, C, device=k.device), diagonal=-1)
+    g_diff = g_c.unsqueeze(-1) - g_c.unsqueeze(-2)
+    g_diff = g_diff * strict_lower
     A = kkt * beta_c.unsqueeze(-1) * torch.exp(g_diff) * strict_lower
     return A.permute(0, 1, 3, 2, 4).reshape(B, T, H, C).to(torch.float32)
 

problems/helion/gated_deltanet_recompute_w_u_py/reference.py

@@ -19,8 +19,9 @@ def _chunk_scaled_dot_kkt_fwd_eager(k, g_cumsum, beta, chunk_size):
     g_c = g_cumsum.float().reshape(B, NT, C, H).permute(0, 1, 3, 2)
     beta_c = beta.float().reshape(B, NT, C, H).permute(0, 1, 3, 2)
     kkt = k_c @ k_c.transpose(-1, -2)
-    g_diff = g_c.unsqueeze(-1) - g_c.unsqueeze(-2)
     strict_lower = torch.tril(torch.ones(C, C, device=k.device), diagonal=-1)
+    g_diff = g_c.unsqueeze(-1) - g_c.unsqueeze(-2)
+    g_diff = g_diff * strict_lower
     A = kkt * beta_c.unsqueeze(-1) * torch.exp(g_diff) * strict_lower
     return A.permute(0, 1, 3, 2, 4).reshape(B, T, H, C).to(torch.float32)