diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 19e1efddf..302aced46 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -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, )