cc908dd556
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
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
import json
|
|
import os
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def _load_jit_cache_module():
|
|
repo_root = Path(__file__).resolve().parents[4]
|
|
module_path = (
|
|
repo_root
|
|
/ "python"
|
|
/ "sglang"
|
|
/ "srt"
|
|
/ "layers"
|
|
/ "deep_gemm_wrapper"
|
|
/ "jit_cache.py"
|
|
)
|
|
spec = importlib.util.spec_from_file_location("deep_gemm_jit_cache_under_test", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader is not None
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_prepare_deep_gemm_jit_env_sets_cache_before_import(tmp_path, monkeypatch):
|
|
prepare_deep_gemm_jit_env = _load_jit_cache_module().prepare_deep_gemm_jit_env
|
|
|
|
cache_dir = tmp_path / "dg-cache"
|
|
monkeypatch.delenv("DG_JIT_CACHE_DIR", raising=False)
|
|
monkeypatch.delenv("DG_JIT_USE_NVRTC", raising=False)
|
|
monkeypatch.setenv("SGLANG_DG_CACHE_DIR", str(cache_dir))
|
|
monkeypatch.setenv("SGLANG_DG_USE_NVRTC", "1")
|
|
|
|
resolved = prepare_deep_gemm_jit_env()
|
|
|
|
assert resolved == str(cache_dir)
|
|
assert os.environ["DG_JIT_CACHE_DIR"] == str(cache_dir)
|
|
assert os.environ["DG_JIT_USE_NVRTC"] == "1"
|
|
assert cache_dir.is_dir()
|
|
|
|
|
|
def test_prepare_deep_gemm_jit_env_respects_explicit_deepgemm_env(
|
|
tmp_path, monkeypatch
|
|
):
|
|
prepare_deep_gemm_jit_env = _load_jit_cache_module().prepare_deep_gemm_jit_env
|
|
|
|
explicit = tmp_path / "explicit"
|
|
ignored = tmp_path / "ignored"
|
|
monkeypatch.setenv("DG_JIT_CACHE_DIR", str(explicit))
|
|
monkeypatch.setenv("SGLANG_DG_CACHE_DIR", str(ignored))
|
|
monkeypatch.setenv("DG_JIT_USE_NVRTC", "0")
|
|
monkeypatch.setenv("SGLANG_DG_USE_NVRTC", "1")
|
|
|
|
resolved = prepare_deep_gemm_jit_env()
|
|
|
|
assert resolved == str(explicit)
|
|
assert os.environ["DG_JIT_CACHE_DIR"] == str(explicit)
|
|
assert os.environ["DG_JIT_USE_NVRTC"] == "0"
|
|
assert explicit.is_dir()
|
|
assert not ignored.exists()
|
|
|
|
|
|
def test_sparse_m_list_is_sparse_but_covers_boundaries():
|
|
build_sparse_m_list = _load_jit_cache_module().build_sparse_m_list
|
|
|
|
m_list = build_sparse_m_list(65536)
|
|
|
|
assert m_list[0] == 1
|
|
assert m_list[-1] == 65536
|
|
assert 1024 in m_list
|
|
assert 2048 in m_list
|
|
assert 4096 in m_list
|
|
assert len(m_list) < 5000
|
|
assert len(m_list) < 65536 // 10
|
|
|
|
|
|
def test_build_warmup_m_lists_is_category_and_cp_aware():
|
|
build_warmup_m_lists = _load_jit_cache_module().build_warmup_m_lists
|
|
|
|
lists = build_warmup_m_lists(
|
|
chunked_prefill_size=65536,
|
|
attn_cp_size=8,
|
|
max_running_requests=200,
|
|
speculative_num_draft_tokens=4,
|
|
)
|
|
|
|
assert lists.per_rank[-1] == 8192
|
|
assert lists.gathered[-1] == 65536
|
|
assert lists.decode[-1] == 800
|
|
assert len(lists.per_rank) < len(lists.gathered)
|
|
assert len(lists.gathered) < 5000
|
|
|
|
|
|
def test_warmup_manifest_requires_cache_key_superset(tmp_path):
|
|
jit_cache = _load_jit_cache_module()
|
|
cache_entries = jit_cache.cache_entries
|
|
make_warmup_manifest_key = jit_cache.make_warmup_manifest_key
|
|
mark_warmup_complete = jit_cache.mark_warmup_complete
|
|
should_skip_warmup = jit_cache.should_skip_warmup
|
|
|
|
cache_root = tmp_path / "deep_gemm"
|
|
(cache_root / "cache" / "a").mkdir(parents=True)
|
|
(cache_root / "cache" / "b").mkdir()
|
|
key = make_warmup_manifest_key(
|
|
kernel_type="GEMM_NT_F8F8BF16",
|
|
n=4096,
|
|
k=8192,
|
|
num_groups=1,
|
|
m_list=[1, 2, 4, 8],
|
|
runtime_fingerprint={"deep_gemm": "0.1.2", "sm": 90},
|
|
)
|
|
|
|
mark_warmup_complete(cache_root, key, cache_entries(cache_root))
|
|
assert should_skip_warmup(cache_root, key)
|
|
|
|
# Deleting a cubin directory must invalidate the SGLang-side manifest.
|
|
(cache_root / "cache" / "b").rmdir()
|
|
assert not should_skip_warmup(cache_root, key)
|
|
|
|
|
|
def test_warmup_manifest_file_is_json_and_atomic(tmp_path):
|
|
jit_cache = _load_jit_cache_module()
|
|
MANIFEST_FILENAME = jit_cache.MANIFEST_FILENAME
|
|
mark_warmup_complete = jit_cache.mark_warmup_complete
|
|
|
|
cache_root = tmp_path / "deep_gemm"
|
|
cache_root.mkdir()
|
|
mark_warmup_complete(cache_root, "shape-key", {"abc"})
|
|
|
|
manifest = json.loads((cache_root / MANIFEST_FILENAME).read_text())
|
|
assert manifest["version"] == 1
|
|
assert manifest["entries"]["shape-key"]["status"] == "complete"
|
|
assert manifest["entries"]["shape-key"]["cache_entries"] == ["abc"]
|