[ROCm][Perf][GLM-5.3-Flash] Stride-aware decode KDA (#57979)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Sep 30, 2026 at 10:16 UTC
3627a6a124896edbc9f802e7c7e120633007e29d
3 files changed
+54
-19
tests/models/glm5next/test_kda_recurrent.py
+3
-4
@@ -5,8 +5,7 @@
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. Both backends must match a pure-PyTorch
8
-recurrence. CUDA reads these slices in place and rejects unsupported layouts;
9
-the ROCm wrapper makes contiguous copies before launching its kernel.
8
+recurrence, read these slices in place and reject unsupported layouts.
9
"""
10
11
import pytest
@@ -166,8 +165,8 @@ def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous(
165
166
167
@pytest.mark.skipif(
169
- not current_platform.is_cuda(),
170
- reason="Only the CUDA implementation requires zero-copy token-strided layouts",
168
+ not current_platform.is_cuda_alike(),
169
+ reason="Both implementations require zero-copy token-strided layouts",
170
)
171
@torch.inference_mode()
172
def test_fused_recurrent_kda_rejects_unaddressable_layouts():
vllm/models/glm5next/amd/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/amd/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,