[Perf] Add SM90 skinny GEMM plans for dense Qwen3 and Llama-3.1 decode (TP=1)

tests/kernels/test_bf16_skinny_gemm.py

@@ -1058,52 +1058,105 @@ def test_skinny_gemm_plans_are_valid(
         assert config.static_k >= 2 * k_tile
 
 
+# One planned (N, K) per capability table, with K = 1024 planned and K = 128 not.
+_PLANNED_SHAPES = {
+    (12, 1): (4096, 1024),
+    (9, 0): (4096, 4096),
+}
+
+
+@pytest.mark.parametrize("capability", sorted(_PLANNED_SHAPES))
 @pytest.mark.parametrize(
-    "dtype,bias,k,ndim,batch_invariant,expected",
+    "dtype,bias,k_planned,ndim,batch_invariant,expected",
     [
-        pytest.param(torch.bfloat16, False, 1024, 2, False, True, id="planned"),
-        pytest.param(torch.float16, False, 1024, 2, False, False, id="fp16"),
-        pytest.param(torch.float32, False, 1024, 2, False, False, id="fp32"),
-        pytest.param(torch.bfloat16, True, 1024, 2, False, False, id="bias"),
-        pytest.param(torch.bfloat16, False, 128, 2, False, False, id="unlisted"),
-        pytest.param(torch.bfloat16, False, 1024, 2, True, False, id="batch-invariant"),
-        pytest.param(torch.bfloat16, False, 1024, 3, False, False, id="conv1d"),
-        pytest.param(torch.bfloat16, False, 1024, 4, False, False, id="conv2d"),
+        pytest.param(torch.bfloat16, False, True, 2, False, True, id="planned"),
+        pytest.param(torch.float16, False, True, 2, False, False, id="fp16"),
+        pytest.param(torch.float32, False, True, 2, False, False, id="fp32"),
+        pytest.param(torch.bfloat16, True, True, 2, False, False, id="bias"),
+        pytest.param(torch.bfloat16, False, False, 2, False, False, id="unlisted"),
+        pytest.param(torch.bfloat16, False, True, 2, True, False, id="batch-invariant"),
+        pytest.param(torch.bfloat16, False, True, 3, False, False, id="conv1d"),
+        pytest.param(torch.bfloat16, False, True, 4, False, False, id="conv2d"),
     ],
 )
 def test_maybe_use_skinny_gemm(
     monkeypatch: pytest.MonkeyPatch,
+    capability: tuple[int, int],
     dtype: torch.dtype,
     bias: bool,
-    k: int,
+    k_planned: bool,
     ndim: int,
     batch_invariant: bool,
     expected: bool,
 ) -> None:
     """Only eligible layers are marked; CPU tensors suffice for selection."""
-    assert (4096, 1024) in skinny_plans.SM121_SKINNY_GEMM_PLANS
-    assert (4096, 128) not in skinny_plans.SM121_SKINNY_GEMM_PLANS
+    n, k = _PLANNED_SHAPES[capability]
+    plans = skinny_plans.SKINNY_GEMM_PLANS_BY_CAPABILITY[capability]
+    assert (n, k) in plans
+    assert (n, 128) not in plans
+    if not k_planned:
+        k = 128
     # Isolate hardware/dependency detection, keeping the real plan lookup.
     monkeypatch.setattr(
-        skinny_plans.current_platform, "get_device_capability", lambda: (12, 1)
+        skinny_plans.current_platform, "get_device_capability", lambda: capability
     )
     monkeypatch.setattr(
         skinny_plans.shape_dynamic_skinny_gemm, "is_available", lambda: True
     )
     monkeypatch.setenv("VLLM_BATCH_INVARIANT", str(int(batch_invariant)))
     layer: nn.Module
     if ndim == 2:
-        layer = nn.Linear(k, 4096, bias=bias, dtype=dtype, device="cpu")
+        layer = nn.Linear(k, n, bias=bias, dtype=dtype, device="cpu")
     elif ndim == 3:
-        layer = nn.Conv1d(k, 4096, 1, bias=bias, dtype=dtype, device="cpu")
+        layer = nn.Conv1d(k, n, 1, bias=bias, dtype=dtype, device="cpu")
     else:
-        layer = nn.Conv2d(k, 4096, 1, bias=bias, dtype=dtype, device="cpu")
+        layer = nn.Conv2d(k, n, 1, bias=bias, dtype=dtype, device="cpu")
 
     maybe_use_skinny_gemm(layer)
 
     assert getattr(layer, "use_skinny_gemm", False) is expected
 
 
+_PLAN_RUN_CASES = [
+    pytest.param(capability, n, k, num_tokens, id=f"{capability}-{n}x{k}-m{num_tokens}")
+    for capability, plans in skinny_plans.SKINNY_GEMM_PLANS_BY_CAPABILITY.items()
+    for (n, k), shape_plans in plans.items()
+    for num_tokens in shape_plans
+]
+
+
+@pytest.mark.parametrize("capability,n,k,num_tokens", _PLAN_RUN_CASES)
+def test_skinny_gemm_plans_run_and_match_linear(
+    capability: tuple[int, int], n: int, k: int, num_tokens: int
+) -> None:
+    """Every registered plan compiles, runs on its device and matches F.linear.
+
+    N is capped: correctness does not depend on it, and the LM-head rows would
+    otherwise allocate a 1.5 GB weight per case.
+    """
+    _require_capability_and_cute(capability)
+    config = skinny_plans.SKINNY_GEMM_PLANS_BY_CAPABILITY[capability][(n, k)][
+        num_tokens
+    ]
+    assert skinny_plans.get_skinny_gemm_plan(n, k)[num_tokens] == config
+    n_test = min(n, 4096)
+    n_test -= n_test % config.outputs_per_block
+    torch.manual_seed(42 + num_tokens)
+    x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(n_test, k, dtype=torch.bfloat16, device="cuda") * 0.02
+
+    # Direct call, so a silent F.linear fallback cannot make this pass.
+    output = skinny_plans.shape_dynamic_skinny_gemm(x, weight, config)
+    reference = torch.nn.functional.linear(x, weight)
+
+    assert (torch.bfloat16, config, False) in (
+        skinny_plans.shape_dynamic_skinny_gemm._compiled
+    )
+    assert output.shape == reference.shape
+    max_err = (output.float() - reference.float()).abs().max().item()
+    assert max_err <= 1e-2 * reference.float().abs().max().item()
+
+
 def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
     _require_sm103_and_cute()
     spec = glm52_gemm.GLM52_Q_B_PROJECTION

vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm_plans.py

@@ -61,9 +61,154 @@
     },
 }
 
