diff --git a/python/sglang/srt/layers/mhc.py b/python/sglang/srt/layers/mhc.py index 11d26cae4..4c56521d1 100644 --- a/python/sglang/srt/layers/mhc.py +++ b/python/sglang/srt/layers/mhc.py @@ -25,9 +25,6 @@ try: 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, @@ -504,7 +501,7 @@ def get_mhc_pre_token_count_representatives( return tuple(sorted(reps.values())) -def _prewarm_mhc_pre( +def prewarm_mhc_pre( residual: torch.Tensor, fn: torch.Tensor, hc_scale: torch.Tensor, @@ -522,7 +519,8 @@ def _prewarm_mhc_pre( """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).""" + cache; later server runs hit the cache. Driven once per process from load_weights. + """ from sglang.srt.server_args import get_global_server_args hc_mult, hidden_size = residual.shape[-2], residual.shape[-1] @@ -534,7 +532,7 @@ def _prewarm_mhc_pre( 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( + mhc_pre( residual.new_zeros(num_tokens, hc_mult, hidden_size), fn, hc_scale, @@ -738,64 +736,6 @@ 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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index d01d17836..a5d50fe71 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1182,121 +1182,6 @@ class DeepseekV4DecoderLayer(nn.Module): self.post_attention_layernorm.weight.data.bfloat16().contiguous() ) - def prewarm_mhc_token_counts( - self, token_counts: Tuple[int, ...], device: torch.device - ) -> None: - paths = ( - ( - "attn", - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.input_layernorm, - ), - ( - "ffn", - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.post_attention_layernorm, - ), - ) - - with torch.inference_mode(): - for num_tokens in token_counts: - for path_name, hc_fn, hc_scale, hc_base, norm in paths: - tic = time.perf_counter() - residual = torch.empty( - (num_tokens, self.hc_mult, self.hidden_size), - dtype=torch.bfloat16, - device=device, - ) - y, post, comb, _ = self.hc_pre( - residual, - hc_fn, - hc_scale, - hc_base, - norm=norm, - ) - del residual, y, post, comb - torch.cuda.synchronize() - logger.info( - "DeepSeek V4 MHC prewarm path=%s num_tokens=%s completed in %.3fs", - path_name, - num_tokens, - time.perf_counter() - tic, - ) - - if self.use_fused_mhc_post_pre: - for num_tokens in token_counts: - for path_name, hc_fn, hc_scale, hc_base, norm in paths: - tic = time.perf_counter() - # Dummy inputs matching the fused kernel's expected shapes. - x = torch.empty( - (num_tokens, self.hidden_size), - dtype=torch.bfloat16, - device=device, - ) - residual = torch.empty( - (num_tokens, self.hc_mult, self.hidden_size), - dtype=torch.bfloat16, - device=device, - ) - post_mix = torch.empty( - (num_tokens, self.hc_mult, 1), - dtype=torch.float32, - device=device, - ) - comb_mix = torch.empty( - (num_tokens, self.hc_mult, self.hc_mult), - dtype=torch.float32, - device=device, - ) - norm_weight = norm.weight.data.bfloat16().contiguous() - mhc_fused_post_pre( - x, - residual, - post_mix, - comb_mix, - hc_fn, - hc_scale, - hc_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - _MHC_POST_MULT_VALUE, - self.hc_sinkhorn_iters, - norm_weight=norm_weight, - norm_eps=norm.variance_epsilon, - ) - del x, residual, post_mix, comb_mix, norm_weight - torch.cuda.synchronize() - logger.info( - "DeepSeek V4 MHC fused prewarm path=%s num_tokens=%s completed in %.3fs", - path_name, - num_tokens, - time.perf_counter() - tic, - ) - - def prewarm_mhc_token_count_buckets( - self, max_num_tokens: int, device: torch.device - ) -> Tuple[int, ...]: - from sglang.srt.layers.mhc import get_mhc_pre_token_count_representatives - - token_counts = get_mhc_pre_token_count_representatives( - max_num_tokens, self.hc_mult * self.hidden_size - ) - if not token_counts: - return token_counts - - logger.info( - "DeepSeek V4 MHC prewarm max_num_tokens=%s representative token counts: %s", - max_num_tokens, - token_counts, - ) - self.prewarm_mhc_token_counts(token_counts, device) - return token_counts - def hc_pre( self, x: torch.Tensor, @@ -1966,6 +1851,11 @@ class DeepseekV4ForCausalLM(nn.Module): self.cp_rank = get_parallel().attn_cp_rank self.cp_size = get_parallel().attn_cp_size + # update_weights_from_disk/_tensor/_distributed re-enter load_weights + # mid-serving (RL refit sends many partial batches); the prewarm and + # its barrier must only run on the first (startup) load. + self._mhc_prewarmed_at_load = False + @property def routed_experts_weights_of_layer(self): return self._routed_experts_weights_of_layer.value @@ -2158,6 +2048,62 @@ class DeepseekV4ForCausalLM(nn.Module): return name + def _prewarm_mhc_pre_kernels(self) -> None: + """One-shot mhc_pre() JIT prewarm at load time, synced across ranks. + + Runs before any forward so the compile burst stays off the serving + path; the barrier keeps ranks from proceeding while a peer is still + compiling. The early returns below must stay rank-uniform. + """ + if self._mhc_prewarmed_at_load: + return + self._mhc_prewarmed_at_load = True + if _is_npu or not ( + envs.SGLANG_DSV4_MHC_PREWARM.get() + and envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() + ): + return + layer = next( + (m for m in self.model.layers if isinstance(m, DeepseekV4DecoderLayer)), + None, + ) + if layer is None: + return + + from sglang.srt.layers.mhc import prewarm_mhc_pre + + tic = time.perf_counter() + prewarm_mhc_pre( + # Template carrying dtype/device; buckets allocate their own sizes. + residual=torch.zeros( + (1, layer.hc_mult, layer.hidden_size), + dtype=torch.bfloat16, + device=layer.hc_attn_fn.device, + ), + fn=layer.hc_attn_fn, + hc_scale=layer.hc_attn_scale, + hc_base=layer.hc_attn_base, + rms_eps=layer.rms_norm_eps, + hc_pre_eps=layer.hc_eps, + hc_sinkhorn_eps=layer.hc_eps, + hc_post_mult_value=_MHC_POST_MULT_VALUE, + sinkhorn_repeat=layer.hc_sinkhorn_iters, + n_splits=1, + n_splits_pre=32, + norm_weight=layer.input_layernorm.weight.data, + norm_eps=layer.input_layernorm.variance_epsilon, + ) + torch.cuda.synchronize() + compile_secs = time.perf_counter() - tic + # Runs before init_memory_pool(); don't let transients skew pool sizing. + torch.cuda.empty_cache() + get_tp_group().barrier() + logger.info( + "DeepSeek V4 MHC prenorm prewarm at load: compile %.1fs, rank sync +%.1fs", + compile_secs, + time.perf_counter() - tic - compile_secs, + ) + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() @@ -2493,6 +2439,9 @@ class DeepseekV4ForCausalLM(nn.Module): self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) + if not is_nextn: + self._prewarm_mhc_pre_kernels() + def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight diff --git a/test/registered/kernels/test_mhc_kernels.py b/test/registered/kernels/test_mhc_kernels.py index 6bddcf867..8f646b785 100644 --- a/test/registered/kernels/test_mhc_kernels.py +++ b/test/registered/kernels/test_mhc_kernels.py @@ -18,7 +18,6 @@ def test_mhc_fused_post_pre_matches_unfused( pytest.skip("CUDA is required for TileLang mHC kernels") monkeypatch.setattr(mhc, "is_dsa_prefill_cp_round_robin_split", lambda: False) - monkeypatch.setattr(mhc, "_mhc_pre_warmed", True) torch.manual_seed(0) device = torch.device("cuda") hc_mult = 4