From 09d8cdeb6c966231d4da54853c3f4fa8def2bd3c Mon Sep 17 00:00:00 2001 From: NolenLiang Date: Sun, 13 Sep 2026 09:40:52 -0700 Subject: [PATCH 1/3] [Bugfix][GLM-5.3-Flash] fused_recurrent KDA/GDN kernels: keep batch x heads out of gridDim.z fused_recurrent_gated_delta_rule_fwd_kernel was launched with grid = (NK, NV, N * HV). CUDA caps gridDim.z at 65535; with GLM-5.3-Flash at TP=1 (every DP rank) the first CUDA-graph capture batch of 1024 times 64 KDA heads is 65536, so cuLaunchKernel fails with 'invalid argument' and the engine cannot start (DP4, DP4+EP, TP1 all crash at capture). Put N * HV in gridDim.x (limit 2^31-1) and swap the program_id mapping in the kernel. Same change mirrored in the AMD copy and the flash_linear_attention copies which share the pattern. Signed-off-by: NolenLiang --- .../glm5next/amd/ops/third_party/kda/fused_recurrent.py | 6 ++++-- vllm/models/glm5next/amd/ops/third_party/kda/kernels.py | 4 +++- .../glm5next/nvidia/ops/third_party/kda/fused_recurrent.py | 6 ++++-- vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py | 4 +++- .../flash_linear_attention/ops/fused_recurrent.py | 6 ++++-- .../flash_linear_attention/ops/fused_sigmoid_gating.py | 6 ++++-- vllm/third_party/flash_linear_attention/ops/kda.py | 4 +++- 7 files changed, 25 insertions(+), 11 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py index d7de85984f85..782f2edcdc58 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py @@ -65,7 +65,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: @@ -238,7 +238,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/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py index be417d6b0526..7c6940fe3cb0 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py @@ -92,7 +92,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, 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 152412b8ca8f..e3f32f721d6e 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 b792f606d4e0..2e22850ff68e 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, diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index eb08b938c2fb..8b400a349976 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -59,7 +59,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( IS_SPEC_DECODING: tl.constexpr, IS_KDA: 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: @@ -214,7 +214,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/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py index 7e0c7e05cab1..b7a207e21de4 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py @@ -59,7 +59,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( IS_SPEC_DECODING: tl.constexpr, IS_KDA: 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: @@ -238,7 +238,9 @@ def fused_sigmoid_gating_delta_rule_update( 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_sigmoid_gating_delta_rule_update_kernel[grid]( A_log=A_log, a=a.contiguous(), diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index 548f600812e2..b63b8216d022 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -68,7 +68,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, From 3a0cf4414da329b357cda786e6dcbea652c75dd8 Mon Sep 17 00:00:00 2001 From: NolenLiang Date: Mon, 14 Sep 2026 22:18:49 -0700 Subject: [PATCH 2/3] [Test][GLM-5.3-Flash] Regression test: fused recurrent KDA with more than 65535 (sequence, head) programs Adds the crashing shape (1024 sequences x 64 heads, and 2048 x 48) to tests/models/glm5next/test_kda_recurrent.py with small head dims; the full batch must be bitwise identical to the same batch run in two halves. Signed-off-by: NolenLiang --- tests/models/glm5next/test_kda_recurrent.py | 69 +++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index 68cf493b443a..b17c3fe9f0ef 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -184,3 +184,72 @@ def test_fused_recurrent_kda_rejects_unaddressable_layouts(): broken["cu_seqlens"] = None if q.shape[0] > 1 else inputs["cu_seqlens"] with pytest.raises(AssertionError, match=r"torch.Size"): run_kernel(broken, state) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@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.cuda.synchronize() + assert torch.equal(full_out, ref_out) + assert torch.equal(full_state, ref_state) From 3bb5fdc99f69fb12599c56b4a462188890eff8ea Mon Sep 17 00:00:00 2001 From: NolenLiang Date: Mon, 21 Sep 2026 22:56:08 -0700 Subject: [PATCH 3/3] [Bugfix][GLM-5.3-Flash] Limit the gridDim.z fix to the NVIDIA GLM-5.3-Flash KDA kernels Revert the launch-grid change in the shared flash-linear-attention copies (fused_recurrent.py, fused_sigmoid_gating.py, kda.py), as requested in review, and also in the AMD GLM copy, to keep this change to the path on which the crash was reproduced and verified (models/glm5next/nvidia/ops). The shared and AMD kernels keep their previous launch grid; their N * HV > 65535 case is left for separate follow-up with its own validation. Signed-off-by: NolenLiang --- .../glm5next/amd/ops/third_party/kda/fused_recurrent.py | 6 ++---- vllm/models/glm5next/amd/ops/third_party/kda/kernels.py | 4 +--- .../flash_linear_attention/ops/fused_recurrent.py | 6 ++---- .../flash_linear_attention/ops/fused_sigmoid_gating.py | 6 ++---- vllm/third_party/flash_linear_attention/ops/kda.py | 4 +--- 5 files changed, 8 insertions(+), 18 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py index 782f2edcdc58..d7de85984f85 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py @@ -65,7 +65,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( SAFE_GATE: tl.constexpr, # bounded gate variant (only branch implemented) LOWER_BOUND: tl.constexpr, ): - i_nh, i_v, i_k = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = 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: @@ -238,9 +238,7 @@ def fused_recurrent_gated_delta_rule_fwd( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - # 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) + grid = (NK, NV, N * HV) fused_recurrent_gated_delta_rule_fwd_kernel[grid]( q=q, k=k, diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py index 7c6940fe3cb0..be417d6b0526 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py @@ -92,9 +92,7 @@ def fused_recurrent_kda_fwd( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - # 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) + grid = (NK, NV, N * HV) fused_recurrent_gated_delta_rule_fwd_kernel[grid]( q=q, k=k, diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index 8b400a349976..eb08b938c2fb 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -59,7 +59,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( IS_SPEC_DECODING: tl.constexpr, IS_KDA: tl.constexpr, ): - i_nh, i_v, i_k = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = 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: @@ -214,9 +214,7 @@ def fused_recurrent_gated_delta_rule_fwd( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - # 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) + grid = (NK, NV, N * HV) fused_recurrent_gated_delta_rule_fwd_kernel[grid]( q=q, k=k, diff --git a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py index b7a207e21de4..7e0c7e05cab1 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py @@ -59,7 +59,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( IS_SPEC_DECODING: tl.constexpr, IS_KDA: tl.constexpr, ): - i_nh, i_v, i_k = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = 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: @@ -238,9 +238,7 @@ def fused_sigmoid_gating_delta_rule_update( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - # 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) + grid = (NK, NV, N * HV) fused_sigmoid_gating_delta_rule_update_kernel[grid]( A_log=A_log, a=a.contiguous(), diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index b63b8216d022..548f600812e2 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -68,9 +68,7 @@ def fused_recurrent_kda_fwd( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - # 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) + grid = (NK, NV, N * HV) fused_recurrent_gated_delta_rule_fwd_kernel[grid]( q=q, k=k,