Reuse DeepGEMM JIT cache across restarts
New sgl-deep-gemm wheels persist cubins, but SGLang was still replaying the full warmup envelope after every restart. The wrapper also translated SGLANG_DG_* env after importing deep_gemm, so the wheel could observe stale cache settings.\n\nMove DeepGEMM JIT env preparation before the first import, bridge the current SGLANG_DG_USE_NVRTC spelling, replace dense M walks with category-specific sparse M grids, and add a SGLang-side warmup manifest under the DeepGEMM cache dir. A manifest hit skips only the SGLang replay loop; missing or stale cache entries still force warmup and refresh the manifest.\n\nConstraint: sgl-deep-gemm 0.1.2 removed the compile-mode API, so dense 1..m_max loops launch real kernels.\nConstraint: DeepGEMM cache keys include compiler flags such as the deep_gemm include path; path-stable environments are still required for cross-container cache hits.\nRejected: SGLANG_JIT_DEEPGEMM_PRECOMPILE=0 | bypasses warmup instead of making cache reuse correct.\nRejected: Enable per-process cache dirs by default | isolates corruption but defeats persistent cross-process cache reuse.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce dense DeepGEMM warmup for new wheels unless compile-mode support is verified on the target wheel.\nTested: Local py_compile for jit_cache.py configurer.py compile_utils.py model_runner.py\nTested: Local pytest test/registered/unit/layers/test_deep_gemm_jit_cache.py (6 passed)\nTested: Remote g0034 cjy-glm5-new py_compile for same files\nTested: Remote g0034 cjy-glm5-new pytest test/registered/unit/layers/test_deep_gemm_jit_cache.py (6 passed)\nTested: Remote env probe verified SGLANG_DG_CACHE_DIR and SGLANG_DG_USE_NVRTC bridge to DG_JIT_* before configurer completes\nNot-tested: Full decode cold-cache/hot-cache restart timing and manifest-hit ETE log validation
This commit is contained in:
laoyao0822
@@ -2,6 +2,7 @@ import logging
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from enum import IntEnum, auto
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import torch
|
||||
@@ -12,6 +13,15 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
||||
restore_symmetric_memory_context,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.deep_gemm_wrapper.jit_cache import (
|
||||
build_warmup_m_lists,
|
||||
build_sparse_m_list,
|
||||
cache_entries,
|
||||
make_warmup_manifest_key,
|
||||
mark_warmup_complete,
|
||||
prepare_deep_gemm_jit_env,
|
||||
should_skip_warmup,
|
||||
)
|
||||
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import ceil_div, get_available_gpu_memory
|
||||
@@ -22,134 +32,51 @@ if ENABLE_JIT_DEEPGEMM:
|
||||
import deep_gemm
|
||||
|
||||
|
||||
_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
|
||||
_M_LIST_PER_RANK: List[int] = build_sparse_m_list(1024 * 16)
|
||||
_M_LIST_GATHERED: List[int] = build_sparse_m_list(1024 * 16)
|
||||
_M_LIST_DECODE: List[int] = build_sparse_m_list(1024)
|
||||
_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()
|
||||
_IN_PRECOMPILE_STAGE = envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get()
|
||||
_FAST_WARMUP = envs.SGLANG_JIT_DEEPGEMM_FAST_WARMUP.get()
|
||||
# New DeepGEMM wheels dropped compile-mode, so dense 1..m_max enumeration would
|
||||
# launch real kernels. Sparse warmup is now always used.
|
||||
|
||||
# Force redirect deep_gemm cache_dir
|
||||
os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
|
||||
"SGLANG_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
|
||||
_DEEP_GEMM_CACHE_DIR = prepare_deep_gemm_jit_env(
|
||||
preload_kernels=_ENABLE_JIT_DEEPGEMM_PRECOMPILE
|
||||
)
|
||||
|
||||
# Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f
|
||||
# NVRTC may have performance loss with some cases.
|
||||
# And NVCC JIT speed is also 9x faster in the ref commit
|
||||
os.environ["DG_JIT_USE_NVRTC"] = os.getenv("SGL_DG_USE_NVRTC", "0")
|
||||
|
||||
# Enable DeepGEMM kernel preloading if precompile is enabled
|
||||
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 _M_LIST_PER_RANK, _M_LIST_GATHERED, _M_LIST_DECODE
|
||||
global _DO_COMPILE_ALL
|
||||
global _IS_FIRST_RANK_ON_NODE
|
||||
|
||||
_BUILTIN_M_LIST = []
|
||||
|
||||
if _FAST_WARMUP:
|
||||
# In fast warmup mode, only compile a small set of typical Ms
|
||||
|
||||
# First cover all the small bs to ensure decode performance
|
||||
_BUILTIN_M_LIST += list(range(1, 1025))
|
||||
|
||||
# Then cover larger batch sizes with gradually increasing steps
|
||||
# For example, when chunekd prefill size is 16384
|
||||
# The sampled Ms would be:
|
||||
# 1024, 1026, ... 2046 (step 2)
|
||||
# 2048, 2052, ... 4092 (step 4)
|
||||
# 4096, 5004, ... 8184 (step 8)
|
||||
# 8192, 9008, ... 16384 (step 16)
|
||||
# Totally 1024 + 1024 / 2 + 2048 / 4 + 4096 / 8 + 8192 / 16 = 3072 kernels
|
||||
next_m, sample_step = 1024, 2
|
||||
max_prefill_bs = (
|
||||
min(server_args.chunked_prefill_size, 32 * 1024)
|
||||
if server_args.chunked_prefill_size >= 1
|
||||
else 16 * 1024
|
||||
)
|
||||
while next_m < max_prefill_bs:
|
||||
_BUILTIN_M_LIST += list(range(next_m, 2 * next_m, sample_step))
|
||||
next_m = next_m * 2
|
||||
sample_step = sample_step * 2
|
||||
_BUILTIN_M_LIST.append(max_prefill_bs)
|
||||
_BUILTIN_M_LIST = sorted(list(set(_BUILTIN_M_LIST)))
|
||||
else:
|
||||
# When fast warmup isn't enabled, generate m_max and compile all the covered Ms.
|
||||
m_max = 1024 * 16
|
||||
if server_args.chunked_prefill_size < 1:
|
||||
m_max = 1024 * 64
|
||||
elif server_args.chunked_prefill_size > 8192:
|
||||
m_max = server_args.chunked_prefill_size * 2
|
||||
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
|
||||
m_lists = build_warmup_m_lists(
|
||||
chunked_prefill_size=server_args.chunked_prefill_size,
|
||||
attn_cp_size=getattr(server_args, "attn_cp_size", 1),
|
||||
max_running_requests=server_args.max_running_requests,
|
||||
speculative_num_draft_tokens=server_args.speculative_num_draft_tokens,
|
||||
)
|
||||
_M_LIST_PER_RANK = m_lists.per_rank
|
||||
_M_LIST_GATHERED = m_lists.gathered
|
||||
_M_LIST_DECODE = m_lists.decode
|
||||
|
||||
_IS_FIRST_RANK_ON_NODE = server_args.base_gpu_id == gpu_id
|
||||
if _IS_FIRST_RANK_ON_NODE:
|
||||
logger.info(
|
||||
"[DeepGEMM precompile] chunked=%s attn_cp=%s -> "
|
||||
"per_rank max_m=%s (%s entries), gathered max_m=%s (%s entries), "
|
||||
"decode max_m=%s (%s entries)",
|
||||
server_args.chunked_prefill_size,
|
||||
getattr(server_args, "attn_cp_size", 1),
|
||||
max(_M_LIST_PER_RANK) if _M_LIST_PER_RANK else 0,
|
||||
len(_M_LIST_PER_RANK),
|
||||
max(_M_LIST_GATHERED) if _M_LIST_GATHERED else 0,
|
||||
len(_M_LIST_GATHERED),
|
||||
max(_M_LIST_DECODE) if _M_LIST_DECODE else 0,
|
||||
len(_M_LIST_DECODE),
|
||||
)
|
||||
|
||||
# Check if is the first rank on node.
|
||||
# Default each rank will try compile all Ms to
|
||||
@@ -167,13 +94,37 @@ 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,
|
||||
)
|
||||
def _m_list_for(kernel_type: DeepGemmKernelType) -> List[int]:
|
||||
if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
|
||||
return _M_LIST_GATHERED
|
||||
if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED:
|
||||
return _M_LIST_DECODE
|
||||
return _M_LIST_PER_RANK
|
||||
|
||||
|
||||
def _deep_gemm_runtime_fingerprint() -> Dict[str, str]:
|
||||
try:
|
||||
deep_gemm_version = version("sgl-deep-gemm")
|
||||
except PackageNotFoundError:
|
||||
try:
|
||||
deep_gemm_version = version("deep-gemm")
|
||||
except PackageNotFoundError:
|
||||
deep_gemm_version = "unknown"
|
||||
|
||||
try:
|
||||
capability = torch.cuda.get_device_capability(0)
|
||||
sm = f"{capability[0]}{capability[1]}"
|
||||
except Exception:
|
||||
sm = "unknown"
|
||||
|
||||
return {
|
||||
"deep_gemm_version": deep_gemm_version,
|
||||
"deep_gemm_file": str(getattr(deep_gemm, "__file__", "")),
|
||||
"torch_version": str(torch.__version__),
|
||||
"cuda_version": str(getattr(torch.version, "cuda", "")),
|
||||
"sm": sm,
|
||||
"dg_jit_use_nvrtc": os.environ.get("DG_JIT_USE_NVRTC", "0"),
|
||||
}
|
||||
|
||||
|
||||
# TODO improve code
|
||||
@@ -184,8 +135,6 @@ def _maybe_compile_deep_gemm_one_type_all(
|
||||
num_groups: int,
|
||||
) -> None:
|
||||
global _INITIALIZATION_DICT
|
||||
global _BUILTIN_M_LIST
|
||||
global _BUILTIN_M_LIST_DENSE
|
||||
|
||||
query_key = (kernel_type, n, k, num_groups)
|
||||
if (
|
||||
@@ -207,13 +156,28 @@ 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
|
||||
m_list = _m_list_for(kernel_type)
|
||||
manifest_key = make_warmup_manifest_key(
|
||||
kernel_type=kernel_type.name,
|
||||
n=n,
|
||||
k=k,
|
||||
num_groups=num_groups,
|
||||
m_list=m_list,
|
||||
runtime_fingerprint=_deep_gemm_runtime_fingerprint(),
|
||||
)
|
||||
if should_skip_warmup(_DEEP_GEMM_CACHE_DIR, manifest_key):
|
||||
logger.info(
|
||||
"[DeepGEMM-cache] warmup manifest hit; skip replay for "
|
||||
"<%s> N=%s K=%s num_groups=%s m_count=%s max_m=%s cache_dir=%s",
|
||||
kernel_type.name,
|
||||
n,
|
||||
k,
|
||||
num_groups,
|
||||
len(m_list),
|
||||
max(m_list) if m_list else 0,
|
||||
_DEEP_GEMM_CACHE_DIR,
|
||||
)
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Try DeepGEMM JIT Compiling for "
|
||||
@@ -229,6 +193,11 @@ def _maybe_compile_deep_gemm_one_type_all(
|
||||
num_groups=num_groups,
|
||||
m_list=m_list,
|
||||
)
|
||||
mark_warmup_complete(
|
||||
_DEEP_GEMM_CACHE_DIR,
|
||||
manifest_key,
|
||||
cache_entries(_DEEP_GEMM_CACHE_DIR),
|
||||
)
|
||||
|
||||
|
||||
# NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import logging
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.deep_gemm_wrapper.jit_cache import prepare_deep_gemm_jit_env
|
||||
from sglang.srt.utils import get_device_sm, is_blackwell_supported
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -11,6 +12,9 @@ def _compute_enable_deep_gemm():
|
||||
if sm_version < 90:
|
||||
return False
|
||||
|
||||
prepare_deep_gemm_jit_env(
|
||||
preload_kernels=envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get()
|
||||
)
|
||||
try:
|
||||
import deep_gemm # noqa: F401
|
||||
except ImportError:
|
||||
|
||||
216
python/sglang/srt/layers/deep_gemm_wrapper/jit_cache.py
Normal file
216
python/sglang/srt/layers/deep_gemm_wrapper/jit_cache.py
Normal file
@@ -0,0 +1,216 @@
|
||||
"""DeepGEMM JIT cache helpers.
|
||||
|
||||
This module is deliberately stdlib-only. It is imported before ``deep_gemm`` in
|
||||
the wrapper startup path, so it must not import torch, SGLang utility modules, or
|
||||
anything that can transitively import ``deep_gemm``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Mapping
|
||||
|
||||
MANIFEST_FILENAME = "sglang_deep_gemm_warmup_manifest.v1.json"
|
||||
MANIFEST_VERSION = 1
|
||||
|
||||
|
||||
def _env_true(value: str | None) -> bool:
|
||||
return value is not None and value.lower() in ("1", "true", "yes")
|
||||
|
||||
|
||||
def default_deep_gemm_cache_dir() -> str:
|
||||
return os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
|
||||
|
||||
|
||||
def prepare_deep_gemm_jit_env(*, preload_kernels: bool = True) -> str:
|
||||
"""Set DeepGEMM JIT env vars before the first ``import deep_gemm``.
|
||||
|
||||
DeepGEMM's JIT cache key depends on its own environment variables and
|
||||
compiler flags. SGLang exposes ``SGLANG_*`` knobs, but the wheel consumes
|
||||
``DG_*`` knobs. This bridge is intentionally idempotent:
|
||||
|
||||
* explicit ``DG_JIT_CACHE_DIR`` / ``DG_JIT_USE_NVRTC`` wins;
|
||||
* otherwise SGLang env aliases are translated once;
|
||||
* cache directory is created eagerly so startup fails at the real cause.
|
||||
"""
|
||||
|
||||
cache_dir = os.environ.get("DG_JIT_CACHE_DIR")
|
||||
if not cache_dir:
|
||||
cache_dir = os.environ.get("SGLANG_DG_CACHE_DIR", default_deep_gemm_cache_dir())
|
||||
if _env_true(os.environ.get("SGLANG_DG_CACHE_DIR_PER_PROCESS")):
|
||||
cache_dir = os.path.join(cache_dir, f"pid_{os.getpid()}")
|
||||
os.environ["DG_JIT_CACHE_DIR"] = cache_dir
|
||||
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
|
||||
if "DG_JIT_USE_NVRTC" not in os.environ:
|
||||
os.environ["DG_JIT_USE_NVRTC"] = os.environ.get(
|
||||
"SGLANG_DG_USE_NVRTC",
|
||||
os.environ.get("SGL_DG_USE_NVRTC", "0"),
|
||||
)
|
||||
|
||||
if preload_kernels and "DG_PRELOAD_KERNELS" not in os.environ:
|
||||
os.environ["DG_PRELOAD_KERNELS"] = "1"
|
||||
|
||||
return cache_dir
|
||||
|
||||
|
||||
def build_sparse_m_list(m_max: int) -> list[int]:
|
||||
"""Sparse M grid matching DeepGEMM v0.1+ block-size breakpoints.
|
||||
|
||||
The old dense 1..m_max walk was only cheap while DeepGEMM exposed
|
||||
compile-mode. New wheels launch real kernels for each M, while cubin
|
||||
selection only changes at coarse breakpoints. Keep 1..1024 dense for
|
||||
decode/small prefill, then double both range and stride.
|
||||
"""
|
||||
|
||||
m_max = max(1, int(m_max))
|
||||
out = list(range(1, min(1024, m_max) + 1))
|
||||
next_m, step = 1024, 2
|
||||
while next_m < m_max:
|
||||
out.extend(range(next_m, min(2 * next_m, m_max), step))
|
||||
next_m *= 2
|
||||
step *= 2
|
||||
out.append(m_max)
|
||||
return sorted(set(out))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WarmupMLists:
|
||||
per_rank: list[int]
|
||||
gathered: list[int]
|
||||
decode: list[int]
|
||||
|
||||
|
||||
def build_warmup_m_lists(
|
||||
*,
|
||||
chunked_prefill_size: int,
|
||||
attn_cp_size: int,
|
||||
max_running_requests: int,
|
||||
speculative_num_draft_tokens: int | None,
|
||||
) -> WarmupMLists:
|
||||
"""Build category-specific DeepGEMM warmup M lists.
|
||||
|
||||
* per_rank: dense/attention/indexer GEMMs after NSA CP split.
|
||||
* gathered: MoE contiguous GEMMs after DeepEP all-to-all.
|
||||
* decode: masked GEMMs, bounded by running speculative decode batch.
|
||||
"""
|
||||
|
||||
chunked = int(chunked_prefill_size)
|
||||
if chunked < 1:
|
||||
chunked = 64 * 1024
|
||||
chunked = min(chunked, 128 * 1024)
|
||||
|
||||
cp_size = max(1, int(attn_cp_size or 1))
|
||||
m_max_per_rank = max(1024, chunked // cp_size)
|
||||
m_max_gathered = chunked
|
||||
|
||||
spec_tokens = max(1, int(speculative_num_draft_tokens or 1))
|
||||
m_max_decode = min(1024, max(256, int(max_running_requests) * spec_tokens))
|
||||
|
||||
return WarmupMLists(
|
||||
per_rank=build_sparse_m_list(m_max_per_rank),
|
||||
gathered=build_sparse_m_list(m_max_gathered),
|
||||
decode=build_sparse_m_list(m_max_decode),
|
||||
)
|
||||
|
||||
|
||||
def cache_entries(cache_root: str | os.PathLike[str]) -> set[str]:
|
||||
"""Return DeepGEMM cubin cache entry directory names under ``cache/``."""
|
||||
|
||||
cache_dir = Path(cache_root) / "cache"
|
||||
if not cache_dir.is_dir():
|
||||
return set()
|
||||
return {p.name for p in cache_dir.iterdir() if p.is_dir()}
|
||||
|
||||
|
||||
def make_warmup_manifest_key(
|
||||
*,
|
||||
kernel_type: str,
|
||||
n: int,
|
||||
k: int,
|
||||
num_groups: int,
|
||||
m_list: Iterable[int],
|
||||
runtime_fingerprint: Mapping[str, object],
|
||||
) -> str:
|
||||
m_values = [int(m) for m in m_list]
|
||||
payload = {
|
||||
"version": MANIFEST_VERSION,
|
||||
"kernel_type": kernel_type,
|
||||
"n": int(n),
|
||||
"k": int(k),
|
||||
"num_groups": int(num_groups),
|
||||
"m_list_count": len(m_values),
|
||||
"m_list_max": max(m_values) if m_values else 0,
|
||||
"m_list_hash": hashlib.sha256(
|
||||
",".join(str(m) for m in m_values).encode("utf-8")
|
||||
).hexdigest(),
|
||||
"runtime": dict(sorted(runtime_fingerprint.items())),
|
||||
}
|
||||
return hashlib.sha256(
|
||||
json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _manifest_path(cache_root: str | os.PathLike[str]) -> Path:
|
||||
return Path(cache_root) / MANIFEST_FILENAME
|
||||
|
||||
|
||||
def _load_manifest(cache_root: str | os.PathLike[str]) -> dict:
|
||||
path = _manifest_path(cache_root)
|
||||
if not path.exists():
|
||||
return {"version": MANIFEST_VERSION, "entries": {}}
|
||||
try:
|
||||
data = json.loads(path.read_text())
|
||||
except Exception:
|
||||
return {"version": MANIFEST_VERSION, "entries": {}}
|
||||
if data.get("version") != MANIFEST_VERSION or not isinstance(
|
||||
data.get("entries"), dict
|
||||
):
|
||||
return {"version": MANIFEST_VERSION, "entries": {}}
|
||||
return data
|
||||
|
||||
|
||||
def should_skip_warmup(cache_root: str | os.PathLike[str], key: str) -> bool:
|
||||
manifest = _load_manifest(cache_root)
|
||||
entry = manifest.get("entries", {}).get(key)
|
||||
if not entry or entry.get("status") != "complete":
|
||||
return False
|
||||
|
||||
required = set(entry.get("cache_entries") or [])
|
||||
if not required:
|
||||
return False
|
||||
return required.issubset(cache_entries(cache_root))
|
||||
|
||||
|
||||
def mark_warmup_complete(
|
||||
cache_root: str | os.PathLike[str],
|
||||
key: str,
|
||||
entries: Iterable[str],
|
||||
) -> None:
|
||||
root = Path(cache_root)
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
manifest = _load_manifest(root)
|
||||
manifest.setdefault("entries", {})[key] = {
|
||||
"status": "complete",
|
||||
"cache_entries": sorted(set(entries)),
|
||||
"updated_at": time.time(),
|
||||
}
|
||||
|
||||
fd, tmp_name = tempfile.mkstemp(
|
||||
prefix=f".{MANIFEST_FILENAME}.", suffix=".tmp", dir=str(root)
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(manifest, f, sort_keys=True)
|
||||
f.write("\n")
|
||||
os.replace(tmp_name, _manifest_path(root))
|
||||
finally:
|
||||
if os.path.exists(tmp_name):
|
||||
os.unlink(tmp_name)
|
||||
@@ -2473,7 +2473,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
|
||||
|
||||
def init_deep_gemm(self):
|
||||
logger.info("[DeepGEMM Debug] Entering init_deep_gemm")
|
||||
logger.info("Entering DeepGEMM init")
|
||||
if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
|
||||
return
|
||||
try:
|
||||
@@ -2528,8 +2528,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
k,
|
||||
num_groups,
|
||||
)
|
||||
# Also compile Dense GEMM kernels for other layers (e.g. Attention) if they use FP8 and DeepGEMM
|
||||
for module in self.model.modules():
|
||||
# Also compile Dense GEMM kernels for other layers (e.g. Attention) if
|
||||
# they use FP8 and DeepGEMM. Skip embed_tokens/lm_head because they do
|
||||
# not dispatch through DeepGEMM at runtime, and restrict BF16 warmup to
|
||||
# NSA indexer weights_proj; warming every BF16 Linear can include
|
||||
# vocab-sized shapes and turn cache-hit startup into minutes of replay.
|
||||
for name, module in self.model.named_modules():
|
||||
if "embed_tokens" in name or "lm_head" in name:
|
||||
continue
|
||||
if hasattr(module, "weight") and module.weight is not None:
|
||||
# Check if it is FP8
|
||||
is_fp8 = module.weight.dtype in [
|
||||
@@ -2544,7 +2550,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
k,
|
||||
1, # num_groups=1 for standard GEMM
|
||||
)
|
||||
if module.weight.dim() == 2 and module.weight.dtype == torch.bfloat16:
|
||||
if (
|
||||
module.weight.dim() == 2
|
||||
and module.weight.dtype == torch.bfloat16
|
||||
and name.endswith("weights_proj")
|
||||
):
|
||||
n, k = module.weight.shape
|
||||
_maybe_compile_deep_gemm_one_type_all(
|
||||
DeepGemmKernelType.GEMM_NT_BF16BF16F32,
|
||||
@@ -2552,7 +2562,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
k,
|
||||
1,
|
||||
)
|
||||
# deep_gemm.preload_kernels()
|
||||
if hasattr(deep_gemm, "preload_all_cached_kernels"):
|
||||
logger.info("Preloading cached DeepGEMM kernels...")
|
||||
deep_gemm.preload_all_cached_kernels()
|
||||
elif hasattr(deep_gemm, "preload_kernels"):
|
||||
logger.info("Preloading cached DeepGEMM kernels...")
|
||||
deep_gemm.preload_kernels()
|
||||
|
||||
def init_threads_binding(self):
|
||||
omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")
|
||||
|
||||
Reference in New Issue
Block a user