[Perf][glm5next] Add SM121 TP=2 low-latency decode GEMM plans for GLM-5.3-Flash

tests/kernels/test_bf16_skinny_gemm.py

@@ -15,6 +15,7 @@
     SkinnyGemmConfig,
 )
 from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_gemm
+from vllm.models.glm5next.nvidia import low_latency_gemm as glm5next_gemm
 from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
 from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
 from vllm.models.qwen4_exp.nvidia import low_latency_gemm as qwen4_exp_gemm
@@ -97,6 +98,12 @@
     for num_tokens, config in plans.items()
 ]
 
+GLM5NEXT_SM121_CASES = [
+    (n, k, num_tokens, config)
+    for (n, k), plans in glm5next_gemm.GLM5NEXT_SM121_GEMM_PLANS.items()
+    for num_tokens, config in plans.items()
+]
+
 EXPECTED_CUTE_CONFIGS = {
     (3072, 7168, 1): (224, 3, 4, 8),
     (3072, 7168, 2): (128, 3, 2, 8),
@@ -972,6 +979,130 @@ def test_qwen4_exp_sm121_selected_shapes(
     assert cosine > 0.999
 
 
+def test_glm5next_sm121_plans_are_valid() -> None:
+    plans = glm5next_gemm.GLM5NEXT_SM121_GEMM_PLANS
+
+    assert len(plans) == 7
+    assert sum(map(len, plans.values())) == 20
+    for (n, k), shape_plans in plans.items():
+        for num_tokens, config in shape_plans.items():
+            assert config.num_rows == num_tokens
+            assert n % config.outputs_per_block == 0
+            assert k % (config.block_size * config.vector_width) == 0
+            assert config.static_k in (None, k)
+
+
+@pytest.mark.parametrize(
+    "capability,expected_plans",
+    [
+        ((12, 1), glm5next_gemm.GLM5NEXT_SM121_GEMM_PLANS),
+        ((12, 0), {}),
+        ((10, 3), {}),
+        ((9, 0), {}),
+    ],
+)
+def test_glm5next_gemm_capability_routing(
+    monkeypatch: pytest.MonkeyPatch,
+    capability: tuple[int, int],
+    expected_plans: dict[tuple[int, int], dict[int, SkinnyGemmConfig]],
+) -> None:
+    monkeypatch.setattr(
+        glm5next_gemm.current_platform,
+        "get_device_capability",
+        lambda: capability,
+    )
+
+    assert (
+        glm5next_gemm.GLM5NEXT_GEMM_PLANS_BY_CAPABILITY.get(
+            glm5next_gemm.current_platform.get_device_capability(), {}
+        )
+        == expected_plans
+    )
+
+
+def test_glm5next_installation_is_shape_specific_and_unquantized(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class FakeLinear(nn.Module):
+        def __init__(self, quant_method: object, n: int, k: int) -> None:
+            super().__init__()
+            self.quant_method = quant_method
+            self.weight = torch.empty(n, k)
+
+    class FakeHead(FakeLinear):
+        pass
+
+    root = nn.Module()
+    root.in_proj = FakeLinear(glm5next_gemm.UnquantizedLinearMethod(), 12576, 4096)
+    # quantized: must be left untouched.
+    quantized_method = object()
+    root.quantized = FakeLinear(quantized_method, 12576, 4096)
+    # shape absent from the table: must be left untouched.
+    root.unlisted = FakeLinear(glm5next_gemm.UnquantizedLinearMethod(), 4096, 128)
+    # A quant config that excludes lm_head (the modelopt NVFP4 checkpoint does)
+    # hands the head an UnquantizedLinearMethod, not an embedding method.
+    root.lm_head = FakeHead(glm5next_gemm.UnquantizedLinearMethod(), 77440, 4096)
+    root.plain_head = FakeHead(glm5next_gemm.UnquantizedEmbeddingMethod(), 77440, 4096)
+
+    monkeypatch.setattr(glm5next_gemm, "LinearBase", FakeLinear)
+    monkeypatch.setattr(glm5next_gemm, "ParallelLMHead", FakeHead)
+    monkeypatch.setattr(
+        glm5next_gemm.current_platform, "get_device_capability", lambda: (12, 1)
+    )
+    monkeypatch.setattr(
+        glm5next_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True
+    )
+    warmup_configs: set[SkinnyGemmConfig] = set()
+    monkeypatch.setattr(
+        glm5next_gemm.shape_dynamic_skinny_gemm,
+        "request_warmup_configs",
+        lambda dtype, configs: warmup_configs.update(configs),
+    )
+
+    glm5next_gemm.enable_glm5next_low_latency_gemm(root, torch.bfloat16)
+
+    assert isinstance(
+        root.in_proj.quant_method, glm5next_gemm.Glm5NextLowLatencyLinearMethod
+    )
+    assert root.quantized.quant_method is quantized_method
+    assert type(root.unlisted.quant_method) is glm5next_gemm.UnquantizedLinearMethod
+    assert isinstance(
+        root.lm_head.quant_method, glm5next_gemm.Glm5NextLowLatencyLinearMethod
+    )
+    assert isinstance(
+        root.plain_head.quant_method,
+        glm5next_gemm.Glm5NextLowLatencyEmbeddingMethod,
+    )
+    assert warmup_configs == {
+        config
+        for key in ((12576, 4096), (77440, 4096))
+        for config in glm5next_gemm.GLM5NEXT_SM121_GEMM_PLANS[key].values()
+    }
+
+
+@pytest.mark.parametrize("n,k,num_tokens,config", GLM5NEXT_SM121_CASES)
+def test_glm5next_sm121_selected_shapes(
+    n: int,
+    k: int,
+    num_tokens: int,
+    config: SkinnyGemmConfig,
+) -> None:
+    _require_capability_and_cute((12, 1))
+    torch.manual_seed(42 + num_tokens)
+    x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+
+    selected = glm5next_gemm.GLM5NEXT_SM121_GEMM_PLANS[(n, k)][num_tokens]
+    assert selected == config
+    output = glm5next_gemm._glm5next_low_latency_gemm(x, weight)
+
+    reference = torch.nn.functional.linear(x, weight)
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.float().flatten(), dim=0
+    ).item()
+    assert cosine > 0.999
+
+
 def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
     _require_sm103_and_cute()
     spec = glm52_gemm.GLM52_Q_B_PROJECTION

vllm/models/glm5next/common/model.py

@@ -100,6 +100,11 @@
     Glm5NextVisionTransformer,
 )
 
+if current_platform.is_cuda():
+    from vllm.models.glm5next.nvidia.low_latency_gemm import (
+        enable_glm5next_low_latency_gemm,
+    )
+
 logger = init_logger(__name__)
 
 
@@ -1069,6 +1074,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.logits_processor = LogitsProcessor(
             self.config.vocab_size, scale=_LOGIT_SCALE
         )
+        if current_platform.is_cuda():
+            enable_glm5next_low_latency_gemm(self, self.model_config.dtype)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)

