From 7e7ad35c186ba2c64e0a04149eb08aa1fed91e9d Mon Sep 17 00:00:00 2001
From: Artem Perevedentsev <aperevedents@nvidia.com>
Date: Fri, 24 Apr 2026 08:55:39 +0300
Subject: [PATCH] fix(gdn): address remaining CodeRabbit feedback from #3001

---
 flashinfer/gdn_prefill.py | 38 ++++++++++++++++++++------------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/flashinfer/gdn_prefill.py b/flashinfer/gdn_prefill.py
index 124784ff22..71f3101ab4 100644
--- a/flashinfer/gdn_prefill.py
+++ b/flashinfer/gdn_prefill.py
@@ -269,23 +269,8 @@ def chunk_gated_delta_rule(
             device=q.device,
         )
 
-    # Allocate output_state if needed
-    if output_final_state and output_state is None:
-        output_state = torch.empty(
-            (num_seqs, num_sab_heads, head_size, head_size),
-            dtype=torch.float32,
-            device=q.device,
-        )
-    elif not output_final_state and output_state is None:
-        # Still need to allocate since kernel always writes state
-        output_state = torch.empty(
-            (num_seqs, num_sab_heads, head_size, head_size),
-            dtype=torch.float32,
-            device=q.device,
-        )
-
     device = q.device
-    _scale = scale if scale is not None else 1.0 / math.sqrt(head_size)
+    _scale = scale if scale is not None and scale != 0.0 else 1.0 / math.sqrt(head_size)
 
     _cuda_major = int(torch.version.cuda.split(".")[0]) if torch.version.cuda else 0
     if _has_blackwell_prefill and is_sm100a_supported(device) and _cuda_major >= 13:
@@ -294,6 +279,16 @@ def chunk_gated_delta_rule(
             f"Blackwell GDN prefill requires head_size=128, got {head_size}"
         )
 
+        # Allocate output_state only when needed
+        if not output_final_state:
+            output_state = None
+        elif output_state is None:
+            output_state = torch.empty(
+                (num_seqs, num_sab_heads, head_size, head_size),
+                dtype=torch.float32,
+                device=device,
+            )
+
         _g = (
             g
             if g is not None
@@ -323,7 +318,7 @@ def chunk_gated_delta_rule(
             output,
             cu_seqlens.to(torch.int32),
             initial_state,
-            output_state if output_final_state else None,
+            output_state,
             _scale,
             checkpoint_every_n_tokens=checkpoint_every_n_tokens,
             cu_checkpoints=_cu_checkpoints,
@@ -331,6 +326,13 @@ def chunk_gated_delta_rule(
         )
     else:
         # SM90 Hopper path (C++ JIT kernel)
+        if output_state is None:
+            output_state = torch.empty(
+                (num_seqs, num_sab_heads, head_size, head_size),
+                dtype=torch.float32,
+                device=device,
+            )
+
         workspace_size = get_device_sm_count(device) * 128
         workspace_buffer = _get_cache_buf(
             "gdn_prefill_workspace", workspace_size, device
@@ -346,7 +348,7 @@ def chunk_gated_delta_rule(
             initial_state,
             g,
             beta,
-            scale if scale is not None else 0.0,
+            _scale,
             workspace_buffer,
             state_checkpoints,
             checkpoint_cu_starts.to(torch.int64)