[Bugfix][Kernel] Keep packed GDN decode beta in FP32 (#53877)
Signed-off-by: CherryLemon <lemon_cn@alumni.sjtu.edu.cn> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Yongji Wu committed
Aug 30, 2026 at 22:15 UTC
56058fd572f6a7fec6899385f4a4ed7f4b964477
2 files changed
+33
-2
tests/kernels/test_fused_recurrent_packed_decode.py
+32
-1
@@ -70,7 +70,7 @@ def test_fused_recurrent_packed_decode_matches_reference(
70
x <= 20.0, torch.log1p(torch.exp(torch.clamp(x, max=20.0))), x
71
)
72
g = (-torch.exp(A_log.float()) * softplus_x).unsqueeze(1)
73
- beta = torch.sigmoid(b.float()).to(dtype).unsqueeze(1)
73
+ beta = torch.sigmoid(b.float()).unsqueeze(1)
74
75
out_ref, state_ref = fused_recurrent_gated_delta_rule(
76
q=q,
@@ -109,6 +109,37 @@ def test_fused_recurrent_packed_decode_matches_reference(
109
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
110
111
112
+def test_packed_decode_keeps_beta_in_fp32():
113
+ device = torch.device(DEVICE)
114
+ dtype = torch.bfloat16
115
+
116
+ mixed_qkv = torch.ones((1, 3), device=device, dtype=dtype)
117
+ a = torch.zeros((1, 1), device=device, dtype=dtype)
118
+ b = torch.full((1, 1), 0.5, device=device, dtype=dtype)
119
+ params = torch.zeros((1,), device=device, dtype=dtype)
120
+ ssm_state_indices = torch.ones((1,), device=device, dtype=torch.int32)
121
+
122
+ state_packed = torch.zeros((2, 1, 1, 1), device=device, dtype=torch.float32)
123
+ out_packed = torch.empty((1, 1, 1, 1), device=device, dtype=dtype)
124
+
125
+ fused_recurrent_gated_delta_rule_packed_decode(
126
+ mixed_qkv=mixed_qkv,
127
+ a=a,
128
+ b=b,
129
+ A_log=params,
130
+ dt_bias=params,
131
+ scale=1.0,
132
+ initial_state=state_packed,
133
+ out=out_packed,
134
+ ssm_state_indices=ssm_state_indices,
135
+ )
136
+
137
+ expected_beta = torch.sigmoid(b.float()).squeeze()
138
+ torch.testing.assert_close(
139
+ state_packed[1, 0, 0, 0], expected_beta, rtol=1e-6, atol=1e-6
140
+ )
141
+
142
+
143
def test_packed_decode_supports_large_batch_head_grid():
144
B, H, HV, K, V = 1024, 8, 64, 1, 1
145
device = torch.device(DEVICE)
vllm/third_party/flash_linear_attention/ops/fused_recurrent.py
+1
-1
@@ -326,7 +326,7 @@ def fused_recurrent_gated_delta_rule_packed_decode_kernel(
326
x = a_val + dt_bias_val
327
softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
328
g_val = -tl.exp(A_log_val) * softplus_x
329
- beta_val = tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32)
329
+ beta_val = tl.sigmoid(b_val)
330
331
b_h *= exp(g_val)
332
b_v -= tl.sum(b_h * b_k[None, :], 1)