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

Deepseek V4 (#23882)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: fzyzcjy <ch271828n@outlook.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: yueming-yuan <yym022502@gmail.com>
Co-authored-by: DarkSharpness <2040703891@qq.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: yhyang201 <yhyang201@users.noreply.github.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: Qiaolin Yu <90088090+qiaolin-yu@users.noreply.github.com>
Co-authored-by: Ethan (Yusheng) Su <11704492+yushengsu-thu@users.noreply.github.com>
Co-authored-by: Mingyi <27337995+wisclmy0611@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Yihao Wang <42559837+againstentropy@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2026-05-07 18:32:21 -07:00
committed by GitHub
co-authored by Baizhou Zhang Claude Opus 4.7 fzyzcjy ispobock Zhiqiang Xie yueming-yuan DarkSharpness Yuhao Yang yhyang201 yhyang201 Qiaolin Yu Ethan Su Mingyi Cheng Wan Yihao Wang
parent 55224fff08
commit 35870d55ac
154 changed files with 24970 additions and 759 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/dsv4/__init__.py
View File
test/manual/dsv4/_common.py
+317
View File
@@ -0,0 +1,317 @@
"""Shared fixture for DeepSeek-V4 cookbook launch-command tests.
Each sibling ``test_<hardware>_<model_size>.py`` declares ONE
``hardware x model_size`` cell from the cookbook (e.g. B200 x Flash)
and contains one ``CustomTestCase`` subclass per recipe
(Low-Latency / Balanced / Max-Throughput / CP, where supported).
Each subclass launches the server with the cookbook's exact flags and
runs two sgl-eval evaluations (https://github.com/sgl-project/sgl-eval):
- ``test_smoke_gsm8k`` — short, cheap GSM8K pass to verify the server
can produce coherent math answers at all (smoke gate).
- ``test_aime25`` — full AIME25 accuracy run (heavy; 16 repeats default).
Cookbook reference:
https://docs.sglang.io/cookbook/autoregressive/DeepSeek/DeepSeek-V4
These are MANUAL tests (not CI). ``sgl-eval`` must be on PATH.
Per-variant defaults (set on the Flash/Pro intermediate base classes):
Flash recipes -> AIME25 score threshold 0.93
Pro recipes -> AIME25 score threshold 0.95
GSM8K smoke threshold (0.93) is shared across Flash and Pro.
AIME25 knobs (env vars):
DSV4_AIME25_NUM_REPEATS (default 16 -> --n-repeats)
DSV4_AIME25_TEMPERATURE (default 1.0 -> --temperature)
DSV4_AIME25_TOP_P (default 1.0 -> --top-p)
DSV4_AIME25_MAX_TOKENS (default 65536 -> --max-tokens)
DSV4_AIME25_NUM_THREADS (default 512 -> --num-threads)
DSV4_AIME25_SCORE_METRIC (default "score"; sgl-eval JSON key under "aggregate")
DSV4_AIME25_SCORE_THRESHOLD (default 0; >0 overrides per-variant default)
GSM8K smoke knobs (env vars):
DSV4_GSM8K_NUM_EXAMPLES (default 50 -> --num-examples)
DSV4_GSM8K_N_REPEATS (default 1 -> --n-repeats)
DSV4_GSM8K_TEMPERATURE (default 0.6 -> --temperature)
DSV4_GSM8K_TOP_P (default 0.95 -> --top-p)
DSV4_GSM8K_MAX_TOKENS (default 8192 -> --max-tokens)
DSV4_GSM8K_NUM_THREADS (default 64 -> --num-threads)
DSV4_GSM8K_SCORE_METRIC (default "score"; sgl-eval JSON key under "aggregate")
DSV4_GSM8K_SCORE_THRESHOLD (default 0.93; set to 0 to skip the assertion)
Shared knobs:
DSV4_SGL_EVAL_OUT_DIR (default /tmp/sgl-eval-out -> --out-dir)
DSV4_SGL_EVAL_BIN (default "sgl-eval"; override path to the CLI)
DSV4_SERVER_LAUNCH_TIMEOUT (default 3600s; the sglang 600s default is
too short for DSV4 model load + DeepGEMM
warmup. 1800s is also tight for the heavier
recipes (DP-attn + DeepEP); 3600s is the
safe default. Bump again for first-run
model downloads if needed.)
Multi-node knobs (only consumed by multi-node test classes; if either
is unset, those classes ``SkipTest``):
DSV4_NODE_RANK (per-node rank for --node-rank)
DSV4_DIST_INIT_ADDR (e.g. 10.0.0.1:20000 for --dist-init-addr)
Always-on env (set by the base class for every recipe; per-recipe EXTRA_ENV
wins on key conflict):
SGLANG_JIT_DEEPGEMM_FAST_WARMUP=1 skip the slow DeepGEMM warmup grid
"""
import json
import os
import shutil
import subprocess
import unittest
from pathlib import Path
from typing import ClassVar, Dict, List, Optional
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
SGL_EVAL_BIN = os.environ.get("DSV4_SGL_EVAL_BIN", "sgl-eval")
SGL_EVAL_OUT_DIR = os.environ.get("DSV4_SGL_EVAL_OUT_DIR", "/tmp/sgl-eval-out")
# DSV4 server launch needs more than the 600s sglang default: model load alone
# can take 5+ min and DeepGEMM warmup another ~5 min. First-run model download
# adds ~10-30 min on top. 1800s covers steady-state; bump via env for downloads.
SERVER_LAUNCH_TIMEOUT = int(os.environ.get("DSV4_SERVER_LAUNCH_TIMEOUT", "3600"))
# Defaults applied to every recipe's EXTRA_ENV. Per-recipe EXTRA_ENV wins on key
# conflict.
BASE_ENV: Dict[str, str] = {
# Skip the slow exhaustive DeepGEMM warmup grid; covers the shapes DSV4
# actually hits and shaves several minutes off server startup.
"SGLANG_JIT_DEEPGEMM_FAST_WARMUP": "1",
}
AIME25_NUM_REPEATS = int(os.environ.get("DSV4_AIME25_NUM_REPEATS", "16"))
AIME25_TEMPERATURE = float(os.environ.get("DSV4_AIME25_TEMPERATURE", "1.0"))
AIME25_TOP_P = float(os.environ.get("DSV4_AIME25_TOP_P", "1.0"))
AIME25_MAX_TOKENS = int(os.environ.get("DSV4_AIME25_MAX_TOKENS", "65536"))
AIME25_NUM_THREADS = int(os.environ.get("DSV4_AIME25_NUM_THREADS", "512"))
AIME25_SCORE_METRIC = os.environ.get("DSV4_AIME25_SCORE_METRIC", "score")
AIME25_SCORE_THRESHOLD = float(os.environ.get("DSV4_AIME25_SCORE_THRESHOLD", "0.0"))
GSM8K_NUM_EXAMPLES = int(os.environ.get("DSV4_GSM8K_NUM_EXAMPLES", "50"))
GSM8K_N_REPEATS = int(os.environ.get("DSV4_GSM8K_N_REPEATS", "1"))
GSM8K_TEMPERATURE = float(os.environ.get("DSV4_GSM8K_TEMPERATURE", "0.6"))
GSM8K_TOP_P = float(os.environ.get("DSV4_GSM8K_TOP_P", "0.95"))
GSM8K_MAX_TOKENS = int(os.environ.get("DSV4_GSM8K_MAX_TOKENS", "8192"))
GSM8K_NUM_THREADS = int(os.environ.get("DSV4_GSM8K_NUM_THREADS", "64"))
GSM8K_SCORE_METRIC = os.environ.get("DSV4_GSM8K_SCORE_METRIC", "score")
GSM8K_SCORE_THRESHOLD = float(os.environ.get("DSV4_GSM8K_SCORE_THRESHOLD", "0.93"))
# DeepEP "large SMS" config — appears as `--deepep-config '{...}'` in every
# DeepEP recipe except multi-node ones (where it is gated off in the JSX).
DEEPEP_LARGE_SMS_CONFIG = (
'{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
)
def multinode_args(nnodes: int) -> List[str]:
"""Return CLI args for a multi-node launch, or skip the test.
Reads DSV4_NODE_RANK and DSV4_DIST_INIT_ADDR from the env. Raises
``unittest.SkipTest`` when either is missing — call from inside
``setUpClass`` so the whole class skips cleanly.
"""
rank = os.environ.get("DSV4_NODE_RANK")
addr = os.environ.get("DSV4_DIST_INIT_ADDR")
if rank is None or addr is None:
raise unittest.SkipTest(
"multi-node test requires DSV4_NODE_RANK and DSV4_DIST_INIT_ADDR"
)
return [
"--nnodes",
str(nnodes),
"--node-rank",
rank,
"--dist-init-addr",
addr,
]
class DSV4Aime25TestBase(CustomTestCase):
"""Subclass via ``DSV4FlashAime25TestBase`` or ``DSV4ProAime25TestBase``,
not directly. Per-recipe subclasses set MODEL / OTHER_ARGS / EXTRA_ENV.
SCORE_THRESHOLD is set by the Flash/Pro intermediate base classes:
Flash 0.93, Pro 0.95.
"""
MODEL: ClassVar[str] = ""
OTHER_ARGS: ClassVar[List[str]] = []
EXTRA_ENV: ClassVar[Dict[str, str]] = {}
SCORE_THRESHOLD: ClassVar[float] = 0.0
_BASE_CLASSES: ClassVar[set] = set()
@classmethod
def setUpClass(cls):
if cls in cls._BASE_CLASSES:
raise unittest.SkipTest("base class; subclass to run")
if not cls.MODEL or not cls.OTHER_ARGS:
raise unittest.SkipTest(f"{cls.__name__}: MODEL and OTHER_ARGS must be set")
cls.base_url = DEFAULT_URL_FOR_TEST
env: Optional[Dict[str, str]] = {**BASE_ENV, **(cls.EXTRA_ENV or {})}
cls.process = popen_launch_server(
cls.MODEL,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=list(cls.OTHER_ARGS),
env=env,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def test_smoke_gsm8k(self):
"""Quick GSM8K pass to verify the server is producing math answers."""
self._run_sgl_eval(
eval_name="gsm8k",
n_repeats=GSM8K_N_REPEATS,
temperature=GSM8K_TEMPERATURE,
top_p=GSM8K_TOP_P,
max_tokens=GSM8K_MAX_TOKENS,
num_threads=GSM8K_NUM_THREADS,
num_examples=GSM8K_NUM_EXAMPLES,
metric=GSM8K_SCORE_METRIC,
threshold=GSM8K_SCORE_THRESHOLD,
)
def test_aime25(self):
"""Full AIME25 accuracy run; threshold gated by Flash vs Pro base."""
threshold = (
AIME25_SCORE_THRESHOLD
if AIME25_SCORE_THRESHOLD > 0
else self.SCORE_THRESHOLD
)
self._run_sgl_eval(
eval_name="aime25",
n_repeats=AIME25_NUM_REPEATS,
temperature=AIME25_TEMPERATURE,
top_p=AIME25_TOP_P,
max_tokens=AIME25_MAX_TOKENS,
num_threads=AIME25_NUM_THREADS,
num_examples=None,
metric=AIME25_SCORE_METRIC,
threshold=threshold,
)
def _run_sgl_eval(
self,
eval_name,
n_repeats,
temperature,
top_p,
max_tokens,
num_threads,
num_examples,
metric,
threshold,
):
if shutil.which(SGL_EVAL_BIN) is None:
self.skipTest(f"{SGL_EVAL_BIN!r} not found on PATH")
out_dir = Path(SGL_EVAL_OUT_DIR)
out_dir.mkdir(parents=True, exist_ok=True)
glob_pattern = f"sgl_eval_{eval_name}_*.json"
before = set(out_dir.glob(glob_pattern))
cmd = [
SGL_EVAL_BIN,
"run",
eval_name,
"--base-url",
f"{self.base_url}/v1",
"--n-repeats",
str(n_repeats),
"--temperature",
str(temperature),
"--top-p",
str(top_p),
"--max-tokens",
str(max_tokens),
"--num-threads",
str(num_threads),
"--out-dir",
str(out_dir),
]
if num_examples is not None:
cmd += ["--num-examples", str(num_examples)]
print(f"[{type(self).__name__}] + {' '.join(cmd)}", flush=True)
subprocess.run(cmd, check=True)
new = sorted(set(out_dir.glob(glob_pattern)) - before)
if not new:
self.fail(f"sgl-eval produced no new {eval_name} JSON in {out_dir}")
result_path = new[-1]
with open(result_path) as f:
result = json.load(f)
print(
f"[{type(self).__name__}] sgl-eval {eval_name} result "
f"({result_path.name}): {json.dumps(result, indent=2)}",
flush=True,
)
score = self._extract_score(result, metric)
if threshold > 0:
self.assertGreaterEqual(
score,
threshold,
f"{eval_name} {metric}={score} below threshold {threshold}",
)
@staticmethod
def _extract_score(result, metric):
"""Find ``metric`` (e.g. "pass@1") anywhere in the sgl-eval JSON tree."""
def walk(o):
if isinstance(o, dict):
if metric in o and isinstance(o[metric], (int, float)):
return float(o[metric])
for v in o.values():
s = walk(v)
if s is not None:
return s
elif isinstance(o, list):
for v in o:
s = walk(v)
if s is not None:
return s
return None
score = walk(result)
if score is None:
raise AssertionError(f"metric {metric!r} not found in sgl-eval result JSON")
return score
class DSV4FlashAime25TestBase(DSV4Aime25TestBase):
"""Base for DeepSeek-V4-Flash recipes: AIME25 threshold 0.93."""
SCORE_THRESHOLD = 0.93
class DSV4ProAime25TestBase(DSV4Aime25TestBase):
"""Base for DeepSeek-V4-Pro recipes: AIME25 threshold 0.95."""
SCORE_THRESHOLD = 0.95
DSV4Aime25TestBase._BASE_CLASSES = {
DSV4Aime25TestBase,
DSV4FlashAime25TestBase,
DSV4ProAime25TestBase,
}
test/manual/dsv4/test_b200_flash.py
+109
View File
@@ -0,0 +1,109 @@
"""B200 (FP4) x DeepSeek-V4-Flash.
Covers the four cookbook recipes for this hardware x model_size cell:
Low-Latency, Balanced, Max-Throughput, Context-Parallel (CP).
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
class TestB200FlashLowLatency(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
]
EXTRA_ENV = {}
class TestB200FlashBalanced(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
class TestB200FlashMaxThroughput(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
class TestB200FlashCP(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.78",
"--max-running-requests",
"1024",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_b200_pro.py
+125
View File
@@ -0,0 +1,125 @@
"""B200 (FP4) x DeepSeek-V4-Pro.
Covers the four cookbook recipes for this hardware x model_size cell:
Low-Latency, Balanced, Max-Throughput, Context-Parallel (CP).
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4ProAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Pro"
class TestB200ProLowLatency(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
"--mem-fraction-static",
"0.88",
]
EXTRA_ENV = {}
class TestB200ProBalanced(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--mem-fraction-static",
"0.82",
"--cuda-graph-max-bs",
"64",
"--max-running-requests",
"128",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
class TestB200ProMaxThroughput(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--mem-fraction-static",
"0.82",
"--cuda-graph-max-bs",
"64",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
class TestB200ProCP(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.78",
"--cuda-graph-max-bs",
"256",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_b300_flash.py
+111
View File
@@ -0,0 +1,111 @@
"""B300 x DeepSeek-V4-Flash.
The cookbook generator aliases B300 to B200, so the launch flags
are identical to the B200(FP4) Flash cell. Kept as a separate file
because the hardware target (and therefore the runtime environment)
is different. Covers Low-Latency, Balanced, Max-Throughput, CP.
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
class TestB300FlashLowLatency(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
]
EXTRA_ENV = {}
class TestB300FlashBalanced(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
class TestB300FlashMaxThroughput(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
class TestB300FlashCP(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.78",
"--max-running-requests",
"1024",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_b300_pro.py
+127
View File
@@ -0,0 +1,127 @@
"""B300 x DeepSeek-V4-Pro.
The cookbook generator aliases B300 to B200, so the launch flags
are identical to the B200(FP4) Pro cell. Kept as a separate file
because the hardware target (and therefore the runtime environment)
is different. Covers Low-Latency, Balanced, Max-Throughput, CP.
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4ProAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Pro"
class TestB300ProLowLatency(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
"--mem-fraction-static",
"0.88",
]
EXTRA_ENV = {}
class TestB300ProBalanced(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--mem-fraction-static",
"0.82",
"--cuda-graph-max-bs",
"64",
"--max-running-requests",
"128",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
class TestB300ProMaxThroughput(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--mem-fraction-static",
"0.82",
"--cuda-graph-max-bs",
"64",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
class TestB300ProCP(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.78",
"--cuda-graph-max-bs",
"256",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_flash_mtp_dp4.py
+175
View File
@@ -0,0 +1,175 @@
"""DSV4 Flash MTP test using EAGLE speculative algorithm.
DSV4 Flash MTP shares the EAGLE wire path: EAGLE algo + NextN head built
into the target model weights. No separate draft model is needed (sglang
auto-falls back `--speculative-draft-model-path` to the target model).
Test matrix mirrors test_eagle_infer_b.TestEAGLEServerBasic to maximize
cuda-graph + buffer-pool coverage on the DSV4 path:
- test_gsm8k (accuracy + spec path full forward)
- test_max_token_one (degenerate spec step, still cuda-graph captured)
- test_request_abort (cuda-graph buffer pool survives abort+restart)
Server launch matches `run_flash_dp4.sh`: tp=4, dp=4, deepep MoE backend,
DSV4 FP8 (FP4 experts disabled).
"""
import random
import threading
import time
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_gsm8k_eval
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_ENV = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
# MTP runs ~num_draft_tokens forward passes per step, so the deepep
# dispatch input size scales by that factor. Default 256 (used by the
# plain server) overflows once cuda-graph-max-bs * num_draft_tokens
# > 256. 1024 covers bs=128 * 4 draft tokens with headroom.
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
PROMPTS = [
"[INST] You are a helpful assistant.\\nWhere are you from? [/INST]",
"[INST] You are a helpful assistant.\\nSummarize gradient descent in 2 sentences. [/INST]",
"[INST] You are a helpful assistant.\\nWhat is 17*23? [/INST]",
"[INST] You are a helpful assistant.\\nList three primary colors. [/INST]",
]
class DSV4FlashMTPServerBase(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DSV4_FLASH_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_CONFIG,
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-fraction-static",
"0.7",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env=DSV4_FLASH_ENV,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def send_request(self):
time.sleep(random.uniform(0, 2))
for prompt in PROMPTS:
resp = requests.post(
self.base_url + "/generate",
json={
"text": prompt,
"sampling_params": {"temperature": 0, "max_new_tokens": 256},
},
)
assert resp.status_code == 200
def send_requests_abort(self):
for prompt in PROMPTS:
try:
time.sleep(random.uniform(0, 2))
requests.post(
self.base_url + "/generate",
json={
"text": prompt,
"sampling_params": {"temperature": 0, "max_new_tokens": 256},
},
timeout=0.5,
)
except requests.exceptions.Timeout:
pass
class TestDSV4FlashMTPBasic(DSV4FlashMTPServerBase):
def test_gsm8k(self):
"""Accuracy + spec path full forward."""
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_gsm8k_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.95)
def test_max_token_one(self):
"""Degenerate spec step (still cuda-graph captured)."""
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=100,
max_new_tokens=1,
parallel=128,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_gsm8k_eval(args)
self.assertGreater(metrics["output_throughput"], 50)
def test_request_abort(self):
"""Cuda-graph buffer pool must survive abort+restart cycles."""
concurrency = 4
threads = [
threading.Thread(target=self.send_request) for _ in range(concurrency)
] + [
threading.Thread(target=self.send_requests_abort)
for _ in range(concurrency)
]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertIsNone(self.process.poll())
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_flash_mtp_tp8.py
+125
View File
@@ -0,0 +1,125 @@
"""DSV4-Flash 285B MTP performance tests on H200 TP=8.
Manual test (8× H200, 285B FP8 weights). Not registered in CI.
"""
import os
import tempfile
import unittest
import requests
from sglang.bench_one_batch_server import BenchArgs as OneBatchBenchArgs
from sglang.bench_one_batch_server import run_benchmark as run_one_batch_benchmark
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_BASE_ENV = {
"SGLANG_ENABLE_SPEC_V2": "1",
"SGLANG_OPT_USE_TOPK_V2": "1",
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_JIT_DEEPGEMM_PRECOMPILE": "0",
}
DSV4_FLASH_SERVER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--max-running-requests",
"8",
]
def _launch_dsv4_flash_server(extra_env=None):
env = dict(DSV4_FLASH_BASE_ENV)
if extra_env:
env.update(extra_env)
return popen_launch_server(
DSV4_FLASH_MODEL_PATH,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 4,
other_args=DSV4_FLASH_SERVER_ARGS,
env=env,
)
class TestDSV4FlashMTPSimulatedAcc(CustomTestCase):
"""bs=1 latency at isl=4096 / 900000 with `SGLANG_SIMULATE_ACC_LEN=3`.
Reference (H200 Flash TP8):
- isl=4096 → output 258.1 tok/s, accept 2.94
- isl=900000 → output 222.9 tok/s, accept 2.90
"""
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = _launch_dsv4_flash_server(
extra_env={"SGLANG_SIMULATE_ACC_LEN": "3"}
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def _run_one_batch(self, input_len):
requests.get(self.base_url + "/flush_cache")
server_args = ServerArgs(model_path=DSV4_FLASH_MODEL_PATH)
bench_args = OneBatchBenchArgs(
run_name=f"dsv4_flash_simacc_isl{input_len}",
batch_size=(1,),
input_len=(input_len,),
output_len=(1024,),
base_url=self.base_url,
skip_warmup=True,
result_filename=os.path.join(
tempfile.gettempdir(), f"dsv4_flash_simacc_isl{input_len}.jsonl"
),
append_to_github_summary=False,
)
results, _ = run_one_batch_benchmark(server_args, bench_args)
self.assertTrue(results, "bench_one_batch_server returned no results")
return results[0]
def test_isl_4096(self):
r = self._run_one_batch(4096)
print(
f"[flash simacc isl=4096] output_throughput={r.output_throughput:.2f} tok/s "
f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
f"acc_length={r.acc_length:.2f}"
)
# Reference 258.1 tok/s / acc=2.94.
self.assertGreater(r.output_throughput, 232.0)
self.assertGreater(r.acc_length, 2.85)
def test_isl_900k(self):
r = self._run_one_batch(900_000)
print(
f"[flash simacc isl=900k] output_throughput={r.output_throughput:.2f} tok/s "
f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
f"acc_length={r.acc_length:.2f}"
)
# Reference 222.9 tok/s / acc=2.90.
self.assertGreater(r.output_throughput, 200.0)
self.assertGreater(r.acc_length, 2.85)
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_flash_sanity_dp4.py
+151
View File
@@ -0,0 +1,151 @@
"""DSV4-Flash 4-GPU server sanity matrix (TP4 variants)."""
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.server_sanity_kit import ServerSanityMixin
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_ENV = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
def _launch(other_args, env_extra=None, timeout_mult=1):
env = dict(DSV4_FLASH_ENV)
if env_extra:
env.update(env_extra)
return popen_launch_server(
DSV4_FLASH_MODEL_PATH,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * timeout_mult,
other_args=other_args,
env=env,
)
_EAGLE_SPEC_ARGS = [
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
class TestDSV4FlashTP4DP4(ServerSanityMixin, CustomTestCase):
"""TP4 + DP4 + deepep + EAGLE MTP."""
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = _launch(
[
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_CONFIG,
"--mem-fraction-static",
"0.7",
*_EAGLE_SPEC_ARGS,
]
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
class TestDSV4FlashTP4EP(ServerSanityMixin, CustomTestCase):
"""TP attn + EP MoE (no DP attn) — exercises the DeepEP + TP-attn path."""
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = _launch(
[
"--trust-remote-code",
"--tp",
"4",
"--ep",
"4",
# No --enable-dp-attention by design: covers TP-attn path.
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"64",
"--deepep-config",
DEEPEP_CONFIG,
"--mem-fraction-static",
"0.7",
*_EAGLE_SPEC_ARGS,
]
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
class TestDSV4FlashTP4DP4ChunkedPrefillLarge(ServerSanityMixin, CustomTestCase):
"""TP4 + DP4 with --chunked-prefill-size 16384 — large chunked prefill."""
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = _launch(
[
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--chunked-prefill-size",
"16384",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_CONFIG,
"--mem-fraction-static",
"0.7",
*_EAGLE_SPEC_ARGS,
]
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_flash_sanity_tp8.py
+50
View File
@@ -0,0 +1,50 @@
"""DSV4-Flash 8-GPU server sanity (TP8, no spec decoding)."""
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.server_sanity_kit import ServerSanityMixin
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_ENV = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
class TestDSV4FlashTP8NoSpec(ServerSanityMixin, CustomTestCase):
"""TP8, no spec decoding."""
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
DSV4_FLASH_MODEL_PATH,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"8",
"--max-running-requests",
"8",
"--mem-fraction-static",
"0.85",
],
env=DSV4_FLASH_ENV,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_pd_disagg_nixl.py
+148
View File
@@ -0,0 +1,148 @@
"""DSV4 Flash PD-disagg with NIXL backend. Both sides run dp-attention
+ deepep + EAGLE MTP so attn_tp_size and the V4 state pool layout are
fully symmetric: same SWA item_len under matching attn_tp, and same
NSA c4/c128 indexer ring buffer size under matching spec status. nixl
`send_state` is page-by-index and has no V4 TP-slice / spec-asymmetric
path, so any layout mismatch would trip the item_len assert in
`nixl/conn.py`."""
import unittest
from types import SimpleNamespace
from sglang.test.few_shot_gsm8k import run_eval as run_gsm8k_eval
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_ENV = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
# MTP num_draft_tokens=4 scales dispatch by ~4x; 256 overflows at bs=128.
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
class TestDSV4FlashPDDisaggNIXL(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.transfer_backend = ["--disaggregation-transfer-backend", "nixl"]
cls.rdma_devices = []
cls.model = DSV4_FLASH_MODEL_PATH
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--base-gpu-id",
"0",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--deepep-config",
DEEPEP_CONFIG,
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--mem-fraction-static",
"0.7",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
*cls.transfer_backend,
*cls.rdma_devices,
]
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
env=DSV4_FLASH_ENV,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--base-gpu-id",
"4",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--deepep-config",
DEEPEP_CONFIG,
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--mem-fraction-static",
"0.7",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
*cls.transfer_backend,
*cls.rdma_devices,
]
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
env=DSV4_FLASH_ENV,
)
def test_gsm8k(self):
"""End-to-end PD-disagg accuracy through the LB."""
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=64,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_gsm8k_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.95)
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_pro_mtp.py
+281
View File
@@ -0,0 +1,281 @@
"""DSV4-Pro 1.6T MTP performance tests on B200 TP=8.
1. TestDSV4ProMTPSimulatedAcc — `SGLANG_SIMULATE_ACC_LEN=3` pins EAGLE accept
length so latency comparisons are apples-to-apples. Runs `bench_one_batch_server`
at bs=1 for isl=4096 and isl=900000 (osl=1024).
2. TestDSV4ProMTPHongloumeng — real EAGLE accept (no SIMULATE) on Chinese
long-context input (`hongloumeng.txt`, ~627k DSV4 tokens). Builds a one-line
custom JSONL dataset on the fly and drives `bench_serving --dataset-name custom`
with one short slice (30k tokens) and the full long prompt.
Manual test (8× B200, 1.6T weights). Not registered in CI.
"""
import json
import os
import tempfile
import unittest
from types import SimpleNamespace
import requests
from sglang.bench_one_batch_server import BenchArgs as OneBatchBenchArgs
from sglang.bench_one_batch_server import run_benchmark as run_one_batch_benchmark
from sglang.bench_serving import run_benchmark as run_serving_benchmark
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DSV4_PRO_MODEL_PATH = "deepseek-ai/DeepSeek-V4-Pro"
HONGLOUMENG_PATH = os.environ.get(
"SGLANG_HONGLOUMENG_PATH",
os.path.join(os.path.dirname(__file__), "hongloumeng.txt"),
)
DSV4_PRO_BASE_ENV = {
"SGLANG_ENABLE_SPEC_V2": "1",
"SGLANG_OPT_USE_TOPK_V2": "1",
"SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2": "1",
"SGLANG_JIT_DEEPGEMM_PRECOMPILE": "0",
}
DSV4_PRO_SERVER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
"--mem-fraction-static",
"0.82",
"--max-running-requests",
"8",
]
def _launch_dsv4_pro_server(extra_env=None):
env = dict(DSV4_PRO_BASE_ENV)
if extra_env:
env.update(extra_env)
return popen_launch_server(
DSV4_PRO_MODEL_PATH,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 4,
other_args=DSV4_PRO_SERVER_ARGS,
env=env,
)
class TestDSV4ProMTPSimulatedAcc(CustomTestCase):
"""bs=1 latency at isl=4096 / 900000 with `SGLANG_SIMULATE_ACC_LEN=3`.
Reference (B200 Pro TP8):
- isl=4096 → output 194.6 tok/s, accept 2.96
- isl=900000 → output 174.6 tok/s, accept 2.93
"""
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = _launch_dsv4_pro_server(
extra_env={"SGLANG_SIMULATE_ACC_LEN": "3"}
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def _run_one_batch(self, input_len):
requests.get(self.base_url + "/flush_cache")
server_args = ServerArgs(model_path=DSV4_PRO_MODEL_PATH)
bench_args = OneBatchBenchArgs(
run_name=f"dsv4_pro_simacc_isl{input_len}",
batch_size=(1,),
input_len=(input_len,),
output_len=(1024,),
base_url=self.base_url,
skip_warmup=True,
result_filename=os.path.join(
tempfile.gettempdir(), f"dsv4_pro_simacc_isl{input_len}.jsonl"
),
append_to_github_summary=False,
)
results, _ = run_one_batch_benchmark(server_args, bench_args)
self.assertTrue(results, "bench_one_batch_server returned no results")
return results[0]
def test_isl_4096(self):
r = self._run_one_batch(4096)
print(
f"[pro simacc isl=4096] output_throughput={r.output_throughput:.2f} tok/s "
f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
f"acc_length={r.acc_length:.2f}"
)
# Reference 194.6 tok/s / acc=2.96 — give 10% throughput margin and a
# generous accept-length floor to absorb run-to-run jitter.
self.assertGreater(r.output_throughput, 175.0)
self.assertGreater(r.acc_length, 2.85)
def test_isl_900k(self):
r = self._run_one_batch(900_000)
print(
f"[pro simacc isl=900k] output_throughput={r.output_throughput:.2f} tok/s "
f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
f"acc_length={r.acc_length:.2f}"
)
# Reference 174.6 tok/s / acc=2.93.
self.assertGreater(r.output_throughput, 155.0)
self.assertGreater(r.acc_length, 2.85)
def _build_hongloumeng_jsonl(num_tokens, tokenizer, out_path):
"""Slice the first `num_tokens` DSV4 tokens of hongloumeng.txt into a
one-line CustomDataset JSONL. Pass num_tokens=None to keep the full text.
"""
with open(HONGLOUMENG_PATH, "r", encoding="utf-8") as f:
text = f.read()
if num_tokens is not None:
ids = tokenizer.encode(text)
text = tokenizer.decode(ids[:num_tokens])
with open(out_path, "w", encoding="utf-8") as f:
f.write(
json.dumps(
{"conversations": [{"value": text}, {"value": "x"}]},
ensure_ascii=False,
)
+ "\n"
)
return out_path
class TestDSV4ProMTPHongloumeng(CustomTestCase):
"""Real EAGLE accept on Chinese long-context (hongloumeng.txt).
Reference (B200 Pro TP8, no SIMULATE):
- isl=30000 → output 124.4 tok/s, decode peak 184 tok/s, accept 2.47
- isl=627059 → output 125.7 tok/s, decode peak 179 tok/s, accept 2.52
"""
SHORT_TOKENS = 30_000
LONG_TOKENS = None # full file (~627k DSV4 tokens)
OUTPUT_TOKENS = 4096
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = _launch_dsv4_pro_server()
# Resolve tokenizer once; the server reports its own tokenizer path so
# on-the-fly token-level slicing matches what the server will see.
info = requests.get(cls.base_url + "/server_info", timeout=60).json()
tokenizer_path = info.get("tokenizer_path") or DSV4_PRO_MODEL_PATH
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
cls.tokenizer = get_tokenizer(tokenizer_path)
cls.tmpdir = tempfile.mkdtemp(prefix="dsv4_hongloumeng_")
cls.short_jsonl = _build_hongloumeng_jsonl(
cls.SHORT_TOKENS,
cls.tokenizer,
os.path.join(cls.tmpdir, "hongloumeng_30k.jsonl"),
)
cls.long_jsonl = _build_hongloumeng_jsonl(
cls.LONG_TOKENS,
cls.tokenizer,
os.path.join(cls.tmpdir, "hongloumeng_full.jsonl"),
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def _run_custom_bench(self, dataset_path):
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
backend="sglang",
base_url=self.base_url,
host=None,
port=None,
dataset_name="custom",
dataset_path=dataset_path,
model=None,
tokenizer=None,
num_prompts=1,
sharegpt_output_len=self.OUTPUT_TOKENS,
sharegpt_context_len=None,
random_input_len=4096,
random_output_len=2048,
random_range_ratio=0.0,
request_rate=float("inf"),
max_concurrency=1,
warmup_requests=0,
flush_cache=True,
multi=None,
output_file=None,
disable_tqdm=False,
disable_stream=False,
return_logprob=False,
return_routed_experts=False,
seed=0,
disable_ignore_eos=False,
extra_request_body=None,
apply_chat_template=False,
profile=None,
lora_name=None,
lora_request_distribution="uniform",
lora_zipf_alpha=1.5,
prompt_suffix="",
device="cuda",
pd_separated=False,
ready_check_timeout_sec=0,
)
return run_serving_benchmark(args)
def test_short_30k(self):
res = self._run_custom_bench(self.short_jsonl)
print(
f"[hongloumeng 30k] output_throughput={res['output_throughput']:.2f} tok/s "
f"accept_length={res['accept_length']:.2f} "
f"mean_ttft_ms={res['mean_ttft_ms']:.0f} "
f"mean_tpot_ms={res['mean_tpot_ms']:.2f}"
)
# Reference 124 tok/s / accept 2.47.
self.assertGreater(res["output_throughput"], 105.0)
self.assertGreater(res["accept_length"], 2.30)
def test_long_full(self):
res = self._run_custom_bench(self.long_jsonl)
print(
f"[hongloumeng full] output_throughput={res['output_throughput']:.2f} tok/s "
f"accept_length={res['accept_length']:.2f} "
f"mean_ttft_ms={res['mean_ttft_ms']:.0f} "
f"mean_tpot_ms={res['mean_tpot_ms']:.2f}"
)
# Reference 125 tok/s / accept 2.52. Cold prefill takes ~85s on 627k
# tokens so the run is dominated by prefill, but decode steady-state
# accept_length is the metric we care about.
self.assertGreater(res["output_throughput"], 105.0)
self.assertGreater(res["accept_length"], 2.30)
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_dsv4_swa_radix_retract.py
+165
View File
@@ -0,0 +1,165 @@
"""DSV4 stress test for SWA radix cache + tombstone + retract interaction.
Reproduces the assert in `swa_radix_cache.cache_unfinished_req`:
assert old_prefix_len <= len(new_indices)
Trip conditions (all required):
1. Fork-only SWA leaf early-release on (`SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW=1`)
2. Multiple requests share a long prefix (so one req's tombstoned leaf
poisons match_prefix for others walking the same radix path).
3. Memory pressure forces retract while at least one req has tombstoned
its leaf (decode_batch_idx >= sliding_window_size at retract time).
After main #19427 changed `old_prefix_len = req.cache_protected_len`
(stable), tombstone-induced shrinks in match's `best_value_len` across
chunked-prefill rounds can make stale `cache_protected_len` exceed
current matchable length -> assert trips.
Test passes iff the scheduler does not crash under this stress workload.
"""
import random
import threading
import time
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
# Long shared prefix forces multi-chunk prefill and ensures cross-request
# prefix-cache hits so one req's tombstone affects later reqs.
SHARED_PREFIX_BLOCK = (
"You are a careful, expert assistant. Answer concisely.\n"
"Context: " + ("the quick brown fox jumps over the lazy dog. " * 600)
)
QUESTION_TAILS = [
" Q: What is 17*23?\n",
" Q: List three primary colors.\n",
" Q: Where is Mount Everest?\n",
" Q: Summarize gradient descent in two sentences.\n",
" Q: Name two bodies of water in Africa.\n",
" Q: What language is spoken in Brazil?\n",
]
class TestDSV4FlashSWARadixRetract(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DSV4_FLASH_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
'{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}',
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
# Tight static memory so SWA pool fills up under load and
# retract is forced.
"--mem-fraction-static",
"0.7",
]
env = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
"SGLANG_OPT_SWA_RADIX_CACHE_COMPACT": "0",
"SGLANG_TEST_RETRACT": "1",
"SGLANG_TEST_RETRACT_INTERVAL": "3",
}
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env=env,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _send_req(self, prompt: str, max_new_tokens: int):
try:
resp = requests.post(
self.base_url + "/generate",
json={
"text": prompt,
"sampling_params": {
# Vary outputs slightly so reqs don't share decode
# paths perfectly; we want some to finish, some to
# be retracted under pressure.
"temperature": 0.7,
"max_new_tokens": max_new_tokens,
},
},
timeout=600,
)
# Per-request success is not the gate; some requests are
# expected to be retracted/aborted under heavy pressure.
return resp.status_code == 200
except Exception:
return False
def test_swa_tombstone_retract_does_not_crash(self):
"""Stress: 64 concurrent long-prompt reqs with long generation force
retract under SWA pool pressure. Reqs share a 30k+ token prefix so
tombstoned leaves from retracted reqs are on the radix path of new
reqs. Scheduler must not crash on the swa_radix_cache assert."""
random.seed(0)
concurrency = 64
# Long enough generation to push past sliding_window_size -> fires
# `dec_swa_lock_only` -> tombstones leaves. Combined with SWA pool
# pressure this guarantees retract while tombstones are live.
max_new_tokens = 1024
threads = []
for i in range(concurrency):
tail = QUESTION_TAILS[i % len(QUESTION_TAILS)]
# Add a small per-req suffix so reqs don't dedup at radix root
# but still share the bulk of the prefix.
prompt = SHARED_PREFIX_BLOCK + tail + f"(seed={i})"
t = threading.Thread(target=self._send_req, args=(prompt, max_new_tokens))
threads.append(t)
t.start()
# Stagger so requests enter prefill in waves; some are still in
# decode (and have tombstoned leaves) when later waves of
# chunked-prefill reqs walk the same radix path.
time.sleep(0.05)
for t in threads:
t.join(timeout=600)
# The only invariant: scheduler survived. Per-request completion is
# best-effort under retract pressure.
self.assertIsNone(self.process.poll())
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_gb300_flash.py
+109
View File
@@ -0,0 +1,109 @@
"""GB300 x DeepSeek-V4-Flash.
Single-node TP=4 path on the deepseek-ai MXFP4 repo. Covers
Low-Latency, Balanced, Max-Throughput, Context-Parallel (CP).
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
class TestGB300FlashLowLatency(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
]
EXTRA_ENV = {}
class TestGB300FlashBalanced(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
class TestGB300FlashMaxThroughput(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
class TestGB300FlashCP(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.78",
"--max-running-requests",
"1024",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_gb300_pro.py
+127
View File
@@ -0,0 +1,127 @@
"""GB300 x DeepSeek-V4-Pro.
Single-node TP=4 path. Note that GB300 Pro CP bumps
mem-fraction-static to 0.88 (1.6T weights at TP=4 on 273 GB don't
fit at the default 0.78). Covers Low-Latency, Balanced,
Max-Throughput, Context-Parallel (CP).
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4ProAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Pro"
class TestGB300ProLowLatency(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"flashinfer_mxfp4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--chunked-prefill-size",
"4096",
"--disable-flashinfer-autotune",
"--mem-fraction-static",
"0.88",
]
EXTRA_ENV = {}
class TestGB300ProBalanced(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--mem-fraction-static",
"0.9",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
class TestGB300ProMaxThroughput(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--mem-fraction-static",
"0.9",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
class TestGB300ProCP(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.88",
"--cuda-graph-max-bs",
"256",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_h200_fp4_flash.py
+71
View File
@@ -0,0 +1,71 @@
"""H200 (FP4 / Marlin) x DeepSeek-V4-Flash.
The cookbook disables Context-Parallel for the H200 FP4 (Marlin)
hardware, so this file only covers Low-Latency, Balanced, and
Max-Throughput.
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DSV4FlashAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
class TestH200Fp4FlashLowLatency(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"marlin",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
EXTRA_ENV = {}
class TestH200Fp4FlashBalanced(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"marlin",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
]
EXTRA_ENV = {}
class TestH200Fp4FlashMaxThroughput(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-runner-backend",
"marlin",
]
EXTRA_ENV = {}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_h200_fp4_pro.py
+77
View File
@@ -0,0 +1,77 @@
"""H200 (FP4 / Marlin) x DeepSeek-V4-Pro.
The cookbook disables Context-Parallel for the H200 FP4 (Marlin)
hardware, so this file only covers Low-Latency, Balanced, and
Max-Throughput.
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DSV4ProAime25TestBase
MODEL = "deepseek-ai/DeepSeek-V4-Pro"
class TestH200Fp4ProLowLatency(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-runner-backend",
"marlin",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-fraction-static",
"0.88",
]
EXTRA_ENV = {}
class TestH200Fp4ProBalanced(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-runner-backend",
"marlin",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--mem-fraction-static",
"0.88",
]
EXTRA_ENV = {}
class TestH200Fp4ProMaxThroughput(DSV4ProAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"8",
"--moe-runner-backend",
"marlin",
"--mem-fraction-static",
"0.88",
]
EXTRA_ENV = {}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_h200_fp8_flash.py
+121
View File
@@ -0,0 +1,121 @@
"""H200 (FP8) x DeepSeek-V4-Flash.
Uses the FP8-repackaged repo (sgl-project/DeepSeek-V4-Flash-FP8) and
the SGLANG_DSV4_FP4_EXPERTS=0 env that the cookbook generator emits
for H200 FP8 cells. Covers Low-Latency, Balanced, Max-Throughput, CP.
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
H200_FP8_ENV = {"SGLANG_DSV4_FP4_EXPERTS": "0"}
class TestH200Fp8FlashLowLatency(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
EXTRA_ENV = dict(H200_FP8_ENV)
class TestH200Fp8FlashBalanced(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"128",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
**H200_FP8_ENV,
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
}
class TestH200Fp8FlashMaxThroughput(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
**H200_FP8_ENV,
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
}
class TestH200Fp8FlashCP(DSV4FlashAime25TestBase):
MODEL = MODEL
OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
"0.78",
"--max-running-requests",
"1024",
"--deepep-config",
DEEPEP_LARGE_SMS_CONFIG,
]
EXTRA_ENV = {
**H200_FP8_ENV,
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_h200_fp8_pro.py
+126
View File
@@ -0,0 +1,126 @@
"""H200 (FP8) x DeepSeek-V4-Pro.
The cookbook ships this cell as a multi-node (2 nodes, TP=16) launch
using the FP8-repackaged repo (sgl-project/DeepSeek-V4-Pro-FP8).
Each test class skips itself unless DSV4_NODE_RANK and
DSV4_DIST_INIT_ADDR are exported. Runtime expectation:
On every node:
DSV4_NODE_RANK=<0 or 1> \\
DSV4_DIST_INIT_ADDR=<head-node-ip>:20000 \\
python test/manual/models/dsv4/test_h200_fp8_pro.py
Context-Parallel is marked TBD in the cookbook for this cell, so it
is intentionally omitted.
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _common import DSV4ProAime25TestBase, multinode_args
MODEL = "sgl-project/DeepSeek-V4-Pro-FP8"
H200_FP8_PRO_ENV = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "128",
}
class TestH200Fp8ProLowLatency(DSV4ProAime25TestBase):
MODEL = MODEL
EXTRA_ENV = dict(H200_FP8_PRO_ENV)
@classmethod
def setUpClass(cls):
cls.OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"16",
"--dp",
"16",
"--enable-dp-attention",
*multinode_args(2),
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"8",
"--max-running-requests",
"32",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-fraction-static",
"0.88",
]
super().setUpClass()
class TestH200Fp8ProBalanced(DSV4ProAime25TestBase):
MODEL = MODEL
EXTRA_ENV = dict(H200_FP8_PRO_ENV)
@classmethod
def setUpClass(cls):
cls.OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"16",
"--dp",
"16",
"--enable-dp-attention",
*multinode_args(2),
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--mem-fraction-static",
"0.88",
"--cuda-graph-max-bs",
"8",
"--max-running-requests",
"32",
]
super().setUpClass()
class TestH200Fp8ProMaxThroughput(DSV4ProAime25TestBase):
MODEL = MODEL
EXTRA_ENV = dict(H200_FP8_PRO_ENV)
@classmethod
def setUpClass(cls):
cls.OTHER_ARGS = [
"--trust-remote-code",
"--tp",
"16",
"--dp",
"16",
"--enable-dp-attention",
*multinode_args(2),
"--moe-a2a-backend",
"deepep",
"--mem-fraction-static",
"0.88",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
]
super().setUpClass()
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_swa_alloc_extend_page_estimation.py
+136
View File
@@ -0,0 +1,136 @@
"""Regression for SWA alloc_extend page estimation.
Old gate in SWATokenToKVPoolAllocator.alloc_extend added one full page_size
per request unconditionally, refusing extends that fit inside the request's
last partial page. Fix replaces with get_num_new_pages-based gating.
"""
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.test.test_utils import CustomTestCase
def _make_self(*, page_size: int, full_available: int, swa_available: int):
full_indices = torch.tensor([10, 11], dtype=torch.int64)
swa_indices = torch.tensor([20, 21], dtype=torch.int64)
return SimpleNamespace(
page_size=page_size,
full_attn_allocator=SimpleNamespace(
available_size=lambda: full_available,
alloc_extend=MagicMock(return_value=full_indices),
),
swa_attn_allocator=SimpleNamespace(
available_size=lambda: swa_available,
alloc_extend=MagicMock(return_value=swa_indices),
),
translate_loc_from_full_to_swa=lambda last_loc: last_loc,
full_to_swa_index_mapping=torch.zeros(64, dtype=torch.int64),
)
def _call(stub, *, prefix_lens_cpu, seq_lens_cpu, extend_num_tokens):
return SWATokenToKVPoolAllocator.alloc_extend(
stub,
prefix_lens=prefix_lens_cpu,
prefix_lens_cpu=prefix_lens_cpu,
seq_lens=seq_lens_cpu,
seq_lens_cpu=seq_lens_cpu,
last_loc=torch.tensor(
[int(p) - 1 for p in prefix_lens_cpu.tolist()], dtype=torch.int64
),
extend_num_tokens=extend_num_tokens,
)
class TestSWAAllocExtendPageEstimation(CustomTestCase):
def test_zero_new_pages_must_succeed(self):
# Old: 2 + 2*8 = 18 > 16 -> would refuse.
# New: prefix 5 -> 6 stays in page 0, 0 new pages.
stub = _make_self(page_size=8, full_available=16, swa_available=16)
result = _call(
stub,
prefix_lens_cpu=torch.tensor([5, 5], dtype=torch.int64),
seq_lens_cpu=torch.tensor([6, 6], dtype=torch.int64),
extend_num_tokens=2,
)
self.assertIsNotNone(result)
stub.full_attn_allocator.alloc_extend.assert_called_once()
stub.swa_attn_allocator.alloc_extend.assert_called_once()
def test_one_new_page_fits(self):
# Old: 6 + 2*8 = 22 > 16. New: 2 new pages == 16 // 8.
stub = _make_self(page_size=8, full_available=16, swa_available=16)
result = _call(
stub,
prefix_lens_cpu=torch.tensor([7, 7], dtype=torch.int64),
seq_lens_cpu=torch.tensor([10, 10], dtype=torch.int64),
extend_num_tokens=6,
)
self.assertIsNotNone(result)
def test_full_pool_genuinely_insufficient(self):
stub = _make_self(page_size=8, full_available=8, swa_available=64)
result = _call(
stub,
prefix_lens_cpu=torch.tensor([8, 8, 8, 8, 8], dtype=torch.int64),
seq_lens_cpu=torch.tensor([9, 9, 9, 9, 9], dtype=torch.int64),
extend_num_tokens=5,
)
self.assertIsNone(result)
stub.full_attn_allocator.alloc_extend.assert_not_called()
def test_swa_pool_genuinely_insufficient(self):
stub = _make_self(page_size=8, full_available=64, swa_available=8)
result = _call(
stub,
prefix_lens_cpu=torch.tensor([8, 8, 8, 8, 8], dtype=torch.int64),
seq_lens_cpu=torch.tensor([9, 9, 9, 9, 9], dtype=torch.int64),
extend_num_tokens=5,
)
self.assertIsNone(result)
stub.swa_attn_allocator.alloc_extend.assert_not_called()
def test_exactly_at_capacity_succeeds(self):
stub = _make_self(page_size=8, full_available=16, swa_available=16)
result = _call(
stub,
prefix_lens_cpu=torch.tensor([8, 8], dtype=torch.int64),
seq_lens_cpu=torch.tensor([9, 9], dtype=torch.int64),
extend_num_tokens=2,
)
self.assertIsNotNone(result)
def test_one_over_capacity_refuses(self):
stub = _make_self(page_size=8, full_available=16, swa_available=16)
result = _call(
stub,
prefix_lens_cpu=torch.tensor([8, 8, 8], dtype=torch.int64),
seq_lens_cpu=torch.tensor([9, 9, 9], dtype=torch.int64),
extend_num_tokens=3,
)
self.assertIsNone(result)
def test_zero_new_pages_across_page_sizes(self):
# Over-estimation gap grows with page_size; sweep to confirm fix
# doesn't depend on the page_size=8 numbers above.
for page_size in (16, 32, 64, 128):
stub = _make_self(
page_size=page_size,
full_available=page_size * 2,
swa_available=page_size * 2,
)
prefix = torch.tensor([page_size - 2] * 4, dtype=torch.int64)
seq = torch.tensor([page_size - 1] * 4, dtype=torch.int64)
result = _call(
stub, prefix_lens_cpu=prefix, seq_lens_cpu=seq, extend_num_tokens=4
)
self.assertIsNotNone(result, f"page_size={page_size}")
if __name__ == "__main__":
unittest.main()
test/manual/dsv4/test_swa_lock_release_lifecycle.py
+463
View File
@@ -0,0 +1,463 @@
"""Regression for SWA lock release lifecycle.
Hybrid-SWA early-release protocol: once a request's decode position passes
the sliding window, drop its prefill SWA lock without touching the full
lock, freeing SWA pages back to LRU.
Covers:
- SWARadixCache.dec_swa_lock_only (leaf tombstone + free, internal protected->evictable)
- SWARadixCache.dec_lock_ref(skip_swa=True)
- SWARadixCache.evict swa branch for leaf with full_lock_ref > 0
- SWARadixCache._delete_leaf skipping swa_evictable_size_ on tombstoned leaves
"""
import unittest
import torch
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.utils import get_device
from sglang.test.test_utils import CustomTestCase
def _build_tree(
*,
sliding_window_size: int = 4,
page_size: int = 1,
kv_size: int = 128,
kv_size_swa: int = 64,
):
head_num, head_dim, num_layers, global_interval = 8, 128, 24, 4
dtype = torch.bfloat16
device = get_device()
full_ids = list(range(0, num_layers, global_interval))
swa_ids = [i for i in range(num_layers) if i not in set(full_ids)]
pool = ReqToTokenPool(
size=8, max_context_len=256, device=device, enable_memory_saver=False
)
kv_pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_ids,
full_attention_layer_ids=full_ids,
enable_kvcache_transpose=False,
device=device,
)
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
need_sort=False,
)
tree = SWARadixCache(
params=CacheInitParams(
req_to_token_pool=pool,
token_to_kv_pool_allocator=allocator,
page_size=page_size,
disable=False,
is_eagle=False,
sliding_window_size=sliding_window_size,
),
)
return tree, allocator, pool
def _swa_alloc(allocator, need_size):
"""Allocate from SWA allocator for any page_size.
SWATokenToKVPoolAllocator.alloc() asserts page_size == 1; for page_size > 1
we drive the underlying paged allocators directly (mirrors the helper in
test_swa_eviction_boundary.py). Required: need_size is a multiple of
page_size when page_size > 1.
"""
if allocator.page_size == 1:
return allocator.alloc(need_size)
assert need_size % allocator.page_size == 0, (
f"page_size > 1 requires page-aligned alloc, got {need_size=} "
f"with {allocator.page_size=}"
)
if need_size > allocator.full_attn_allocator.available_size():
return None
if need_size > allocator.swa_attn_allocator.available_size():
return None
full_indices = allocator.full_attn_allocator.alloc(need_size)
swa_indices = allocator.swa_attn_allocator.alloc(need_size)
assert full_indices is not None and swa_indices is not None
allocator.full_to_swa_index_mapping[full_indices] = swa_indices
return full_indices
def _insert_chain(tree, allocator, token_ids):
indices = _swa_alloc(allocator, len(token_ids))
assert indices is not None
tree.insert(InsertParams(key=RadixKey(token_ids), value=indices))
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids)))
return match.last_device_node
def _release_swa_lock_chain_in_place(tree, leaf, swa_uuid_for_lock):
# Mirrors dec_swa_lock_only's non-tombstone arm (protected->evictable on
# internal nodes) but skips the leaf-free + tombstone step, to construct
# the post-revival state where SWA was already early-released yet the
# leaf is back in swa_lru_list with full_lock_ref still > 0.
node = leaf
while node is not tree.root_node:
if node.swa_lock_ref > 0:
if node.swa_lock_ref == 1:
tree.swa_protected_size_ -= len(node.value)
tree.swa_evictable_size_ += len(node.value)
node.swa_lock_ref -= 1
if swa_uuid_for_lock and node.swa_uuid == swa_uuid_for_lock:
break
node = node.parent
class TestSWALockReleaseLifecycle(CustomTestCase):
"""Each test pins one component of the early-release fix; method names
are prefixed with the API surface they exercise so pytest output groups
them naturally."""
def test_dec_swa_lock_only_leaf_tombstones_and_frees(self):
tree, allocator, _ = _build_tree(sliding_window_size=4)
leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8])
self.assertEqual(len(leaf.value), 8)
inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock
self.assertIsNotNone(swa_uuid)
swa_avail_before = allocator.swa_available_size()
full_avail_before = allocator.full_available_size()
self.assertEqual(leaf.swa_lock_ref, 1)
self.assertEqual(leaf.full_lock_ref, 1)
self.assertFalse(leaf.swa_tombstone)
self.assertTrue(tree.swa_lru_list.in_list(leaf))
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid)
self.assertTrue(leaf.swa_tombstone)
self.assertFalse(tree.swa_lru_list.in_list(leaf))
self.assertEqual(leaf.swa_lock_ref, 0)
self.assertEqual(
allocator.swa_available_size(), swa_avail_before + len(leaf.value)
)
self.assertEqual(leaf.full_lock_ref, 1)
self.assertEqual(allocator.full_available_size(), full_avail_before)
# sanity_check forbids live locks; release the full half before checking.
tree.dec_lock_ref(
leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True
)
tree.sanity_check()
def test_dec_swa_lock_only_internal_no_tombstone_no_free(self):
# Two siblings force an internal node at the shared prefix.
tree, allocator, _ = _build_tree(sliding_window_size=4)
leaf_a = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8])
_insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 9])
# Post-split: leaf_a now carries [8] only, parent holds the shared 7.
self.assertEqual(len(leaf_a.value), 1)
internal = leaf_a.parent
self.assertGreater(len(internal.children), 1)
self.assertEqual(len(internal.value), 7)
inc_res = tree.inc_lock_ref(leaf_a)
swa_uuid = inc_res.swa_uuid_for_lock
# window=4, value 1 (leaf) + 7 (internal): swa lock chain ends at internal.
self.assertEqual(swa_uuid, internal.swa_uuid)
swa_protected_before = tree.swa_protected_size_
swa_evictable_before = tree.swa_evictable_size_
swa_avail_before = allocator.swa_available_size()
tree.dec_swa_lock_only(leaf_a, swa_uuid_for_lock=swa_uuid)
self.assertFalse(internal.swa_tombstone)
self.assertTrue(tree.swa_lru_list.in_list(internal))
self.assertEqual(internal.swa_lock_ref, 0)
self.assertEqual(
tree.swa_protected_size_, swa_protected_before - (len(leaf_a.value) + 7)
)
self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + 7)
self.assertEqual(
allocator.swa_available_size(), swa_avail_before + len(leaf_a.value)
)
tree.dec_lock_ref(
leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True
)
tree.sanity_check()
def test_dec_lock_ref_skip_swa_true_drops_full_only(self):
tree, allocator, _ = _build_tree(sliding_window_size=4)
leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8])
inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid)
self.assertTrue(leaf.swa_tombstone)
self.assertEqual(leaf.full_lock_ref, 1)
swa_avail_after_release = allocator.swa_available_size()
swa_protected_after_release = tree.swa_protected_size_
# Without skip_swa, dec_lock_ref would assert on the swa_tombstone leaf.
tree.dec_lock_ref(
leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True
)
self.assertEqual(leaf.full_lock_ref, 0)
self.assertEqual(allocator.swa_available_size(), swa_avail_after_release)
self.assertEqual(tree.swa_protected_size_, swa_protected_after_release)
tree.sanity_check()
def test_dec_lock_ref_skip_swa_false_drops_both(self):
# Default skip_swa=False must keep legacy behavior intact.
tree, allocator, _ = _build_tree(sliding_window_size=4)
leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8])
inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock
full_avail_before = allocator.full_available_size()
swa_avail_before = allocator.swa_available_size()
tree.dec_lock_ref(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid))
self.assertEqual(leaf.full_lock_ref, 0)
self.assertEqual(leaf.swa_lock_ref, 0)
self.assertEqual(tree.full_protected_size_, 0)
self.assertEqual(tree.swa_protected_size_, 0)
# dec_lock_ref releases locks but doesn't free; eviction does.
self.assertEqual(allocator.full_available_size(), full_avail_before)
self.assertEqual(allocator.swa_available_size(), swa_avail_before)
tree.sanity_check()
def test_evict_swa_leaf_with_full_lock_tombstones_in_place(self):
# Large window so inc_lock_ref locks the entire SWA chain.
tree, allocator, _ = _build_tree(sliding_window_size=64)
leaf = _insert_chain(tree, allocator, [1, 2, 3, 4])
self.assertEqual(len(leaf.value), 4)
inc_res = tree.inc_lock_ref(leaf)
_release_swa_lock_chain_in_place(tree, leaf, inc_res.swa_uuid_for_lock)
self.assertEqual(leaf.full_lock_ref, 1)
self.assertEqual(leaf.swa_lock_ref, 0)
self.assertFalse(leaf.swa_tombstone)
self.assertTrue(tree.swa_lru_list.in_list(leaf))
swa_avail_before = allocator.swa_available_size()
swa_evictable_before = tree.swa_evictable_size_
# num_tokens=0 skips the full eviction loop; swa loop hits the new branch.
evict_res = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=4))
self.assertGreaterEqual(evict_res.swa_num_tokens_evicted, 4)
self.assertTrue(leaf.swa_tombstone)
self.assertFalse(tree.swa_lru_list.in_list(leaf))
self.assertEqual(leaf.full_lock_ref, 1)
self.assertEqual(
allocator.swa_available_size(), swa_avail_before + len(leaf.value)
)
# Full lock prevents _delete_leaf, so the node stays attached.
self.assertIs(leaf.parent.children[leaf.key.child_key(tree.page_size)], leaf)
self.assertEqual(
tree.swa_evictable_size_, swa_evictable_before - len(leaf.value)
)
tree.dec_lock_ref(
leaf,
DecLockRefParams(swa_uuid_for_lock=inc_res.swa_uuid_for_lock),
skip_swa=True,
)
tree.sanity_check()
def test_delete_leaf_skips_swa_size_on_tombstone(self):
# Tombstone removes the count once; _delete_leaf must not subtract again.
tree, allocator, _ = _build_tree(sliding_window_size=4)
leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8])
inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid)
self.assertTrue(leaf.swa_tombstone)
swa_evictable_before_delete = tree.swa_evictable_size_
tree.full_lru_list.remove_node(leaf)
tree._delete_leaf(leaf)
self.assertEqual(tree.swa_evictable_size_, swa_evictable_before_delete)
def test_dec_swa_lock_only_leaf_page_size_variants(self):
"""Single-leaf tombstone+free across all (page_size, window) regimes.
Sweep covers:
- window multiple of page_size (page_size=2, window=4)
- page_size > window (page_size=8, window=4)
- window not multiple of page (page_size=4, window=6)
With page_size > 1, _swa_alloc routes through the paged allocators;
free_swa(leaf.value) must release exactly len(leaf.value) tokens
(page-aligned) regardless of how page_size relates to the window.
"""
for page_size, window in [(2, 4), (8, 4), (4, 6)]:
with self.subTest(page_size=page_size, window=window):
tree, allocator, _ = _build_tree(
sliding_window_size=window,
page_size=page_size,
kv_size=max(128, 32 * page_size),
kv_size_swa=max(64, 16 * page_size),
)
n_tokens = max(window, 2 * page_size)
n_tokens = (n_tokens + page_size - 1) // page_size * page_size
leaf = _insert_chain(tree, allocator, list(range(1, n_tokens + 1)))
self.assertEqual(len(leaf.value), n_tokens)
self.assertEqual(len(leaf.value) % page_size, 0)
inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock
self.assertIsNotNone(
swa_uuid,
f"inc_lock_ref must reach the window with leaf.value="
f"{len(leaf.value)} >= window={window}",
)
swa_avail_before = allocator.swa_available_size()
full_avail_before = allocator.full_available_size()
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid)
self.assertTrue(leaf.swa_tombstone)
self.assertFalse(tree.swa_lru_list.in_list(leaf))
self.assertEqual(leaf.swa_lock_ref, 0)
self.assertEqual(
allocator.swa_available_size(),
swa_avail_before + len(leaf.value),
"free_swa must release the leaf's full page-aligned slot count",
)
self.assertEqual(leaf.full_lock_ref, 1)
self.assertEqual(allocator.full_available_size(), full_avail_before)
tree.dec_lock_ref(
leaf,
DecLockRefParams(swa_uuid_for_lock=swa_uuid),
skip_swa=True,
)
tree.sanity_check()
def test_dec_swa_lock_only_internal_page_size_gt_1(self):
"""Internal-node chain release with page_size > 1.
Two siblings sharing a page-aligned prefix force a radix split on a
page boundary. The swa lock chain therefore spans leaf -> internal,
and dec_swa_lock_only must:
- tombstone the leaf and free len(leaf.value) SWA tokens
- flip the internal node from protected -> evictable (no free,
no tombstone)
"""
page_size, window = 2, 6
tree, allocator, _ = _build_tree(
sliding_window_size=window, page_size=page_size
)
# Shared prefix len 4 (2 pages); divergent suffix len 2 (1 page each).
leaf_a = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6])
_insert_chain(tree, allocator, [1, 2, 3, 4, 7, 8])
self.assertEqual(len(leaf_a.value), 2)
internal = leaf_a.parent
self.assertGreater(len(internal.children), 1)
self.assertEqual(len(internal.value), 4)
inc_res = tree.inc_lock_ref(leaf_a)
swa_uuid = inc_res.swa_uuid_for_lock
# leaf_a (2) + internal (4) = 6 >= window=6, so uuid stops at internal.
self.assertEqual(swa_uuid, internal.swa_uuid)
swa_protected_before = tree.swa_protected_size_
swa_evictable_before = tree.swa_evictable_size_
swa_avail_before = allocator.swa_available_size()
tree.dec_swa_lock_only(leaf_a, swa_uuid_for_lock=swa_uuid)
# Leaf side: tombstoned and pages freed.
self.assertTrue(leaf_a.swa_tombstone)
self.assertFalse(tree.swa_lru_list.in_list(leaf_a))
self.assertEqual(
allocator.swa_available_size(),
swa_avail_before + len(leaf_a.value),
)
# Internal side: protected -> evictable, still in lru, no free.
self.assertFalse(internal.swa_tombstone)
self.assertTrue(tree.swa_lru_list.in_list(internal))
self.assertEqual(internal.swa_lock_ref, 0)
self.assertEqual(
tree.swa_protected_size_,
swa_protected_before - (len(leaf_a.value) + len(internal.value)),
)
self.assertEqual(
tree.swa_evictable_size_,
swa_evictable_before + len(internal.value),
)
tree.dec_lock_ref(
leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True
)
tree.sanity_check()
def test_full_lifecycle_inc_dec_swa_dec_lock_balances(self):
tree, allocator, _ = _build_tree(sliding_window_size=4)
leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8])
full_protected0 = tree.full_protected_size_
swa_protected0 = tree.swa_protected_size_
full_avail0 = allocator.full_available_size()
swa_avail0 = allocator.swa_available_size()
inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock
self.assertGreater(tree.full_protected_size_, full_protected0)
self.assertGreater(tree.swa_protected_size_, swa_protected0)
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid)
self.assertEqual(tree.swa_protected_size_, swa_protected0)
self.assertGreater(tree.full_protected_size_, full_protected0)
tree.dec_lock_ref(
leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True
)
self.assertEqual(tree.full_protected_size_, full_protected0)
self.assertEqual(tree.swa_protected_size_, swa_protected0)
self.assertEqual(allocator.full_available_size(), full_avail0)
self.assertEqual(allocator.swa_available_size(), swa_avail0 + len(leaf.value))
tree.sanity_check()
if __name__ == "__main__":
unittest.main()