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

[dsv4] Prewarm MHC prenorm kernel at startup (#27986)

Browse Source
This commit is contained in:
YAMY
2026-06-15 13:26:26 -07:00
committed by GitHub
parent 7e629a2f8c
commit f870bf1ed0
4 changed files with 125 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -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)
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
-38
View File
@@ -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
):
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
+2 -9
View File
@@ -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):
python/sglang/srt/layers/mhc.py
+122
View File
@@ -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