[Kernel][Qwen] Add Hopper LL-GEMM tuning table for Qwen4Exp

tests/kernels/test_bf16_skinny_gemm.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM90/SM100/SM103 selectors."""
+"""Tests for BF16 skinny GEMMs and model-specific selectors."""
 
 from pathlib import Path
 from types import SimpleNamespace
@@ -17,6 +17,7 @@
 from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_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
 
 # Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is
 # the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16).
@@ -77,6 +78,12 @@
     for _, config in spec.cute_configs
 ]
 
+QWEN4_EXP_SM90_CASES = [
+    (n, k, num_tokens, config)
+    for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM90_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),
@@ -528,6 +535,43 @@ def test_low_latency_table_capability_routing(
     assert k3_gemm._low_latency_table() is None
 
 
+def test_qwen4_exp_hopper_plans_are_valid() -> None:
+    plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS
+
+    assert len(plans) == 9
+    assert sum(map(len, plans.values())) == 31
+    assert (320, 10240) in plans
+    assert (10240, 320) not in plans
+    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",
+    [
+        ((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
+        ((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
+        ((8, 0), {}),
+    ],
+)
+def test_qwen4_exp_gemm_capability_routing(
+    monkeypatch: pytest.MonkeyPatch,
+    capability: tuple[int, int],
+    expected_plans: dict[tuple[int, int], dict[int, SkinnyGemmConfig]],
+) -> None:
+    monkeypatch.setattr(
+        qwen4_exp_gemm.current_platform,
+        "is_device_capability",
+        lambda target: capability == target,
+    )
+
+    assert qwen4_exp_gemm._gemm_plans() == expected_plans
+
+
 def test_installation_is_shape_specific_and_unquantized(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -643,7 +687,7 @@ def _require_capability_and_cute(capability: tuple[int, int]) -> None:
         not torch.cuda.is_available()
         or torch.cuda.get_device_capability() != capability
     ):
-        pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}")
+        pytest.skip(f"CuTe DSL selection requires SM{capability[0]}{capability[1]}")
     if not k3_gemm.shape_dynamic_skinny_gemm.is_available():
         pytest.skip("CuTe DSL is not available")
 
@@ -679,6 +723,29 @@ def test_glm_cute_selected_shapes(
     torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
 
 
+@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM90_CASES)
+def test_qwen4_exp_sm90_selected_shapes(
+    n: int,
+    k: int,
+    num_tokens: int,
+    config: SkinnyGemmConfig,
+) -> None:
+    _require_capability_and_cute((9, 0))
+    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 = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS[(n, k)][num_tokens]
+    assert selected == config
+    output = qwen4_exp_gemm._qwen4_exp_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/qwen4_exp/nvidia/low_latency_gemm.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Qwen4Exp decode GEMM selection on Blackwell.
+"""Qwen4Exp decode GEMM selection on Hopper and Blackwell.
 
 Dispatch follows Kimi-K3 and uses the local ``(N, K)`` shape and token count.
 Plans contain measured CUDA graph capture sizes; other token counts use the
@@ -78,11 +78,88 @@
     },
 }
 
+# H200 plans selected by exhaustive CUDA graph replay measurements over
+# M={1, 2, 4, 8, 16}. Only points that beat the standard linear implementation
+# in both hot-cache and L2-flush measurements are retained; other token counts
+# keep the standard implementation and its GEMM heuristics.
+QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
+    # GDN fused QKVZ projection, TP=4.
+    (4096, 2560): {
+        1: SkinnyGemmConfig(1, 128, 2, vector_width=4, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=2560),
+    },
+    # GDN and QSA output projections, TP=4.
+    (2560, 1536): {
+        1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=1536),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=1536),
+        4: SkinnyGemmConfig(4, 64, 4, k_unroll=6, vector_width=4),
+    },
+    # GDN fused B/A projection, TP=4.
+    (24, 2560): {
+        1: SkinnyGemmConfig(1, 128, 3, vector_width=4, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 2, vector_width=4, static_k=2560),
+        4: SkinnyGemmConfig(4, 64, 1, static_k=2560),
+        8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
+        16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
+    },
+    # QSA fused QKV/gate projection, TP=4.
+    (3584, 2560): {
+        1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 4, k_unroll=5),
+    },
+    # QSA indexer Q/K projection, replicated in a TP=4 deployment.
+    (640, 2560): {
+        1: SkinnyGemmConfig(1, 256, 2, vector_width=2, static_k=2560),
+        2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
+        4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560),
+        8: SkinnyGemmConfig(8, 128, 2, vector_width=4, static_k=2560),
+    },
+    # Shared-expert fused gate/up projection, TP=4.
+    (320, 2560): {
+        1: SkinnyGemmConfig(1, 64, 2, vector_width=4, static_k=2560),
+        2: SkinnyGemmConfig(2, 128, 4, k_unroll=5, vector_width=4),
+        4: SkinnyGemmConfig(4, 160, 1, k_unroll=2),
+        8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
+        16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
+    },
+    # LM head, TP=4.
+    (62080, 2560): {
+        1: SkinnyGemmConfig(1, 64, 2, vector_width=2, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 2, vector_width=2, static_k=2560),
+    },
+    # HC merged down/injection projection, replicated in a TP=4 deployment.
+    (336, 10240): {
+        1: SkinnyGemmConfig(1, 256, 1, k_unroll=5),
+        2: SkinnyGemmConfig(2, 256, 3, static_k=10240),
+        4: SkinnyGemmConfig(4, 256, 3, static_k=10240),
+        8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
+    },
+    # Final HC down projection, replicated in a TP=4 deployment.
+    (320, 10240): {
+        1: SkinnyGemmConfig(1, 256, 1, static_k=10240),
+        2: SkinnyGemmConfig(2, 128, 1, k_unroll=10),
+        4: SkinnyGemmConfig(4, 128, 1, k_unroll=10),
+        8: SkinnyGemmConfig(8, 128, 1, k_unroll=10),
+    },
+}
+
 
 def _is_sm103() -> bool:
     return current_platform.is_device_capability((10, 3))
 
 
+def _is_sm90() -> bool:
+    return current_platform.is_device_capability((9, 0))
+
+
+def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
+    if _is_sm103():
+        return QWEN4_EXP_GEMM_PLANS
+    if _is_sm90():
+        return QWEN4_EXP_SM90_GEMM_PLANS
+    return {}
+
+
 def _is_packed_row_major(tensor: torch.Tensor) -> bool:
     return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
 
@@ -124,7 +201,7 @@ class Qwen4ExpLowLatencyEmbeddingMethod(
 
 
 def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
-    plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1]))
+    plan = _gemm_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
@@ -152,7 +229,8 @@ def enable_qwen4_exp_low_latency_gemm(
     module: nn.Module,
     dtype: torch.dtype,
 ) -> None:
-    if dtype != torch.bfloat16 or not _is_sm103():
+    plans = _gemm_plans()
+    if dtype != torch.bfloat16 or not plans:
         return
     if not shape_dynamic_skinny_gemm.is_available():
         return
@@ -172,7 +250,7 @@ def enable_qwen4_exp_low_latency_gemm(
         weight = getattr(child, "weight", None)
         if weight is None or weight.dim() != 2:
             continue
-        plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1]))
+        plan = plans.get((weight.shape[0], weight.shape[1]))
         if plan is None:
             continue
         if is_linear: