chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

perf(deepgemm): CP-aware dense warmup M-grid for NSA prefill CP

Browse Source 
Shrink the non-grouped (dense/attention) DeepGEMM warmup M grid by attn_cp_size under NSA prefill in-seq CP (per-rank M = tokens/cp), while keeping the grouped MoE GEMM grid full (deepep all-to-all re-gathers all tokens; topk==ep_size keeps MoE M ~= chunked). Gated by _cp_dense_warmup_divisor.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 672ef3a6609abd100e5e95b42016d17d0a2966e5)
This commit is contained in:
孙远航
2026-06-08 13:22:34 +00:00
committed by laoyao0822
parent 134b88b9b4
commit 81a0191331
1 changed files with 82 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
View File

@@ -23,6 +23,9 @@ if ENABLE_JIT_DEEPGEMM:
_BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1))
# Separate, possibly-shrunk M grid for NON-grouped (dense/attention) shapes under
# CP prefill -- see _cp_dense_warmup_divisor / update_deep_gemm_config.
_BUILTIN_M_LIST_DENSE = _BUILTIN_M_LIST
_ENABLE_JIT_DEEPGEMM_PRECOMPILE = envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get()
_DO_COMPILE_ALL = True
_IS_FIRST_RANK_ON_NODE = envs.SGLANG_IS_FIRST_RANK_ON_NODE.get()
@@ -44,8 +47,46 @@ if _ENABLE_JIT_DEEPGEMM_PRECOMPILE:
os.environ["DG_PRELOAD_KERNELS"] = "1"
def _cp_dense_warmup_divisor(server_args: ServerArgs) -> int:
"""CP divisor for the NON-grouped (dense/attention) DeepGEMM warmup M grid.
Under NSA prefill context-parallel ``in-seq-split`` the sequence is split
across ``attn_cp_size`` ranks BEFORE the transformer layers
(deepseek_v2.py cp_split_and_rebuild_data), so every dense ``num_groups==1``
GEMM -- attention q/k/v/o projections, the dense/shared-expert MLP, the NSA
indexer ``weights_proj`` -- runs on only ~tokens/attn_cp_size per rank.
Warming those for the full non-CP M range wastes ~cp_size x.
The MoE GROUPED GEMM is deliberately NOT shrunk: deepep all-to-all re-gathers
every token, so its M = sum(num_recv_tokens_per_expert) ~= chunked * topk /
ep_size (== chunked for GLM-5.1 where topk==ep_size==8), independent of CP --
it keeps the full grid.
Shrink only when EVERY dense-extend path is CP-split: the main prefill extend
always is; the EAGLE draft extend is CP-split only when
``SGLANG_CP_DRAFT_SHARED_KV`` is set (``_is_cp_shared_kv_draft_extend``,
``include_v2=True``). If a draft is configured without that env, the draft
extend runs non-CP at full tokens, so the dense shapes still need the full
grid and we return 1.
"""
cp_size = int(getattr(server_args, "attn_cp_size", 1) or 1)
if cp_size <= 1:
return 1
prefill_cp_on = bool(
getattr(server_args, "enable_nsa_prefill_context_parallel", False)
) and getattr(server_args, "nsa_prefill_cp_mode", None) == "in-seq-split"
if not prefill_cp_on:
return 1
has_draft = getattr(server_args, "speculative_algorithm", None) is not None
if has_draft and not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
return 1
return cp_size
def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs):
global _BUILTIN_M_LIST
global _BUILTIN_M_LIST_DENSE
global _DO_COMPILE_ALL
global _IS_FIRST_RANK_ON_NODE
@@ -87,6 +128,27 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs):
m_max = min(1024 * 128, m_max)
_BUILTIN_M_LIST += list(range(1, m_max + 1))
# Dense (non-grouped) shapes run at ~tokens/cp_size under CP prefill in-seq-split;
# shrink their M grid by the (gated) CP divisor while the MoE grouped grid stays
# full. Keep an absolute floor so degenerate cp_size/chunk combos still cover a
# reasonable dense M range.
cp_dense_div = _cp_dense_warmup_divisor(server_args)
if cp_dense_div > 1 and _BUILTIN_M_LIST:
dense_m_max = max(ceil_div(max(_BUILTIN_M_LIST), cp_dense_div), 2048)
_BUILTIN_M_LIST_DENSE = [m for m in _BUILTIN_M_LIST if m <= dense_m_max]
logger.info(
"DeepGEMM warmup: CP in-seq-split active (attn_cp_size divisor=%s); "
"dense (non-grouped) shapes warmed up to M=%s (%s Ms) vs full grouped "
"grid M=%s (%s Ms).",
cp_dense_div,
dense_m_max,
len(_BUILTIN_M_LIST_DENSE),
max(_BUILTIN_M_LIST),
len(_BUILTIN_M_LIST),
)
else:
_BUILTIN_M_LIST_DENSE = _BUILTIN_M_LIST
_IS_FIRST_RANK_ON_NODE = server_args.base_gpu_id == gpu_id
# Check if is the first rank on node.
@@ -105,6 +167,14 @@ class DeepGemmKernelType(IntEnum):
_INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict()
# Grouped (MoE) GEMMs are fed by the deepep all-to-all and run at M ~= chunked
# regardless of CP, so they always use the full M grid; non-grouped (dense) GEMMs
# are CP-per-rank and use the (possibly shrunk) dense grid.
_GROUPED_GEMM_KERNEL_TYPES = (
DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED,
DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG,
)
# TODO improve code
def _maybe_compile_deep_gemm_one_type_all(
@@ -115,6 +185,7 @@ def _maybe_compile_deep_gemm_one_type_all(
) -> None:
global _INITIALIZATION_DICT
global _BUILTIN_M_LIST
global _BUILTIN_M_LIST_DENSE
query_key = (kernel_type, n, k, num_groups)
if (
@@ -136,9 +207,18 @@ def _maybe_compile_deep_gemm_one_type_all(
"`python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code`"
)
# Grouped (MoE) shapes need the full M grid (deepep re-gathers all tokens);
# dense shapes are CP-per-rank and use the shrunk dense grid.
m_list = (
_BUILTIN_M_LIST
if kernel_type in _GROUPED_GEMM_KERNEL_TYPES
else _BUILTIN_M_LIST_DENSE
)
logger.info(
f"Try DeepGEMM JIT Compiling for "
f"<{kernel_type.name}> N={n}, K={k}, num_groups={num_groups} with all Ms."
f"<{kernel_type.name}> N={n}, K={k}, num_groups={num_groups} "
f"with {len(m_list)} Ms (max M={max(m_list) if m_list else 0})."
f"{' It only takes a little time (typically 1 sec) if you have run `python3 -m sglang.compile_deep_gemm`. ' if not _IN_PRECOMPILE_STAGE else ''}"
)
@@ -147,7 +227,7 @@ def _maybe_compile_deep_gemm_one_type_all(
n=n,
k=k,
num_groups=num_groups,
m_list=_BUILTIN_M_LIST,
m_list=m_list,
)