Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ZJY0516 merged 10 commits into from
Oct 5, 2026
Merged

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
09d8cde
[Bugfix][GLM-5.3-Flash] fused_recurrent KDA/GDN kernels: keep batch x…
NolenLiang Sep 13, 2026
3a0cf44
[Test][GLM-5.3-Flash] Regression test: fused recurrent KDA with more …
NolenLiang Sep 15, 2026
3bb5fdc
[Bugfix][GLM-5.3-Flash] Limit the gridDim.z fix to the NVIDIA GLM-5.3…
NolenLiang Sep 22, 2026
b281d66
Merge branch 'main' into kda-grid-z-fix
mergify[bot] Sep 28, 2026
f2c5651
Merge main and fix GLM KDA regression CI
NolenLiang Sep 29, 2026
8d96809
Merge branch 'main' into kda-grid-z-fix
mergify[bot] Sep 30, 2026
99bab4b
Merge branch 'main' into kda-grid-z-fix
mergify[bot] Sep 30, 2026
b2b249e
Merge branch 'main' into kda-grid-z-fix
mergify[bot] Oct 1, 2026
3201f0f
Merge branch 'main' into kda-grid-z-fix
mergify[bot] Oct 2, 2026
1fe2ac9
Merge branch 'main' into kda-grid-z-fix
ZJY0516 Oct 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 69 additions & 0 deletions tests/models/glm5next/test_kda_recurrent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
6 changes: 4 additions & 2 deletions vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion vllm/models/glm5next/amd/ops/third_party/kda/kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions vllm/third_party/flash_linear_attention/ops/fused_recurrent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@ZJY0516 ZJY0516 Sep 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we don't change these unrelated files

@NolenLiang NolenLiang Sep 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reverted in the latest commit — the PR now only touches the GLM-5.3-Flash NVIDIA copy under models/glm5next/nvidia/ops and its test. Details in the thread above.

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:
Expand Down Expand Up @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(),
Expand Down
4 changes: 3 additions & 1 deletion vllm/third_party/flash_linear_attention/ops/kda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@simondanielsson simondanielsson Sep 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: Do we see this error only for GLM-5.3-Flash or for all models using KDA?

@JaredforReal JaredforReal Sep 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same question here, I have encoutered this issue with GLM-5.3-Flash, but did't try or see any other model with this issue
@NolenLiang

@NolenLiang NolenLiang Sep 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have reproduced the end-to-end failure only with GLM-5.3-Flash, in the DP configuration we tested (TP=1 per rank, default max_num_seqs 1024 on ≥160 GB GPUs, so the largest capture batch is 1024 × 64 heads = 65536). The underlying limit is not model-specific: every launch site changed here puts N * HV in gridDim.z, where N is the sequence count of that kernel call and HV the per-rank value-head count, and CUDA caps gridDim.z at 65535.

Direct kernel tests on the 2026-09-11 nightly (synthetic inputs, one GB200, wrappers called directly): the stock fused_sigmoid_gating_delta_rule_update, fused_recurrent_gated_delta_rule and fused_recurrent_kda wrappers each fail with Triton Error [CUDA]: invalid argument at 1024 × 64, 2048 × 32 and 2048 × 48 (sequences × heads) and pass at 512 × 64; with this PR's files all calls succeed and the full batch is bitwise identical to the same batch run as two halves. These are kernel-level tests, not end-to-end tests of other models.

Whether another model exposes the issue depends on which kernel its dispatch path reaches and on the sequence and head counts of those calls. From code reading only (not tested): the Qwen GDN layer sends uniform non-spec decode batches to fused_recurrent_gated_delta_rule_packed_decode, which already guards B * HV > 65535, and reaches the changed kernel for spec-decode batches on the Triton path, for the decode tokens of mixed prefill+decode steps, and for all decode when VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=0; the Kimi-K3 / Kimi-Linear KDA copy under models/kimi_k3 is separate and not changed here — its recurrent kernel uses a 1-D grid, while its packed-decode kernel puts B * H in gridDim.y, which has the same 65535 limit. A direct call to that packed-decode kernel fails with the same invalid argument at 2048 × 32 and 1024 × 64 (gridDim.y = 65536) and passes at 1024 × 32, so it is the same class of limit; it is outside this PR and we have not hit it end to end. We have not verified any of this end to end for other models.

Happy to retitle the PR to say "shared KDA/GDN fused recurrent kernels" with GLM-5.3-Flash as the verified reproducer, if you prefer.

@simondanielsson simondanielsson Sep 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One option is to do one of the following:

  1. Keep current changes but also validate the blast radius is small by testing this on another model like k3 and TP>1
  2. Narrow the change so it only affects glm-5.3-flash tp1 (or models with the same kernel shapes).
    WDYT?

@NolenLiang NolenLiang Sep 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Went with option 2. The PR is now limited to the NVIDIA GLM-5.3-Flash copy: vllm/models/glm5next/nvidia/ops/third_party/kda/{kernels.py,fused_recurrent.py} plus the regression test. The shared third_party/flash_linear_attention files are reverted (per @ZJY0516's comment), and I also took the AMD GLM copy out — not something the review asked for, but we have no AMD reproduction or hardware to validate it, and it is not needed for the failure we confirmed. Both are left for separate follow-up with their own validation; the kernel-level results from the earlier comment stand as the pointer for whoever picks that up.

Re-ran the kernel regression test on the 2026-09-11 nightly image (4× GB200) with only the two NVIDIA GLM kernel files copied over the installed package (the shared flash-linear-attention wrappers left as installed): stock 2 failed / 6 passed in tests/models/glm5next/test_kda_recurrent.py (the two new > 65535-program cases fail with Triton Error [CUDA]: invalid argument), with the two files 8 passed. The end-to-end numbers in the description (TP4 with the patch matching stock TP4 throughput, DP4+EP serving, gsm8k) are from before the narrowing and were not repeated; the NVIDIA GLM kernel files are the same in both versions. Scope note: this copy serves every GLM-5.3-Flash TP configuration, not only TP=1; the failure itself needs TP=1 per rank at the default capture sizes.

Done as an additional commit rather than a rewrite so the review history stays intact. PR title and description updated accordingly.

@simondanielsson simondanielsson Sep 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am happy to try this PR on AMD MI350 so we can keep the rocm code. Just LMK once it's re-added, thanks.

Or if you rather have this merged ASAP I can open an AMD version of the PR

fused_recurrent_gated_delta_rule_fwd_kernel[grid](
q=q,
k=k,
Expand Down
Loading