[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>
Roberto L. Castro committed
Jul 27, 2026 at 16:25 UTC
56c96b0d91f05140b61a2005c222d83e9ec042db
2 files changed
+93
-20
vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
+61
-9
@@ -31,20 +31,71 @@ def is_available() -> bool:
31
return _cutedsl_available
32
33
34
+# Default configs
35
_DEFAULT_DOTPROD_BS = 128
36
_DEFAULT_DOTPROD_MAX_M = 4
37
_DEFAULT_SPLITK_CONFIG = (6, 4)
37
-_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = {
38
- (7168, 256): 6,
38
+
39
+# ll_bf16 router shapes covered by warmup/tuning:
40
+# (4096, 256) # DSV4-Flash
41
+# (6144, 256) # GLM5.2
42
+# (7168, 256) # DSV3.2
43
+# (7168, 384) # DSV4-Pro
44
+
45
+# SM100f-specific tuned configs
46
+_SM100F_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = {
47
+ (6144, 256): {M: 256 for M in (1, 3, 4)},
48
+}
49
+_SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
50
+ (4096, 256): {
51
+ **{M: (8, 5) for M in (5, 8)},
52
+ 9: (8, 2),
53
+ },
54
+ (7168, 256): {14: (8, 2)},
55
+ (6144, 256): {M: (8, 2) for M in (9, 12, 16)},
56
+ (7168, 384): {M: (7, 5) for M in (13, 16)},
57
}
40
-_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
58
+
59
+# SM90-specific tuned configs
60
+_SM90_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = {
61
+ (4096, 256): {M: 256 for M in (1, 3)},
62
+ (7168, 384): {M: 256 for M in (1, 2)},
63
+}
64
+_SM90_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
65
+ (4096, 256): {
66
+ **{M: (8, 2) for M in range(5, 8)},
67
+ **{M: (8, 5) for M in range(10, 12)},
68
+ **{M: (8, 2) for M in (13, 16)},
69
+ **{M: (8, 5) for M in (14, 15)},
70
+ },
71
+ (7168, 256): {8: (6, 5)},
72
+ (6144, 256): {M: (8, 2) for M in (9, 11)},
73
(7168, 384): {
42
- 5: (4, 4),
43
- **{M: (5, 4) for M in range(6, 17)},
74
+ **{M: (8, 2) for M in (6, 8, 12)},
75
+ **{M: (7, 5) for M in (7, 9, 10, 11, 13, 14, 15, 16)},
76
},
77
}
78
79
80
+def _arch_tuned_configs() -> tuple[
81
+ dict[tuple[int, int], dict[int, int]],
82
+ dict[tuple[int, int], dict[int, tuple[int, int]]],
83
+]:
84
+ from vllm.platforms import current_platform
85
+
86
+ if current_platform.is_device_capability_family(100):
87
+ return (
88
+ _SM100F_TUNED_DOTPROD_BS,
89
+ _SM100F_TUNED_SPLITK_CONFIGS,
90
+ )
91
+ if current_platform.is_device_capability(90):
92
+ return (
93
+ _SM90_TUNED_DOTPROD_BS,
94
+ _SM90_TUNED_SPLITK_CONFIGS,
95
+ )
96
+ return {}, {}
97
+
98
+
99
_cute_ctx = None
100
101
@@ -89,11 +140,12 @@ class LLBf16Gemm:
140
self._splitk_cache: dict[tuple[int, int], Any] = {}
141
142
def dispatch(self, *, M: int, K: int, N: int) -> CompileKey:
92
- dotprod_max_m = _TUNED_DOTPROD_MAX_M.get((K, N), _DEFAULT_DOTPROD_MAX_M)
93
- if dotprod_max_m >= M or K < 2048:
94
- return self.CompileKey(backend="dotprod", M=M, K=K, bs=_DEFAULT_DOTPROD_BS)
143
+ tuned_bs, tuned_splitk = _arch_tuned_configs()
144
+ if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048:
145
+ bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS)
146
+ return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs)
147
96
- split_k, num_stages = _TUNED_CONFIGS.get((K, N), {}).get(
148
+ split_k, num_stages = tuned_splitk.get((K, N), {}).get(
149
M, _DEFAULT_SPLITK_CONFIG
150
)
151
return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages)
vllm/model_executor/warmup/kernel_warmup.py
+32
-11
@@ -45,16 +45,30 @@ if TYPE_CHECKING:
45
46
logger = init_logger(__name__)
47
48
-_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = (
49
- (6144, 264), # Inkling
50
- (7168, 256), # DSV3
51
- (7168, 384), # DSV4-Pro
52
- (14400, 256), # DSV4-Flash
53
-)
48
_LL_BF16_WARMUP_M_RANGE = range(1, 17)
49
50
57
-def _warmup_ll_bf16_router_gemm() -> None:
51
+def _ll_bf16_router_shapes_from_model(
52
+ model: torch.nn.Module,
53
+) -> tuple[tuple[int, int], ...]:
54
+ from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
55
+
56
+ shapes: set[tuple[int, int]] = set()
57
+ for module in model.modules():
58
+ if not isinstance(module, GateLinear):
59
+ continue
60
+ weight = getattr(module, "weight", None)
61
+ if not isinstance(weight, torch.Tensor):
62
+ continue
63
+ if weight.dim() != 2 or weight.dtype != torch.bfloat16:
64
+ continue
65
+ n, k = weight.shape
66
+ if k % 8 == 0:
67
+ shapes.add((int(k), int(n)))
68
+ return tuple(sorted(shapes))
69
+
70
+
71
+def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None:
72
from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
73
is_available as is_ll_bf16_gemm_available,
74
)
@@ -65,9 +79,16 @@ def _warmup_ll_bf16_router_gemm() -> None:
79
if not is_ll_bf16_gemm_available():
80
return
81
68
- logger.info("Warming up ll_bf16 router GEMM kernels.")
82
+ shapes = _ll_bf16_router_shapes_from_model(model)
83
+ if not shapes:
84
+ logger.info(
85
+ "Skipping ll_bf16 router GEMM warmup: no bf16 GateLinear shapes found."
86
+ )
87
+ return
88
+
89
+ logger.info("Warming up ll_bf16 router GEMM kernels for shapes: %s.", shapes)
90
ll_bf16_gemm_kernel.warmup(
70
- shapes=_LL_BF16_WARMUP_MODEL_SHAPES,
91
+ shapes=shapes,
92
m_values=_LL_BF16_WARMUP_M_RANGE,
93
)
94
@@ -122,8 +143,8 @@ def kernel_warmup(worker: "Worker"):
143
elif has_flashinfer() and current_platform.has_device_capability(90):
144
flashinfer_autotune(worker.model_runner)
145
125
- if current_platform.has_device_capability(90) and worker.model_config.is_moe:
126
- _warmup_ll_bf16_router_gemm()
146
+ if current_platform.has_device_capability(90):
147
+ _warmup_ll_bf16_router_gemm(worker.get_model())
148
149
# FlashInfer attention warmup
150
# Only warmup if the model has FlashInfer attention groups