fix(gdn): use block-end decay for SM100 state updates

flashinfer/gdn_kernels/blackwell/gated_delta_net_chunked.py

@@ -4387,8 +4387,6 @@ def compute_group_1_chunk(
         tCsO = thr_o_r2s.partition_D(sO)
 
         sub_tile_size = 32
-        max_coord = tTR_tCcShared[cute.size(tTR_tCcShared) - 1]
-
         sV_vt_view = utils.gemm.sm100.transform_partitioned_tensor_layout(sV)
         tCsV = thr_v_s2r.partition_S(sV_vt_view)
 
@@ -4409,7 +4407,10 @@ def compute_group_1_chunk(
         # the per-row gate work until the previous state has been published
         # and decayed, keeping its register fragment in one contiguous region.
         gate_handle = load_gate_consumer.wait_and_advance()
-        cumprod_total = sCumprod[max_coord[1], 0, gate_handle.index]
+        # valid_len = max(0, min(self.b_t, seqlen_b - chunk_iter * self.b_t))
+        # gamma_end = 1.0 if valid_len == 0 else sCumprod[valid_len - 1, 0, gate_handle.index]
+        # OOB alpha is padded with 1 before the prefix scan, so the last physical slot equals gamma_end.
+        cumprod_total = sCumprod[self.b_t - 1, 0, gate_handle.index]
 
         kv_prev_handle = kv_acc_consumer.current_handle()
         if valid_state:

tests/gdn/test_prefill_delta_rule.py

@@ -171,6 +171,67 @@ def _test_prefill_kernel(
     torch.testing.assert_close(our_state, ref_state, atol=atol_kv, rtol=rtol_kv)
 
 
+@torch.inference_mode()
+def test_prefill_block_end_decay(qkv_factory, seed=0):
+    _skip_if_unsupported()
+    random.seed(seed)
+    torch.random.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+    seq_lens = [64, 111, 192]
+    total_seqlen = sum(seq_lens)
+    num_heads = 1
+    head_size = 128
+    dtype = torch.float16
+    device = torch.device("cuda")
+
+    with device:
+        q, k, v = qkv_factory(
+            seq_lens, num_heads, num_heads, num_heads, head_size, dtype
+        )
+        k = torch.nn.functional.normalize(k, p=2.0, dim=-1)
+        alpha = 0.99 + 0.01 * torch.rand(total_seqlen, num_heads)
+        beta = 0.99 + 0.01 * torch.rand(total_seqlen, num_heads)
+        cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int64)
+
+    our_o = torch.empty_like(q)
+    our_state = torch.empty(
+        (len(seq_lens), num_heads, head_size, head_size),
+        dtype=torch.float32,
+        device=device,
+    )
+    chunk_gated_delta_rule(
+        q,
+        k,
+        v,
+        alpha,
+        beta,
+        1.0,
+        None,
+        True,
+        cu_seqlens,
+        True,
+        output=our_o,
+        output_state=our_state,
+        use_cp=False,
+    )
+
+    ref_o, ref_state = blockwise_delta_rule(
+        q.float(),
+        k.float(),
+        v.float(),
+        seq_lens,
+        alpha=alpha,
+        beta=beta,
+        block_size=64,
+        state_dtype=torch.float32,
+    )
+    torch.testing.assert_close(our_o, ref_o.to(dtype), atol=2e-3, rtol=1e-3)
+    torch.testing.assert_close(
+        our_state.transpose(-1, -2), ref_state, atol=1e-3, rtol=1e-4
+    )
+
+
 @pytest.mark.parametrize("beta", [False, True])
 @pytest.mark.parametrize("alpha", [False, True])
 @pytest.mark.parametrize("scale", [1.0, "auto"])