[Bugfix] Fix packed GDN decode launch for large batch-head grids (#52030)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 13, 2026 at 10:19 UTC
2d24355eb87b716fc1169e66731dc0386ed1a3a2
2 files changed
+33
-3
tests/kernels/test_fused_recurrent_packed_decode.py
+23
@@ -100,3 +100,26 @@ def test_fused_recurrent_packed_decode_matches_reference(
100
valid = ssm_state_indices > 0
101
torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
102
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
103
+
104
+
105
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")
106
+def test_packed_decode_supports_large_batch_head_grid():
107
+ B, H, HV, K, V = 1024, 8, 64, 1, 1
108
+ device = torch.device("cuda")
109
+ gates = torch.empty((B, HV), device=device)
110
+ params = torch.empty((HV,), device=device)
111
+ out = torch.empty((B, 1, HV, V), device=device)
112
+
113
+ fused_recurrent_gated_delta_rule_packed_decode(
114
+ mixed_qkv=torch.empty((B, 2 * H * K + HV * V), device=device),
115
+ a=gates,
116
+ b=gates,
117
+ A_log=params,
118
+ dt_bias=params,
119
+ scale=1.0,
120
+ initial_state=torch.empty((1, HV, V, K), device=device),
121
+ out=out,
122
+ ssm_state_indices=torch.zeros((B,), device=device, dtype=torch.int32),
123
+ )
124
+
125
+ assert torch.count_nonzero(out).item() == 0
vllm/third_party/flash_linear_attention/ops/fused_recurrent.py
+10
-3
@@ -278,9 +278,13 @@ def fused_recurrent_gated_delta_rule_packed_decode_kernel(
278
BV: tl.constexpr,
279
SOFTPLUS_THRESHOLD: tl.constexpr,
280
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
281
+ SPLIT_BATCH_HEAD_GRID: tl.constexpr,
282
):
282
- i_v, i_nh = tl.program_id(0), tl.program_id(1)
283
- i_n, i_hv = i_nh // HV, i_nh % HV
283
+ if SPLIT_BATCH_HEAD_GRID:
284
+ i_v, i_hv, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2)
285
+ else:
286
+ i_v, i_nh = tl.program_id(0), tl.program_id(1)
287
+ i_n, i_hv = i_nh // HV, i_nh % HV
288
i_h = i_hv // (HV // H)
289
290
o_k = tl.arange(0, BK)
@@ -446,7 +450,9 @@ def fused_recurrent_gated_delta_rule_packed_decode(
450
stride_indices_seq = ssm_state_indices.stride(0)
451
452
NV = triton.cdiv(V, BV)
449
- grid = (NV, B * HV)
453
+ # CUDA limits grid Y/Z dimensions to 65535.
454
+ split_batch_head_grid = B * HV > 65535
455
+ grid = (NV, HV, B) if split_batch_head_grid else (NV, B * HV)
456
fused_recurrent_gated_delta_rule_packed_decode_kernel[grid](
457
mixed_qkv=mixed_qkv,
458
a=a,
@@ -472,6 +478,7 @@ def fused_recurrent_gated_delta_rule_packed_decode(
478
BV=BV,
479
SOFTPLUS_THRESHOLD=20.0,
480
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
481
+ SPLIT_BATCH_HEAD_GRID=split_batch_head_grid,
482
num_warps=num_warps,
483
num_stages=num_stages,
484
)