+# SM90, keyed by capability (9, 0) so it also selects H100 / H800 / H20; every
+# entry was measured on an H200 (132 SMs, 4.8 TB/s) only. TP=1 local shapes of
+# dense Qwen3 / Llama-3.1 models. Every 4-token entry was also checked inside
+# a running model (nsys, 4 concurrent requests); 4-token cells that win in
+# isolation but lose there are left out. static_k entries have no k_unroll:
+# the static path is fully unrolled.
+SM90_SKINNY_GEMM_PLANS: dict[tuple[int, int], SkinnyGemmPlan] = {
+    # Qwen3-30B-A3B attention output projection.
+    (2048, 4096): {
+        1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4, static_k=4096),
+        2: SkinnyGemmConfig(2, 128, 4, k_unroll=2, vector_width=4, static_k=4096),
+        4: SkinnyGemmConfig(4, 128, 4, vector_width=4, static_k=4096),
+    },
+    # Qwen3-30B-A3B fused QKV projection. M=4 is left to cuBLAS (1.03x).
+    (5120, 2048): {
+        1: SkinnyGemmConfig(1, 64, 4, k_unroll=4, vector_width=4),
+        2: SkinnyGemmConfig(2, 64, 4, k_unroll=4, vector_width=4),
+    },
+    # Qwen3-30B-A3B LM head. M=4 is left to cuBLAS (1.01x).
+    (151936, 2048): {
+        1: SkinnyGemmConfig(1, 128, 4),
+        2: SkinnyGemmConfig(2, 64, 4, k_unroll=4, vector_width=4),
+    },
+    # Qwen3-4B attention output projection.
+    (2560, 4096): {
+        1: SkinnyGemmConfig(1, 128, 1, vector_width=4, static_k=4096),
+        2: SkinnyGemmConfig(2, 128, 4, k_unroll=4, vector_width=4),
+        4: SkinnyGemmConfig(4, 128, 4),
+    },
+    # Qwen3-4B down projection.
+    (2560, 9728): {
+        1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=9728),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=9728),
+    },
+    # Qwen3-8B / Llama-3.1-8B attention output projection.
+    (4096, 4096): {
+        1: SkinnyGemmConfig(1, 128, 1, vector_width=2, static_k=4096),
+        2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=4096),
+        4: SkinnyGemmConfig(4, 64, 4, vector_width=4, static_k=4096),
+    },
+    # Qwen3-8B down projection.
+    (4096, 12288): {
+        1: SkinnyGemmConfig(1, 256, 4, vector_width=4, static_k=12288),
+        2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=12288),
+        4: SkinnyGemmConfig(4, 128, 4, vector_width=4, static_k=12288),
+    },
+    # Llama-3.1-8B down projection. M=4 is left to cuBLAS: it wins in isolation
+    # but loses inside the model (measured with nsys at 4 concurrent requests).
+    (4096, 14336): {
+        1: SkinnyGemmConfig(1, 224, 1, vector_width=4, static_k=14336),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=14336),
+    },
+    # Qwen3-14B attention output projection.
+    (5120, 5120): {
+        1: SkinnyGemmConfig(1, 64, 4, static_k=5120),
+        2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=5120),
+    },
+    # Qwen3-32B attention output projection.
+    (5120, 8192): {
+        1: SkinnyGemmConfig(1, 256, 2, vector_width=2, static_k=8192),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=8192),
+        4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4),
+    },
+    # Qwen3-14B down projection.
+    (5120, 17408): {
+        1: SkinnyGemmConfig(1, 64, 4, static_k=17408),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=17408),
+        4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4),
+    },
+    # Qwen3-32B down projection.
+    (5120, 25600): {
+        1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=25600),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=25600),
+        4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4),
+    },
+    # Qwen3-4B fused QKV projection.
+    (6144, 2560): {
+        1: SkinnyGemmConfig(1, 32, 4, static_k=2560),
+        2: SkinnyGemmConfig(2, 32, 4, static_k=2560),
+    },
+    # Qwen3-8B / Llama-3.1-8B fused QKV projection.
+    (6144, 4096): {
+        1: SkinnyGemmConfig(1, 128, 1, vector_width=2, static_k=4096),
+        2: SkinnyGemmConfig(2, 128, 3, vector_width=2, static_k=4096),
+    },
+    # Qwen3-14B fused QKV projection.
+    (7168, 5120): {
+        1: SkinnyGemmConfig(1, 64, 2, static_k=5120),
+        2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=5120),
+    },
+    # Qwen3-32B fused QKV projection.
+    (10240, 5120): {
+        1: SkinnyGemmConfig(1, 128, 2, vector_width=2, static_k=5120),
+        2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=5120),
+    },
+    # Qwen3-4B fused gate/up projection.
+    (19456, 2560): {
+        1: SkinnyGemmConfig(1, 64, 4, vector_width=2, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 2, vector_width=2, static_k=2560),
+        4: SkinnyGemmConfig(4, 32, 2, vector_width=4, static_k=2560),
+    },
+    # Qwen3-8B fused gate/up projection.
+    (24576, 4096): {
+        1: SkinnyGemmConfig(1, 256, 2, k_unroll=2),
+        2: SkinnyGemmConfig(2, 128, 2, vector_width=2, static_k=4096),
+        4: SkinnyGemmConfig(4, 32, 4, vector_width=4, static_k=4096),
+    },
+    # Llama-3.1-8B fused gate/up projection.
+    (28672, 4096): {
+        1: SkinnyGemmConfig(1, 256, 2, k_unroll=2),
+        2: SkinnyGemmConfig(2, 128, 2, vector_width=2, static_k=4096),
+    },
+    # Qwen3-14B fused gate/up projection.
+    (34816, 5120): {
+        1: SkinnyGemmConfig(1, 64, 1, vector_width=2, static_k=5120),
+        2: SkinnyGemmConfig(2, 64, 2, vector_width=4, static_k=5120),
+    },
+    # Qwen3-32B fused gate/up projection.
+    (51200, 5120): {
+        1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=5120),
+        2: SkinnyGemmConfig(2, 64, 2, vector_width=4, static_k=5120),
+    },
+    # Llama-3.1-8B LM head.
+    (128256, 4096): {
+        1: SkinnyGemmConfig(1, 256, 4, k_unroll=2, vector_width=4),
+        2: SkinnyGemmConfig(2, 256, 4, vector_width=4, static_k=4096),
+    },
+    # Qwen3-4B LM head.
+    (151936, 2560): {
+        1: SkinnyGemmConfig(1, 64, 4, vector_width=4, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 4),
+    },
+    # Qwen3-8B LM head.
+    (151936, 4096): {
+        1: SkinnyGemmConfig(1, 256, 4, k_unroll=2, vector_width=4),
+        2: SkinnyGemmConfig(2, 256, 4, vector_width=4, static_k=4096),
+    },
+    # Qwen3-14B / Qwen3-32B LM head.
+    (151936, 5120): {
+        1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=5120),
+        2: SkinnyGemmConfig(2, 64, 2, vector_width=4, static_k=5120),
+    },
+}
+
 SKINNY_GEMM_PLANS_BY_CAPABILITY: dict[
     tuple[int, int] | None, dict[tuple[int, int], SkinnyGemmPlan]
 ] = {
+    (9, 0): SM90_SKINNY_GEMM_PLANS,
     (12, 1): SM121_SKINNY_GEMM_PLANS,
 }