vllm/models/glm5next/common/mtp.py

@@ -41,6 +41,11 @@
     get_spec_layer_idx_from_weight_name,
 )
 
+if current_platform.is_cuda():
+    from vllm.models.glm5next.nvidia.low_latency_gemm import (
+        enable_glm5next_low_latency_gemm,
+    )
+
 
 class Glm5NextMultiTokenPredictorLayer(nn.Module):
     def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
@@ -211,6 +216,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.model = Glm5NextMultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
+        if current_platform.is_cuda():
+            enable_glm5next_low_latency_gemm(self, vllm_config.model_config.dtype)
         self.set_moe_parameters()
         self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
             self.model.layers.values(), Glm5NextMoE, "mtp_block.mlp"

vllm/models/glm5next/nvidia/low_latency_gemm.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-5.3-Flash decode GEMM selection on NVIDIA GPUs.
+
+Dispatch follows Qwen4Exp and uses the local ``(N, K)`` shape and token count.
+Plans contain measured CUDA graph capture sizes; other token counts use the
+standard linear implementation.
+"""
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+    SkinnyGemmConfig,
+    row_stride_ok,
+    shape_dynamic_skinny_gemm,
+)
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    UnquantizedEmbeddingMethod,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
+
+# DGX Spark (GB10) plans for TP=2 local shapes.
+GLM5NEXT_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
+    # Shared-expert fused gate/up and MLA fused q_a/kv_a projections, TP=2.
+    (2048, 4096): {
+        1: SkinnyGemmConfig(1, 32, 1, k_unroll=4, static_k=4096),
+        2: SkinnyGemmConfig(2, 64, 1, k_unroll=4, vector_width=4, static_k=4096),
+        4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=4096),
+    },
+    # Shared-expert down projection, TP=2.
+    (4096, 1024): {
+        1: SkinnyGemmConfig(1, 64, 1, vector_width=2, static_k=1024),
+        2: SkinnyGemmConfig(2, 64, 1, vector_width=2, static_k=1024),
+        4: SkinnyGemmConfig(4, 64, 1, vector_width=2, static_k=1024),
+    },
+    # KDA output projection, TP=2.
+    (4096, 4096): {
+        1: SkinnyGemmConfig(1, 32, 2, k_unroll=4, vector_width=4, static_k=4096),
+        2: SkinnyGemmConfig(2, 64, 1, k_unroll=4, vector_width=2, static_k=4096),
+        4: SkinnyGemmConfig(4, 64, 1, vector_width=2, static_k=4096),
+    },
+    # MLA output projection, TP=2.
+    (4096, 8192): {
+        2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, vector_width=2, static_k=8192),
+        4: SkinnyGemmConfig(4, 256, 1, k_unroll=2, static_k=8192),
+    },
+    # MLA q_b projection, TP=2.
+    (8192, 1536): {
+        1: SkinnyGemmConfig(1, 32, 1, k_unroll=4, vector_width=2, static_k=1536),
+        2: SkinnyGemmConfig(2, 32, 1, k_unroll=2, vector_width=2, static_k=1536),
+        4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=1536),
+    },
+    # KDA fused QKV/beta/f_a/g_a input projection, TP=2.
+    (12576, 4096): {
+        1: SkinnyGemmConfig(1, 256, 1, k_unroll=2, vector_width=2),
+        2: SkinnyGemmConfig(2, 256, 1),
+        4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4, static_k=4096),
+    },
+    # LM head, TP=2.
+    (77440, 4096): {
+        1: SkinnyGemmConfig(1, 256, 1, static_k=4096),
+        2: SkinnyGemmConfig(2, 256, 1, k_unroll=2),
+        4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
+    },
+}
+
+
+GLM5NEXT_GEMM_PLANS_BY_CAPABILITY: dict[
+    tuple[int, int] | None, dict[tuple[int, int], dict[int, SkinnyGemmConfig]]
+] = {
+    (12, 1): GLM5NEXT_SM121_GEMM_PLANS,
+}
+
+
+def _is_packed_row_major(tensor: torch.Tensor) -> bool:
+    return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
+
+
+def _runtime_ok(
+    x: torch.Tensor, weight: torch.Tensor, config: SkinnyGemmConfig
+) -> bool:
+    return (
+        not envs.VLLM_BATCH_INVARIANT
+        and x.dim() == 2
+        and row_stride_ok(x, config)
+        and _is_packed_row_major(weight)
+        and x.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and x.is_cuda
+        and weight.is_cuda
+        and x.device == weight.device
+        and x.shape[1] == weight.shape[1]
+    )
+
+
+class _Glm5NextLowLatencyApply:
+    def apply(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if bias is None and not envs.VLLM_BATCH_INVARIANT:
+            return torch.ops.vllm.glm5next_low_latency_gemm(x, layer.weight)
+        return super().apply(layer, x, bias)  # type: ignore[misc]
+
+
+class Glm5NextLowLatencyLinearMethod(_Glm5NextLowLatencyApply, UnquantizedLinearMethod):
+    pass
+
+
+class Glm5NextLowLatencyEmbeddingMethod(
+    _Glm5NextLowLatencyApply, UnquantizedEmbeddingMethod
+):
+    pass
+
+
+# A quant config that excludes ``lm_head`` (the modelopt NVFP4 checkpoints do)
+# hands ParallelLMHead an UnquantizedLinearMethod rather than an
+# UnquantizedEmbeddingMethod, so the replacement is chosen by the method the
+# layer actually carries.
+_LOW_LATENCY_METHODS: dict[type, type] = {
+    UnquantizedLinearMethod: Glm5NextLowLatencyLinearMethod,
+    UnquantizedEmbeddingMethod: Glm5NextLowLatencyEmbeddingMethod,
+}
+
+
+def _glm5next_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    plans = GLM5NEXT_GEMM_PLANS_BY_CAPABILITY.get(
+        current_platform.get_device_capability(), {}
+    )
+    plan = plans.get((weight.shape[0], weight.shape[1]))
+    config = None if plan is None else plan.get(x.shape[0])
+    if (
+        config is not None
+        and _runtime_ok(x, weight, config)
+        and shape_dynamic_skinny_gemm.is_available()
+    ):
+        return shape_dynamic_skinny_gemm(x, weight, config)
+    return torch.nn.functional.linear(x, weight)
+
+
+def _glm5next_low_latency_gemm_fake(
+    x: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+    return x.new_empty((*x.shape[:-1], weight.shape[0]))
+
+
+direct_register_custom_op(
+    op_name="glm5next_low_latency_gemm",
+    op_func=_glm5next_low_latency_gemm,
+    fake_impl=_glm5next_low_latency_gemm_fake,
+)
+
+
+def enable_glm5next_low_latency_gemm(
+    module: nn.Module,
+    dtype: torch.dtype,
+) -> None:
+    plans = GLM5NEXT_GEMM_PLANS_BY_CAPABILITY.get(
+        current_platform.get_device_capability(), {}
+    )
+    if dtype != torch.bfloat16 or not plans:
+        return
+    if not shape_dynamic_skinny_gemm.is_available():
+        return
+
+    warmup_configs: set[SkinnyGemmConfig] = set()
+    for child in module.modules():
+        if not isinstance(child, (LinearBase, ParallelLMHead)):
+            continue
+        method_cls = _LOW_LATENCY_METHODS.get(type(child.quant_method))
+        if method_cls is None:
+            continue
+        weight = getattr(child, "weight", None)
+        if weight is None or weight.dim() != 2:
+            continue
+        plan = plans.get((weight.shape[0], weight.shape[1]))
+        if plan is None:
+            continue
+        child.quant_method = method_cls()
+        warmup_configs.update(plan.values())
+
+    if warmup_configs:
+        shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs)