From ae53b06898bcd4ae13553eb019f7381dab90edb2 Mon Sep 17 00:00:00 2001 From: bcsdhjew Date: Mon, 5 Oct 2026 17:31:45 +0800 Subject: [PATCH] [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 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Codex --- tests/models/glm5next/test_kda_recurrent.py | 71 +++++++++++++++++++ .../ops/third_party/kda/fused_recurrent.py | 6 +- .../nvidia/ops/third_party/kda/kernels.py | 4 +- 3 files changed, 78 insertions(+), 3 deletions(-) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index fe017e4131..06b2eafac6 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/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]) diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py index a9d0fba476..767b822170 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py +++ b/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, diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py index 9805b37471..8f662f68aa 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py +++ b/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,