From 56c96b0d91f05140b61a2005c222d83e9ec042db Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:25:37 +0200 Subject: [PATCH] [Perf] Tune LL BF16 Router GEMM (#48774) Signed-off-by: LopezCastroRoberto Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> --- .../kernels/linear/cute_dsl/ll_bf16.py | 70 ++++++++++++++++--- vllm/model_executor/warmup/kernel_warmup.py | 43 +++++++++--- 2 files changed, 93 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 16bf963296..b0471b5657 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/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) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index fae88eb7f4..e461dae0bb 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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