From f870bf1ed0d7c95775ee5908d616e3692197563d Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Mon, 15 Jun 2026 13:26:26 -0700 Subject: [PATCH] [dsv4] Prewarm MHC prenorm kernel at startup (#27986) --- python/sglang/srt/environ.py | 1 + .../layers/deep_gemm_wrapper/compile_utils.py | 38 ------ .../layers/deep_gemm_wrapper/entrypoint.py | 11 +- python/sglang/srt/layers/mhc.py | 122 ++++++++++++++++++ 4 files changed, 125 insertions(+), 47 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 209c56fd0..b0f1923e8 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -775,6 +775,7 @@ class Envs: SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True) + SGLANG_DSV4_MHC_PREWARM = EnvBool(True) SGLANG_OPT_USE_TRITON_FUSED_MHC = EnvBool(True) SGLANG_OPT_FUSE_MHC_POST_PRE = EnvBool(False) SGLANG_OPT_USE_TILELANG_INDEXER = EnvBool(False) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 5e8139f6b..e46f29980 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -103,7 +103,6 @@ class DeepGemmKernelType(IntEnum): GROUPED_GEMM_NT_BF16_CONTIG = auto() GEMM_NT_F8F8BF16 = auto() GEMM_NT_BF16BF16F32 = auto() - TF32_HC_PRENORM_GEMM = auto() _INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict() @@ -236,7 +235,6 @@ class _BaseWarmupExecutor: DeepGemmKernelType.GEMM_NT_BF16BF16F32: _BF16F32WarmupExecutor, DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG: _BF16GroupedContWarmupExecutor, DeepGemmKernelType.GROUPED_GEMM_NT_BF16_MASKED: _BF16GroupedMaskedWarmupExecutor, - DeepGemmKernelType.TF32_HC_PRENORM_GEMM: _TF32HcPrenormWarmupExecutor, }[kernel_type](**kwargs) @staticmethod @@ -270,11 +268,6 @@ class _BaseWarmupExecutor: + num_groups * 4 + num_groups * max_m * n * 2 ) / _GB - elif kernel_type == DeepGemmKernelType.TF32_HC_PRENORM_GEMM: - # The generic hook's fourth dimension is num_splits for MHC. - # A value of 0 represents DeepGEMM's unsplit num_splits=None path. - num_splits = num_groups if num_groups > 0 else 1 - return (max_m * k * 2 + n * k * 4 + num_splits * max_m * (n + 1) * 4) / _GB else: raise ValueError(f"Invalid kernel type: {kernel_type}") @@ -405,37 +398,6 @@ class _BF16GroupedMaskedWarmupExecutor(_BaseWarmupExecutor): ) -class _TF32HcPrenormWarmupExecutor(_BaseWarmupExecutor): - def __init__(self, max_m: int, n: int, k: int, num_groups: int): - self.x = torch.empty((max_m, k), device="cuda", dtype=torch.bfloat16) - self.fn = torch.empty((n, k), device="cuda", dtype=torch.float32) - self.n = n - # The generic warmup executor's num_groups argument is num_splits here. - # A value of 0 represents DeepGEMM's unsplit num_splits=None path. - self.num_splits = num_groups if num_groups > 0 else None - - def execute(self, m): - if self.num_splits is None: - out = torch.empty((m, self.n), device="cuda", dtype=torch.float32) - sqrsum = torch.empty((m,), device="cuda", dtype=torch.float32) - else: - # Slicing the middle dimension of a preallocated - # (num_splits, max_m, n) output would create a strided view. - out = torch.empty( - (self.num_splits, m, self.n), device="cuda", dtype=torch.float32 - ) - sqrsum = torch.empty( - (self.num_splits, m), device="cuda", dtype=torch.float32 - ) - deep_gemm.tf32_hc_prenorm_gemm( - self.x[:m], - self.fn, - out, - sqrsum, - num_splits=self.num_splits, - ) - - def deep_gemm_execution_hook( m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType ): diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 3d967f1e1..003a716a7 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -192,16 +192,9 @@ def tf32_hc_prenorm_gemm( sqrsum: torch.Tensor, num_splits: Optional[int], ): - m, k = x.shape - n, _ = fn.shape - num_splits_key = num_splits if num_splits is not None else 0 - kernel_type = compile_utils.DeepGemmKernelType.TF32_HC_PRENORM_GEMM - - if m == 0: + if x.shape[0] == 0: return - - with compile_utils.deep_gemm_execution_hook(m, n, k, num_splits_key, kernel_type): - deep_gemm.tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_splits=num_splits) + deep_gemm.tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_splits=num_splits) def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): diff --git a/python/sglang/srt/layers/mhc.py b/python/sglang/srt/layers/mhc.py index 76bf69557..6fdb0def7 100644 --- a/python/sglang/srt/layers/mhc.py +++ b/python/sglang/srt/layers/mhc.py @@ -1,4 +1,5 @@ import functools +import logging import math from typing import Tuple @@ -11,8 +12,13 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_round_robin_split from sglang.srt.layers.utils.common import strict_contiguous +logger = logging.getLogger(__name__) + tilelang.set_log_level("WARNING") +# Set once mhc_pre() has compiled every n_splits bucket at startup. +_mhc_pre_warmed = False + pass_configs = { tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, @@ -450,6 +456,65 @@ def _compute_num_split_for_mhc_pre(num_tokens: int, hc_hidden_size: int) -> int: return max(1, min(n_sms // max(grid_size, 1), num_block_k // 4)) +def get_mhc_pre_token_count_representatives( + max_num_tokens: int, hc_hidden_size: int +) -> Tuple[int, ...]: + """One representative token count per distinct mhc_pre n_splits bucket over + [1, max_num_tokens] (the kernel is specialized only by n_splits).""" + reps = {} + for grid in range(1, (max(1, max_num_tokens) + 63) // 64 + 1): + num_tokens = min(grid * 64, max_num_tokens) + reps[_compute_num_split_for_mhc_pre(num_tokens, hc_hidden_size)] = num_tokens + return tuple(sorted(reps.values())) + + +def _prewarm_mhc_pre( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int, + n_splits_pre: int, + norm_weight: torch.Tensor | None, + norm_eps: float | None, +): + """Compile the prenorm kernel for every n_splits bucket by replaying the + prenorm with the call's real weights. The compiled kernels are written to + the TileLang/DeepGEMM on-disk JIT cache, so this cost is paid only on a cold + cache; later server runs hit the cache. Runs once (gated in mhc_pre).""" + from sglang.srt.server_args import get_global_server_args + + hc_mult, hidden_size = residual.shape[-2], residual.shape[-1] + max_num_tokens = get_global_server_args().chunked_prefill_size + buckets = get_mhc_pre_token_count_representatives( + max_num_tokens, hc_mult * hidden_size + ) + + logger.info("DeepSeek V4 MHC prenorm prewarm: %d n_splits buckets", len(buckets)) + with torch.inference_mode(): + for num_tokens in buckets: + _mhc_pre_impl( + residual.new_zeros(num_tokens, hc_mult, hidden_size), + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + n_splits_pre, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + + @tilelang.jit( pass_configs={ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, @@ -638,7 +703,64 @@ def mhc_pre( norm_weight: torch.Tensor | None = None, norm_eps: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # One-shot startup prewarm: on the first non-capturing call, compile every + # n_splits bucket up front so it isn't JIT-compiled lazily on the first + # prefill. Replays the prenorm via _mhc_pre_impl (no re-entry into mhc_pre). + global _mhc_pre_warmed + if ( + not _mhc_pre_warmed + and envs.SGLANG_DSV4_MHC_PREWARM.get() + and not torch.cuda.is_current_stream_capturing() + ): + _mhc_pre_warmed = True + _prewarm_mhc_pre( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + n_splits_pre, + norm_weight, + norm_eps, + ) + return _mhc_pre_impl( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + n_splits_pre, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + +def _mhc_pre_impl( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + n_splits_pre: int = 32, + *, + norm_weight: torch.Tensor | None = None, + norm_eps: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: assert residual.dtype == torch.bfloat16 assert fn.dtype == torch.float32 assert hc_scale.dtype == torch.float32