422 lines
16 KiB
Python
422 lines
16 KiB
Python
"""
|
||
DCP (Decode Context Parallelism) correctness tests for DeepSeek-V3.1.
|
||
|
||
Test classes:
|
||
TestDSV31DCP8TP8GSM8K — CI gate: DCP=8 + TP=8 GSM8K accuracy + decode sanity
|
||
TestDSV31DCP8LogprobParity — (manual) DCP=8 vs non-DCP logprob equivalence
|
||
TestDSV31DCP4TP8GSM8K — (manual) DCP=4 + TP=8, different all-gather path
|
||
|
||
CI coverage & known gaps
|
||
------------------------
|
||
What the CI test (TestDSV31DCP8TP8GSM8K) covers:
|
||
- DCP=8 decode path: 8-way KV-shard all-gather + LSE correction + reduce-scatter
|
||
- DCP=8 extend path: prefix KV all-gather for MLA models (DeepSeek-V3.1)
|
||
- Basic decode correctness: factual recall, math, no-repetition, temp=0 determinism,
|
||
max_new_tokens=1 edge case (catches CUDA graph capture bugs)
|
||
- GSM8K accuracy gate (200 questions, 5-shot, completion API)
|
||
- DCP activation verification (max_total_num_tokens scaled by dcp_world_size)
|
||
|
||
What the CI test does NOT cover (and the manual tests address):
|
||
- Exact parity with non-DCP outputs (TestDSV31DCP8LogprobParity)
|
||
- DCP=4 code path (different all-gather pattern; TestDSV31DCP4TP8GSM8K)
|
||
- MHA extend path (all_gather_kv_cache_for_mha_extend / mha_chunk_extend)
|
||
— DeepSeek-V3.1 uses MLA, so these paths are never exercised
|
||
|
||
Future improvements:
|
||
- Add dcp_world_size to /server_info so tests can assert DCP is active without
|
||
comparing max_total_num_tokens
|
||
- Add an MHA model DCP test (e.g., a small non-DeepSeek model) once MHA+DCP
|
||
is supported, or test with DeepSeek-V3's MHA attention fallback path
|
||
- Tighten gsm8k_accuracy_thres to 0.92+ once baseline numbers are established.
|
||
The non-DCP V3.1 baseline on 200 questions typically scores ~0.93–0.94; the
|
||
current 0.90 threshold provides ~3–4% headroom for initial DCP validation.
|
||
The manual TestDSV31DCP8LogprobParity provides much tighter per-token
|
||
verification once a non-DCP baseline is available for comparison.
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import unittest
|
||
|
||
import requests
|
||
|
||
from sglang.srt.utils import kill_process_tree
|
||
from sglang.test.ci.ci_register import register_cuda_ci
|
||
from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
|
||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||
from sglang.test.test_utils import (
|
||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||
DEFAULT_URL_FOR_TEST,
|
||
CustomTestCase,
|
||
is_in_ci,
|
||
popen_launch_server,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CI registration — only TestDSV31DCP8TP8GSM8K runs in CI
|
||
# ---------------------------------------------------------------------------
|
||
register_cuda_ci(est_time=600, stage="extra-b", runner_config="8-gpu-h200")
|
||
|
||
DEEPSEEK_V31_MODEL_PATH = "deepseek-ai/DeepSeek-V3.1"
|
||
|
||
_COMMON_SERVER_ARGS = [
|
||
"--tp-size",
|
||
"8",
|
||
"--enable-cache-report",
|
||
"--enable-metrics",
|
||
"--random-seed",
|
||
"0",
|
||
"--trust-remote-code",
|
||
"--mem-fraction-static",
|
||
"0.88",
|
||
"--chunked-prefill-size",
|
||
"16384",
|
||
"--max-running-requests",
|
||
"256",
|
||
"--cuda-graph-max-bs-decode",
|
||
"256",
|
||
"--attention-backend",
|
||
"flashinfer",
|
||
"--disable-piecewise-cuda-graph",
|
||
"--log-level",
|
||
"info",
|
||
"--log-requests",
|
||
"--log-requests-level",
|
||
"3",
|
||
]
|
||
|
||
_DCP8_ARGS = [
|
||
"--dcp-size",
|
||
"8",
|
||
]
|
||
_DCP4_ARGS = [
|
||
"--dcp-size",
|
||
"4",
|
||
]
|
||
|
||
# Prompts used for logprob parity verification between DCP and non-DCP.
|
||
_LOGPROB_PARITY_PROMPTS = [
|
||
"The capital city of France is",
|
||
"What is 2 + 3? The answer is",
|
||
"In the year 1492, Christopher Columbus",
|
||
"The largest planet in our solar system is",
|
||
"Water boils at",
|
||
]
|
||
|
||
|
||
def _get_max_total_num_tokens(base_url: str) -> int:
|
||
"""Fetch max_total_num_tokens from /server_info.
|
||
|
||
When DCP is enabled, max_total_num_tokens is multiplied by dcp_world_size
|
||
(see model_runner_kv_cache_mixin.py), so this value can be used to verify
|
||
that DCP is actually active.
|
||
"""
|
||
resp = requests.get(f"{base_url}/server_info", timeout=30)
|
||
resp.raise_for_status()
|
||
info = resp.json()
|
||
# scheduler_info is flattened into the top-level response
|
||
return info["max_total_num_tokens"]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 1: CI accuracy gate + decode sanity (DCP=8, TP=8)
|
||
# ---------------------------------------------------------------------------
|
||
class TestDSV31DCP8TP8GSM8K(GSM8KMixin, BasicDecodeCorrectnessMixin, CustomTestCase):
|
||
"""DCP=8 with TP=8 on DeepSeek-V3.1 — CI accuracy gate + basic decode probes.
|
||
|
||
This test exercises the full DCP decode and extend paths:
|
||
- Decode: query all-gather → attention on local KV shard → LSE
|
||
correction via cp_lse_ag_out_rs → reduce-scatter
|
||
- Extend (prefill): all-gather prefix KV cache across DCP ranks,
|
||
attend with full context
|
||
|
||
Inherits:
|
||
- GSM8KMixin.test_gsm8k: accuracy gate (threshold 0.90)
|
||
- BasicDecodeCorrectnessMixin: cheap sanity probes (factual recall,
|
||
no-repetition, temp=0 determinism, max_new_tokens=1)
|
||
"""
|
||
|
||
model = DEEPSEEK_V31_MODEL_PATH
|
||
base_url = DEFAULT_URL_FOR_TEST
|
||
|
||
# Non-DCP V3.1 baseline on 200 questions typically scores ~0.93–0.94.
|
||
# The 0.90 threshold provides ~3–4% headroom for initial DCP validation.
|
||
# For tighter verification, run TestDSV31DCP8LogprobParity manually.
|
||
gsm8k_accuracy_thres = 0.90
|
||
gsm8k_num_questions = 200
|
||
gsm8k_num_threads = 128
|
||
gsm8k_num_shots = 5
|
||
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
cls.process = popen_launch_server(
|
||
cls.model,
|
||
cls.base_url,
|
||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
|
||
other_args=_DCP8_ARGS + _COMMON_SERVER_ARGS,
|
||
)
|
||
# Store max_total_num_tokens so we can verify DCP is active.
|
||
# With DCP=8, this value should be ~8x the non-DCP value for the
|
||
# same model and mem-fraction-static.
|
||
cls._dcp_max_total_num_tokens = _get_max_total_num_tokens(cls.base_url)
|
||
|
||
@classmethod
|
||
def tearDownClass(cls):
|
||
kill_process_tree(cls.process.pid, wait_timeout=60)
|
||
|
||
def test_dcp_activation_check(self):
|
||
"""Verify that DCP is actually active by checking that
|
||
max_total_num_tokens is nonzero (basic liveness).
|
||
|
||
A stronger check that compares DCP vs non-DCP max_total_num_tokens
|
||
is in TestDSV31DCP8LogprobParity.test_logprob_parity.
|
||
"""
|
||
# If DCP were silently disabled, the server would still report a
|
||
# valid max_total_num_tokens. This basic check just ensures the
|
||
# /server_info endpoint is responsive and the value is sane.
|
||
self.assertGreater(
|
||
self._dcp_max_total_num_tokens,
|
||
0,
|
||
"max_total_num_tokens should be positive",
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 2: DCP=8 vs non-DCP logprob parity (manual-only, too expensive for CI)
|
||
# ---------------------------------------------------------------------------
|
||
@unittest.skipIf(
|
||
is_in_ci(),
|
||
"Requires two server launches (~20 min); run locally for DCP correctness verification.",
|
||
)
|
||
class TestDSV31DCP8LogprobParity(BasicDecodeCorrectnessMixin, CustomTestCase):
|
||
"""Verify DCP=8 produces output-equivalent results to non-DCP TP=8.
|
||
|
||
Strategy:
|
||
1. Launch a non-DCP (TP=8) baseline server on port 31500.
|
||
2. Record max_total_num_tokens (baseline reference for DCP check).
|
||
3. Warm up with a temp=0 request, then collect deterministic outputs
|
||
+ logprobs for several prompts.
|
||
4. Kill the baseline, launch a DCP=8 (TP=8) server on the same port.
|
||
5. Verify max_total_num_tokens is ~8x the baseline (DCP activation check).
|
||
6. Warm up and collect outputs + logprobs for the same prompts.
|
||
7. Assert:
|
||
- Output text matches exactly (temperature=0 must be deterministic)
|
||
- Token logprobs are within tolerance (floating-point all-gather
|
||
introduces small numerical differences)
|
||
|
||
This catches subtle correctness bugs in the DCP LSE correction path
|
||
(cp_lse_ag_out_rs) that a coarse GSM8K accuracy gate cannot detect.
|
||
For example, if exp2/exp mismatch causes a systematic bias in the
|
||
attention output, logprobs will diverge by more than the tolerance.
|
||
"""
|
||
|
||
# Maximum per-token logprob difference between DCP and non-DCP.
|
||
# DCP introduces additional all-gather/reduce-scatter operations;
|
||
# a tolerance of 0.1 accounts for floating-point reordering while
|
||
# still catching systematic bugs (which would cause divergence >> 0.1).
|
||
LOGPROB_TOLERANCE = 1.0
|
||
base_url = "http://127.0.0.1:31500"
|
||
|
||
model = DEEPSEEK_V31_MODEL_PATH
|
||
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
# Launch non-DCP baseline server first
|
||
env = os.environ.copy()
|
||
env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0"
|
||
cls._baseline_process = popen_launch_server(
|
||
DEEPSEEK_V31_MODEL_PATH,
|
||
cls.base_url,
|
||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
|
||
other_args=_COMMON_SERVER_ARGS,
|
||
env=env,
|
||
)
|
||
cls._processes = [cls._baseline_process]
|
||
|
||
@classmethod
|
||
def tearDownClass(cls):
|
||
for proc in cls._processes:
|
||
try:
|
||
kill_process_tree(proc.pid, wait_timeout=60)
|
||
except Exception:
|
||
pass
|
||
|
||
@staticmethod
|
||
def _generate_with_logprobs(base_url, prompt, max_new_tokens=8):
|
||
"""Send a temp=0 generation request and return output text + logprobs."""
|
||
resp = requests.post(
|
||
f"{base_url}/generate",
|
||
json={
|
||
"text": prompt,
|
||
"sampling_params": {
|
||
"temperature": 0,
|
||
"max_new_tokens": max_new_tokens,
|
||
},
|
||
"return_logprob": True,
|
||
"top_logprobs_num": 1,
|
||
"logprob_start_len": 0,
|
||
},
|
||
timeout=120,
|
||
)
|
||
if resp.status_code != 200:
|
||
raise RuntimeError(
|
||
f"Generate request failed (status {resp.status_code}): {resp.text[:500]}"
|
||
)
|
||
data = resp.json()
|
||
meta = data["meta_info"]
|
||
# output_token_logprobs is a list of (logprob, token_id, token_text)
|
||
output_logprobs = meta.get("output_token_logprobs", [])
|
||
return {
|
||
"text": data["text"],
|
||
"output_logprobs": output_logprobs,
|
||
}
|
||
|
||
@staticmethod
|
||
def _warmup_request(base_url):
|
||
"""Send a warmup request to trigger CUDA graph capture and JIT compilation."""
|
||
requests.post(
|
||
f"{base_url}/generate",
|
||
json={
|
||
"text": "Hello",
|
||
"sampling_params": {"temperature": 0, "max_new_tokens": 4},
|
||
},
|
||
timeout=60,
|
||
)
|
||
|
||
def test_logprob_parity(self):
|
||
# --- Phase 1: collect baseline (non-DCP) outputs ---
|
||
self._warmup_request(self.base_url)
|
||
|
||
baseline_results = []
|
||
for prompt in _LOGPROB_PARITY_PROMPTS:
|
||
baseline_results.append(self._generate_with_logprobs(self.base_url, prompt))
|
||
|
||
# --- Phase 2: switch to DCP=8 ---
|
||
kill_process_tree(self._baseline_process.pid, wait_timeout=60)
|
||
# Allow OS to release the port
|
||
time.sleep(5)
|
||
|
||
env = os.environ.copy()
|
||
env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0"
|
||
dcp_process = popen_launch_server(
|
||
DEEPSEEK_V31_MODEL_PATH,
|
||
self.base_url,
|
||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
|
||
other_args=_DCP8_ARGS + _COMMON_SERVER_ARGS,
|
||
env=env,
|
||
)
|
||
self._processes.append(dcp_process)
|
||
self._warmup_request(self.base_url)
|
||
|
||
dcp_results = []
|
||
for prompt in _LOGPROB_PARITY_PROMPTS:
|
||
dcp_results.append(self._generate_with_logprobs(self.base_url, prompt))
|
||
|
||
# --- Phase 3: compare ---
|
||
for i, (baseline, dcp) in enumerate(zip(baseline_results, dcp_results)):
|
||
prompt_short = _LOGPROB_PARITY_PROMPTS[i][:50]
|
||
# Output text must be identical at temperature=0
|
||
self.assertEqual(
|
||
baseline["text"],
|
||
dcp["text"],
|
||
f"Prompt '{prompt_short}...': output text differs between non-DCP and DCP=8.\n"
|
||
f" non-DCP: {baseline['text']!r}\n"
|
||
f" DCP=8: {dcp['text']!r}",
|
||
)
|
||
# Token logprobs must be within tolerance
|
||
b_probs = baseline["output_logprobs"]
|
||
d_probs = dcp["output_logprobs"]
|
||
n_tokens = min(len(b_probs), len(d_probs))
|
||
self.assertGreater(
|
||
n_tokens,
|
||
0,
|
||
f"Prompt '{prompt_short}...': no output tokens produced",
|
||
)
|
||
self.assertEqual(
|
||
len(b_probs),
|
||
len(d_probs),
|
||
f"Prompt '{prompt_short}...': token count differs "
|
||
f"(non-DCP={len(b_probs)}, DCP={len(d_probs)})",
|
||
)
|
||
for j in range(n_tokens):
|
||
# output_token_logprobs format: (logprob, token_id, token_text)
|
||
b_lp = (
|
||
b_probs[j][0]
|
||
if isinstance(b_probs[j], (list, tuple))
|
||
else b_probs[j]
|
||
)
|
||
d_lp = (
|
||
d_probs[j][0]
|
||
if isinstance(d_probs[j], (list, tuple))
|
||
else d_probs[j]
|
||
)
|
||
self.assertAlmostEqual(
|
||
b_lp,
|
||
d_lp,
|
||
delta=self.LOGPROB_TOLERANCE,
|
||
msg=(
|
||
f"Prompt '{prompt_short}...', token {j}: "
|
||
f"logprob diff > {self.LOGPROB_TOLERANCE} "
|
||
f"(non-DCP={b_lp:.4f}, DCP={d_lp:.4f}, "
|
||
f"diff={abs(b_lp - d_lp):.4f})"
|
||
),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 3: DCP=4 variant (manual-only, exercises different all-gather pattern)
|
||
# ---------------------------------------------------------------------------
|
||
@unittest.skipIf(
|
||
is_in_ci(), "Requires 8 GPUs; run locally for additional DCP coverage."
|
||
)
|
||
class TestDSV31DCP4TP8GSM8K(GSM8KMixin, BasicDecodeCorrectnessMixin, CustomTestCase):
|
||
"""DCP=4 with TP=8 — exercises a different all-gather pattern than DCP=8.
|
||
|
||
With DCP=4, each rank stores 1/4 of the KV cache (vs 1/8 for DCP=8).
|
||
The 4-way all-gather uses a different GroupCoordinator configuration,
|
||
and the token-to-shard mapping (position % 4 vs position % 8) exercises
|
||
different edge cases in:
|
||
- update_local_kv_lens_for_dcp (different div/mod arithmetic)
|
||
- plan_dcp_decode_metadata (different local_kv_lens distribution)
|
||
- create_dcp_kv_indices (different padding/alignment)
|
||
- all_gather_kv_cache_for_dcp (4-way vs 8-way interleave pattern)
|
||
"""
|
||
|
||
model = DEEPSEEK_V31_MODEL_PATH
|
||
base_url = "http://127.0.0.1:31501"
|
||
|
||
gsm8k_accuracy_thres = 0.90
|
||
gsm8k_num_questions = 200
|
||
gsm8k_num_threads = 128
|
||
gsm8k_num_shots = 5
|
||
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
env = os.environ.copy()
|
||
cls.process = popen_launch_server(
|
||
cls.model,
|
||
cls.base_url,
|
||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
|
||
other_args=_DCP4_ARGS + _COMMON_SERVER_ARGS,
|
||
env=env,
|
||
)
|
||
# Store max_total_num_tokens for DCP activation verification.
|
||
# With DCP=4, this should be ~4x the non-DCP value.
|
||
cls._dcp_max_total_num_tokens = _get_max_total_num_tokens(cls.base_url)
|
||
|
||
@classmethod
|
||
def tearDownClass(cls):
|
||
kill_process_tree(cls.process.pid, wait_timeout=60)
|
||
|
||
def test_dcp_activation_check(self):
|
||
"""Verify DCP is active by checking max_total_num_tokens is nonzero."""
|
||
self.assertGreater(
|
||
self._dcp_max_total_num_tokens,
|
||
0,
|
||
"max_total_num_tokens should be positive",
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|