| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Warmup deep_gemm kernels. |
| 4 | DeepGEMM JIT's the kernels. The warmup aims to JIT all the kernels that would |
| 5 | be used during model execution beforehand. |
| 6 | """ |
| 7 | |
| 8 | import torch |
| 9 | from tqdm import tqdm |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank |
| 13 | from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( |
| 14 | DeepGemmFp8BlockScaledMMKernel, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 17 | from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( |
| 18 | compute_aligned_M_and_alignment, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts |
| 21 | from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( |
| 22 | TritonOrDeepGemmExperts, |
| 23 | ) |
| 24 | from vllm.tracing import instrument |
| 25 | from vllm.utils.deep_gemm import ( |
| 26 | fp8_gemm_nt, |
| 27 | get_mk_alignment_for_contiguous_layout, |
| 28 | m_grouped_fp8_gemm_nt_contiguous, |
| 29 | mk_alignment_scope, |
| 30 | ) |
| 31 | from vllm.utils.math_utils import cdiv |
| 32 | from vllm.utils.platform_utils import num_compute_units |
| 33 | from vllm.v1.worker.workspace import current_workspace_manager |
| 34 | |
| 35 | |
| 36 | def _generate_optimal_warmup_m_values( |
| 37 | max_tokens: int, n: int, device: torch.device |
| 38 | ) -> list[int]: |
| 39 | """Generate M values that cover all possible DeepGEMM kernel configurations. |
| 40 | Reference: https://github.com/deepseek-ai/DeepGEMM/blob/79f48ee15a82dd5fad5cd9beaa393c1f755e6b55/csrc/jit_kernels/heuristics/common.hpp |
| 41 | |
| 42 | Args: |
| 43 | max_tokens: Maximum number of tokens to warmup for |
| 44 | n: The actual N dimension from the weight tensor |
| 45 | device: The torch device to get properties from. |
| 46 | |
| 47 | """ |
| 48 | # DeepGEMM's possible block sizes |
| 49 | block_ms = [64, 128, 256] |
| 50 | block_ns = list(range(16, min(257, n + 1), 16)) |
| 51 | num_sms = num_compute_units(device.index) |
| 52 | |
| 53 | m_values = set() |
| 54 | |
| 55 | # Always include small cases |
| 56 | m_values.update([1, 2, 4] + [i for i in range(8, 65, 8)]) |
| 57 | |
| 58 | # Collect M values where different wave patterns occur |
| 59 | for block_m in block_ms: |
| 60 | for block_n in block_ns: |
| 61 | if block_n > n: |
| 62 | continue |
| 63 | |
| 64 | # Add key M boundaries for this block combination |
| 65 | for wave in range(1, 11): # Up to 10 waves |
| 66 | # M where this block config transitions to next wave |
| 67 | target_blocks = wave * num_sms |
| 68 | m = target_blocks * block_m // cdiv(n, block_n) |
| 69 | if 1 <= m <= max_tokens: |
| 70 | m_values.add(m) |
| 71 | |
| 72 | # Add block_m boundaries |
| 73 | for multiple in range(1, max_tokens // block_m + 1): |
| 74 | m = multiple * block_m |
| 75 | if m <= max_tokens: |
| 76 | m_values.add(m) |
| 77 | |
| 78 | return sorted(m_values) |
| 79 | |
| 80 | |
| 81 | def _extract_data_from_fused_moe_module( |
| 82 | m_: torch.nn.Module, |
| 83 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: |
| 84 | """Extract weights, weight scales and num_topk from MoERunner module.""" |
| 85 | assert isinstance(m_, MoERunner) |
| 86 | m = m_.routed_experts |
| 87 | w13 = m.w13_weight |
| 88 | w13_s = ( |
| 89 | m.w13_weight_scale_inv |
| 90 | if hasattr(m, "w13_weight_scale_inv") |
| 91 | else m.w13_weight_scale |
| 92 | ) |
| 93 | w2 = m.w2_weight |
| 94 | w2_s = ( |
| 95 | m.w2_weight_scale_inv |
| 96 | if hasattr(m, "w2_weight_scale_inv") |
| 97 | else m.w2_weight_scale |
| 98 | ) |
| 99 | num_topk = m.top_k |
| 100 | |
| 101 | assert isinstance(w13, torch.Tensor) |
| 102 | assert isinstance(w13_s, torch.Tensor) |
| 103 | assert isinstance(w2, torch.Tensor) |
| 104 | assert isinstance(w2_s, torch.Tensor) |
| 105 | return w13, w13_s, w2, w2_s, num_topk |
| 106 | |
| 107 | |
| 108 | def _deep_gemm_linear_data( |
| 109 | module: torch.nn.Module, |
| 110 | ) -> tuple[torch.Tensor, torch.Tensor] | None: |
| 111 | """Return (weight, weight_scale) for a layer whose kernel stamped itself as |
| 112 | ``deep_gemm_warmup_provider`` in ``process_weights_after_loading``, else None. |
| 113 | """ |
| 114 | provider = getattr(module, "deep_gemm_warmup_provider", None) |
| 115 | if not isinstance(provider, DeepGemmFp8BlockScaledMMKernel): |
| 116 | return None |
| 117 | if module.weight_block_size != get_mk_alignment_for_contiguous_layout(): |
| 118 | return None |
| 119 | return provider.get_deep_gemm_warmup_weights(module) |
| 120 | |
| 121 | |
| 122 | def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool: |
| 123 | if not (envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM): |
| 124 | return False |
| 125 | |
| 126 | if not isinstance(module, MoERunner): |
| 127 | return False |
| 128 | |
| 129 | quant_method = module._quant_method |
| 130 | moe_quant_config = quant_method.get_fused_moe_quant_config(module.routed_experts) |
| 131 | |
| 132 | if ( |
| 133 | moe_quant_config is None |
| 134 | or moe_quant_config.quant_dtype != torch.float8_e4m3fn |
| 135 | or moe_quant_config.block_shape != get_mk_alignment_for_contiguous_layout() |
| 136 | ): |
| 137 | return False |
| 138 | |
| 139 | moe_kernel = getattr(quant_method, "moe_kernel", None) |
| 140 | if moe_kernel is None: |
| 141 | return False |
| 142 | |
| 143 | fused_experts = moe_kernel.impl.fused_experts |
| 144 | return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts)) |
| 145 | |
| 146 | |
| 147 | FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set() |
| 148 | |
| 149 | |
| 150 | def _get_fp8_gemm_nt_m_values(w: torch.Tensor, max_tokens: int) -> list[int]: |
| 151 | """Get the M values to warmup for a given weight tensor.""" |
| 152 | n, _ = w.size() |
| 153 | device = w.device |
| 154 | |
| 155 | # Use optimal M values only if VLLM_DEEP_GEMM_WARMUP is set to "relax". |
| 156 | # Otherwise warmup all token sizes to avoid JIT compilation in hotpath |
| 157 | if envs.VLLM_DEEP_GEMM_WARMUP == "relax": |
| 158 | return _generate_optimal_warmup_m_values(max_tokens, n, device) |
| 159 | else: |
| 160 | assert envs.VLLM_DEEP_GEMM_WARMUP == "full", ( |
| 161 | "Expected " |
| 162 | 'VLLM_DEEP_GEMM_WARMUP env to be set to "full" but got ' |
| 163 | f"{envs.VLLM_DEEP_GEMM_WARMUP}" |
| 164 | ) |
| 165 | return list(range(1, max_tokens + 1)) |
| 166 | |
| 167 | |
| 168 | def _deepgemm_fp8_gemm_nt_warmup( |
| 169 | w: torch.Tensor, |
| 170 | ws: torch.Tensor, |
| 171 | max_tokens: int, |
| 172 | pbar: tqdm | None = None, |
| 173 | ): |
| 174 | if w.size() in FP8_GEMM_NT_WARMUP_CACHE: |
| 175 | return |
| 176 | |
| 177 | n, k = w.size() |
| 178 | block_m = get_mk_alignment_for_contiguous_layout()[0] |
| 179 | |
| 180 | device = w.device |
| 181 | a1q = torch.empty((max_tokens, k), device=device, dtype=torch.float8_e4m3fn) |
| 182 | a1q_scales = torch.zeros( |
| 183 | (max_tokens, k // block_m), device=device, dtype=torch.float32 |
| 184 | ) |
| 185 | out = torch.empty((max_tokens, n), device=device, dtype=torch.bfloat16) |
| 186 | |
| 187 | m_values = _get_fp8_gemm_nt_m_values(w, max_tokens) |
| 188 | |
| 189 | for num_tokens in m_values: |
| 190 | fp8_gemm_nt( |
| 191 | (a1q[:num_tokens], a1q_scales[:num_tokens]), (w, ws), out[:num_tokens] |
| 192 | ) |
| 193 | if pbar is not None: |
| 194 | pbar.update(1) |
| 195 | |
| 196 | FP8_GEMM_NT_WARMUP_CACHE.add(w.size()) |
| 197 | |
| 198 | |
| 199 | GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE: set[torch.Size] = set() |
| 200 | |
| 201 | |
| 202 | def _get_grouped_gemm_params( |
| 203 | w1: torch.Tensor, |
| 204 | w2: torch.Tensor, |
| 205 | num_topk: int, |
| 206 | max_tokens: int, |
| 207 | ) -> tuple[int, int, list[tuple[int, int, torch.Tensor]]]: |
| 208 | assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts" |
| 209 | |
| 210 | block_m = get_mk_alignment_for_contiguous_layout()[0] |
| 211 | num_experts = w1.size(0) |
| 212 | device = w1.device |
| 213 | |
| 214 | # Assumes all ranks have the same max_num_batched_tokens |
| 215 | max_tokens = get_dp_group().world_size * max_tokens |
| 216 | |
| 217 | request_m_values = _generate_optimal_warmup_m_values( |
| 218 | max_tokens, |
| 219 | max(w1.size(1), w2.size(1)), |
| 220 | device, |
| 221 | ) |
| 222 | request_m_values = sorted({m for m in (*request_m_values, max_tokens) if m > 0}) |
| 223 | if not request_m_values: |
| 224 | return 0, block_m, [] |
| 225 | |
| 226 | cases_by_shape: dict[tuple[int, int], torch.Tensor] = {} |
| 227 | for request_m in request_m_values: |
| 228 | M_sum, align_used = compute_aligned_M_and_alignment( |
| 229 | M=request_m, |
| 230 | num_topk=num_topk, |
| 231 | local_num_experts=num_experts, |
| 232 | alignment=block_m, |
| 233 | expert_tokens_meta=None, |
| 234 | ) |
| 235 | if (M_sum, align_used) in cases_by_shape: |
| 236 | continue |
| 237 | |
| 238 | num_blocks = M_sum // align_used |
| 239 | expert_ids_block = torch.randint( |
| 240 | low=0, |
| 241 | high=num_experts, |
| 242 | size=(num_blocks,), |
| 243 | device=device, |
| 244 | dtype=torch.int32, |
| 245 | ) |
| 246 | cases_by_shape[(M_sum, align_used)] = torch.repeat_interleave( |
| 247 | expert_ids_block, align_used, dim=0 |
| 248 | ) |
| 249 | |
| 250 | max_m = max(M_sum for M_sum, _ in cases_by_shape) |
| 251 | warmup_cases = [ |
| 252 | (M_sum, align_used, expert_ids) |
| 253 | for (M_sum, align_used), expert_ids in sorted(cases_by_shape.items()) |
| 254 | ] |
| 255 | |
| 256 | return max_m, block_m, warmup_cases |
| 257 | |
| 258 | |
| 259 | def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( |
| 260 | w1: torch.Tensor, |
| 261 | w2: torch.Tensor, |
| 262 | w1_scale: torch.Tensor, |
| 263 | w2_scale: torch.Tensor, |
| 264 | num_topk: int, |
| 265 | max_tokens: int, |
| 266 | pbar: tqdm | None = None, |
| 267 | ): |
| 268 | if ( |
| 269 | w1.size() in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE |
| 270 | and w2.size() in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE |
| 271 | ): |
| 272 | return |
| 273 | |
| 274 | MAX_M, block_m, warmup_cases = _get_grouped_gemm_params( |
| 275 | w1, w2, num_topk, max_tokens |
| 276 | ) |
| 277 | if not warmup_cases: |
| 278 | return |
| 279 | |
| 280 | def _warmup(w: torch.Tensor, w_scale: torch.Tensor): |
| 281 | _, n, k = w.size() |
| 282 | a1q, a1q_scales, out = current_workspace_manager().get_simultaneous( |
| 283 | ((MAX_M, k), torch.float8_e4m3fn), |
| 284 | ((MAX_M, k // block_m), torch.float32), |
| 285 | ((MAX_M, n), torch.bfloat16), |
| 286 | ) |
| 287 | a1q_scales.zero_() |
| 288 | |
| 289 | for num_tokens, align_used, expert_ids in warmup_cases: |
| 290 | with mk_alignment_scope(align_used): |
| 291 | m_grouped_fp8_gemm_nt_contiguous( |
| 292 | (a1q[:num_tokens], a1q_scales[:num_tokens]), |
| 293 | (w, w_scale), |
| 294 | out[:num_tokens], |
| 295 | expert_ids, |
| 296 | ) |
| 297 | if pbar is not None: |
| 298 | pbar.update(1) |
| 299 | |
| 300 | for w, ws in [(w1, w1_scale), (w2, w2_scale)]: |
| 301 | if w.size() not in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE: |
| 302 | _warmup(w, ws) |
| 303 | GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE.add(w.size()) |
| 304 | |
| 305 | |
| 306 | def deepgemm_fp8_gemm_nt_warmup( |
| 307 | model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None |
| 308 | ): |
| 309 | for m in model.modules(): |
| 310 | if (data := _deep_gemm_linear_data(m)) is not None: |
| 311 | w, ws = data |
| 312 | _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar) |
| 313 | |
| 314 | |
| 315 | def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( |
| 316 | model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None |
| 317 | ): |
| 318 | dg_modules = [ |
| 319 | m for m in model.modules() if _fused_moe_grouped_gemm_may_use_deep_gemm(m) |
| 320 | ] |
| 321 | |
| 322 | for dgm in dg_modules: |
| 323 | w13, w13_scale, w2, w2_scale, num_topk = _extract_data_from_fused_moe_module( |
| 324 | dgm |
| 325 | ) |
| 326 | _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( |
| 327 | w13, w2, w13_scale, w2_scale, num_topk, max_tokens, pbar=pbar |
| 328 | ) |
| 329 | |
| 330 | |
| 331 | def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: |
| 332 | seen_fp8_sizes: set[torch.Size] = set(FP8_GEMM_NT_WARMUP_CACHE) |
| 333 | seen_grouped_sizes: set[torch.Size] = set( |
| 334 | GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE |
| 335 | ) |
| 336 | |
| 337 | total = 0 |
| 338 | for m in model.modules(): |
| 339 | if (data := _deep_gemm_linear_data(m)) is not None: |
| 340 | w, _ = data |
| 341 | if w.size() not in seen_fp8_sizes: |
| 342 | total += len(_get_fp8_gemm_nt_m_values(w, max_tokens)) |
| 343 | seen_fp8_sizes.add(w.size()) |
| 344 | elif _fused_moe_grouped_gemm_may_use_deep_gemm(m): |
| 345 | w13, _, w2, _, num_topk = _extract_data_from_fused_moe_module(m) |
| 346 | if w13.size() in seen_grouped_sizes and w2.size() in seen_grouped_sizes: |
| 347 | continue |
| 348 | _, _, warmup_cases = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens) |
| 349 | n_values = len(warmup_cases) |
| 350 | if w13.size() not in seen_grouped_sizes: |
| 351 | total += n_values |
| 352 | seen_grouped_sizes.add(w13.size()) |
| 353 | if w2.size() not in seen_grouped_sizes: |
| 354 | total += n_values |
| 355 | seen_grouped_sizes.add(w2.size()) |
| 356 | return total |
| 357 | |
| 358 | |
| 359 | @instrument(span_name="DeepGemm warmup") |
| 360 | def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int): |
| 361 | total = _count_warmup_iterations(model, max_tokens) |
| 362 | if total == 0: |
| 363 | return |
| 364 | |
| 365 | # Only show progress bar on rank 0 to avoid cluttered output |
| 366 | if is_global_first_rank(): |
| 367 | with tqdm(total=total, desc="DeepGEMM warmup") as pbar: |
| 368 | deepgemm_fp8_gemm_nt_warmup(model, max_tokens, pbar) |
| 369 | deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, pbar) |
| 370 | else: |
| 371 | deepgemm_fp8_gemm_nt_warmup(model, max_tokens, None) |
| 372 | deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, None) |