[Bugfix][GLM-5.3-Flash] Keep batch x heads out of gridDim.z in the GLM-5.3-Flash fused recurrent KDA kernel (#56974)

Signed-off-by: NolenLiang <nliang@nvidia.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>

tests/models/glm5next/test_kda_recurrent.py

@@ -194,6 +194,77 @@ def test_fused_recurrent_kda_rejects_unaddressable_layouts():
             run_kernel(broken, state)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="The large-grid fix is NVIDIA-only"
+)
+@pytest.mark.parametrize("num_seqs,num_heads", [(1024, 64), (2048, 48)])
+def test_fused_recurrent_kda_more_than_65535_programs(num_seqs: int, num_heads: int):
+    """Regression test for the launch grid.
+
+    GLM-5.3-Flash has 64 KDA value heads; at TP=1 (every DP rank) a decode
+    batch of 1024 sequences needs 1024 * 64 = 65536 (sequence, head)
+    programs. The kernel used to put that product into gridDim.z, whose CUDA
+    limit is 65535, so the launch failed with ``Triton Error [CUDA]: invalid
+    argument`` at CUDA-graph capture and DP deployments could not start. The
+    grid now carries the product in gridDim.x. Small head dims keep the test
+    cheap; the shape that matters is the number of programs. Every (sequence,
+    head) pair is independent, so the full batch must be bitwise identical to
+    the same batch run in two halves.
+    """
+    torch.manual_seed(0)
+    dev = torch.device("cuda")
+    head_dim = 16
+    pool = num_seqs + 8
+    q = torch.randn(1, num_seqs, num_heads, head_dim, device=dev, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    g = torch.randn_like(q)
+    beta = torch.randn(1, num_seqs, num_heads, device=dev, dtype=torch.bfloat16)
+    state = torch.randn(
+        pool, num_heads, head_dim, head_dim, device=dev, dtype=torch.float32
+    )
+    cu = torch.arange(0, num_seqs + 1, device=dev, dtype=torch.int32)
+    # slot 0 is NULL_BLOCK_ID and is skipped by the kernel: use slots 1..
+    idx = (torch.randperm(pool - 1, device=dev)[:num_seqs] + 1).to(torch.int32)
+    a_log = torch.randn(num_heads, device=dev, dtype=torch.float32)
+    g_bias = torch.randn(num_heads, head_dim, device=dev, dtype=torch.float32)
+
+    def run(sl: slice, cu_: torch.Tensor, st: torch.Tensor, out: torch.Tensor):
+        fused_recurrent_kda(
+            q=q[:, sl],
+            k=k[:, sl],
+            v=v[:, sl],
+            g=g[:, sl],
+            beta=beta[:, sl],
+            initial_state=st,
+            use_qk_l2norm_in_kernel=True,
+            cu_seqlens=cu_,
+            ssm_state_indices=idx[sl],
+            sigmoid_beta=True,
+            a_log=a_log,
+            g_bias=g_bias,
+            compute_gate=True,
+            lower_bound=-5.0,
+            out=out,
+        )
+
+    full_state = state.clone()
+    full_out = torch.zeros_like(k)
+    run(slice(0, num_seqs), cu, full_state, full_out)
+
+    half = num_seqs // 2
+    ref_state = state.clone()
+    ref_out = torch.zeros_like(k)
+    for s, e in ((0, half), (half, num_seqs)):
+        cu_h = torch.arange(0, e - s + 1, device=dev, dtype=torch.int32)
+        part = torch.zeros_like(k[:, s:e])
+        run(slice(s, e), cu_h, ref_state, part)
+        ref_out[:, s:e] = part
+    torch.accelerator.synchronize()
+    assert torch.equal(full_out, ref_out)
+    assert torch.equal(full_state, ref_state)
+
+
 @pytest.mark.parametrize("offset", [16, 144, 256])
 @pytest.mark.parametrize("dim_first", [False, True])
 @pytest.mark.parametrize("num_spec", [0, 3])

vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py

@@ -87,7 +87,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
     SAFE_GATE: tl.constexpr,  # bounded gate variant (only branch implemented)
     LOWER_BOUND: tl.constexpr,
 ):
-    i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_nh, i_v, i_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
     i_n, i_hv = i_nh // HV, i_nh % HV
     i_h = i_hv // (HV // H)
     if IS_VARLEN:
@@ -260,7 +260,9 @@ def fused_recurrent_gated_delta_rule_fwd(
     else:
         stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
 
-    grid = (NK, NV, N * HV)
+    # N * HV goes in gridDim.x: gridDim.z is capped at 65535 and batch x heads exceeds it
+    # (e.g. GLM-5.3-Flash at TP=1: 1024 x 64 = 65536 -> "invalid argument" at CUDA-graph capture).
+    grid = (N * HV, NV, NK)
     fused_recurrent_gated_delta_rule_fwd_kernel[grid](
         q=q,
         k=k,

vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py

@@ -95,7 +95,9 @@ def fused_recurrent_kda_fwd(
     else:
         stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
 
-    grid = (NK, NV, N * HV)
+    # N * HV goes in gridDim.x: gridDim.z is capped at 65535 and batch x heads exceeds it
+    # (e.g. GLM-5.3-Flash at TP=1: 1024 x 64 = 65536 -> "invalid argument" at CUDA-graph capture).
+    grid = (N * HV, NV, NK)
     fused_recurrent_gated_delta_rule_fwd_kernel[grid](
         q=q,
         k=k,