[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>

.buildkite/test_areas/models_basic.yaml

@@ -86,6 +86,16 @@ steps:
     # The native NVIDIA Kimi K3 kernels require the SM100 family.
     - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
 
+- label: ":nvidia: (H200) GLM5Next"
+  key: glm5next-unit-tests
+  timeout_in_minutes: 20
+  device: h200_35gb
+  source_file_dependencies:
+  - vllm/models/glm5next/
+  - tests/models/glm5next/
+  commands:
+    - pytest -v -s models/glm5next
+
 - label: ":nvidia: (H200) Qwen4 Exp"
   key: qwen4-exp-unit-tests
   timeout_in_minutes: 40

tests/models/glm5next/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/models/glm5next/test_kda_recurrent.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-5.3-Flash KDA recurrent (decode) kernel.
+
+The decode path hands the kernel column slices of the merged ``q|k|v`` conv
+output and of the fused ``qkvbfg_a`` projection (beta), so q/k/v/beta are
+token-strided rather than contiguous. The kernel must read them in place,
+match a pure-PyTorch recurrence, and reject layouts it cannot address.
+"""
+
+import pytest
+import torch
+
+from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="CUDA-only Triton kernel"
+)
+
+H, D = 16, 128
+LOWER_BOUND = -5.0
+
+
+def naive_recurrent_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    a_log: torch.Tensor,
+    g_bias: torch.Tensor,
+    state: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """fp32 reference for one sequence: ``[T, H, D]`` inputs, ``[H, D, D]``
+    (v-major) state; mirrors the kernel's in-kernel gate, beta sigmoid and
+    q/k l2norm.
+    """
+    q, k, v, raw_g, raw_beta = (x.float() for x in (q, k, v, raw_g, raw_beta))
+    q = q / torch.sqrt(q.square().sum(-1, keepdim=True) + 1e-6) * D**-0.5
+    k = k / torch.sqrt(k.square().sum(-1, keepdim=True) + 1e-6)
+    gate = LOWER_BOUND * torch.sigmoid(a_log.exp()[:, None] * (raw_g + g_bias))
+    beta = torch.sigmoid(raw_beta)
+    s = state.clone()
+    out = torch.empty_like(v)
+    for t in range(q.shape[0]):
+        s = s * gate[t].exp()[:, None, :]
+        u = beta[t][:, None] * (v[t] - torch.einsum("hvk,hk->hv", s, k[t]))
+        s = s + u[:, :, None] * k[t][:, None, :]
+        out[t] = torch.einsum("hvk,hk->hv", s, q[t])
+    return out, s
+
+
+def make_inputs(num_seqs: int, query_len: int, device: torch.device):
+    """Token-strided q/k/v/beta as the decode path produces them: column
+    slices of a merged ``[T, q|k|v]`` conv output and of the fused
+    ``[T, qkv|beta|f_a|g_a]`` projection.
+    """
+    T, proj = num_seqs * query_len, H * D
+    qkv = torch.randn(T, 3 * proj, dtype=torch.bfloat16, device=device)
+    projected = torch.randn(
+        T, 3 * proj + H + 2 * D, dtype=torch.bfloat16, device=device
+    )
+    q, k, v = (x.view(1, T, H, D) for x in qkv.split(proj, dim=-1))
+    beta = projected[:, 3 * proj : 3 * proj + H].unsqueeze(0)
+    # (A size-1 token dim gets an arbitrary stride from `view`.)
+    assert T == 1 or (q.stride(1) == 3 * proj and beta.stride(1) == projected.stride(0))
+    inputs = dict(
+        q=q,
+        k=k,
+        v=v,
+        g=torch.randn(1, T, H, D, dtype=torch.bfloat16, device=device),
+        beta=beta,
+        a_log=0.5 * torch.randn(H, dtype=torch.float32, device=device),
+        g_bias=0.1 * torch.randn(H * D, dtype=torch.float32, device=device),
+        cu_seqlens=torch.arange(0, T + 1, query_len, dtype=torch.int32, device=device),
+    )
+    # Slot 0 is NULL_BLOCK_ID; sequences own random distinct slots (one per
+    # token in the spec-decode layout).
+    slots = torch.randperm(T, device=device).to(torch.int32) + 1
+    if query_len == 1:
+        inputs["ssm_state_indices"] = slots
+    else:
+        inputs["ssm_state_indices"] = slots.view(num_seqs, query_len)
+        inputs["num_accepted_tokens"] = torch.randint(
+            1, query_len + 1, (num_seqs,), dtype=torch.int32, device=device
+        )
+    state = torch.randn(T + 1, H, D, D, dtype=torch.float32, device=device)
+    return inputs, state
+
+
+def run_kernel(inputs: dict, state: torch.Tensor) -> torch.Tensor:
+    out, _ = fused_recurrent_kda(
+        **inputs,
+        initial_state=state,
+        use_qk_l2norm_in_kernel=True,
+        sigmoid_beta=True,
+        compute_gate=True,
+        lower_bound=LOWER_BOUND,
+    )
+    return out
+
+
+@pytest.mark.parametrize(
+    ("num_seqs", "query_len"), [(1, 1), (7, 1), (3, 3)], ids=["1x1", "7x1", "3x3"]
+)
+@torch.inference_mode()
+def test_fused_recurrent_kda_matches_reference(num_seqs: int, query_len: int):
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    inputs, state = make_inputs(num_seqs, query_len, device)
+    expected_state = state.clone()
+    out = run_kernel(inputs, state)
+
+    indices = inputs["ssm_state_indices"].view(num_seqs, query_len)
+    accepted = inputs.get("num_accepted_tokens")
+    expected = torch.empty_like(out[0], dtype=torch.float32)
+    for n in range(num_seqs):
+        first = indices[n, 0 if accepted is None else accepted[n] - 1]
+        s = expected_state[first]
+        for t in range(query_len):
+            tok = slice(n * query_len + t, n * query_len + t + 1)
+            expected[tok], s = naive_recurrent_kda(
+                inputs["q"][0, tok],
+                inputs["k"][0, tok],
+                inputs["v"][0, tok],
+                inputs["g"][0, tok],
+                inputs["beta"][0, tok],
+                inputs["a_log"],
+                inputs["g_bias"].view(H, D),
+                s,
+            )
+            expected_state[indices[n, t]] = s
+
+    torch.testing.assert_close(out[0].float(), expected, rtol=1e-2, atol=1e-3)
+    used = indices.flatten().long()
+    torch.testing.assert_close(state[used], expected_state[used], rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("num_seqs", "query_len"), [(7, 1), (3, 3)], ids=["7x1", "3x3"]
+)
+@torch.inference_mode()
+def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous(
+    num_seqs: int, query_len: int
+):
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    inputs, state = make_inputs(num_seqs, query_len, device)
+    for name in ("q", "k", "v", "beta"):
+        assert not inputs[name].is_contiguous()
+    contiguous = {
+        name: x.contiguous() if name in ("q", "k", "v", "beta") else x
+        for name, x in inputs.items()
+    }
+    state_ref = state.clone()
+    out_ref = run_kernel(contiguous, state_ref)
+    out = run_kernel(inputs, state)
+    torch.testing.assert_close(out, out_ref, rtol=0, atol=0)
+    torch.testing.assert_close(state, state_ref, rtol=0, atol=0)
+
+
+@torch.inference_mode()
+def test_fused_recurrent_kda_rejects_unaddressable_layouts():
+    """Layouts the token-stride addressing cannot express must fail loudly
+    rather than read the wrong tokens: a batch slice of a wider buffer
+    (``stride(0) != T * stride(1)``), overlapping tokens, and a head-strided
+    (transposed) block.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    inputs, state = make_inputs(1, 1, device)
+    T = 4
+    base = torch.randn(4, T, 2 * H * D, dtype=torch.bfloat16, device=device)
+    bad_q = {
+        "batch slice": base[::2, :, : H * D].view(2, T, H, D),
+        "overlapping tokens": base[:1, :, : H * D].as_strided(
+            (1, T, H, D), (0, D, D, 1)
+        ),
+        "head-strided": base[:1, :, : H * D].view(1, T, D, H).transpose(2, 3),
+    }
+    for q in bad_q.values():
+        broken = dict(inputs, q=q, k=q, v=q)
+        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)

