[dsv4] Trigger MHC prenorm prewarm at weight-load time with rank sync (#29988)
This commit is contained in:
YAMY
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user