minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Prewarm DSV4 MHC post kernel at model load (#30741)

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-08-03 23:41:56 -07:00
committed by GitHub
co-authored by weireweire
parent 157401f050
commit 23ea7b6481
1 changed files with 25 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v4.py
+25 -10
View File
@@ -2752,8 +2752,8 @@ class DeepseekV4ForCausalLM(nn.Module):
return name return name
def _prewarm_mhc_pre_kernels(self) -> None: def _prewarm_mhc_kernels(self) -> None:
"""One-shot mhc_pre() JIT prewarm at load time, synced across ranks. """One-shot MHC JIT prewarm at load time, synced across ranks.
Runs before any forward so the compile burst stays off the serving Runs before any forward so the compile burst stays off the serving
path; the barrier keeps ranks from proceeding while a peer is still path; the barrier keeps ranks from proceeding while a peer is still
@@ -2774,16 +2774,17 @@ class DeepseekV4ForCausalLM(nn.Module):
if layer is None: if layer is None:
return return
from sglang.kernels.ops.layernorm.mhc import prewarm_mhc_pre from sglang.kernels.ops.layernorm.mhc import mhc_post, prewarm_mhc_pre
tic = time.perf_counter() tic = time.perf_counter()
residual = torch.zeros(
(1, layer.hc_mult, layer.hidden_size),
dtype=torch.bfloat16,
device=layer.hc_attn_fn.device,
)
prewarm_mhc_pre( prewarm_mhc_pre(
# Template carrying dtype/device; buckets allocate their own sizes. # Template carrying dtype/device; buckets allocate their own sizes.
residual=torch.zeros( residual=residual,
(1, layer.hc_mult, layer.hidden_size),
dtype=torch.bfloat16,
device=layer.hc_attn_fn.device,
),
fn=layer.hc_attn_fn, fn=layer.hc_attn_fn,
hc_scale=layer.hc_attn_scale, hc_scale=layer.hc_attn_scale,
hc_base=layer.hc_attn_base, hc_base=layer.hc_attn_base,
@@ -2797,13 +2798,27 @@ class DeepseekV4ForCausalLM(nn.Module):
norm_weight=layer.input_layernorm.weight.data, norm_weight=layer.input_layernorm.weight.data,
norm_eps=layer.input_layernorm.variance_epsilon, norm_eps=layer.input_layernorm.variance_epsilon,
) )
mhc_post(
x=residual.new_zeros((1, layer.hidden_size)),
residual=residual,
post_layer_mix=torch.zeros(
(1, layer.hc_mult, 1),
dtype=torch.float32,
device=residual.device,
),
comb_res_mix=torch.zeros(
(1, layer.hc_mult, layer.hc_mult),
dtype=torch.float32,
device=residual.device,
),
)
torch.cuda.synchronize() torch.cuda.synchronize()
compile_secs = time.perf_counter() - tic compile_secs = time.perf_counter() - tic
# Runs before init_memory_pool(); don't let transients skew pool sizing. # Runs before init_memory_pool(); don't let transients skew pool sizing.
torch.cuda.empty_cache() torch.cuda.empty_cache()
get_tp_group().barrier() get_tp_group().barrier()
logger.info( logger.info(
"DeepSeek V4 MHC prenorm prewarm at load: compile %.1fs, rank sync +%.1fs", "DeepSeek V4 MHC prewarm at load: compile %.1fs, rank sync +%.1fs",
compile_secs, compile_secs,
time.perf_counter() - tic - compile_secs, time.perf_counter() - tic - compile_secs,
) )
@@ -3153,7 +3168,7 @@ class DeepseekV4ForCausalLM(nn.Module):
self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
if not is_nextn: if not is_nextn:
self._prewarm_mhc_pre_kernels() self._prewarm_mhc_kernels()
def get_embed_and_head(self): def get_embed_and_head(self):
return self.model.embed_tokens.weight, self.lm_head.weight return self.model.embed_tokens.weight, self.lm_head.weight