vllm/model_executor/layers/attention/mla_attention.py

@@ -998,14 +998,17 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 if self.q_pad_num_heads is not None:
                     mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
                     mqa_ql_nope.resize_((N, B, L))
+                    # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
+                    torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)
+                    # Convert from (N, B, L) to (B, N, L)
+                    mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
                 else:
-                    mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))
-
-                # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
-                torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)
-
-                # Convert from (N, B, L) to (B, N, L)
-                mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
+                    # Write the (N, B, L) bmm result straight into a
+                    # token-major (B, N, L) buffer so the MQA query is already
+                    # contiguous; a NoPE model (qk_rope_head_dim == 0) then
+                    # needs no concat at all.
+                    mqa_ql_nope = mqa_q_nope.new_empty((B, N, L))
+                    torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope.transpose(0, 1))
 
             if fp8_attention and self.impl.supports_quant_query_input:
                 assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]

vllm/models/glm5next/nvidia/model.py

@@ -259,11 +259,11 @@ class Glm5NextMoE(nn.Module):
         if self.is_sequence_parallel and not already_sequence_parallel:
             hidden_states = sequence_parallel_chunk(hidden_states)
 
-        # The router is always external (self.gate); main's MoERunner expects
-        # pre-computed router_logits, so compute them here unconditionally.
-        router_logits, _ = self.gate(hidden_states)
+        # MoERunner holds the gate (passed to FusedMoEFactory) and computes
+        # the router logits itself, so nothing is precomputed here (matches
+        # DeepseekV2MoE; `router_logits` is a placeholder).
         final_hidden_states = self.experts(
-            hidden_states=hidden_states, router_logits=router_logits
+            hidden_states=hidden_states, router_logits=hidden_states
         )
 
         if self.is_sequence_parallel and not already_sequence_parallel:

vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py

@@ -15,6 +15,23 @@ from vllm.third_party.flash_linear_attention.ops.op import exp
 from vllm.triton_utils import tl, triton
 
 
+def token_stride(x: torch.Tensor) -> int:
+    """Token stride (elements) of a ``[B, T, H, D]`` or ``[B, T, H]`` tensor.
+
+    The recurrent kernel walks tokens with this stride and addresses heads
+    densely inside a token, so each token's ``[H, D]`` (or ``[H]``) block must
+    be contiguous, tokens must not overlap, and with ``B > 1`` sequence ``n``
+    must start at token ``n * T`` (dense batch). Column slices of a wider
+    per-token projection buffer satisfy this and are consumed in place.
+    """
+    st = x.stride()
+    assert x.dim() in (3, 4) and st[-1] == 1, (x.shape, st)
+    assert x.dim() == 3 or st[2] == x.shape[3], (x.shape, st)
+    assert st[1] >= x.shape[2] * (x.shape[3] if x.dim() == 4 else 1), (x.shape, st)
+    assert x.shape[0] == 1 or st[0] == x.shape[1] * st[1], (x.shape, st)
+    return st[1]
+
+
 @triton.heuristics(
     {
         "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
@@ -52,6 +69,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
     stride_final_state_token: tl.constexpr,
     stride_indices_seq: tl.constexpr,
     stride_indices_tok: tl.constexpr,
+    # Token strides of q/k/v/beta (elements), see `token_stride`.
+    stride_q_t,
+    stride_k_t,
+    stride_v_t,
+    stride_beta_t,
     USE_INITIAL_STATE: tl.constexpr,  # whether to use initial state
     INPLACE_FINAL_STATE: tl.constexpr,  # whether to store final state inplace
     IS_BETA_HEADWISE: tl.constexpr,  # whether beta is headwise vector or scalar,
@@ -86,13 +108,13 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
     o_k = i_k * BK + tl.arange(0, BK)
     o_v = i_v * BV + tl.arange(0, BV)
 
-    p_q = q + (bos * H + i_h) * K + o_k
-    p_k = k + (bos * H + i_h) * K + o_k
-    p_v = v + (bos * HV + i_hv) * V + o_v
+    p_q = q + bos * stride_q_t + i_h * K + o_k
+    p_k = k + bos * stride_k_t + i_h * K + o_k
+    p_v = v + bos * stride_v_t + i_hv * V + o_v
     if IS_BETA_HEADWISE:
-        p_beta = beta + (bos * HV + i_hv) * V + o_v
+        p_beta = beta + bos * stride_beta_t + i_hv * V + o_v
     else:
-        p_beta = beta + bos * HV + i_hv
+        p_beta = beta + bos * stride_beta_t + i_hv
 
     if not IS_KDA:
         p_g = g + bos * HV + i_hv
@@ -188,15 +210,15 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
             p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
             tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
 
-        p_q += H * K
-        p_k += H * K
+        p_q += stride_q_t
+        p_k += stride_k_t
         p_o += HV * V
-        p_v += HV * V
+        p_v += stride_v_t
         if not IS_KDA:
             p_g += HV
         else:
             p_gk += HV * K
-        p_beta += HV * (V if IS_BETA_HEADWISE else 1)
+        p_beta += stride_beta_t
 
 
 def fused_recurrent_gated_delta_rule_fwd(
@@ -265,6 +287,10 @@ def fused_recurrent_gated_delta_rule_fwd(
         stride_final_state_token=stride_final_state_token,
         stride_indices_seq=stride_indices_seq,
         stride_indices_tok=stride_indices_tok,
+        stride_q_t=token_stride(q),
+        stride_k_t=token_stride(k),
+        stride_v_t=token_stride(v),
+        stride_beta_t=token_stride(beta),
         IS_BETA_HEADWISE=beta.ndim == v.ndim,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
         INPLACE_FINAL_STATE=inplace_final_state,

vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py

@@ -24,7 +24,10 @@ from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
 
-from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel
+from .fused_recurrent import (
+    fused_recurrent_gated_delta_rule_fwd_kernel,
+    token_stride,
+)
 
 BT_LIST_AUTOTUNE = [32, 64, 128]
 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
@@ -70,7 +73,7 @@ def fused_recurrent_kda_fwd(
         g_bias = g_bias.reshape(-1).contiguous()
 
     if out is None:
-        o = torch.empty_like(k)
+        o = torch.empty(k.shape, dtype=k.dtype, device=k.device)
     else:
         # Caller-provided output buffer; must be layout-compatible with the
         # tensor the kernel indexes (contiguous, same shape/dtype as k).
@@ -119,6 +122,10 @@ def fused_recurrent_kda_fwd(
         stride_final_state_token=stride_final_state_token,
         stride_indices_seq=stride_indices_seq,
         stride_indices_tok=stride_indices_tok,
+        stride_q_t=token_stride(q),
+        stride_k_t=token_stride(k),
+        stride_v_t=token_stride(v),
+        stride_beta_t=token_stride(beta),
         IS_BETA_HEADWISE=beta.ndim == v.ndim,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
         INPLACE_FINAL_STATE=inplace_final_state,
@@ -165,12 +172,15 @@ def fused_recurrent_kda(
     if scale is None:
         scale = k.shape[-1] ** -0.5
 
+    # q/k/v/beta are consumed in place with an explicit token stride, so
+    # column slices of the fused projection buffer need no copy; layouts the
+    # kernel cannot address fail loudly in `token_stride`.
     o, final_state = fused_recurrent_kda_fwd(
-        q=q.contiguous(),
-        k=k.contiguous(),
-        v=v.contiguous(),
+        q=q,
+        k=k,
+        v=v,
         g=g.contiguous(),
-        beta=beta.contiguous(),
+        beta=beta,
         scale=scale,
         initial_state=initial_state,
         inplace_final_state=inplace_final_state,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -391,7 +391,11 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
         layer: AttentionLayer,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         if isinstance(q, tuple):
-            q = torch.cat(q, dim=-1)
+            ql_nope, q_pe = q
+            if q_pe.shape[-1] == 0 and ql_nope.is_contiguous():
+                q = ql_nope
+            else:
+                q = torch.cat(q, dim=-1)
 
         num_actual_toks = q.shape[0]