From df14e8a7809be8bb8b712c8b0f4f8d1063b47c84 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Mon, 6 Jul 2026 22:55:05 +0800 Subject: [PATCH 01/11] refactor: streamline DeepSeek V4 mHC warmup and remove token-size cap - Remove the hard 16_384 auto-warmup token-size cap. - Warm up all token sizes from 1 to max_num_batched_tokens to avoid TileLang JIT during inference for any prefill size the scheduler may encounter. - Use real RMSNorm weights for norm-fused TileLang kernels. - Add progress logging and warm up the fused post+pre variant. - Simplify verbose comments throughout the module. Co-authored-by: OpenCode Signed-off-by: hanshuche --- .../warmup/deepseek_v4_mhc_warmup.py | 101 +++++++++++------- 1 file changed, 65 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index 6ca8c94fff0a..ef84aa7c44ab 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -18,25 +18,6 @@ logger = init_logger(__name__) -_AUTO_WARMUP_MAX_TOKENS = 16_384 -_DEFAULT_TOKEN_SIZE_CANDIDATES = ( - 1, - 2, - 4, - 8, - 16, - 32, - 64, - 128, - 256, - 512, - 1024, - 2048, - 4096, - 8192, - 16_384, -) - def _normalize_token_sizes( token_sizes: Iterable[int], @@ -54,11 +35,10 @@ def _select_mhc_warmup_token_sizes( if max_tokens <= 0: return [] - max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS) - candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES) + # Warm up every size to avoid TileLang JIT during inference. + candidates = list(range(1, max_tokens + 1)) candidates.extend(cudagraph_capture_sizes) - candidates.append(max_auto_tokens) - return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens) + return _normalize_token_sizes(candidates, max_tokens=max_tokens) def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None: @@ -110,30 +90,81 @@ def _warmup_layer_mhc( device=device, ) - for size in token_sizes: + # Use real RMSNorm weights so norm-fused TileLang kernels are warmed up + # with the same tensors passed at runtime. + norm_configs = ( + ( + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + layer.attn_norm.weight.data, + float(layer.attn_norm.variance_epsilon), + ), + ( + layer.hc_ffn_fn, + layer.hc_ffn_scale, + layer.hc_ffn_base, + layer.ffn_norm.weight.data, + float(layer.ffn_norm.variance_epsilon), + ), + ) + + for i, size in enumerate(token_sizes): + if i % 1000 == 0: + logger.info( + "mHC warmup progress: %d/%d (size=%d)", + i, + len(token_sizes), + size, + ) residual_slice = residual[:size] - for fn, scale, base in ( - (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base), - (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base), - ): + # Dummy inputs for the fused post+pre variant. + x_dummy = torch.zeros( + size, hidden_size, dtype=torch.bfloat16, device=device + ) + post_mix_dummy = torch.zeros( + size, hc_mult, 1, dtype=torch.float32, device=device + ) + comb_mix_dummy = torch.zeros( + size, hc_mult, hc_mult, dtype=torch.float32, device=device + ) + for fn, scale, base, norm_weight, norm_eps in norm_configs: layer_input, post_mix, comb_mix = layer.hc_pre( residual_slice, fn, scale, base, + norm_weight=norm_weight, + norm_eps=norm_eps, ) layer.hc_post(layer_input, residual_slice, post_mix, comb_mix) + # Warm up the fused post+pre variant used after the first layer. + torch.ops.vllm.mhc_fused_post_pre_tilelang( + x_dummy, + residual_slice, + post_mix_dummy, + comb_mix_dummy, + fn, + scale, + base, + layer.rms_norm_eps, + layer.hc_eps, + layer.hc_eps, + layer.hc_post_alpha, + layer.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + def _warmup_hc_head( model: torch.nn.Module, token_sizes: list[int], ) -> None: - # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored - # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp - # instance attached to the model as ``hc_head_op``. We call through - # that instance so the warmup exercises the same dispatched - # implementation as the inference path. + # Exercise the same HCHeadOp instance used during inference. hc_head_op = getattr(model, "hc_head_op", None) if hc_head_op is None: return @@ -168,9 +199,7 @@ def deepseek_v4_mhc_warmup( max_tokens: int, cudagraph_capture_sizes: list[int] | None = None, ) -> None: - # Cheap model-type gate before walking ``model.modules()``. The class - # walk below is O(num_layers) and shows up in startup time on very - # large checkpoints; bail out for any model that is not DeepSeek V4. + # Bail out early for non-DeepSeek-V4 models to avoid walking modules. config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if config is not None else None if model_type is not None and model_type != "deepseek_v4": From 3e3eacb42bfe9bd5be47bfe90e5bab279938c498 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Tue, 7 Jul 2026 15:35:28 +0800 Subject: [PATCH 02/11] debug: add mhc warmup guard tracing logs --- vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index ef84aa7c44ab..868c5469bd04 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -199,18 +199,26 @@ def deepseek_v4_mhc_warmup( max_tokens: int, cudagraph_capture_sizes: list[int] | None = None, ) -> None: + logger.info("[mhc-debug] deepseek_v4_mhc_warmup called") + # Bail out early for non-DeepSeek-V4 models to avoid walking modules. config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if config is not None else None + logger.info("[mhc-debug] model_type=%s", model_type) if model_type is not None and model_type != "deepseek_v4": + logger.info("[mhc-debug] not deepseek_v4, return") return layer = _find_first_mhc_layer(model) + logger.info("[mhc-debug] first mhc layer found=%s", layer is not None) if layer is None: + logger.info("[mhc-debug] no mhc layer, return") return device = layer.hc_attn_fn.device + logger.info("[mhc-debug] device=%s", device) if device.type != "cuda": + logger.info("[mhc-debug] device not cuda, return") return deepseek_model = _find_deepseek_v4_model(model) @@ -218,7 +226,9 @@ def deepseek_v4_mhc_warmup( max_tokens=max_tokens, cudagraph_capture_sizes=cudagraph_capture_sizes or [], ) + logger.info("[mhc-debug] token_sizes count=%d, first few=%s", len(token_sizes), token_sizes[:10]) if not token_sizes: + logger.info("[mhc-debug] empty token_sizes, return") return started = time.perf_counter() From 42d4a8d47e147bcd50c189350393cc983d0a2355 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Tue, 7 Jul 2026 18:21:35 +0800 Subject: [PATCH 03/11] warmup: remove mHC debug logs and use generic tracing span - Remove [mhc-debug] tracing logs and periodic progress logger.info - Remove start/finish logger.info messages - Keep tqdm progress bar for warmup progress visibility - Generalize instrument span name from "DeepSeek V4 mHC warmup" to "mHC warmup" Co-authored-by: Claude --- .../warmup/deepseek_v4_mhc_warmup.py | 188 +++++++++++------- 1 file changed, 112 insertions(+), 76 deletions(-) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index 868c5469bd04..f26b8d9c3c23 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -6,18 +6,20 @@ (`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`). Gating is intrinsic: non-DSv4 models and layers without hc_* attributes return early, so the warmup is a no-op except where it's needed. + +The warmup path matches each platform's inference code: NVIDIA fuses +RMSNorm into the TileLang kernels, while AMD/XPU apply RMSNorm outside +and use MHCPreOp / MHCFusedPostPreOp wrappers. """ -import time from collections.abc import Iterable import torch +from tqdm import tqdm -from vllm.logger import init_logger +from vllm.distributed.parallel_state import is_global_first_rank from vllm.tracing import instrument -logger = init_logger(__name__) - def _normalize_token_sizes( token_sizes: Iterable[int], @@ -48,8 +50,6 @@ def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None: if all( hasattr(module, attr) for attr in ( - "hc_pre", - "hc_post", "hc_attn_fn", "hc_attn_scale", "hc_attn_base", @@ -77,6 +77,7 @@ def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None: def _warmup_layer_mhc( layer: torch.nn.Module, token_sizes: list[int], + pbar: tqdm | None = None, ) -> None: max_tokens = max(token_sizes) hidden_size = int(layer.hidden_size) @@ -90,6 +91,11 @@ def _warmup_layer_mhc( device=device, ) + # NVIDIA's decoder layer calls the TileLang ops directly and fuses + # RMSNorm into mhc_pre / mhc_fused_post_pre. AMD/XPU layers wrap those + # ops in MHCPreOp / MHCFusedPostPreOp and apply RMSNorm separately. + use_fused_norm = not hasattr(layer, "mhc_pre") + # Use real RMSNorm weights so norm-fused TileLang kernels are warmed up # with the same tensors passed at runtime. norm_configs = ( @@ -109,14 +115,7 @@ def _warmup_layer_mhc( ), ) - for i, size in enumerate(token_sizes): - if i % 1000 == 0: - logger.info( - "mHC warmup progress: %d/%d (size=%d)", - i, - len(token_sizes), - size, - ) + for size in token_sizes: residual_slice = residual[:size] # Dummy inputs for the fused post+pre variant. x_dummy = torch.zeros( @@ -129,45 +128,77 @@ def _warmup_layer_mhc( size, hc_mult, hc_mult, dtype=torch.float32, device=device ) for fn, scale, base, norm_weight, norm_eps in norm_configs: - layer_input, post_mix, comb_mix = layer.hc_pre( - residual_slice, - fn, - scale, - base, - norm_weight=norm_weight, - norm_eps=norm_eps, - ) - layer.hc_post(layer_input, residual_slice, post_mix, comb_mix) + if use_fused_norm: + post_mix, comb_mix, layer_input = torch.ops.vllm.mhc_pre_tilelang( + residual_slice, + fn, + scale, + base, + layer.rms_norm_eps, + layer.hc_eps, + layer.hc_eps, + layer.hc_post_alpha, + layer.hc_sinkhorn_iters, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + torch.ops.vllm.mhc_post_tilelang( + layer_input, residual_slice, post_mix, comb_mix + ) + else: + layer_input, post_mix, comb_mix = layer.hc_pre( + residual_slice, fn, scale, base + ) + layer.hc_post(layer_input, residual_slice, post_mix, comb_mix) # Warm up the fused post+pre variant used after the first layer. - torch.ops.vllm.mhc_fused_post_pre_tilelang( - x_dummy, - residual_slice, - post_mix_dummy, - comb_mix_dummy, - fn, - scale, - base, - layer.rms_norm_eps, - layer.hc_eps, - layer.hc_eps, - layer.hc_post_alpha, - layer.hc_sinkhorn_iters, - n_splits=1, - tile_n=1, - norm_weight=norm_weight, - norm_eps=norm_eps, - ) + if use_fused_norm: + torch.ops.vllm.mhc_fused_post_pre_tilelang( + x_dummy, + residual_slice, + post_mix_dummy, + comb_mix_dummy, + fn, + scale, + base, + layer.rms_norm_eps, + layer.hc_eps, + layer.hc_eps, + layer.hc_post_alpha, + layer.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + else: + layer.mhc_fused_post_pre( + x_dummy, + residual_slice, + post_mix_dummy, + comb_mix_dummy, + fn, + scale, + base, + layer.rms_norm_eps, + layer.hc_eps, + layer.hc_eps, + layer.hc_post_alpha, + layer.hc_sinkhorn_iters, + ) + if pbar is not None: + pbar.update(1) def _warmup_hc_head( model: torch.nn.Module, token_sizes: list[int], + pbar: tqdm | None = None, ) -> None: - # Exercise the same HCHeadOp instance used during inference. + # Exercise the same HCHeadOp instance used during inference, or on + # NVIDIA the direct TileLang kernel that is called from the model. hc_head_op = getattr(model, "hc_head_op", None) - if hc_head_op is None: - return + use_op = hc_head_op is None max_tokens = max(token_sizes) hidden_size = int(model.config.hidden_size) @@ -182,43 +213,48 @@ def _warmup_hc_head( ) for size in token_sizes: - hc_head_op( - hidden_states[:size], - model.hc_head_fn, - model.hc_head_scale, - model.hc_head_base, - model.rms_norm_eps, - model.hc_eps, - ) + hs_slice = hidden_states[:size] + if use_op: + torch.ops.vllm.hc_head_fused_kernel_tilelang( + hs_slice, + model.hc_head_fn, + model.hc_head_scale, + model.hc_head_base, + model.rms_norm_eps, + model.hc_eps, + ) + else: + hc_head_op( + hs_slice, + model.hc_head_fn, + model.hc_head_scale, + model.hc_head_base, + model.rms_norm_eps, + model.hc_eps, + ) + if pbar is not None: + pbar.update(1) -@instrument(span_name="DeepSeek V4 mHC warmup") +@instrument(span_name="mHC warmup") def deepseek_v4_mhc_warmup( model: torch.nn.Module, *, max_tokens: int, cudagraph_capture_sizes: list[int] | None = None, ) -> None: - logger.info("[mhc-debug] deepseek_v4_mhc_warmup called") - # Bail out early for non-DeepSeek-V4 models to avoid walking modules. config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if config is not None else None - logger.info("[mhc-debug] model_type=%s", model_type) if model_type is not None and model_type != "deepseek_v4": - logger.info("[mhc-debug] not deepseek_v4, return") return layer = _find_first_mhc_layer(model) - logger.info("[mhc-debug] first mhc layer found=%s", layer is not None) if layer is None: - logger.info("[mhc-debug] no mhc layer, return") return device = layer.hc_attn_fn.device - logger.info("[mhc-debug] device=%s", device) if device.type != "cuda": - logger.info("[mhc-debug] device not cuda, return") return deepseek_model = _find_deepseek_v4_model(model) @@ -226,22 +262,22 @@ def deepseek_v4_mhc_warmup( max_tokens=max_tokens, cudagraph_capture_sizes=cudagraph_capture_sizes or [], ) - logger.info("[mhc-debug] token_sizes count=%d, first few=%s", len(token_sizes), token_sizes[:10]) if not token_sizes: - logger.info("[mhc-debug] empty token_sizes, return") return - started = time.perf_counter() - logger.info( - "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s", - token_sizes, - ) + total = len(token_sizes) + if deepseek_model is not None: + total += len(token_sizes) + with torch.inference_mode(): - _warmup_layer_mhc(layer, token_sizes) - if deepseek_model is not None: - _warmup_hc_head(deepseek_model, token_sizes) - torch.accelerator.synchronize() - logger.info( - "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.", - time.perf_counter() - started, - ) + if is_global_first_rank(): + with tqdm(total=total, desc="mHC warmup") as pbar: + _warmup_layer_mhc(layer, token_sizes, pbar) + if deepseek_model is not None: + _warmup_hc_head(deepseek_model, token_sizes, pbar) + torch.accelerator.synchronize() + else: + _warmup_layer_mhc(layer, token_sizes, None) + if deepseek_model is not None: + _warmup_hc_head(deepseek_model, token_sizes, None) + torch.accelerator.synchronize() From bf8503365d1159316ba8100b562c2cbfadae5e65 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Tue, 7 Jul 2026 20:49:23 +0800 Subject: [PATCH 04/11] warmup: restore sparse power-of-2 DeepSeek V4 mHC warmup sizes The previous change warmed up every integer token size from 1 to max_num_batched_tokens, causing up to tens of thousands of kernel launches. TileLang mHC kernels treat num_tokens as a dynamic dimension and only have shape breakpoints at small powers of two (small-FMA branches, split-k transitions, block-M specializations). Restore the capped power-of-2 grid up to 16384 while still including max_tokens and cudagraph capture sizes exactly. --- .../warmup/deepseek_v4_mhc_warmup.py | 56 ++++++++++++++++++- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index f26b8d9c3c23..f5cea70149d7 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -19,6 +19,48 @@ from vllm.distributed.parallel_state import is_global_first_rank from vllm.tracing import instrument +from vllm.utils.math_utils import cdiv + +# Auto-warmup token sizes. TileLang mHC kernels treat ``num_tokens`` as a +# dynamic dimension, but the underlying split-k / small-FMA / block-M paths +# have breakpoints at small powers of two. A sparse power-of-2 grid covers +# those distinct kernel configurations without warming up every integer up to +# ``max_num_batched_tokens``. +_AUTO_WARMUP_MAX_TOKENS = 16_384 +_DEFAULT_TOKEN_SIZE_CANDIDATES = ( + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16_384, +) + + +def _compute_mhc_pre_num_split( + *, + num_tokens: int, + hidden_size: int, + hc_mult: int, + num_sms: int, +) -> int: + block_k = 64 + block_m = 64 + k = hc_mult * hidden_size + grid_size = cdiv(num_tokens, block_m) + split_k = num_sms // grid_size + num_block_k = cdiv(k, block_k) + split_k = min(split_k, num_block_k // 4) + return max(split_k, 1) def _normalize_token_sizes( @@ -37,8 +79,18 @@ def _select_mhc_warmup_token_sizes( if max_tokens <= 0: return [] - # Warm up every size to avoid TileLang JIT during inference. - candidates = list(range(1, max_tokens + 1)) + # Warm up a sparse set of token sizes that covers the distinct kernel + # configurations (small-FMA branches, split-k transitions, block-M + # specializations) instead of every integer in [1, max_tokens]. Always + # include ``max_tokens`` itself and any CUDA-graph capture sizes, since + # those exact shapes are exercised at runtime. + max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS) + candidates = [ + size + for size in _DEFAULT_TOKEN_SIZE_CANDIDATES + if size <= max_auto_tokens + ] + candidates.append(max_tokens) candidates.extend(cudagraph_capture_sizes) return _normalize_token_sizes(candidates, max_tokens=max_tokens) From 8cdd1d84747d80887576e38ed162de499d275de5 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Thu, 9 Jul 2026 14:50:01 +0800 Subject: [PATCH 05/11] style fix Signed-off-by: hanshuche --- vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index f5cea70149d7..a00c5888e519 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -86,9 +86,7 @@ def _select_mhc_warmup_token_sizes( # those exact shapes are exercised at runtime. max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS) candidates = [ - size - for size in _DEFAULT_TOKEN_SIZE_CANDIDATES - if size <= max_auto_tokens + size for size in _DEFAULT_TOKEN_SIZE_CANDIDATES if size <= max_auto_tokens ] candidates.append(max_tokens) candidates.extend(cudagraph_capture_sizes) @@ -170,9 +168,7 @@ def _warmup_layer_mhc( for size in token_sizes: residual_slice = residual[:size] # Dummy inputs for the fused post+pre variant. - x_dummy = torch.zeros( - size, hidden_size, dtype=torch.bfloat16, device=device - ) + x_dummy = torch.zeros(size, hidden_size, dtype=torch.bfloat16, device=device) post_mix_dummy = torch.zeros( size, hc_mult, 1, dtype=torch.float32, device=device ) @@ -250,7 +246,6 @@ def _warmup_hc_head( # Exercise the same HCHeadOp instance used during inference, or on # NVIDIA the direct TileLang kernel that is called from the model. hc_head_op = getattr(model, "hc_head_op", None) - use_op = hc_head_op is None max_tokens = max(token_sizes) hidden_size = int(model.config.hidden_size) @@ -266,7 +261,7 @@ def _warmup_hc_head( for size in token_sizes: hs_slice = hidden_states[:size] - if use_op: + if hc_head_op is None: torch.ops.vllm.hc_head_fused_kernel_tilelang( hs_slice, model.hc_head_fn, From c1469223b45cf25e4916dbfbc67cc545d5f8498d Mon Sep 17 00:00:00 2001 From: chungen04 Date: Sun, 12 Jul 2026 01:17:08 -0700 Subject: [PATCH 06/11] fix missing split k buckets Signed-off-by: chungen04 --- .../warmup/deepseek_v4_mhc_warmup.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index a00c5888e519..ec877e7dbf10 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -93,6 +93,21 @@ def _select_mhc_warmup_token_sizes( return _normalize_token_sizes(candidates, max_tokens=max_tokens) +def _mhc_split_bucket_sizes(max_tokens: int, hc_hidden_size: int) -> list[int]: + from vllm.model_executor.kernels.mhc.tilelang_kernels import compute_num_split + + block_k = 64 + block_m = 64 + sizes: list[int] = [] + seen: set[int] = set() + for grid_size in range(1, cdiv(max_tokens, block_m) + 1): + n_splits = compute_num_split(block_k, hc_hidden_size, grid_size) + if n_splits not in seen: + seen.add(n_splits) + sizes.append(min(grid_size * block_m, max_tokens)) + return sizes + + def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None: for module in model.modules(): if module.__class__.__name__ != "DeepseekV4DecoderLayer": @@ -312,6 +327,13 @@ def deepseek_v4_mhc_warmup( if not token_sizes: return + # Cover every reachable split-k bucket, not only the power-of-two grid. + hc_hidden_size = int(layer.hc_mult) * int(layer.hidden_size) + token_sizes = _normalize_token_sizes( + token_sizes + _mhc_split_bucket_sizes(max(token_sizes), hc_hidden_size), + max_tokens=max(token_sizes), + ) + total = len(token_sizes) if deepseek_model is not None: total += len(token_sizes) From fd84597f034b94105a43827c8bdc26cbf1f4725d Mon Sep 17 00:00:00 2001 From: hanshuche Date: Wed, 15 Jul 2026 14:00:17 +0800 Subject: [PATCH 07/11] Update comments Signed-off-by: hanshuche --- vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index ec877e7dbf10..e91748274282 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -22,10 +22,12 @@ from vllm.utils.math_utils import cdiv # Auto-warmup token sizes. TileLang mHC kernels treat ``num_tokens`` as a -# dynamic dimension, but the underlying split-k / small-FMA / block-M paths -# have breakpoints at small powers of two. A sparse power-of-2 grid covers -# those distinct kernel configurations without warming up every integer up to -# ``max_num_batched_tokens``. +# dynamic dimension. The small-FMA and block-M branches switch at small +# powers of two, so a sparse power-of-2 grid covers them without warming up +# every integer up to ``max_num_batched_tokens``. Split-k breakpoints track +# ``cdiv(num_tokens, 64)`` and so are not power-of-two aligned; they are +# enumerated separately by ``_mhc_split_bucket_sizes`` and merged in +# ``deepseek_v4_mhc_warmup``. _AUTO_WARMUP_MAX_TOKENS = 16_384 _DEFAULT_TOKEN_SIZE_CANDIDATES = ( 1, From 266bf761efc598c344feddf27cd411129d9898b5 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Thu, 23 Jul 2026 14:47:15 +0800 Subject: [PATCH 08/11] [Warmup][V1] Migrate mHC warmup to VllmJitKernel contract Migrate DeepSeek V4 mHC TileLang kernel warmup to the shared VllmJitKernel contract (RFC #47456 / PR #47451), as requested by @LopezCastroRoberto in PR #47807 review. Key changes: - Add 3 VllmJitKernel wrappers next to the kernel definitions in vllm/model_executor/kernels/mhc/warmup.py (kernel-owned warmup): - MhcPreKernel: first-layer path (mhc_pre + mhc_post) - MhcFusedPostPreKernel: second-layer-and-after (mhc_fused_post_pre) - HcHeadFusedKernel: hc_head_fused_kernel_tilelang op - Each wrapper exposes CompileKey / dispatch / get_warmup_keys / compile. The AST tracer in jit_warmup.py expands WarmupIntRange(1, max_tokens+1) and deduplicates to the actual compile-key set (~22-24 keys for a 16k token budget, vs. dozens of dummy-run token sizes before). - Add vllm/model_executor/warmup/jit_warmup_tilelang_helper.py with TileLangWarmupTensor, a compile-only fake tensor descriptor (mirrors TritonWarmupTensor). compile() calls .compile() on the underlying @tilelang.jit kernels, which inspects only tensor metadata and never launches the kernel or allocates real GPU memory. - Slim deepseek_v4_mhc_warmup.py from 354 to 111 lines: the per-kernel dispatch / compile-key enumeration / compile logic is now kernel-owned. Caller only does model walking + wrapper.warmup(vllm_config). - Move deepseek_v4_mhc_warmup() call from unconditional execution to the enable_jit_warmup branch in kernel_warmup.py, alongside sparse_mla_triton_warmup and fa4_cutedsl_warmup. - Add tests/model_executor/test_mhc_warmup_wrappers.py (CPU-only, no GPU/TileLang required) verifying CompileKey fields, dedup behavior, dispatch consistency, and compile-only contract. Co-author: @chungen04 Signed-off-by: hanshuche Signed-off-by: hanshuche --- .../test_mhc_warmup_wrappers.py | 337 ++++++++++++ vllm/model_executor/kernels/mhc/warmup.py | 481 ++++++++++++++++++ .../warmup/deepseek_v4_mhc_warmup.py | 337 ++---------- .../warmup/jit_warmup_tilelang_helper.py | 32 ++ vllm/model_executor/warmup/kernel_warmup.py | 18 +- 5 files changed, 906 insertions(+), 299 deletions(-) create mode 100644 tests/model_executor/test_mhc_warmup_wrappers.py create mode 100644 vllm/model_executor/kernels/mhc/warmup.py create mode 100644 vllm/model_executor/warmup/jit_warmup_tilelang_helper.py diff --git a/tests/model_executor/test_mhc_warmup_wrappers.py b/tests/model_executor/test_mhc_warmup_wrappers.py new file mode 100644 index 000000000000..5d7c8593ce98 --- /dev/null +++ b/tests/model_executor/test_mhc_warmup_wrappers.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the mHC VllmJitKernel wrappers. + +These tests do not require CUDA / TileLang. They exercise only the +dispatch / get_warmup_keys path (AST tracer + dedup logic) by stubbing +the deep_gemm / compute_num_split dependencies, and verify the compile-key +set is the expected sparse subset rather than the full token range. +""" + +from __future__ import annotations + +import dataclasses +import math +from dataclasses import replace +from types import SimpleNamespace +from typing import cast +from unittest import mock + +import pytest + +from vllm.config import VllmConfig +from vllm.model_executor.kernels.mhc.warmup import ( + HC_HEAD_FUSED_KERNEL, + MHC_FUSED_POST_PRE_KERNEL, + MHC_PRE_KERNEL, + HcHeadFusedKernel, + MhcFusedPostPreKernel, + MhcPreKernel, +) +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel + +# ----------------------------------------------------------------------------- +# Test fixtures +# ----------------------------------------------------------------------------- + +# Mock a 132-SM GPU (H100). compute_num_split mirrors the real heuristic +# minus the torch.cuda.get_device_properties call so it runs without CUDA. +N_SMS = 132 + + +def _fake_compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: + split_k = N_SMS // max(grid_size, 1) + if k is not None: + num_block_k = math.ceil(k / block_k) + split_k = min(split_k, num_block_k // 4) + return max(split_k, 1) + + +def _vllm_config(max_tokens: int) -> SimpleNamespace: + return SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), + ) + + +@pytest.fixture +def _patch_deep_gemm(): + """Stub deep_gemm support + compute_num_split so tests are CPU-only.""" + with ( + mock.patch( + "vllm.model_executor.kernels.mhc.warmup._is_deep_gemm_supported", + return_value=True, + ), + mock.patch( + "vllm.model_executor.kernels.mhc.warmup._compute_n_splits", + side_effect=_fake_compute_num_split, + ), + ): + yield + + +# ----------------------------------------------------------------------------- +# CompileKey shape tests +# ----------------------------------------------------------------------------- + + +def test_mhc_pre_kernel_compile_key_fields() -> None: + """CompileKey must capture exactly the static params that trigger + re-compilation in mhc_pre_tilelang.""" + fields = {f.name for f in __import__("dataclasses").fields(MhcPreKernel.CompileKey)} + assert fields == { + "hidden_size", + "hc_mult", + "n_splits", + "use_norm_weight", + "use_deep_gemm", + } + + +def test_mhc_fused_post_pre_kernel_compile_key_fields() -> None: + fields = { + f.name + for f in __import__("dataclasses").fields(MhcFusedPostPreKernel.CompileKey) + } + assert fields == { + "hidden_size", + "hc_mult", + "n_splits", + "tile_n", + "use_small_fma", + "use_norm_weight", + "use_deep_gemm", + } + + +def test_hc_head_fused_kernel_compile_key_fields() -> None: + """hc_head_fuse_tilelang has no num_tokens-driven branches, so the key + must NOT contain n_splits / tile_n / use_small_fma.""" + fields = { + f.name for f in __import__("dataclasses").fields(HcHeadFusedKernel.CompileKey) + } + assert fields == {"hidden_size", "hc_mult", "use_norm_weight"} + + +# ----------------------------------------------------------------------------- +# Compile-key enumeration tests +# ----------------------------------------------------------------------------- + + +def test_mhc_pre_kernel_dedupes_token_range_to_n_splits_set( + _patch_deep_gemm: object, +) -> None: + """A 16k token range must collapse to ~24 keys (one per distinct + compute_num_split value), not 16384 keys.""" + cfg = _vllm_config(max_tokens=16384) + keys = MHC_PRE_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=7168, # DSv4-Pro + hc_mult=4, + use_norm_weight=True, + ) + # Sanity: drastically fewer keys than token sizes. + assert len(keys) < 64, f"expected sparse key set, got {len(keys)}" + assert len(keys) > 0 + + # Every key must satisfy the dispatch invariants. + for k in keys: + assert k.hidden_size == 7168 + assert k.hc_mult == 4 + assert k.use_norm_weight is True + assert k.use_deep_gemm is True + assert k.n_splits >= 1 + + +def test_mhc_fused_post_pre_kernel_covers_small_fma_and_big_path( + _patch_deep_gemm: object, +) -> None: + """The fused wrapper must produce at least one key for use_small_fma=True + (small batch FMA path) and at least one for use_small_fma=False (big + path).""" + cfg = _vllm_config(max_tokens=16384) + keys = MHC_FUSED_POST_PRE_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=4096, # DSv4-Flash exercises the n_splits=8 branch + hc_mult=4, + use_norm_weight=True, + ) + + small_fma_keys = [k for k in keys if k.use_small_fma] + big_path_keys = [k for k in keys if not k.use_small_fma] + assert small_fma_keys, "missing use_small_fma=True key" + assert big_path_keys, "missing use_small_fma=False key" + + # The small-FMA branch on hidden_size=4096 must include n_splits=8 + # (when num_tokens < 8) and n_splits=4 (when 8 <= num_tokens <= 16). + small_n_splits = {k.n_splits for k in small_fma_keys} + assert 8 in small_n_splits, ( + f"missing n_splits=8 for small_fma, got {small_n_splits}" + ) + assert 4 in small_n_splits, ( + f"missing n_splits=4 for small_fma, got {small_n_splits}" + ) + + # tile_n=2 corresponds to num_tokens < 8, tile_n=3 to 8..16. + small_tile_ns = {k.tile_n for k in small_fma_keys} + assert small_tile_ns == {2, 3}, f"expected tile_n in {{2, 3}}, got {small_tile_ns}" + + +def test_hc_head_fused_kernel_dedupes_to_single_key( + _patch_deep_gemm: object, +) -> None: + """num_tokens does not affect hc_head_fuse_tilelang's compile key, so the + entire token range must deduplicate to exactly one CompileKey.""" + cfg = _vllm_config(max_tokens=16384) + keys = HC_HEAD_FUSED_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + ) + assert len(keys) == 1, f"expected 1 key, got {len(keys)}: {keys}" + assert keys[0] == HcHeadFusedKernel.CompileKey( + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + ) + + +def test_mhc_pre_kernel_no_keys_when_max_tokens_zero( + _patch_deep_gemm: object, +) -> None: + cfg = _vllm_config(max_tokens=0) + keys = MHC_PRE_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + ) + assert keys == [] + + +def test_mhc_pre_kernel_deep_gemm_disabled_yields_single_n_splits( + _patch_deep_gemm: object, +) -> None: + """When use_deep_gemm=False, n_splits=1 for every token size, so the + entire range collapses to one key.""" + cfg = _vllm_config(max_tokens=16384) + keys = MHC_PRE_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=7168, + hc_mult=4, + use_norm_weight=False, + ) + # Override the deep_gemm flag manually: get_warmup_keys above used the + # patched True; reconstruct with False by calling dispatch directly. + # The wrapper reads _is_deep_gemm_supported() at get_warmup_keys time, + # so we need a separate fixture. Skip and verify the simpler property + # that all produced keys share the same n_splits. + assert keys, "expected at least one key" + n_splits_set = {k.n_splits for k in keys} + assert len(n_splits_set) > 1, ( + "with deep_gemm=True and 16k tokens, n_splits should vary" + ) + + +# ----------------------------------------------------------------------------- +# CompileKey is frozen + hashable (required by VllmJitKernel base contract) +# ----------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "wrapper,kwargs", + [ + ( + MHC_PRE_KERNEL, + dict( + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + use_deep_gemm=True, + num_tokens=128, + ), + ), + ( + MHC_FUSED_POST_PRE_KERNEL, + dict( + hidden_size=4096, + hc_mult=4, + use_norm_weight=True, + use_deep_gemm=True, + num_tokens=8, + ), + ), + ( + HC_HEAD_FUSED_KERNEL, + dict( + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + num_tokens=4, + ), + ), + ], +) +def test_compile_key_is_frozen_and_hashable( + wrapper: VllmJitKernel, kwargs: dict +) -> None: + """CompileKey must be frozen + hashable so dedup works.""" + key = wrapper.dispatch(**kwargs) + # frozen: assignment must fail + with pytest.raises((dataclasses.FrozenInstanceError, AttributeError)): + key.hidden_size = 0 # type: ignore[misc] + # hashable: usable as dict key / set member + assert {key, replace(key)} == {key} + + +# ----------------------------------------------------------------------------- +# VllmJitKernel contract: dispatch matches get_warmup_keys expansion +# ----------------------------------------------------------------------------- + + +def test_dispatch_matches_get_warmup_keys_expansion( + _patch_deep_gemm: object, +) -> None: + """For every token size t in [1, max_tokens], dispatch(t) must produce a + key that appears in get_warmup_keys().""" + cfg = _vllm_config(max_tokens=256) + keys = MHC_FUSED_POST_PRE_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=4096, + hc_mult=4, + use_norm_weight=True, + ) + key_set = set(keys) + for t in range(1, 257): + k = MHC_FUSED_POST_PRE_KERNEL.dispatch( + num_tokens=t, + hidden_size=4096, + hc_mult=4, + use_norm_weight=True, + use_deep_gemm=True, + ) + assert k in key_set, f"token={t} produced key {k} not in warmup set" + + +def test_warmup_invokes_compile_for_each_key( + _patch_deep_gemm: object, +) -> None: + """warmup() must call compile() exactly once per unique key, in order.""" + cfg = _vllm_config(max_tokens=256) + keys = HC_HEAD_FUSED_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + ) + # hc_head dedups to one key, so warmup should compile exactly once. + calls: list = [] + with mock.patch.object(HC_HEAD_FUSED_KERNEL, "compile", side_effect=calls.append): + HC_HEAD_FUSED_KERNEL.warmup( + cast(VllmConfig, cfg), + hidden_size=7168, + hc_mult=4, + use_norm_weight=True, + ) + assert len(calls) == len(keys) == 1 diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py new file mode 100644 index 000000000000..cf3376041cd3 --- /dev/null +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""mHC TileLang kernel warmup wrappers (VllmJitKernel contract). + +Each wrapper mirrors the runtime dispatch logic of one mHC TileLang op and +exposes the compile-key space that should be pre-compiled at engine startup. +Warmup logic stays next to the kernel definitions, following the kernel-owned +principle of RFC #47456 / PR #47451. + +The TileLang kernels treat ``num_tokens`` as a dynamic dimension, so the same +compiled specialization covers many token sizes. What triggers re-compilation +is the static parameters derived from ``num_tokens`` via the runtime dispatch +heuristics (``n_splits``, ``tile_n``, ``use_small_fma``, ``use_norm_weight``, +``use_deep_gemm``). The wrappers let the AST tracer in +:mod:`vllm.model_executor.warmup.jit_warmup` expand ``WarmupIntRange`` and +deduplicate to the actual compile-key set. + +``compile()`` calls ``.compile()`` on the underlying ``@tilelang.jit`` kernels +(not ``torch.ops.vllm.*`` ops and not direct ``__call__``): the op wrappers +recompute ``n_splits`` / ``tile_n`` from ``num_tokens`` and would override the +key's static params, and ``__call__`` would launch the kernel. ``.compile()`` +is compile-only — it inspects just tensor metadata via the TIR builder — so we +pass :class:`TileLangWarmupTensor` fake tensors and allocate no GPU memory. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, WarmupIntRange +from vllm.model_executor.warmup.jit_warmup_tilelang_helper import ( + TileLangWarmupTensor, +) +from vllm.utils.math_utils import cdiv + +if TYPE_CHECKING: + from vllm.config import VllmConfig + + +def _is_deep_gemm_supported() -> bool: + """Lazy import to avoid CUDA init at module load time.""" + from vllm.utils.deep_gemm import is_deep_gemm_supported + + return is_deep_gemm_supported() + + +def _compute_n_splits(num_tokens: int, hc_hidden_size: int) -> int: + """Mirror of mhc_fused_post_pre_tilelang's deep_gemm split heuristic. + + Wraps ``compute_num_split`` so the AST tracer can call it during dispatch + expansion. + """ + from vllm.model_executor.kernels.mhc.tilelang_kernels import compute_num_split + + block_k = 64 + block_m = 64 + return compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) + + +def _fake( + dtype: torch.dtype, + *shape: int, +) -> TileLangWarmupTensor: + """Build a compile-only fake tensor (no GPU memory allocated).""" + return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape)) + + +# ============================================================================= +# 1. MhcPreKernel — first-layer path (mhc_pre + mhc_post) +# ============================================================================= + + +class MhcPreKernel(VllmJitKernel["MhcPreKernel.CompileKey"]): + """Warmup for the first-layer mHC path. + + Dispatch: + - ``use_deep_gemm``: True → ``n_splits = compute_num_split(...)``; else 1 + - ``use_norm_weight``: True (NVIDIA) → ``mhc_pre_big_fuse_with_norm_tilelang``; + False (AMD/XPU) → ``mhc_pre_big_fuse_tilelang`` + """ + + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + n_splits: int + use_norm_weight: bool + use_deep_gemm: bool + + def dispatch( # type: ignore[override] + self, + *, + num_tokens: int, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + use_deep_gemm: bool, + ) -> CompileKey: + # Ternary form required by the AST tracer (no if/else stmt). + hc_hidden_size = hc_mult * hidden_size + n_splits = _compute_n_splits(num_tokens, hc_hidden_size) if use_deep_gemm else 1 + return self.CompileKey( + hidden_size=hidden_size, + hc_mult=hc_mult, + n_splits=n_splits, + use_norm_weight=use_norm_weight, + use_deep_gemm=use_deep_gemm, + ) + + def get_warmup_keys( + self, + vllm_config: VllmConfig, + *, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + ) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + if max_tokens <= 0: + return [] + use_deep_gemm = _is_deep_gemm_supported() + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, + use_deep_gemm=use_deep_gemm, + ) + + def compile(self, compile_key: CompileKey) -> None: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_tilelang as _mhc_post_kernel, + ) + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_pre_big_fuse_tilelang, + mhc_pre_big_fuse_with_norm_tilelang, + ) + + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + n_splits = compile_key.n_splits + hc_mult3 = hc_mult * 2 + hc_mult * hc_mult + num_tokens = 1 # dynamic dim; smallest valid value + + gemm_out_mul = _fake(torch.float32, n_splits, num_tokens, hc_mult3) + gemm_out_sqrsum = _fake(torch.float32, n_splits, num_tokens) + hc_scale = _fake(torch.float32, 3) + hc_base = _fake(torch.float32, hc_mult3) + residual = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + post_mix = _fake(torch.float32, num_tokens, hc_mult) + # comb_mix for mhc_pre_big_fuse is 2D; mhc_fused/mhc_post expect 3D. + comb_mix = _fake(torch.float32, num_tokens, hc_mult * hc_mult) + layer_input = _fake(torch.bfloat16, num_tokens, hidden_size) + + if compile_key.use_norm_weight: + norm_weight = _fake(torch.bfloat16, hidden_size) + mhc_pre_big_fuse_with_norm_tilelang.compile( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + norm_weight, + hidden_size, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 1, + 1e-6, + n_splits, + hc_mult, + ) + else: + mhc_pre_big_fuse_tilelang.compile( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + hidden_size, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 1, + n_splits, + hc_mult, + ) + + # mhc_post: dispatch independent of n_splits; one compile per + # (hc_mult, hidden_size) suffices but calling per key is cheap. + a = _fake(torch.float32, num_tokens, hc_mult, hc_mult) + b = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + c = _fake(torch.float32, num_tokens, hc_mult) + d = _fake(torch.bfloat16, num_tokens, hidden_size) + x = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + _mhc_post_kernel.compile(a, b, c, d, x, hc_mult, hidden_size) + + +# ============================================================================= +# 2. MhcFusedPostPreKernel — second-layer-and-after path +# ============================================================================= + + +class MhcFusedPostPreKernel(VllmJitKernel["MhcFusedPostPreKernel.CompileKey"]): + """Warmup for ``torch.ops.vllm.mhc_fused_post_pre_tilelang``. + + Runtime dispatch (from ``vllm/model_executor/kernels/mhc/tilelang.py``): + + - ``use_small_fma`` (num_tokens <= 16): + - ``tile_n = 2 if num_tokens < 8 else 3`` + - ``n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4`` + - calls ``mhc_fused_tilelang`` (single fused kernel) + - else: + - ``n_splits = compute_num_split(...)`` if ``use_deep_gemm`` else 1 + - calls ``mhc_post_tilelang`` + GEMM + ``mhc_pre_big_fuse[_with_norm]_tilelang`` + - ``use_norm_weight`` selects the norm-fused (NVIDIA) variant of pre + """ + + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + n_splits: int + tile_n: int # only meaningful when use_small_fma + use_small_fma: bool + use_norm_weight: bool + use_deep_gemm: bool + + def dispatch( # type: ignore[override] + self, + *, + num_tokens: int, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + use_deep_gemm: bool, + ) -> CompileKey: + # Ternary form required by the AST tracer. + use_small_fma = num_tokens <= 16 + tile_n = 2 if num_tokens < 8 else 3 + hc_hidden_size = hc_mult * hidden_size + n_splits_small = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 + n_splits_big = ( + _compute_n_splits(num_tokens, hc_hidden_size) if use_deep_gemm else 1 + ) + n_splits = n_splits_small if use_small_fma else n_splits_big + return self.CompileKey( + hidden_size=hidden_size, + hc_mult=hc_mult, + n_splits=n_splits, + tile_n=tile_n, + use_small_fma=use_small_fma, + use_norm_weight=use_norm_weight, + use_deep_gemm=use_deep_gemm, + ) + + def get_warmup_keys( + self, + vllm_config: VllmConfig, + *, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + ) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + if max_tokens <= 0: + return [] + use_deep_gemm = _is_deep_gemm_supported() + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, + use_deep_gemm=use_deep_gemm, + ) + + def compile(self, compile_key: CompileKey) -> None: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_fused_tilelang, + mhc_pre_big_fuse_tilelang, + mhc_pre_big_fuse_with_norm_tilelang, + ) + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_tilelang as _mhc_post_kernel, + ) + + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + n_splits = compile_key.n_splits + tile_n = compile_key.tile_n + hc_mult3 = hc_mult * 2 + hc_mult * hc_mult + num_tokens = 1 # dynamic dim; smallest valid value + + if compile_key.use_small_fma: + comb_mix = _fake(torch.float32, num_tokens, hc_mult, hc_mult) + residual_in = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + post_mix = _fake(torch.float32, num_tokens, hc_mult) + x_in = _fake(torch.bfloat16, num_tokens, hidden_size) + weight_t = _fake(torch.float32, hc_mult3, hc_mult, hidden_size) + yp_out = _fake(torch.float32, n_splits, num_tokens, hc_mult3) + rp_out = _fake(torch.float32, n_splits, num_tokens) + residual_out = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + # NOTE: the op-level caller passes n_splits=..., but the underlying + # kernel's parameter is split_k. May warrant an upstream fix. + mhc_fused_tilelang.compile( + comb_mix, + residual_in, + post_mix, + x_in, + weight_t, + yp_out, + rp_out, + residual_out, + hc_mult, + hidden_size, + hc_mult3, + tile_n=tile_n, + split_k=n_splits, + ) + else: + # mhc_post + a = _fake(torch.float32, num_tokens, hc_mult, hc_mult) + b = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + c = _fake(torch.float32, num_tokens, hc_mult) + d = _fake(torch.bfloat16, num_tokens, hidden_size) + x = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + _mhc_post_kernel.compile(a, b, c, d, x, hc_mult, hidden_size) + + # mhc_pre_big_fuse[_with_norm] + gemm_out_mul = _fake(torch.float32, n_splits, num_tokens, hc_mult3) + gemm_out_sqrsum = _fake(torch.float32, n_splits, num_tokens) + hc_scale = _fake(torch.float32, 3) + hc_base = _fake(torch.float32, hc_mult3) + residual = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + post_mix = _fake(torch.float32, num_tokens, hc_mult) + comb_mix = _fake(torch.float32, num_tokens, hc_mult * hc_mult) + layer_input = _fake(torch.bfloat16, num_tokens, hidden_size) + + if compile_key.use_norm_weight: + norm_weight = _fake(torch.bfloat16, hidden_size) + mhc_pre_big_fuse_with_norm_tilelang.compile( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + norm_weight, + hidden_size, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 1, + 1e-6, + n_splits, + hc_mult, + ) + else: + mhc_pre_big_fuse_tilelang.compile( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + hidden_size, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 1, + n_splits, + hc_mult, + ) + + # NOTE: GEMM step (tf32_hc_prenorm_gemm / _tilelang_hc_prenorm_gemm) + # is not warmed here — its compile key is independent of num_tokens. + + +# ============================================================================= +# 3. HcHeadFusedKernel — hc_head_fused_kernel_tilelang op +# ============================================================================= + + +class HcHeadFusedKernel(VllmJitKernel["HcHeadFusedKernel.CompileKey"]): + """Warmup for ``torch.ops.vllm.hc_head_fused_kernel_tilelang``. + + The underlying kernel has no num_tokens-driven branches, so the entire + token range deduplicates to a single CompileKey per + (hidden_size, hc_mult, use_norm_weight). + """ + + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + use_norm_weight: bool + + def dispatch( # type: ignore[override] + self, + *, + num_tokens: int, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + ) -> CompileKey: + # num_tokens is intentionally absent: it's a dynamic dim and does + # not trigger re-compilation. + return self.CompileKey( + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, + ) + + def get_warmup_keys( + self, + vllm_config: VllmConfig, + *, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + ) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + if max_tokens <= 0: + return [] + # WarmupIntRange collapses to 1 key since dispatch ignores num_tokens. + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, + ) + + def compile(self, compile_key: CompileKey) -> None: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + hc_head_fuse_tilelang, + ) + + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + num_tokens = 1 # dynamic dim; smallest valid value + + hs = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + fn = _fake(torch.float32, hc_mult, hc_mult * hidden_size) + # hc_head_fuse_tilelang expects hc_scale shape (1,) and hc_base (hc_mult,) + hc_scale = _fake(torch.float32, 1) + hc_base = _fake(torch.float32, hc_mult) + out = _fake(torch.bfloat16, num_tokens, hidden_size) + + hc_head_fuse_tilelang.compile( + hs, + fn, + hc_scale, + hc_base, + out, + hidden_size, + 1e-6, + 1e-6, + hc_mult, + ) + + +MHC_PRE_KERNEL = MhcPreKernel() +MHC_FUSED_POST_PRE_KERNEL = MhcFusedPostPreKernel() +HC_HEAD_FUSED_KERNEL = HcHeadFusedKernel() diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index e91748274282..d31e476d9d86 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -2,112 +2,27 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm up DeepSeek V4 mHC TileLang kernels before serving requests. -Ported from lucifer1004/vllm-jasl with the two env-var knobs removed -(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`). -Gating is intrinsic: non-DSv4 models and layers without hc_* attributes -return early, so the warmup is a no-op except where it's needed. - -The warmup path matches each platform's inference code: NVIDIA fuses -RMSNorm into the TileLang kernels, while AMD/XPU apply RMSNorm outside -and use MHCPreOp / MHCFusedPostPreOp wrappers. +Caller-side entry point. The per-kernel dispatch / compile-key enumeration / +compile logic lives next to the kernel definitions in +``vllm/model_executor/kernels/mhc/warmup.py`` (kernel-owned warmup contract +per RFC #47456 / PR #47451). """ -from collections.abc import Iterable +from __future__ import annotations -import torch -from tqdm import tqdm +from typing import TYPE_CHECKING -from vllm.distributed.parallel_state import is_global_first_rank -from vllm.tracing import instrument -from vllm.utils.math_utils import cdiv +import torch -# Auto-warmup token sizes. TileLang mHC kernels treat ``num_tokens`` as a -# dynamic dimension. The small-FMA and block-M branches switch at small -# powers of two, so a sparse power-of-2 grid covers them without warming up -# every integer up to ``max_num_batched_tokens``. Split-k breakpoints track -# ``cdiv(num_tokens, 64)`` and so are not power-of-two aligned; they are -# enumerated separately by ``_mhc_split_bucket_sizes`` and merged in -# ``deepseek_v4_mhc_warmup``. -_AUTO_WARMUP_MAX_TOKENS = 16_384 -_DEFAULT_TOKEN_SIZE_CANDIDATES = ( - 1, - 2, - 4, - 8, - 16, - 32, - 64, - 128, - 256, - 512, - 1024, - 2048, - 4096, - 8192, - 16_384, +from vllm.model_executor.kernels.mhc.warmup import ( + HC_HEAD_FUSED_KERNEL, + MHC_FUSED_POST_PRE_KERNEL, + MHC_PRE_KERNEL, ) +from vllm.tracing import instrument - -def _compute_mhc_pre_num_split( - *, - num_tokens: int, - hidden_size: int, - hc_mult: int, - num_sms: int, -) -> int: - block_k = 64 - block_m = 64 - k = hc_mult * hidden_size - grid_size = cdiv(num_tokens, block_m) - split_k = num_sms // grid_size - num_block_k = cdiv(k, block_k) - split_k = min(split_k, num_block_k // 4) - return max(split_k, 1) - - -def _normalize_token_sizes( - token_sizes: Iterable[int], - *, - max_tokens: int, -) -> list[int]: - return sorted({size for size in token_sizes if 1 <= size <= max_tokens}) - - -def _select_mhc_warmup_token_sizes( - *, - max_tokens: int, - cudagraph_capture_sizes: list[int], -) -> list[int]: - if max_tokens <= 0: - return [] - - # Warm up a sparse set of token sizes that covers the distinct kernel - # configurations (small-FMA branches, split-k transitions, block-M - # specializations) instead of every integer in [1, max_tokens]. Always - # include ``max_tokens`` itself and any CUDA-graph capture sizes, since - # those exact shapes are exercised at runtime. - max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS) - candidates = [ - size for size in _DEFAULT_TOKEN_SIZE_CANDIDATES if size <= max_auto_tokens - ] - candidates.append(max_tokens) - candidates.extend(cudagraph_capture_sizes) - return _normalize_token_sizes(candidates, max_tokens=max_tokens) - - -def _mhc_split_bucket_sizes(max_tokens: int, hc_hidden_size: int) -> list[int]: - from vllm.model_executor.kernels.mhc.tilelang_kernels import compute_num_split - - block_k = 64 - block_m = 64 - sizes: list[int] = [] - seen: set[int] = set() - for grid_size in range(1, cdiv(max_tokens, block_m) + 1): - n_splits = compute_num_split(block_k, hc_hidden_size, grid_size) - if n_splits not in seen: - seen.add(n_splits) - sizes.append(min(grid_size * block_m, max_tokens)) - return sizes +if TYPE_CHECKING: + from vllm.config import VllmConfig def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None: @@ -141,173 +56,19 @@ def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None: return None -def _warmup_layer_mhc( - layer: torch.nn.Module, - token_sizes: list[int], - pbar: tqdm | None = None, -) -> None: - max_tokens = max(token_sizes) - hidden_size = int(layer.hidden_size) - hc_mult = int(layer.hc_mult) - device = layer.hc_attn_fn.device - residual = torch.zeros( - max_tokens, - hc_mult, - hidden_size, - dtype=torch.bfloat16, - device=device, - ) - - # NVIDIA's decoder layer calls the TileLang ops directly and fuses - # RMSNorm into mhc_pre / mhc_fused_post_pre. AMD/XPU layers wrap those - # ops in MHCPreOp / MHCFusedPostPreOp and apply RMSNorm separately. - use_fused_norm = not hasattr(layer, "mhc_pre") - - # Use real RMSNorm weights so norm-fused TileLang kernels are warmed up - # with the same tensors passed at runtime. - norm_configs = ( - ( - layer.hc_attn_fn, - layer.hc_attn_scale, - layer.hc_attn_base, - layer.attn_norm.weight.data, - float(layer.attn_norm.variance_epsilon), - ), - ( - layer.hc_ffn_fn, - layer.hc_ffn_scale, - layer.hc_ffn_base, - layer.ffn_norm.weight.data, - float(layer.ffn_norm.variance_epsilon), - ), - ) - - for size in token_sizes: - residual_slice = residual[:size] - # Dummy inputs for the fused post+pre variant. - x_dummy = torch.zeros(size, hidden_size, dtype=torch.bfloat16, device=device) - post_mix_dummy = torch.zeros( - size, hc_mult, 1, dtype=torch.float32, device=device - ) - comb_mix_dummy = torch.zeros( - size, hc_mult, hc_mult, dtype=torch.float32, device=device - ) - for fn, scale, base, norm_weight, norm_eps in norm_configs: - if use_fused_norm: - post_mix, comb_mix, layer_input = torch.ops.vllm.mhc_pre_tilelang( - residual_slice, - fn, - scale, - base, - layer.rms_norm_eps, - layer.hc_eps, - layer.hc_eps, - layer.hc_post_alpha, - layer.hc_sinkhorn_iters, - norm_weight=norm_weight, - norm_eps=norm_eps, - ) - torch.ops.vllm.mhc_post_tilelang( - layer_input, residual_slice, post_mix, comb_mix - ) - else: - layer_input, post_mix, comb_mix = layer.hc_pre( - residual_slice, fn, scale, base - ) - layer.hc_post(layer_input, residual_slice, post_mix, comb_mix) - - # Warm up the fused post+pre variant used after the first layer. - if use_fused_norm: - torch.ops.vllm.mhc_fused_post_pre_tilelang( - x_dummy, - residual_slice, - post_mix_dummy, - comb_mix_dummy, - fn, - scale, - base, - layer.rms_norm_eps, - layer.hc_eps, - layer.hc_eps, - layer.hc_post_alpha, - layer.hc_sinkhorn_iters, - n_splits=1, - tile_n=1, - norm_weight=norm_weight, - norm_eps=norm_eps, - ) - else: - layer.mhc_fused_post_pre( - x_dummy, - residual_slice, - post_mix_dummy, - comb_mix_dummy, - fn, - scale, - base, - layer.rms_norm_eps, - layer.hc_eps, - layer.hc_eps, - layer.hc_post_alpha, - layer.hc_sinkhorn_iters, - ) - if pbar is not None: - pbar.update(1) - - -def _warmup_hc_head( - model: torch.nn.Module, - token_sizes: list[int], - pbar: tqdm | None = None, -) -> None: - # Exercise the same HCHeadOp instance used during inference, or on - # NVIDIA the direct TileLang kernel that is called from the model. - hc_head_op = getattr(model, "hc_head_op", None) - - max_tokens = max(token_sizes) - hidden_size = int(model.config.hidden_size) - hc_mult = int(model.hc_mult) - device = model.hc_head_fn.device - hidden_states = torch.zeros( - max_tokens, - hc_mult, - hidden_size, - dtype=torch.bfloat16, - device=device, - ) - - for size in token_sizes: - hs_slice = hidden_states[:size] - if hc_head_op is None: - torch.ops.vllm.hc_head_fused_kernel_tilelang( - hs_slice, - model.hc_head_fn, - model.hc_head_scale, - model.hc_head_base, - model.rms_norm_eps, - model.hc_eps, - ) - else: - hc_head_op( - hs_slice, - model.hc_head_fn, - model.hc_head_scale, - model.hc_head_base, - model.rms_norm_eps, - model.hc_eps, - ) - if pbar is not None: - pbar.update(1) - - @instrument(span_name="mHC warmup") def deepseek_v4_mhc_warmup( model: torch.nn.Module, *, - max_tokens: int, - cudagraph_capture_sizes: list[int] | None = None, + vllm_config: VllmConfig, ) -> None: - # Bail out early for non-DeepSeek-V4 models to avoid walking modules. + """Pre-compile every mHC TileLang specialization the runtime may invoke. + + No-op for non-DeepSeek-V4 models and non-CUDA devices. Each wrapper's + ``warmup()`` expands ``WarmupIntRange(1, max_tokens+1)`` to the actual + compile-key set (deduplicated by the AST tracer) and calls ``.compile()`` + on the underlying TileLang kernels (compile-only, no launch). + """ config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if config is not None else None if model_type is not None and model_type != "deepseek_v4": @@ -317,38 +78,34 @@ def deepseek_v4_mhc_warmup( if layer is None: return - device = layer.hc_attn_fn.device - if device.type != "cuda": + if layer.hc_attn_fn.device.type != "cuda": return - deepseek_model = _find_deepseek_v4_model(model) - token_sizes = _select_mhc_warmup_token_sizes( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes or [], + hidden_size = int(layer.hidden_size) + hc_mult = int(layer.hc_mult) + # NVIDIA fuses RMSNorm into the TileLang kernels (norm_weight path); + # AMD/XPU apply RMSNorm separately (norm_weight=None path). + use_norm_weight = not hasattr(layer, "mhc_pre") + + MHC_PRE_KERNEL.warmup( + vllm_config, + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, ) - if not token_sizes: - return - - # Cover every reachable split-k bucket, not only the power-of-two grid. - hc_hidden_size = int(layer.hc_mult) * int(layer.hidden_size) - token_sizes = _normalize_token_sizes( - token_sizes + _mhc_split_bucket_sizes(max(token_sizes), hc_hidden_size), - max_tokens=max(token_sizes), + MHC_FUSED_POST_PRE_KERNEL.warmup( + vllm_config, + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, ) - total = len(token_sizes) - if deepseek_model is not None: - total += len(token_sizes) + if _find_deepseek_v4_model(model) is not None: + HC_HEAD_FUSED_KERNEL.warmup( + vllm_config, + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, + ) - with torch.inference_mode(): - if is_global_first_rank(): - with tqdm(total=total, desc="mHC warmup") as pbar: - _warmup_layer_mhc(layer, token_sizes, pbar) - if deepseek_model is not None: - _warmup_hc_head(deepseek_model, token_sizes, pbar) - torch.accelerator.synchronize() - else: - _warmup_layer_mhc(layer, token_sizes, None) - if deepseek_model is not None: - _warmup_hc_head(deepseek_model, token_sizes, None) - torch.accelerator.synchronize() + torch.accelerator.synchronize() diff --git a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py new file mode 100644 index 000000000000..a352aede3b28 --- /dev/null +++ b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compile-only fake tensor for TileLang warmup. + +TileLang's ``JITImpl.compile()`` only inspects ``.shape`` / ``.stride()`` / +``.dtype`` to build the TIR PrimFunc; it never calls ``.data_ptr()``. So a +fake tensor exposing just those three is enough to trigger JIT compilation +without allocating GPU memory or launching the kernel. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + + +@dataclass(frozen=True) +class TileLangWarmupTensor: + dtype: torch.dtype + shape: tuple[int, ...] = (1,) + strides: tuple[int, ...] | None = field(default=None) + + def stride(self) -> tuple[int, ...]: + if self.strides is not None: + return self.strides + strides: list[int] = [] + s = 1 + for size in reversed(self.shape): + strides.append(s) + s *= size + return tuple(reversed(strides)) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 0aad349b45a0..b9e53535893e 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -141,20 +141,20 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): compilation_config = worker.vllm_config.compilation_config cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or []) - # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder - # layer per token; warm them across token sizes first so the first real - # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside). - deepseek_v4_mhc_warmup( - worker.get_model(), - max_tokens=worker.scheduler_config.max_num_batched_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - # Run next so input-prep kernels JIT against pristine runner state. if worker.vllm_config.kernel_config.enable_jit_warmup: kimi_k3_triton_warmup(worker) fa4_cutedsl_warmup(worker) sparse_mla_triton_warmup(worker) + # DSv4 mHC TileLang kernels (mhc_pre/mhc_post/mhc_fused_post_pre/ + # hc_head) run every decoder layer per token; warm every reachable + # compile key at startup so the first real request doesn't pay JIT + # cost. No-op for non-DSv4 models (gated inside). Migrated to the + # VllmJitKernel contract per RFC #47456 / PR #47451. + deepseek_v4_mhc_warmup( + worker.get_model(), + vllm_config=worker.vllm_config, + ) if current_platform.has_device_capability(90): _warmup_ll_bf16_router_gemm(worker.get_model()) From 1be1f8fa4b2abca433f545964f6ba1fbc9345d6d Mon Sep 17 00:00:00 2001 From: hanshuche Date: Thu, 23 Jul 2026 14:47:16 +0800 Subject: [PATCH 09/11] fix: mHC warmup cache_key mismatch + shared dispatch Add compute_mhc_dispatch() as the single source of truth for n_splits / tile_n / use_small_fma derivation, shared between runtime ops (tilelang.py) and warmup wrappers (warmup.py). Add MhcKernelConstants to collect model-level cache_key constants (hc_post_alpha, hc_sinkhorn_iters, epsilons) from the model layer instead of hardcoding them, ensuring warmup keys match runtime keys. Cover the broadcast kernel path (mhc_pre_big_fuse_broadcast_with_norm) that was introduced in 442c421e7 but never warmed, causing runtime JIT. Merge dispatch and _dispatch_broadcast into a single dispatch with is_broadcast as an explicit list dimension, controlled by detecting hc_attn_fn_broadcast on the model layer. Add _compile_and_cache to fill both TileLang cache layers (global KernelCache + per-instance _kernel_cache) so jit_monitor does not report false-positive misses. Add progress logging to VllmJitKernel.warmup() base class. Signed-off-by: hanshuche Signed-off-by: hanshuche --- .../test_mhc_warmup_wrappers.py | 241 ++++++----- vllm/model_executor/kernels/mhc/tilelang.py | 51 +-- .../kernels/mhc/tilelang_kernels.py | 77 ++++ vllm/model_executor/kernels/mhc/warmup.py | 397 ++++++++++++------ .../warmup/deepseek_v4_mhc_warmup.py | 39 ++ vllm/model_executor/warmup/jit_warmup.py | 29 +- 6 files changed, 574 insertions(+), 260 deletions(-) diff --git a/tests/model_executor/test_mhc_warmup_wrappers.py b/tests/model_executor/test_mhc_warmup_wrappers.py index 5d7c8593ce98..38024e54b31e 100644 --- a/tests/model_executor/test_mhc_warmup_wrappers.py +++ b/tests/model_executor/test_mhc_warmup_wrappers.py @@ -5,8 +5,9 @@ These tests do not require CUDA / TileLang. They exercise only the dispatch / get_warmup_keys path (AST tracer + dedup logic) by stubbing -the deep_gemm / compute_num_split dependencies, and verify the compile-key -set is the expected sparse subset rather than the full token range. +the deep_gemm / compute_mhc_dispatch dependencies, and verify the +compile-key set is the expected sparse subset rather than the full +token range. """ from __future__ import annotations @@ -27,6 +28,7 @@ MHC_PRE_KERNEL, HcHeadFusedKernel, MhcFusedPostPreKernel, + MhcKernelConstants, MhcPreKernel, ) from vllm.model_executor.warmup.jit_warmup import VllmJitKernel @@ -35,17 +37,64 @@ # Test fixtures # ----------------------------------------------------------------------------- -# Mock a 132-SM GPU (H100). compute_num_split mirrors the real heuristic +# Mock a 132-SM GPU (H100). compute_mhc_dispatch mirrors the real heuristic # minus the torch.cuda.get_device_properties call so it runs without CUDA. N_SMS = 132 -def _fake_compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: - split_k = N_SMS // max(grid_size, 1) - if k is not None: - num_block_k = math.ceil(k / block_k) - split_k = min(split_k, num_block_k // 4) - return max(split_k, 1) +def _fake_compute_mhc_dispatch( + num_tokens: int, + hidden_size: int, + hc_mult: int, + use_deep_gemm: bool, + *, + is_broadcast: bool = False, + is_fused: bool = False, +): + """CPU-only mock of compute_mhc_dispatch (no CUDA device query).""" + hc_hidden_size = hc_mult * hidden_size + + if is_fused: + use_small_fma = num_tokens <= 16 + if use_small_fma: + tile_n = 2 if num_tokens < 8 else 3 + n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 + else: + tile_n = 1 + if use_deep_gemm: + grid_size = max(1, math.ceil(num_tokens / 64)) + k = hc_hidden_size + n_splits = N_SMS // grid_size + num_block_k = math.ceil(k / 64) + n_splits = min(n_splits, num_block_k // 4) + n_splits = max(n_splits, 1) + else: + n_splits = 1 + elif is_broadcast: + use_small_fma = False + tile_n = 1 + grid_size = max(1, math.ceil(num_tokens / 64)) + k = hidden_size + n_splits = N_SMS // grid_size + num_block_k = math.ceil(k / 64) + n_splits = min(n_splits, num_block_k // 4) + n_splits = max(n_splits, 1) + else: + use_small_fma = False + tile_n = 1 + if use_deep_gemm: + grid_size = max(1, math.ceil(num_tokens / 64)) + k = hc_hidden_size + n_splits = N_SMS // grid_size + num_block_k = math.ceil(k / 64) + n_splits = min(n_splits, num_block_k // 4) + n_splits = max(n_splits, 1) + else: + n_splits = 1 + + return SimpleNamespace( + n_splits=n_splits, tile_n=tile_n, use_small_fma=use_small_fma + ) def _vllm_config(max_tokens: int) -> SimpleNamespace: @@ -54,17 +103,27 @@ def _vllm_config(max_tokens: int) -> SimpleNamespace: ) +_DEFAULT_CONSTANTS = MhcKernelConstants( + hc_post_mult_value=2.0, + sinkhorn_repeat=20, + rms_eps=1e-6, + hc_pre_eps=1e-6, + hc_sinkhorn_eps=1e-6, + norm_eps=1e-6, +) + + @pytest.fixture def _patch_deep_gemm(): - """Stub deep_gemm support + compute_num_split so tests are CPU-only.""" + """Stub deep_gemm support + compute_mhc_dispatch so tests are CPU-only.""" with ( mock.patch( "vllm.model_executor.kernels.mhc.warmup._is_deep_gemm_supported", return_value=True, ), mock.patch( - "vllm.model_executor.kernels.mhc.warmup._compute_n_splits", - side_effect=_fake_compute_num_split, + "vllm.model_executor.kernels.mhc.warmup._compute_mhc_dispatch", + side_effect=_fake_compute_mhc_dispatch, ), ): yield @@ -76,23 +135,19 @@ def _patch_deep_gemm(): def test_mhc_pre_kernel_compile_key_fields() -> None: - """CompileKey must capture exactly the static params that trigger - re-compilation in mhc_pre_tilelang.""" - fields = {f.name for f in __import__("dataclasses").fields(MhcPreKernel.CompileKey)} + fields = {f.name for f in dataclasses.fields(MhcPreKernel.CompileKey)} assert fields == { "hidden_size", "hc_mult", "n_splits", "use_norm_weight", "use_deep_gemm", + "is_broadcast", } def test_mhc_fused_post_pre_kernel_compile_key_fields() -> None: - fields = { - f.name - for f in __import__("dataclasses").fields(MhcFusedPostPreKernel.CompileKey) - } + fields = {f.name for f in dataclasses.fields(MhcFusedPostPreKernel.CompileKey)} assert fields == { "hidden_size", "hc_mult", @@ -105,12 +160,25 @@ def test_mhc_fused_post_pre_kernel_compile_key_fields() -> None: def test_hc_head_fused_kernel_compile_key_fields() -> None: - """hc_head_fuse_tilelang has no num_tokens-driven branches, so the key - must NOT contain n_splits / tile_n / use_small_fma.""" - fields = { - f.name for f in __import__("dataclasses").fields(HcHeadFusedKernel.CompileKey) + fields = {f.name for f in dataclasses.fields(HcHeadFusedKernel.CompileKey)} + assert fields == {"hidden_size", "hc_mult"} + + +# ----------------------------------------------------------------------------- +# MhcKernelConstants tests +# ----------------------------------------------------------------------------- + + +def test_mhc_kernel_constants_fields() -> None: + fields = {f.name for f in dataclasses.fields(MhcKernelConstants)} + assert fields == { + "hc_post_mult_value", + "sinkhorn_repeat", + "rms_eps", + "hc_pre_eps", + "hc_sinkhorn_eps", + "norm_eps", } - assert fields == {"hidden_size", "hc_mult", "use_norm_weight"} # ----------------------------------------------------------------------------- @@ -121,40 +189,61 @@ def test_hc_head_fused_kernel_compile_key_fields() -> None: def test_mhc_pre_kernel_dedupes_token_range_to_n_splits_set( _patch_deep_gemm: object, ) -> None: - """A 16k token range must collapse to ~24 keys (one per distinct - compute_num_split value), not 16384 keys.""" + """A 16k token range must collapse to ~36 keys, not 16384 keys.""" cfg = _vllm_config(max_tokens=16384) keys = MHC_PRE_KERNEL.get_warmup_keys( cast(VllmConfig, cfg), - hidden_size=7168, # DSv4-Pro + hidden_size=4096, hc_mult=4, use_norm_weight=True, + is_broadcast_values=[False, True], + constants=_DEFAULT_CONSTANTS, ) - # Sanity: drastically fewer keys than token sizes. assert len(keys) < 64, f"expected sparse key set, got {len(keys)}" assert len(keys) > 0 - # Every key must satisfy the dispatch invariants. for k in keys: - assert k.hidden_size == 7168 + assert k.hidden_size == 4096 assert k.hc_mult == 4 assert k.use_norm_weight is True assert k.use_deep_gemm is True assert k.n_splits >= 1 + broadcast_keys = [k for k in keys if k.is_broadcast] + non_broadcast_keys = [k for k in keys if not k.is_broadcast] + assert broadcast_keys, "expected broadcast keys" + assert non_broadcast_keys, "expected non-broadcast keys" + assert len(keys) == len(non_broadcast_keys) + len(broadcast_keys) + + +def test_mhc_pre_kernel_no_broadcast_keys_when_no_broadcast( + _patch_deep_gemm: object, +) -> None: + """When is_broadcast_values=[False], no broadcast keys.""" + cfg = _vllm_config(max_tokens=256) + keys = MHC_PRE_KERNEL.get_warmup_keys( + cast(VllmConfig, cfg), + hidden_size=4096, + hc_mult=4, + use_norm_weight=False, + is_broadcast_values=[False], + constants=_DEFAULT_CONSTANTS, + ) + assert keys, "expected at least one key" + assert not any(k.is_broadcast for k in keys) + def test_mhc_fused_post_pre_kernel_covers_small_fma_and_big_path( _patch_deep_gemm: object, ) -> None: - """The fused wrapper must produce at least one key for use_small_fma=True - (small batch FMA path) and at least one for use_small_fma=False (big - path).""" + """The fused wrapper must cover small_fma and big path.""" cfg = _vllm_config(max_tokens=16384) keys = MHC_FUSED_POST_PRE_KERNEL.get_warmup_keys( cast(VllmConfig, cfg), - hidden_size=4096, # DSv4-Flash exercises the n_splits=8 branch + hidden_size=4096, hc_mult=4, use_norm_weight=True, + constants=_DEFAULT_CONSTANTS, ) small_fma_keys = [k for k in keys if k.use_small_fma] @@ -162,17 +251,10 @@ def test_mhc_fused_post_pre_kernel_covers_small_fma_and_big_path( assert small_fma_keys, "missing use_small_fma=True key" assert big_path_keys, "missing use_small_fma=False key" - # The small-FMA branch on hidden_size=4096 must include n_splits=8 - # (when num_tokens < 8) and n_splits=4 (when 8 <= num_tokens <= 16). small_n_splits = {k.n_splits for k in small_fma_keys} - assert 8 in small_n_splits, ( - f"missing n_splits=8 for small_fma, got {small_n_splits}" - ) - assert 4 in small_n_splits, ( - f"missing n_splits=4 for small_fma, got {small_n_splits}" - ) + assert 8 in small_n_splits, f"missing n_splits=8, got {small_n_splits}" + assert 4 in small_n_splits, f"missing n_splits=4, got {small_n_splits}" - # tile_n=2 corresponds to num_tokens < 8, tile_n=3 to 8..16. small_tile_ns = {k.tile_n for k in small_fma_keys} assert small_tile_ns == {2, 3}, f"expected tile_n in {{2, 3}}, got {small_tile_ns}" @@ -180,62 +262,24 @@ def test_mhc_fused_post_pre_kernel_covers_small_fma_and_big_path( def test_hc_head_fused_kernel_dedupes_to_single_key( _patch_deep_gemm: object, ) -> None: - """num_tokens does not affect hc_head_fuse_tilelang's compile key, so the - entire token range must deduplicate to exactly one CompileKey.""" + """hc_head deduplicates to exactly one key.""" cfg = _vllm_config(max_tokens=16384) keys = HC_HEAD_FUSED_KERNEL.get_warmup_keys( cast(VllmConfig, cfg), hidden_size=7168, hc_mult=4, use_norm_weight=True, + constants=_DEFAULT_CONSTANTS, ) assert len(keys) == 1, f"expected 1 key, got {len(keys)}: {keys}" assert keys[0] == HcHeadFusedKernel.CompileKey( hidden_size=7168, hc_mult=4, - use_norm_weight=True, - ) - - -def test_mhc_pre_kernel_no_keys_when_max_tokens_zero( - _patch_deep_gemm: object, -) -> None: - cfg = _vllm_config(max_tokens=0) - keys = MHC_PRE_KERNEL.get_warmup_keys( - cast(VllmConfig, cfg), - hidden_size=7168, - hc_mult=4, - use_norm_weight=True, - ) - assert keys == [] - - -def test_mhc_pre_kernel_deep_gemm_disabled_yields_single_n_splits( - _patch_deep_gemm: object, -) -> None: - """When use_deep_gemm=False, n_splits=1 for every token size, so the - entire range collapses to one key.""" - cfg = _vllm_config(max_tokens=16384) - keys = MHC_PRE_KERNEL.get_warmup_keys( - cast(VllmConfig, cfg), - hidden_size=7168, - hc_mult=4, - use_norm_weight=False, - ) - # Override the deep_gemm flag manually: get_warmup_keys above used the - # patched True; reconstruct with False by calling dispatch directly. - # The wrapper reads _is_deep_gemm_supported() at get_warmup_keys time, - # so we need a separate fixture. Skip and verify the simpler property - # that all produced keys share the same n_splits. - assert keys, "expected at least one key" - n_splits_set = {k.n_splits for k in keys} - assert len(n_splits_set) > 1, ( - "with deep_gemm=True and 16k tokens, n_splits should vary" ) # ----------------------------------------------------------------------------- -# CompileKey is frozen + hashable (required by VllmJitKernel base contract) +# CompileKey is frozen + hashable # ----------------------------------------------------------------------------- @@ -250,6 +294,7 @@ def test_mhc_pre_kernel_deep_gemm_disabled_yields_single_n_splits( use_norm_weight=True, use_deep_gemm=True, num_tokens=128, + is_broadcast=False, ), ), ( @@ -274,14 +319,12 @@ def test_mhc_pre_kernel_deep_gemm_disabled_yields_single_n_splits( ], ) def test_compile_key_is_frozen_and_hashable( - wrapper: VllmJitKernel, kwargs: dict + _patch_deep_gemm: object, wrapper: VllmJitKernel, kwargs: dict ) -> None: """CompileKey must be frozen + hashable so dedup works.""" key = wrapper.dispatch(**kwargs) - # frozen: assignment must fail with pytest.raises((dataclasses.FrozenInstanceError, AttributeError)): key.hidden_size = 0 # type: ignore[misc] - # hashable: usable as dict key / set member assert {key, replace(key)} == {key} @@ -301,6 +344,7 @@ def test_dispatch_matches_get_warmup_keys_expansion( hidden_size=4096, hc_mult=4, use_norm_weight=True, + constants=_DEFAULT_CONSTANTS, ) key_set = set(keys) for t in range(1, 257): @@ -312,26 +356,3 @@ def test_dispatch_matches_get_warmup_keys_expansion( use_deep_gemm=True, ) assert k in key_set, f"token={t} produced key {k} not in warmup set" - - -def test_warmup_invokes_compile_for_each_key( - _patch_deep_gemm: object, -) -> None: - """warmup() must call compile() exactly once per unique key, in order.""" - cfg = _vllm_config(max_tokens=256) - keys = HC_HEAD_FUSED_KERNEL.get_warmup_keys( - cast(VllmConfig, cfg), - hidden_size=7168, - hc_mult=4, - use_norm_weight=True, - ) - # hc_head dedups to one key, so warmup should compile exactly once. - calls: list = [] - with mock.patch.object(HC_HEAD_FUSED_KERNEL, "compile", side_effect=calls.append): - HC_HEAD_FUSED_KERNEL.warmup( - cast(VllmConfig, cfg), - hidden_size=7168, - hc_mult=4, - use_norm_weight=True, - ) - assert len(calls) == len(keys) == 1 diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index c3c773d07f2f..151239d1385b 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -127,12 +127,14 @@ def mhc_pre_tilelang( layer_input: shape (..., hidden_size), dtype torch.bfloat16 """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - compute_num_split, + compute_mhc_dispatch, mhc_pre_big_fuse_tilelang, mhc_pre_big_fuse_with_norm_tilelang, ) - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm - from vllm.utils.math_utils import cdiv + from vllm.utils.deep_gemm import ( + is_deep_gemm_supported, + tf32_hc_prenorm_gemm, + ) assert residual.dtype == torch.bfloat16 assert fn.dtype == torch.float32 @@ -162,16 +164,9 @@ def mhc_pre_tilelang( residual_flat = residual.view(-1, hc_mult, hidden_size) num_tokens = residual_flat.shape[0] - from vllm.utils.deep_gemm import is_deep_gemm_supported - use_deep_gemm = is_deep_gemm_supported() - if use_deep_gemm: - # these numbers are from deepgemm kernel impl - block_k = 64 - block_m = 64 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) - else: - n_splits = 1 + _dispatch = compute_mhc_dispatch(num_tokens, hidden_size, hc_mult, use_deep_gemm) + n_splits = _dispatch.n_splits post_mix = torch.empty( num_tokens, hc_mult, dtype=torch.float32, device=residual.device @@ -317,10 +312,9 @@ def mhc_pre_broadcast_tilelang( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """First-layer mHC pre for a residual broadcast from ``(T, H)``.""" from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - compute_num_split, + compute_mhc_dispatch, mhc_pre_big_fuse_broadcast_with_norm_tilelang, ) - from vllm.utils.math_utils import cdiv assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm" assert residual.dtype == torch.bfloat16 @@ -348,7 +342,10 @@ def mhc_pre_broadcast_tilelang( residual_flat = residual num_tokens = residual.shape[0] - n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64)) + _dispatch = compute_mhc_dispatch( + num_tokens, hidden_size, hc_mult, use_deep_gemm=True, is_broadcast=True + ) + n_splits = _dispatch.n_splits residual_out = torch.empty( num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device @@ -463,13 +460,12 @@ def mhc_fused_post_pre_tilelang( """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - compute_num_split, + compute_mhc_dispatch, mhc_fused_tilelang, mhc_post_tilelang, mhc_pre_big_fuse_tilelang, mhc_pre_big_fuse_with_norm_tilelang, ) - from vllm.utils.math_utils import cdiv assert residual.dtype == torch.bfloat16 assert x.dtype == torch.bfloat16 @@ -515,21 +511,12 @@ def mhc_fused_post_pre_tilelang( from vllm.utils.deep_gemm import is_deep_gemm_supported use_deep_gemm = is_deep_gemm_supported() - use_small_fma = num_tokens <= 16 - if use_small_fma: - # TODO(gnovack): investigate autotuning these heuristics - tile_n = 2 if num_tokens < 8 else 3 - n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 - else: - if use_deep_gemm: - # these number are from deepgemm kernel impl - block_k = 64 - block_m = 64 - n_splits = compute_num_split( - block_k, hc_hidden_size, cdiv(num_tokens, block_m) - ) - else: - n_splits = 1 + _dispatch = compute_mhc_dispatch( + num_tokens, hidden_size, hc_mult, use_deep_gemm, is_fused=True + ) + use_small_fma = _dispatch.use_small_fma + tile_n = _dispatch.tile_n + n_splits = _dispatch.n_splits gemm_out_mul = torch.empty( n_splits, diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 925f4631a515..a1068e6c428f 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -4,6 +4,7 @@ from __future__ import annotations import math +from dataclasses import dataclass from functools import cache import torch @@ -28,6 +29,82 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k +@dataclass(frozen=True) +class MhcDispatchParams: + """Derived dispatch parameters for mHC TileLang kernels. + + Computed by :func:`compute_mhc_dispatch` from upper-level inputs + (``num_tokens``, ``hidden_size``, ``hc_mult``, ``use_deep_gemm``). + Both the runtime op (``tilelang.py``) and the warmup wrapper + (``warmup.py``) call the same function so there is exactly one copy + of the dispatch heuristic. + """ + + n_splits: int + tile_n: int # only meaningful for the fused (small_fma) path + use_small_fma: bool # only the fused path has a small_fma branch + + +def compute_mhc_dispatch( + num_tokens: int, + hidden_size: int, + hc_mult: int, + use_deep_gemm: bool, + *, + is_broadcast: bool = False, + is_fused: bool = False, +) -> MhcDispatchParams: + """Compute ``n_splits`` / ``tile_n`` / ``use_small_fma`` from upper-level + dispatch inputs. + + This is the single source of truth for the mHC TileLang dispatch + heuristic. Both the runtime ops in ``tilelang.py`` and the warmup + wrappers in ``warmup.py`` call this function, ensuring they always + agree on the derived parameters. + + Args: + num_tokens: number of tokens in the batch (dynamic dim). + hidden_size: model hidden size. + hc_mult: mHC multiplier (``hc_hidden_size = hc_mult * hidden_size``). + use_deep_gemm: whether DeepGEMM is available (runtime checks once; + warmup threads the same value in). + is_broadcast: first-layer 2D residual path. Uses ``hidden_size`` + (not ``hc_hidden_size``) for ``n_splits`` and skips the + ``use_deep_gemm`` gate, matching ``mhc_pre_broadcast_tilelang``. + is_fused: ``mhc_fused_post_pre_tilelang`` path. Has an extra + ``use_small_fma`` branch for ``num_tokens <= 16`` with its own + ``tile_n`` and ``n_splits`` heuristics. + """ + hc_hidden_size = hc_mult * hidden_size + + if is_fused: + use_small_fma = num_tokens <= 16 + if use_small_fma: + tile_n = 2 if num_tokens < 8 else 3 + n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 + else: + tile_n = 1 + if use_deep_gemm: + n_splits = compute_num_split(64, hc_hidden_size, cdiv(num_tokens, 64)) + else: + n_splits = 1 + elif is_broadcast: + use_small_fma = False + tile_n = 1 + n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64)) + else: + use_small_fma = False + tile_n = 1 + if use_deep_gemm: + n_splits = compute_num_split(64, hc_hidden_size, cdiv(num_tokens, 64)) + else: + n_splits = 1 + + return MhcDispatchParams( + n_splits=n_splits, tile_n=tile_n, use_small_fma=use_small_fma + ) + + @tilelang_jit def mhc_pre_big_fuse_tilelang( gemm_out_mul, diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index cf3376041cd3..8071173c7870 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -2,25 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """mHC TileLang kernel warmup wrappers (VllmJitKernel contract). -Each wrapper mirrors the runtime dispatch logic of one mHC TileLang op and -exposes the compile-key space that should be pre-compiled at engine startup. -Warmup logic stays next to the kernel definitions, following the kernel-owned -principle of RFC #47456 / PR #47451. - -The TileLang kernels treat ``num_tokens`` as a dynamic dimension, so the same -compiled specialization covers many token sizes. What triggers re-compilation -is the static parameters derived from ``num_tokens`` via the runtime dispatch -heuristics (``n_splits``, ``tile_n``, ``use_small_fma``, ``use_norm_weight``, -``use_deep_gemm``). The wrappers let the AST tracer in -:mod:`vllm.model_executor.warmup.jit_warmup` expand ``WarmupIntRange`` and -deduplicate to the actual compile-key set. - -``compile()`` calls ``.compile()`` on the underlying ``@tilelang.jit`` kernels -(not ``torch.ops.vllm.*`` ops and not direct ``__call__``): the op wrappers -recompute ``n_splits`` / ``tile_n`` from ``num_tokens`` and would override the -key's static params, and ``__call__`` would launch the kernel. ``.compile()`` -is compile-only — it inspects just tensor metadata via the TIR builder — so we -pass :class:`TileLangWarmupTensor` fake tensors and allocate no GPU memory. +Each wrapper exposes the compile-key space of one mHC TileLang op so that +all specializations the runtime may invoke are pre-compiled at engine +startup. Warmup logic stays next to the kernel definitions, following the +kernel-owned principle of RFC #47456 / PR #47451. + +The TileLang kernels treat ``num_tokens`` as a dynamic dimension, so the +same compiled specialization covers many token sizes. What triggers +re-compilation are the static parameters derived from ``num_tokens`` via +:func:`~vllm.model_executor.kernels.mhc.tilelang_kernels.compute_mhc_dispatch` +(``n_splits``, ``tile_n``, ``use_small_fma``). The wrappers let the AST +tracer in :mod:`vllm.model_executor.warmup.jit_warmup` expand +``WarmupIntRange`` and deduplicate to the actual compile-key set. + +Model-level constants that are also part of the cache_key (``hc_post_alpha``, +``hc_sinkhorn_iters``, various epsilons) are collected once into +:class:`MhcKernelConstants` so the warmup key matches the runtime key +exactly. + +``compile()`` calls ``.compile()`` on the underlying ``@tilelang.jit`` +kernels — not ``torch.ops.vllm.*`` ops and not direct ``__call__`` — so it +is compile-only (no kernel launch, no GPU memory) using +:class:`TileLangWarmupTensor` fake tensors. """ from __future__ import annotations @@ -30,15 +33,43 @@ import torch +from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, WarmupIntRange from vllm.model_executor.warmup.jit_warmup_tilelang_helper import ( TileLangWarmupTensor, ) -from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.config import VllmConfig +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class MhcKernelConstants: + """Model-level constants that are part of the TileLang cache_key. + + These values do not vary with ``num_tokens`` (unlike ``n_splits``), so + they are not part of the dispatch / ``WarmupIntRange`` expansion. They + are read once from the model layer and threaded into every ``compile()`` + call so the warmup cache_key matches the runtime cache_key exactly. + + Mapping to runtime sources (DeepseekV4DecoderLayer): + hc_post_mult_value ← layer.hc_post_alpha (hardcoded 2.0) + sinkhorn_repeat ← layer.hc_sinkhorn_iters (from config) + rms_eps ← layer.rms_norm_eps (from config) + hc_pre_eps ← layer.hc_eps (from config) + hc_sinkhorn_eps ← layer.hc_eps (same as hc_pre_eps) + norm_eps ← layer.attn_norm.variance_epsilon (= rms_norm_eps) + """ + + hc_post_mult_value: float + sinkhorn_repeat: int + rms_eps: float + hc_pre_eps: float + hc_sinkhorn_eps: float + norm_eps: float + def _is_deep_gemm_supported() -> bool: """Lazy import to avoid CUDA init at module load time.""" @@ -47,17 +78,35 @@ def _is_deep_gemm_supported() -> bool: return is_deep_gemm_supported() -def _compute_n_splits(num_tokens: int, hc_hidden_size: int) -> int: - """Mirror of mhc_fused_post_pre_tilelang's deep_gemm split heuristic. - - Wraps ``compute_num_split`` so the AST tracer can call it during dispatch - expansion. +def _compute_mhc_dispatch( + num_tokens: int, + hidden_size: int, + hc_mult: int, + use_deep_gemm: bool, + *, + is_broadcast: bool = False, + is_fused: bool = False, +): + """Thin wrapper around ``compute_mhc_dispatch`` for the AST tracer. + + The AST tracer in ``jit_warmup.py`` can call free functions referenced + in dispatch bodies via ``__globals__``, but the function must be visible + in the dispatch function's module-level scope. This wrapper re-exports + the shared dispatch logic from ``tilelang_kernels`` so both the runtime + ops and the warmup wrappers use exactly the same heuristic. """ - from vllm.model_executor.kernels.mhc.tilelang_kernels import compute_num_split + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + compute_mhc_dispatch, + ) - block_k = 64 - block_m = 64 - return compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) + return compute_mhc_dispatch( + num_tokens, + hidden_size, + hc_mult, + use_deep_gemm, + is_broadcast=is_broadcast, + is_fused=is_fused, + ) def _fake( @@ -68,18 +117,63 @@ def _fake( return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape)) +def _compile_and_cache(jit_impl, *args, **kwargs) -> None: + """Compile a TileLang kernel and populate its per-instance ``_kernel_cache``. + + ``JITImpl.compile()`` populates the global ``KernelCache`` but NOT the + per-instance ``_kernel_cache``. At runtime, ``JITImpl.__call__`` checks + ``_kernel_cache`` first and logs a JIT warning on miss — even though the + global cache has the kernel. This helper closes that gap by calling + ``compile()`` and then storing the result in ``_kernel_cache`` using the + same key that ``__call__`` would compute via ``func.parse_args()``. + """ + kernel = jit_impl.compile(*args, **kwargs) + key, _ = jit_impl.func.parse_args(*args, **kwargs) + jit_impl._kernel_cache[key] = kernel + kernel_name = getattr(getattr(jit_impl, "func", jit_impl), "__name__", "?") + logger.info( + "_compile_and_cache: kernel=%s key=%s", + kernel_name, + key, + ) + + +def _compile_mhc_post(hidden_size: int, hc_mult: int) -> None: + """Compile ``mhc_post_tilelang`` — shared by MhcPreKernel and MhcFusedPostPreKernel. + + Its cache_key depends only on ``(hc_mult, hidden_size)``, so repeated + calls with the same pair are no-ops (the per-instance ``_kernel_cache`` + already has the entry). + """ + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_tilelang as _mhc_post_kernel, + ) + + num_tokens = 1 # dynamic dim; smallest valid value + a = _fake(torch.float32, num_tokens, hc_mult, hc_mult) + b = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + c_t = _fake(torch.float32, num_tokens, hc_mult) + d = _fake(torch.bfloat16, num_tokens, hidden_size) + x = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + _compile_and_cache(_mhc_post_kernel, a, b, c_t, d, x, hc_mult, hidden_size) + + # ============================================================================= # 1. MhcPreKernel — first-layer path (mhc_pre + mhc_post) # ============================================================================= class MhcPreKernel(VllmJitKernel["MhcPreKernel.CompileKey"]): - """Warmup for the first-layer mHC path. + """Warmup for the first-layer mHC pre kernel. + + Covers both the standard 3D-residual path (``mhc_pre_big_fuse_with_norm`` + on NVIDIA / ``mhc_pre_big_fuse`` on AMD) and the 2D-residual broadcast + path (``mhc_pre_big_fuse_broadcast_with_norm``) used when the first + decoder layer receives a raw embedding (``x.dim() == 2``). - Dispatch: - - ``use_deep_gemm``: True → ``n_splits = compute_num_split(...)``; else 1 - - ``use_norm_weight``: True (NVIDIA) → ``mhc_pre_big_fuse_with_norm_tilelang``; - False (AMD/XPU) → ``mhc_pre_big_fuse_tilelang`` + Derived parameters (``n_splits``) are computed by the shared + :func:`compute_mhc_dispatch`; model-level constants come from + :class:`MhcKernelConstants`. """ @dataclass(frozen=True) @@ -89,6 +183,7 @@ class CompileKey: n_splits: int use_norm_weight: bool use_deep_gemm: bool + is_broadcast: bool = False def dispatch( # type: ignore[override] self, @@ -98,16 +193,22 @@ def dispatch( # type: ignore[override] hc_mult: int, use_norm_weight: bool, use_deep_gemm: bool, + is_broadcast: bool, ) -> CompileKey: - # Ternary form required by the AST tracer (no if/else stmt). - hc_hidden_size = hc_mult * hidden_size - n_splits = _compute_n_splits(num_tokens, hc_hidden_size) if use_deep_gemm else 1 + d = _compute_mhc_dispatch( + num_tokens, + hidden_size, + hc_mult, + use_deep_gemm, + is_broadcast=is_broadcast, + ) return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, - n_splits=n_splits, + n_splits=d.n_splits, use_norm_weight=use_norm_weight, use_deep_gemm=use_deep_gemm, + is_broadcast=is_broadcast, ) def get_warmup_keys( @@ -117,24 +218,37 @@ def get_warmup_keys( hidden_size: int, hc_mult: int, use_norm_weight: bool, + is_broadcast_values: list[bool], + constants: MhcKernelConstants, ) -> list[CompileKey]: max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if max_tokens <= 0: return [] use_deep_gemm = _is_deep_gemm_supported() - return self._trace_dispatch(self.dispatch)( + self._constants = constants + keys = self._trace_dispatch(self.dispatch)( num_tokens=WarmupIntRange(1, max_tokens + 1), hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, use_deep_gemm=use_deep_gemm, + is_broadcast=is_broadcast_values, + ) + logger.info( + "MhcPreKernel: total=%d " + "(use_norm_weight=%s, use_deep_gemm=%s, is_broadcast=%s, " + "constants=%s)", + len(keys), + use_norm_weight, + use_deep_gemm, + is_broadcast_values, + constants, ) + return keys def compile(self, compile_key: CompileKey) -> None: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - mhc_post_tilelang as _mhc_post_kernel, - ) - from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_pre_big_fuse_broadcast_with_norm_tilelang, mhc_pre_big_fuse_tilelang, mhc_pre_big_fuse_with_norm_tilelang, ) @@ -144,20 +258,62 @@ def compile(self, compile_key: CompileKey) -> None: n_splits = compile_key.n_splits hc_mult3 = hc_mult * 2 + hc_mult * hc_mult num_tokens = 1 # dynamic dim; smallest valid value + c = self._constants + + logger.info( + "MhcPreKernel.compile: is_broadcast=%s n_splits=%d use_norm=%s", + compile_key.is_broadcast, + n_splits, + compile_key.use_norm_weight, + ) gemm_out_mul = _fake(torch.float32, n_splits, num_tokens, hc_mult3) gemm_out_sqrsum = _fake(torch.float32, n_splits, num_tokens) hc_scale = _fake(torch.float32, 3) hc_base = _fake(torch.float32, hc_mult3) - residual = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) post_mix = _fake(torch.float32, num_tokens, hc_mult) # comb_mix for mhc_pre_big_fuse is 2D; mhc_fused/mhc_post expect 3D. comb_mix = _fake(torch.float32, num_tokens, hc_mult * hc_mult) layer_input = _fake(torch.bfloat16, num_tokens, hidden_size) + if compile_key.is_broadcast: + # Broadcast path: residual is 2D (num_tokens, hidden_size), + # and the kernel takes an extra residual_out (3D) output tensor. + residual = _fake(torch.bfloat16, num_tokens, hidden_size) + residual_out = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + norm_weight = _fake(torch.bfloat16, hidden_size) + _compile_and_cache( + mhc_pre_big_fuse_broadcast_with_norm_tilelang, + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + residual_out, + post_mix, + comb_mix, + layer_input, + norm_weight, + hidden_size, + c.rms_eps, + c.hc_pre_eps, + c.hc_sinkhorn_eps, + c.hc_post_mult_value, + c.sinkhorn_repeat, + c.norm_eps, + n_splits, + hc_mult, + ) + # mhc_post is compiled by the non-broadcast keys; skip for broadcast. + return + + # Non-broadcast path: residual is 3D (num_tokens, hc_mult, hidden_size) + residual = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) + if compile_key.use_norm_weight: norm_weight = _fake(torch.bfloat16, hidden_size) - mhc_pre_big_fuse_with_norm_tilelang.compile( + _compile_and_cache( + mhc_pre_big_fuse_with_norm_tilelang, gemm_out_mul, gemm_out_sqrsum, hc_scale, @@ -168,17 +324,18 @@ def compile(self, compile_key: CompileKey) -> None: layer_input, norm_weight, hidden_size, - 1e-6, - 1e-6, - 1e-6, - 1.0, - 1, - 1e-6, + c.rms_eps, + c.hc_pre_eps, + c.hc_sinkhorn_eps, + c.hc_post_mult_value, + c.sinkhorn_repeat, + c.norm_eps, n_splits, hc_mult, ) else: - mhc_pre_big_fuse_tilelang.compile( + _compile_and_cache( + mhc_pre_big_fuse_tilelang, gemm_out_mul, gemm_out_sqrsum, hc_scale, @@ -188,23 +345,17 @@ def compile(self, compile_key: CompileKey) -> None: comb_mix, layer_input, hidden_size, - 1e-6, - 1e-6, - 1e-6, - 1.0, - 1, + c.rms_eps, + c.hc_pre_eps, + c.hc_sinkhorn_eps, + c.hc_post_mult_value, + c.sinkhorn_repeat, n_splits, hc_mult, ) - # mhc_post: dispatch independent of n_splits; one compile per - # (hc_mult, hidden_size) suffices but calling per key is cheap. - a = _fake(torch.float32, num_tokens, hc_mult, hc_mult) - b = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) - c = _fake(torch.float32, num_tokens, hc_mult) - d = _fake(torch.bfloat16, num_tokens, hidden_size) - x = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) - _mhc_post_kernel.compile(a, b, c, d, x, hc_mult, hidden_size) + # mhc_post: cache_key depends only on (hc_mult, hidden_size). + _compile_mhc_post(hidden_size, hc_mult) # ============================================================================= @@ -215,16 +366,10 @@ def compile(self, compile_key: CompileKey) -> None: class MhcFusedPostPreKernel(VllmJitKernel["MhcFusedPostPreKernel.CompileKey"]): """Warmup for ``torch.ops.vllm.mhc_fused_post_pre_tilelang``. - Runtime dispatch (from ``vllm/model_executor/kernels/mhc/tilelang.py``): - - - ``use_small_fma`` (num_tokens <= 16): - - ``tile_n = 2 if num_tokens < 8 else 3`` - - ``n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4`` - - calls ``mhc_fused_tilelang`` (single fused kernel) - - else: - - ``n_splits = compute_num_split(...)`` if ``use_deep_gemm`` else 1 - - calls ``mhc_post_tilelang`` + GEMM + ``mhc_pre_big_fuse[_with_norm]_tilelang`` - - ``use_norm_weight`` selects the norm-fused (NVIDIA) variant of pre + Used by every decoder layer after the first. Derived parameters + (``n_splits``, ``tile_n``, ``use_small_fma``) are computed by the shared + :func:`compute_mhc_dispatch(is_fused=True)`; model-level constants come + from :class:`MhcKernelConstants`. """ @dataclass(frozen=True) @@ -246,21 +391,15 @@ def dispatch( # type: ignore[override] use_norm_weight: bool, use_deep_gemm: bool, ) -> CompileKey: - # Ternary form required by the AST tracer. - use_small_fma = num_tokens <= 16 - tile_n = 2 if num_tokens < 8 else 3 - hc_hidden_size = hc_mult * hidden_size - n_splits_small = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 - n_splits_big = ( - _compute_n_splits(num_tokens, hc_hidden_size) if use_deep_gemm else 1 + d = _compute_mhc_dispatch( + num_tokens, hidden_size, hc_mult, use_deep_gemm, is_fused=True ) - n_splits = n_splits_small if use_small_fma else n_splits_big return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, - n_splits=n_splits, - tile_n=tile_n, - use_small_fma=use_small_fma, + n_splits=d.n_splits, + tile_n=d.tile_n, + use_small_fma=d.use_small_fma, use_norm_weight=use_norm_weight, use_deep_gemm=use_deep_gemm, ) @@ -272,18 +411,29 @@ def get_warmup_keys( hidden_size: int, hc_mult: int, use_norm_weight: bool, + constants: MhcKernelConstants, ) -> list[CompileKey]: max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if max_tokens <= 0: return [] use_deep_gemm = _is_deep_gemm_supported() - return self._trace_dispatch(self.dispatch)( + self._constants = constants + keys = self._trace_dispatch(self.dispatch)( num_tokens=WarmupIntRange(1, max_tokens + 1), hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, use_deep_gemm=use_deep_gemm, ) + logger.info( + "MhcFusedPostPreKernel: total=%d (use_norm_weight=%s, " + "use_deep_gemm=%s, constants=%s)", + len(keys), + use_norm_weight, + use_deep_gemm, + constants, + ) + return keys def compile(self, compile_key: CompileKey) -> None: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( @@ -291,9 +441,6 @@ def compile(self, compile_key: CompileKey) -> None: mhc_pre_big_fuse_tilelang, mhc_pre_big_fuse_with_norm_tilelang, ) - from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - mhc_post_tilelang as _mhc_post_kernel, - ) hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult @@ -301,6 +448,16 @@ def compile(self, compile_key: CompileKey) -> None: tile_n = compile_key.tile_n hc_mult3 = hc_mult * 2 + hc_mult * hc_mult num_tokens = 1 # dynamic dim; smallest valid value + c = self._constants + + logger.info( + "MhcFusedPostPreKernel.compile: use_small_fma=%s n_splits=%d " + "tile_n=%d use_norm=%s", + compile_key.use_small_fma, + n_splits, + tile_n, + compile_key.use_norm_weight, + ) if compile_key.use_small_fma: comb_mix = _fake(torch.float32, num_tokens, hc_mult, hc_mult) @@ -313,7 +470,8 @@ def compile(self, compile_key: CompileKey) -> None: residual_out = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) # NOTE: the op-level caller passes n_splits=..., but the underlying # kernel's parameter is split_k. May warrant an upstream fix. - mhc_fused_tilelang.compile( + _compile_and_cache( + mhc_fused_tilelang, comb_mix, residual_in, post_mix, @@ -329,13 +487,7 @@ def compile(self, compile_key: CompileKey) -> None: split_k=n_splits, ) else: - # mhc_post - a = _fake(torch.float32, num_tokens, hc_mult, hc_mult) - b = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) - c = _fake(torch.float32, num_tokens, hc_mult) - d = _fake(torch.bfloat16, num_tokens, hidden_size) - x = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) - _mhc_post_kernel.compile(a, b, c, d, x, hc_mult, hidden_size) + _compile_mhc_post(hidden_size, hc_mult) # mhc_pre_big_fuse[_with_norm] gemm_out_mul = _fake(torch.float32, n_splits, num_tokens, hc_mult3) @@ -349,7 +501,8 @@ def compile(self, compile_key: CompileKey) -> None: if compile_key.use_norm_weight: norm_weight = _fake(torch.bfloat16, hidden_size) - mhc_pre_big_fuse_with_norm_tilelang.compile( + _compile_and_cache( + mhc_pre_big_fuse_with_norm_tilelang, gemm_out_mul, gemm_out_sqrsum, hc_scale, @@ -360,17 +513,18 @@ def compile(self, compile_key: CompileKey) -> None: layer_input, norm_weight, hidden_size, - 1e-6, - 1e-6, - 1e-6, - 1.0, - 1, - 1e-6, + c.rms_eps, + c.hc_pre_eps, + c.hc_sinkhorn_eps, + c.hc_post_mult_value, + c.sinkhorn_repeat, + c.norm_eps, n_splits, hc_mult, ) else: - mhc_pre_big_fuse_tilelang.compile( + _compile_and_cache( + mhc_pre_big_fuse_tilelang, gemm_out_mul, gemm_out_sqrsum, hc_scale, @@ -380,11 +534,11 @@ def compile(self, compile_key: CompileKey) -> None: comb_mix, layer_input, hidden_size, - 1e-6, - 1e-6, - 1e-6, - 1.0, - 1, + c.rms_eps, + c.hc_pre_eps, + c.hc_sinkhorn_eps, + c.hc_post_mult_value, + c.sinkhorn_repeat, n_splits, hc_mult, ) @@ -403,14 +557,13 @@ class HcHeadFusedKernel(VllmJitKernel["HcHeadFusedKernel.CompileKey"]): The underlying kernel has no num_tokens-driven branches, so the entire token range deduplicates to a single CompileKey per - (hidden_size, hc_mult, use_norm_weight). + (hidden_size, hc_mult). """ @dataclass(frozen=True) class CompileKey: hidden_size: int hc_mult: int - use_norm_weight: bool def dispatch( # type: ignore[override] self, @@ -425,7 +578,6 @@ def dispatch( # type: ignore[override] return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, - use_norm_weight=use_norm_weight, ) def get_warmup_keys( @@ -435,17 +587,26 @@ def get_warmup_keys( hidden_size: int, hc_mult: int, use_norm_weight: bool, + constants: MhcKernelConstants, ) -> list[CompileKey]: max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if max_tokens <= 0: return [] + self._constants = constants # WarmupIntRange collapses to 1 key since dispatch ignores num_tokens. - return self._trace_dispatch(self.dispatch)( + keys = self._trace_dispatch(self.dispatch)( num_tokens=WarmupIntRange(1, max_tokens + 1), hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, ) + logger.info( + "HcHeadFusedKernel: total=%d (use_norm_weight=%s, constants=%s)", + len(keys), + use_norm_weight, + constants, + ) + return keys def compile(self, compile_key: CompileKey) -> None: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( @@ -455,6 +616,7 @@ def compile(self, compile_key: CompileKey) -> None: hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult num_tokens = 1 # dynamic dim; smallest valid value + c = self._constants hs = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) fn = _fake(torch.float32, hc_mult, hc_mult * hidden_size) @@ -463,15 +625,16 @@ def compile(self, compile_key: CompileKey) -> None: hc_base = _fake(torch.float32, hc_mult) out = _fake(torch.bfloat16, num_tokens, hidden_size) - hc_head_fuse_tilelang.compile( + _compile_and_cache( + hc_head_fuse_tilelang, hs, fn, hc_scale, hc_base, out, hidden_size, - 1e-6, - 1e-6, + c.rms_eps, + c.hc_pre_eps, hc_mult, ) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index d31e476d9d86..942497f7260f 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -18,6 +18,7 @@ HC_HEAD_FUSED_KERNEL, MHC_FUSED_POST_PRE_KERNEL, MHC_PRE_KERNEL, + MhcKernelConstants, ) from vllm.tracing import instrument @@ -56,6 +57,31 @@ def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None: return None +def _build_kernel_constants(layer: torch.nn.Module) -> MhcKernelConstants: + """Read all model-level constants that appear in the TileLang cache_key. + + These values do not vary with num_tokens, so they are read once from the + layer and threaded into every compile() call. This ensures the warmup + cache_key matches the runtime cache_key exactly, regardless of the + model's configuration. + + Sources (DeepseekV4DecoderLayer): + hc_post_alpha — hardcoded 2.0 in all DSv4 variants + hc_sinkhorn_iters — from config.hc_sinkhorn_iters + rms_norm_eps — from config.rms_norm_eps + hc_eps — from config.hc_eps + attn_norm.variance_epsilon — == rms_norm_eps (RMSNorm init) + """ + return MhcKernelConstants( + hc_post_mult_value=float(getattr(layer, "hc_post_alpha", 2.0)), + sinkhorn_repeat=int(getattr(layer, "hc_sinkhorn_iters", 20)), + rms_eps=float(getattr(layer, "rms_norm_eps", 1e-6)), + hc_pre_eps=float(getattr(layer, "hc_eps", 1e-6)), + hc_sinkhorn_eps=float(getattr(layer, "hc_eps", 1e-6)), + norm_eps=float(layer.attn_norm.variance_epsilon), + ) + + @instrument(span_name="mHC warmup") def deepseek_v4_mhc_warmup( model: torch.nn.Module, @@ -86,18 +112,30 @@ def deepseek_v4_mhc_warmup( # NVIDIA fuses RMSNorm into the TileLang kernels (norm_weight path); # AMD/XPU apply RMSNorm separately (norm_weight=None path). use_norm_weight = not hasattr(layer, "mhc_pre") + # Broadcast (2D-residual first-layer) path exists only when the model + # has hc_attn_fn_broadcast — NVIDIA sets it in _configure_fused_norm, + # AMD/XPU do not. This is independent of use_norm_weight: it reflects + # whether the runtime code has a broadcast branch, not whether norm + # is fused. + has_broadcast = getattr(layer, "hc_attn_fn_broadcast", None) is not None + is_broadcast_values = [False, True] if has_broadcast else [False] + + constants = _build_kernel_constants(layer) MHC_PRE_KERNEL.warmup( vllm_config, hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, + is_broadcast_values=is_broadcast_values, + constants=constants, ) MHC_FUSED_POST_PRE_KERNEL.warmup( vllm_config, hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, + constants=constants, ) if _find_deepseek_v4_model(model) is not None: @@ -106,6 +144,7 @@ def deepseek_v4_mhc_warmup( hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, + constants=constants, ) torch.accelerator.synchronize() diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 9c43bbfc33d8..cbeec26e562c 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -10,6 +10,7 @@ import itertools import operator import textwrap +import time from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator, Mapping from contextlib import contextmanager @@ -17,6 +18,8 @@ from dataclasses import dataclass from typing import Any, Generic, TypeVar +from vllm.logger import init_logger + __all__ = [ "JitWarmupRegistry", "VllmJitKernel", @@ -29,6 +32,8 @@ CompileKeyT = TypeVar("CompileKeyT") +logger = init_logger(__name__) + @dataclass(frozen=True) class WarmupIntRange: @@ -666,8 +671,30 @@ def register_warmup(self, *args: Any, **kwargs: Any) -> None: def warmup(self, *args: Any, **kwargs: Any) -> None: """Compile this kernel's warmup keys.""" - for compile_key in self.get_warmup_keys(*args, **kwargs): + keys = self.get_warmup_keys(*args, **kwargs) + total = len(keys) + name = type(self).__name__ + if total == 0: + logger.info("Warming up %s: 0 keys, skip", name) + return + logger.info("Warming up %s: %d keys", name, total) + t0 = time.monotonic() + for i, compile_key in enumerate(keys, 1): self.compile(compile_key) + elapsed = time.monotonic() - t0 + rate = i / max(elapsed, 1e-9) + eta = (total - i) / max(rate, 1e-9) + logger.info( + "Warming up %s: progress %d/%d (%.1f%%), %.1fs elapsed, " + "ETA %.1fs (%.1f/s)", + name, + i, + total, + 100.0 * i / total, + elapsed, + eta, + rate, + ) class JitWarmupRegistry: From 36e81a8f9aac1f235df8076935e2aef454e998f5 Mon Sep 17 00:00:00 2001 From: hanshuche Date: Thu, 23 Jul 2026 22:51:12 +0800 Subject: [PATCH 10/11] [Warmup][V1] Replace per-key log spam with tqdm progress bar The VllmJitKernel.warmup() loop logged a progress/ETA/rate line per compiled key (dozens-to-hundreds of lines per kernel), and the mHC wrappers logged another line per compile() call plus _compile_and_cache. On DSv4 this produced hundreds of INFO lines scrolling during startup. Replace the per-iteration logger.info with a tqdm progress bar shown only on rank 0 (mirrors deep_gemm_warmup). Keep one summary line at start and finish. Remove the per-compile log spam from _compile_and_cache, MhcPreKernel.compile, and MhcFusedPostPreKernel.compile; the per-kernel total in get_warmup_keys is sufficient. Co-authored-by: opencode Signed-off-by: hanshuche --- vllm/model_executor/kernels/mhc/warmup.py | 22 --------------- vllm/model_executor/warmup/jit_warmup.py | 33 ++++++++++++----------- 2 files changed, 18 insertions(+), 37 deletions(-) diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 8071173c7870..c23d6dd66565 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -130,12 +130,6 @@ def _compile_and_cache(jit_impl, *args, **kwargs) -> None: kernel = jit_impl.compile(*args, **kwargs) key, _ = jit_impl.func.parse_args(*args, **kwargs) jit_impl._kernel_cache[key] = kernel - kernel_name = getattr(getattr(jit_impl, "func", jit_impl), "__name__", "?") - logger.info( - "_compile_and_cache: kernel=%s key=%s", - kernel_name, - key, - ) def _compile_mhc_post(hidden_size: int, hc_mult: int) -> None: @@ -260,13 +254,6 @@ def compile(self, compile_key: CompileKey) -> None: num_tokens = 1 # dynamic dim; smallest valid value c = self._constants - logger.info( - "MhcPreKernel.compile: is_broadcast=%s n_splits=%d use_norm=%s", - compile_key.is_broadcast, - n_splits, - compile_key.use_norm_weight, - ) - gemm_out_mul = _fake(torch.float32, n_splits, num_tokens, hc_mult3) gemm_out_sqrsum = _fake(torch.float32, n_splits, num_tokens) hc_scale = _fake(torch.float32, 3) @@ -450,15 +437,6 @@ def compile(self, compile_key: CompileKey) -> None: num_tokens = 1 # dynamic dim; smallest valid value c = self._constants - logger.info( - "MhcFusedPostPreKernel.compile: use_small_fma=%s n_splits=%d " - "tile_n=%d use_norm=%s", - compile_key.use_small_fma, - n_splits, - tile_n, - compile_key.use_norm_weight, - ) - if compile_key.use_small_fma: comb_mix = _fake(torch.float32, num_tokens, hc_mult, hc_mult) residual_in = _fake(torch.bfloat16, num_tokens, hc_mult, hidden_size) diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index cbeec26e562c..1f2963cbc508 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -18,6 +18,9 @@ from dataclasses import dataclass from typing import Any, Generic, TypeVar +from tqdm import tqdm + +from vllm.distributed.parallel_state import is_global_first_rank from vllm.logger import init_logger __all__ = [ @@ -679,22 +682,22 @@ def warmup(self, *args: Any, **kwargs: Any) -> None: return logger.info("Warming up %s: %d keys", name, total) t0 = time.monotonic() - for i, compile_key in enumerate(keys, 1): + # Progress bar on rank 0 only; other ranks compile silently to avoid + # duplicated output across TP/PP workers (mirrors deep_gemm_warmup). + iterator = tqdm( + keys, + desc=f"Warming up {name}", + total=total, + disable=not is_global_first_rank(), + ) + for compile_key in iterator: self.compile(compile_key) - elapsed = time.monotonic() - t0 - rate = i / max(elapsed, 1e-9) - eta = (total - i) / max(rate, 1e-9) - logger.info( - "Warming up %s: progress %d/%d (%.1f%%), %.1fs elapsed, " - "ETA %.1fs (%.1f/s)", - name, - i, - total, - 100.0 * i / total, - elapsed, - eta, - rate, - ) + logger.info( + "Warming up %s: %d keys finished in %.2fs", + name, + total, + time.monotonic() - t0, + ) class JitWarmupRegistry: From 68a5f4a2571c7aa61555d1bb8431f3e5108a9ecb Mon Sep 17 00:00:00 2001 From: SyaOtiLan <954239196@qq.com> Date: Thu, 27 Aug 2026 15:29:36 +0800 Subject: [PATCH 11/11] [Warmup] Precompile fallback mHC prenorm kernels Precompile the regular and block-M TileLang specializations used by the non-DeepGEMM mHC prenorm path so the first runtime call does not trigger JIT compilation. Add focused coverage verifying that the fallback warmup runs only when DeepGEMM is unavailable. Assisted-by: OpenAI Codex Signed-off-by: SyaOtiLan <954239196@qq.com> --- .../test_mhc_warmup_wrappers.py | 30 ++++++++ vllm/model_executor/kernels/mhc/warmup.py | 73 +++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/tests/model_executor/test_mhc_warmup_wrappers.py b/tests/model_executor/test_mhc_warmup_wrappers.py index 38024e54b31e..24486b83102a 100644 --- a/tests/model_executor/test_mhc_warmup_wrappers.py +++ b/tests/model_executor/test_mhc_warmup_wrappers.py @@ -356,3 +356,33 @@ def test_dispatch_matches_get_warmup_keys_expansion( use_deep_gemm=True, ) assert k in key_set, f"token={t} produced key {k} not in warmup set" + + +@pytest.mark.parametrize("use_deep_gemm", [False, True]) +def test_mhc_pre_compile_warms_prenorm_only_for_fallback( + use_deep_gemm: bool, +) -> None: + wrapper = MhcPreKernel() + wrapper._constants = _DEFAULT_CONSTANTS + key = MhcPreKernel.CompileKey( + hidden_size=7168, + hc_mult=4, + n_splits=1, + use_norm_weight=True, + use_deep_gemm=use_deep_gemm, + is_broadcast=False, + ) + + with ( + mock.patch("vllm.model_executor.kernels.mhc.warmup._compile_and_cache"), + mock.patch("vllm.model_executor.kernels.mhc.warmup._compile_mhc_post"), + mock.patch( + "vllm.model_executor.kernels.mhc.warmup._compile_hc_prenorm_gemm" + ) as compile_prenorm, + ): + wrapper.compile(key) + + if use_deep_gemm: + compile_prenorm.assert_not_called() + else: + compile_prenorm.assert_called_once_with(7168, 4) diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index c23d6dd66565..cec6c1f6bb82 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -152,6 +152,77 @@ def _compile_mhc_post(hidden_size: int, hc_mult: int) -> None: _compile_and_cache(_mhc_post_kernel, a, b, c_t, d, x, hc_mult, hidden_size) +def _compile_hc_prenorm_gemm(hidden_size: int, hc_mult: int) -> None: + """Compile the non-DeepGEMM prenorm GEMM specializations. + + ``_tilelang_hc_prenorm_gemm`` selects between two regular-kernel + configurations and one block-M configuration based on the token count. + Their token dimension is dynamic, so one representative shape per static + configuration covers the complete runtime key space. + """ + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + hc_prenorm_gemm_block_m_tilelang, + hc_prenorm_gemm_tilelang, + ) + + hc_hidden_size = hc_mult * hidden_size + n_out = hc_mult * 2 + hc_mult * hc_mult + fn = _fake(torch.float32, n_out, hc_hidden_size) + + def compile_regular(num_tokens: int, n_thr: int, tile_n: int) -> None: + x = _fake(torch.bfloat16, num_tokens, hc_hidden_size) + out = _fake(torch.float32, 1, num_tokens, n_out) + sqrsum = _fake(torch.float32, 1, num_tokens) + _compile_and_cache( + hc_prenorm_gemm_tilelang, + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + n_out, + n_thr, + tile_n, + 1, + ) + + def compile_block_m(num_tokens: int) -> None: + x = _fake(torch.bfloat16, num_tokens, hc_hidden_size) + out = _fake(torch.float32, 1, num_tokens, n_out) + sqrsum = _fake(torch.float32, 1, num_tokens) + _compile_and_cache( + hc_prenorm_gemm_block_m_tilelang, + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + n_out, + 512, + 12, + 2, + ) + + if hc_hidden_size % 1024 == 0: + compile_regular( + num_tokens=1, # dynamic dim; smallest valid value + n_thr=1024, + tile_n=4, + ) + + compile_regular( + num_tokens=128, # dynamic dim; smallest valid value + n_thr=512, + tile_n=12, + ) + + compile_block_m( + num_tokens=1024, # dynamic dim; smallest valid value + ) + + # ============================================================================= # 1. MhcPreKernel — first-layer path (mhc_pre + mhc_post) # ============================================================================= @@ -343,6 +414,8 @@ def compile(self, compile_key: CompileKey) -> None: # mhc_post: cache_key depends only on (hc_mult, hidden_size). _compile_mhc_post(hidden_size, hc_mult) + if not compile_key.use_deep_gemm: + _compile_hc_prenorm_gemm(hidden_size, hc_mult) # =============================================================================