[Perf][GLM-5.3-Flash] Decode hot-path cleanups: strided KDA recurrent inputs, NoPE MQA query without concat, no duplicate router GEMM (#55736)
Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Jared Wen committed
Sep 10, 2026 at 16:28 UTC
1768273c13d8e083aa07d48f0fb79c05cec39df5
8 files changed
+268
-27
.buildkite/test_areas/models_basic.yaml
+10
@@ -86,6 +86,16 @@ steps:
86
# The native NVIDIA Kimi K3 kernels require the SM100 family.
87
- pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
88
89
+- label: ":nvidia: (H200) GLM5Next"
90
+ key: glm5next-unit-tests
91
+ timeout_in_minutes: 20
92
+ device: h200_35gb
93
+ source_file_dependencies:
94
+ - vllm/models/glm5next/
95
+ - tests/models/glm5next/
96
+ commands:
97
+ - pytest -v -s models/glm5next
98
+
99
- label: ":nvidia: (H200) Qwen4 Exp"
100
key: qwen4-exp-unit-tests
101
timeout_in_minutes: 40
tests/models/glm5next/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/models/glm5next/test_kda_recurrent.py
new
+186
@@ -0,0 +1,186 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""GLM-5.3-Flash KDA recurrent (decode) kernel.
4
+
5
+The decode path hands the kernel column slices of the merged ``q|k|v`` conv
6
+output and of the fused ``qkvbfg_a`` projection (beta), so q/k/v/beta are
7
+token-strided rather than contiguous. The kernel must read them in place,
8
+match a pure-PyTorch recurrence, and reject layouts it cannot address.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+
14
+from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
15
+from vllm.platforms import current_platform
16
+
17
+pytestmark = pytest.mark.skipif(
18
+ not current_platform.is_cuda(), reason="CUDA-only Triton kernel"
19
+)
20
+
21
+H, D = 16, 128
22
+LOWER_BOUND = -5.0
23
+
24
+
25
+def naive_recurrent_kda(
26
+ q: torch.Tensor,
27
+ k: torch.Tensor,
28
+ v: torch.Tensor,
29
+ raw_g: torch.Tensor,
30
+ raw_beta: torch.Tensor,
31
+ a_log: torch.Tensor,
32
+ g_bias: torch.Tensor,
33
+ state: torch.Tensor,
34
+) -> tuple[torch.Tensor, torch.Tensor]:
35
+ """fp32 reference for one sequence: ``[T, H, D]`` inputs, ``[H, D, D]``
36
+ (v-major) state; mirrors the kernel's in-kernel gate, beta sigmoid and
37
+ q/k l2norm.
38
+ """
39
+ q, k, v, raw_g, raw_beta = (x.float() for x in (q, k, v, raw_g, raw_beta))
40
+ q = q / torch.sqrt(q.square().sum(-1, keepdim=True) + 1e-6) * D**-0.5
41
+ k = k / torch.sqrt(k.square().sum(-1, keepdim=True) + 1e-6)
42
+ gate = LOWER_BOUND * torch.sigmoid(a_log.exp()[:, None] * (raw_g + g_bias))
43
+ beta = torch.sigmoid(raw_beta)
44
+ s = state.clone()
45
+ out = torch.empty_like(v)
46
+ for t in range(q.shape[0]):
47
+ s = s * gate[t].exp()[:, None, :]
48
+ u = beta[t][:, None] * (v[t] - torch.einsum("hvk,hk->hv", s, k[t]))
49
+ s = s + u[:, :, None] * k[t][:, None, :]
50
+ out[t] = torch.einsum("hvk,hk->hv", s, q[t])
51
+ return out, s
52
+
53
+
54
+def make_inputs(num_seqs: int, query_len: int, device: torch.device):
55
+ """Token-strided q/k/v/beta as the decode path produces them: column
56
+ slices of a merged ``[T, q|k|v]`` conv output and of the fused
57
+ ``[T, qkv|beta|f_a|g_a]`` projection.
58
+ """
59
+ T, proj = num_seqs * query_len, H * D
60
+ qkv = torch.randn(T, 3 * proj, dtype=torch.bfloat16, device=device)
61
+ projected = torch.randn(
62
+ T, 3 * proj + H + 2 * D, dtype=torch.bfloat16, device=device
63
+ )
64
+ q, k, v = (x.view(1, T, H, D) for x in qkv.split(proj, dim=-1))
65
+ beta = projected[:, 3 * proj : 3 * proj + H].unsqueeze(0)
66
+ # (A size-1 token dim gets an arbitrary stride from `view`.)
67
+ assert T == 1 or (q.stride(1) == 3 * proj and beta.stride(1) == projected.stride(0))
68
+ inputs = dict(
69
+ q=q,
70
+ k=k,
71
+ v=v,
72
+ g=torch.randn(1, T, H, D, dtype=torch.bfloat16, device=device),
73
+ beta=beta,
74
+ a_log=0.5 * torch.randn(H, dtype=torch.float32, device=device),
75
+ g_bias=0.1 * torch.randn(H * D, dtype=torch.float32, device=device),
76
+ cu_seqlens=torch.arange(0, T + 1, query_len, dtype=torch.int32, device=device),
77
+ )
78
+ # Slot 0 is NULL_BLOCK_ID; sequences own random distinct slots (one per
79
+ # token in the spec-decode layout).
80
+ slots = torch.randperm(T, device=device).to(torch.int32) + 1
81
+ if query_len == 1:
82
+ inputs["ssm_state_indices"] = slots
83
+ else:
84
+ inputs["ssm_state_indices"] = slots.view(num_seqs, query_len)
85
+ inputs["num_accepted_tokens"] = torch.randint(
86
+ 1, query_len + 1, (num_seqs,), dtype=torch.int32, device=device
87
+ )
88
+ state = torch.randn(T + 1, H, D, D, dtype=torch.float32, device=device)
89
+ return inputs, state
90
+
91
+
92
+def run_kernel(inputs: dict, state: torch.Tensor) -> torch.Tensor:
93
+ out, _ = fused_recurrent_kda(
94
+ **inputs,
95
+ initial_state=state,
96
+ use_qk_l2norm_in_kernel=True,
97
+ sigmoid_beta=True,
98
+ compute_gate=True,
99
+ lower_bound=LOWER_BOUND,
100
+ )
101
+ return out
102
+
103
+
104
+@pytest.mark.parametrize(
105
+ ("num_seqs", "query_len"), [(1, 1), (7, 1), (3, 3)], ids=["1x1", "7x1", "3x3"]
106
+)
107
+@torch.inference_mode()
108
+def test_fused_recurrent_kda_matches_reference(num_seqs: int, query_len: int):
109
+ torch.manual_seed(0)
110
+ device = torch.device("cuda")
111
+ inputs, state = make_inputs(num_seqs, query_len, device)
112
+ expected_state = state.clone()
113
+ out = run_kernel(inputs, state)
114
+
115
+ indices = inputs["ssm_state_indices"].view(num_seqs, query_len)
116
+ accepted = inputs.get("num_accepted_tokens")
117
+ expected = torch.empty_like(out[0], dtype=torch.float32)
118
+ for n in range(num_seqs):
119
+ first = indices[n, 0 if accepted is None else accepted[n] - 1]
120
+ s = expected_state[first]
121
+ for t in range(query_len):
122
+ tok = slice(n * query_len + t, n * query_len + t + 1)
123
+ expected[tok], s = naive_recurrent_kda(
124
+ inputs["q"][0, tok],
125
+ inputs["k"][0, tok],
126
+ inputs["v"][0, tok],
127
+ inputs["g"][0, tok],
128
+ inputs["beta"][0, tok],
129
+ inputs["a_log"],
130
+ inputs["g_bias"].view(H, D),
131
+ s,
132
+ )
133
+ expected_state[indices[n, t]] = s
134
+
135
+ torch.testing.assert_close(out[0].float(), expected, rtol=1e-2, atol=1e-3)
136
+ used = indices.flatten().long()
137
+ torch.testing.assert_close(state[used], expected_state[used], rtol=1e-5, atol=1e-5)
138
+
139
+
140
+@pytest.mark.parametrize(
141
+ ("num_seqs", "query_len"), [(7, 1), (3, 3)], ids=["7x1", "3x3"]
142
+)
143
+@torch.inference_mode()
144
+def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous(
145
+ num_seqs: int, query_len: int
146
+):
147
+ torch.manual_seed(0)
148
+ device = torch.device("cuda")
149
+ inputs, state = make_inputs(num_seqs, query_len, device)
150
+ for name in ("q", "k", "v", "beta"):
151
+ assert not inputs[name].is_contiguous()
152
+ contiguous = {
153
+ name: x.contiguous() if name in ("q", "k", "v", "beta") else x
154
+ for name, x in inputs.items()
155
+ }
156
+ state_ref = state.clone()
157
+ out_ref = run_kernel(contiguous, state_ref)
158
+ out = run_kernel(inputs, state)
159
+ torch.testing.assert_close(out, out_ref, rtol=0, atol=0)
160
+ torch.testing.assert_close(state, state_ref, rtol=0, atol=0)
161
+
162
+
163
+@torch.inference_mode()
164
+def test_fused_recurrent_kda_rejects_unaddressable_layouts():
165
+ """Layouts the token-stride addressing cannot express must fail loudly
166
+ rather than read the wrong tokens: a batch slice of a wider buffer
167
+ (``stride(0) != T * stride(1)``), overlapping tokens, and a head-strided
168
+ (transposed) block.
169
+ """
170
+ torch.manual_seed(0)
171
+ device = torch.device("cuda")
172
+ inputs, state = make_inputs(1, 1, device)
173
+ T = 4
174
+ base = torch.randn(4, T, 2 * H * D, dtype=torch.bfloat16, device=device)
175
+ bad_q = {
176
+ "batch slice": base[::2, :, : H * D].view(2, T, H, D),
177
+ "overlapping tokens": base[:1, :, : H * D].as_strided(
178
+ (1, T, H, D), (0, D, D, 1)
179
+ ),
180
+ "head-strided": base[:1, :, : H * D].view(1, T, D, H).transpose(2, 3),
181
+ }
182
+ for q in bad_q.values():
183
+ broken = dict(inputs, q=q, k=q, v=q)
184
+ broken["cu_seqlens"] = None if q.shape[0] > 1 else inputs["cu_seqlens"]
185
+ with pytest.raises(AssertionError, match=r"torch.Size"):
186
+ run_kernel(broken, state)
vllm/model_executor/layers/attention/mla_attention.py
+10
-7
@@ -998,14 +998,17 @@ class MLAAttention(nn.Module, AttentionLayerBase):
998
if self.q_pad_num_heads is not None:
999
mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
1000
mqa_ql_nope.resize_((N, B, L))
1001
+ # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1002
+ torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)
1003
+ # Convert from (N, B, L) to (B, N, L)
1004
+ mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
1005
else:
1002
- mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))
1003
-
1004
- # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1005
- torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)
1006
-
1007
- # Convert from (N, B, L) to (B, N, L)
1008
- mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
1006
+ # Write the (N, B, L) bmm result straight into a
1007
+ # token-major (B, N, L) buffer so the MQA query is already
1008
+ # contiguous; a NoPE model (qk_rope_head_dim == 0) then
1009
+ # needs no concat at all.
1010
+ mqa_ql_nope = mqa_q_nope.new_empty((B, N, L))
1011
+ torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope.transpose(0, 1))
1012
1013
if fp8_attention and self.impl.supports_quant_query_input:
1014
assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
vllm/models/glm5next/nvidia/model.py
+4
-4
@@ -259,11 +259,11 @@ class Glm5NextMoE(nn.Module):
259
if self.is_sequence_parallel and not already_sequence_parallel:
260
hidden_states = sequence_parallel_chunk(hidden_states)
261
262
- # The router is always external (self.gate); main's MoERunner expects
263
- # pre-computed router_logits, so compute them here unconditionally.
264
- router_logits, _ = self.gate(hidden_states)
262
+ # MoERunner holds the gate (passed to FusedMoEFactory) and computes
263
+ # the router logits itself, so nothing is precomputed here (matches
264
+ # DeepseekV2MoE; `router_logits` is a placeholder).
265
final_hidden_states = self.experts(
266
- hidden_states=hidden_states, router_logits=router_logits
266
+ hidden_states=hidden_states, router_logits=hidden_states
267
)
268
269
if self.is_sequence_parallel and not already_sequence_parallel:
vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py
+35
-9
@@ -15,6 +15,23 @@ from vllm.third_party.flash_linear_attention.ops.op import exp
15
from vllm.triton_utils import tl, triton
16
17
18
+def token_stride(x: torch.Tensor) -> int:
19
+ """Token stride (elements) of a ``[B, T, H, D]`` or ``[B, T, H]`` tensor.
20
+
21
+ The recurrent kernel walks tokens with this stride and addresses heads
22
+ densely inside a token, so each token's ``[H, D]`` (or ``[H]``) block must
23
+ be contiguous, tokens must not overlap, and with ``B > 1`` sequence ``n``
24
+ must start at token ``n * T`` (dense batch). Column slices of a wider
25
+ per-token projection buffer satisfy this and are consumed in place.
26
+ """
27
+ st = x.stride()
28
+ assert x.dim() in (3, 4) and st[-1] == 1, (x.shape, st)
29
+ assert x.dim() == 3 or st[2] == x.shape[3], (x.shape, st)
30
+ assert st[1] >= x.shape[2] * (x.shape[3] if x.dim() == 4 else 1), (x.shape, st)
31
+ assert x.shape[0] == 1 or st[0] == x.shape[1] * st[1], (x.shape, st)
32
+ return st[1]
33
+
34
+
35
@triton.heuristics(
36
{
37
"USE_INITIAL_STATE": lambda args: args["h0"] is not None,
@@ -52,6 +69,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
69
stride_final_state_token: tl.constexpr,
70
stride_indices_seq: tl.constexpr,
71
stride_indices_tok: tl.constexpr,
72
+ # Token strides of q/k/v/beta (elements), see `token_stride`.
73
+ stride_q_t,
74
+ stride_k_t,
75
+ stride_v_t,
76
+ stride_beta_t,
77
USE_INITIAL_STATE: tl.constexpr, # whether to use initial state
78
INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace
79
IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar,
@@ -86,13 +108,13 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
108
o_k = i_k * BK + tl.arange(0, BK)
109
o_v = i_v * BV + tl.arange(0, BV)
110
89
- p_q = q + (bos * H + i_h) * K + o_k
90
- p_k = k + (bos * H + i_h) * K + o_k
91
- p_v = v + (bos * HV + i_hv) * V + o_v
111
+ p_q = q + bos * stride_q_t + i_h * K + o_k
112
+ p_k = k + bos * stride_k_t + i_h * K + o_k
113
+ p_v = v + bos * stride_v_t + i_hv * V + o_v
114
if IS_BETA_HEADWISE:
93
- p_beta = beta + (bos * HV + i_hv) * V + o_v
115
+ p_beta = beta + bos * stride_beta_t + i_hv * V + o_v
116
else:
95
- p_beta = beta + bos * HV + i_hv
117
+ p_beta = beta + bos * stride_beta_t + i_hv
118
119
if not IS_KDA:
120
p_g = g + bos * HV + i_hv
@@ -188,15 +210,15 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
210
p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
211
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
212
191
- p_q += H * K
192
- p_k += H * K
213
+ p_q += stride_q_t
214
+ p_k += stride_k_t
215
p_o += HV * V
194
- p_v += HV * V
216
+ p_v += stride_v_t
217
if not IS_KDA:
218
p_g += HV
219
else:
220
p_gk += HV * K
199
- p_beta += HV * (V if IS_BETA_HEADWISE else 1)
221
+ p_beta += stride_beta_t
222
223
224
def fused_recurrent_gated_delta_rule_fwd(
@@ -265,6 +287,10 @@ def fused_recurrent_gated_delta_rule_fwd(
287
stride_final_state_token=stride_final_state_token,
288
stride_indices_seq=stride_indices_seq,
289
stride_indices_tok=stride_indices_tok,
290
+ stride_q_t=token_stride(q),
291
+ stride_k_t=token_stride(k),
292
+ stride_v_t=token_stride(v),
293
+ stride_beta_t=token_stride(beta),
294
IS_BETA_HEADWISE=beta.ndim == v.ndim,
295
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
296
INPLACE_FINAL_STATE=inplace_final_state,
vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
+16
-6
@@ -24,7 +24,10 @@ from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is
24
from vllm.triton_utils import tl, triton
25
from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
26
27
-from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel
27
+from .fused_recurrent import (
28
+ fused_recurrent_gated_delta_rule_fwd_kernel,
29
+ token_stride,
30
+)
31
32
BT_LIST_AUTOTUNE = [32, 64, 128]
33
NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
@@ -70,7 +73,7 @@ def fused_recurrent_kda_fwd(
73
g_bias = g_bias.reshape(-1).contiguous()
74
75
if out is None:
73
- o = torch.empty_like(k)
76
+ o = torch.empty(k.shape, dtype=k.dtype, device=k.device)
77
else:
78
# Caller-provided output buffer; must be layout-compatible with the
79
# tensor the kernel indexes (contiguous, same shape/dtype as k).
@@ -119,6 +122,10 @@ def fused_recurrent_kda_fwd(
122
stride_final_state_token=stride_final_state_token,
123
stride_indices_seq=stride_indices_seq,
124
stride_indices_tok=stride_indices_tok,
125
+ stride_q_t=token_stride(q),
126
+ stride_k_t=token_stride(k),
127
+ stride_v_t=token_stride(v),
128
+ stride_beta_t=token_stride(beta),
129
IS_BETA_HEADWISE=beta.ndim == v.ndim,
130
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
131
INPLACE_FINAL_STATE=inplace_final_state,
@@ -165,12 +172,15 @@ def fused_recurrent_kda(
172
if scale is None:
173
scale = k.shape[-1] ** -0.5
174
175
+ # q/k/v/beta are consumed in place with an explicit token stride, so
176
+ # column slices of the fused projection buffer need no copy; layouts the
177
+ # kernel cannot address fail loudly in `token_stride`.
178
o, final_state = fused_recurrent_kda_fwd(
169
- q=q.contiguous(),
170
- k=k.contiguous(),
171
- v=v.contiguous(),
179
+ q=q,
180
+ k=k,
181
+ v=v,
182
g=g.contiguous(),
173
- beta=beta.contiguous(),
183
+ beta=beta,
184
scale=scale,
185
initial_state=initial_state,
186
inplace_final_state=inplace_final_state,
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+5
-1
@@ -391,7 +391,11 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
391
layer: AttentionLayer,
392
) -> tuple[torch.Tensor, torch.Tensor | None]:
393
if isinstance(q, tuple):
394
- q = torch.cat(q, dim=-1)
394
+ ql_nope, q_pe = q
395
+ if q_pe.shape[-1] == 0 and ql_nope.is_contiguous():
396
+ q = ql_nope
397
+ else:
398
+ q = torch.cat(q, dim=-1)
399
400
num_actual_toks = q.shape[0]
401