[Perf] Tune LL BF16 Router GEMM (#48774)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>

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

@@ -31,20 +31,71 @@ def is_available() -> bool:
     return _cutedsl_available
 
 
+# Default configs
 _DEFAULT_DOTPROD_BS = 128
 _DEFAULT_DOTPROD_MAX_M = 4
 _DEFAULT_SPLITK_CONFIG = (6, 4)
-_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = {
-    (7168, 256): 6,
+
+# ll_bf16 router shapes covered by warmup/tuning:
+# (4096, 256)  # DSV4-Flash
+# (6144, 256)  # GLM5.2
+# (7168, 256)  # DSV3.2
+# (7168, 384)  # DSV4-Pro
+
+# SM100f-specific tuned configs
+_SM100F_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = {
+    (6144, 256): {M: 256 for M in (1, 3, 4)},
 }
-_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
+_SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
+    (4096, 256): {
+        **{M: (8, 5) for M in (5, 8)},
+        9: (8, 2),
+    },
+    (7168, 256): {14: (8, 2)},
+    (6144, 256): {M: (8, 2) for M in (9, 12, 16)},
+    (7168, 384): {M: (7, 5) for M in (13, 16)},
+}
+
+# SM90-specific tuned configs
+_SM90_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = {
+    (4096, 256): {M: 256 for M in (1, 3)},
+    (7168, 384): {M: 256 for M in (1, 2)},
+}
+_SM90_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
+    (4096, 256): {
+        **{M: (8, 2) for M in range(5, 8)},
+        **{M: (8, 5) for M in range(10, 12)},
+        **{M: (8, 2) for M in (13, 16)},
+        **{M: (8, 5) for M in (14, 15)},
+    },
+    (7168, 256): {8: (6, 5)},
+    (6144, 256): {M: (8, 2) for M in (9, 11)},
     (7168, 384): {
-        5: (4, 4),
-        **{M: (5, 4) for M in range(6, 17)},
+        **{M: (8, 2) for M in (6, 8, 12)},
+        **{M: (7, 5) for M in (7, 9, 10, 11, 13, 14, 15, 16)},
     },
 }
 
 
+def _arch_tuned_configs() -> tuple[
+    dict[tuple[int, int], dict[int, int]],
+    dict[tuple[int, int], dict[int, tuple[int, int]]],
+]:
+    from vllm.platforms import current_platform
+
+    if current_platform.is_device_capability_family(100):
+        return (
+            _SM100F_TUNED_DOTPROD_BS,
+            _SM100F_TUNED_SPLITK_CONFIGS,
+        )
+    if current_platform.is_device_capability(90):
+        return (
+            _SM90_TUNED_DOTPROD_BS,
+            _SM90_TUNED_SPLITK_CONFIGS,
+        )
+    return {}, {}
+
+
 _cute_ctx = None
 
 
@@ -89,11 +140,12 @@ class LLBf16Gemm:
         self._splitk_cache: dict[tuple[int, int], Any] = {}
 
     def dispatch(self, *, M: int, K: int, N: int) -> CompileKey:
-        dotprod_max_m = _TUNED_DOTPROD_MAX_M.get((K, N), _DEFAULT_DOTPROD_MAX_M)
-        if dotprod_max_m >= M or K < 2048:
-            return self.CompileKey(backend="dotprod", M=M, K=K, bs=_DEFAULT_DOTPROD_BS)
+        tuned_bs, tuned_splitk = _arch_tuned_configs()
+        if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048:
+            bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS)
+            return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs)
 
-        split_k, num_stages = _TUNED_CONFIGS.get((K, N), {}).get(
+        split_k, num_stages = tuned_splitk.get((K, N), {}).get(
             M, _DEFAULT_SPLITK_CONFIG
         )
         return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages)

vllm/model_executor/warmup/kernel_warmup.py

@@ -45,16 +45,30 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = (
-    (6144, 264),  # Inkling
-    (7168, 256),  # DSV3
-    (7168, 384),  # DSV4-Pro
-    (14400, 256),  # DSV4-Flash
-)
 _LL_BF16_WARMUP_M_RANGE = range(1, 17)
 
 
-def _warmup_ll_bf16_router_gemm() -> None:
+def _ll_bf16_router_shapes_from_model(
+    model: torch.nn.Module,
+) -> tuple[tuple[int, int], ...]:
+    from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
+
+    shapes: set[tuple[int, int]] = set()
+    for module in model.modules():
+        if not isinstance(module, GateLinear):
+            continue
+        weight = getattr(module, "weight", None)
+        if not isinstance(weight, torch.Tensor):
+            continue
+        if weight.dim() != 2 or weight.dtype != torch.bfloat16:
+            continue
+        n, k = weight.shape
+        if k % 8 == 0:
+            shapes.add((int(k), int(n)))
+    return tuple(sorted(shapes))
+
+
+def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None:
     from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
         is_available as is_ll_bf16_gemm_available,
     )
@@ -65,9 +79,16 @@ def _warmup_ll_bf16_router_gemm() -> None:
     if not is_ll_bf16_gemm_available():
         return
 
-    logger.info("Warming up ll_bf16 router GEMM kernels.")
+    shapes = _ll_bf16_router_shapes_from_model(model)
+    if not shapes:
+        logger.info(
+            "Skipping ll_bf16 router GEMM warmup: no bf16 GateLinear shapes found."
+        )
+        return
+
+    logger.info("Warming up ll_bf16 router GEMM kernels for shapes: %s.", shapes)
     ll_bf16_gemm_kernel.warmup(
-        shapes=_LL_BF16_WARMUP_MODEL_SHAPES,
+        shapes=shapes,
         m_values=_LL_BF16_WARMUP_M_RANGE,
     )
 
@@ -122,8 +143,8 @@ def kernel_warmup(worker: "Worker"):
     elif has_flashinfer() and current_platform.has_device_capability(90):
         flashinfer_autotune(worker.model_runner)
 
-    if current_platform.has_device_capability(90) and worker.model_config.is_moe:
-        _warmup_ll_bf16_router_gemm()
+    if current_platform.has_device_capability(90):
+        _warmup_ll_bf16_router_gemm(worker.get_model())
 
     # FlashInfer attention warmup
     # Only warmup if the model has FlashInfer attention groups