diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 15d6667126..93d6962868 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.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 diff --git a/tests/models/glm5next/__init__.py b/tests/models/glm5next/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/models/glm5next/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py new file mode 100644 index 0000000000..68cf493b44 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index a572d94522..34df61fc82 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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] diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index 5d632df137..8444c81416 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/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: 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 d7de85984f..152412b8ca 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 @@ -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, 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 be417d6b05..b792f606d4 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 636278d425..b34c983752 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/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]