262 lines
8.3 KiB
Python
Executable File
262 lines
8.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Profile DSpark inference with Nsight Systems and/or PyTorch profiler.
|
|
|
|
Usage:
|
|
python scripts/profile_dspark.py --config dspark-st3 --concurrency 64 --tool nsys
|
|
python scripts/profile_dspark.py --config nospec --concurrency 64 --tool nsys
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
ROOT = Path("/data/user1/yy")
|
|
RESULT_DIR = ROOT / "bench_results" / f"dspark_profile_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
|
RESULT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
DATASET = "/data/user1/yy/datasets/ShareGPT_V4.3_unfiltered_cleaned_split.json"
|
|
NUM_PROMPTS = 200
|
|
SEED = 42
|
|
OUTPUT_LEN = 256
|
|
HOST = "127.0.0.1"
|
|
|
|
CONFIGS = {
|
|
"dspark-st3": {
|
|
"model": "/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"port": 30004,
|
|
"spec_tokens": 3,
|
|
"cmd": [
|
|
"/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
|
|
"/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"--trust-remote-code",
|
|
"--tensor-parallel-size", "8",
|
|
"--kv-cache-dtype", "fp8",
|
|
"--block-size", "256",
|
|
"--max-model-len", "auto",
|
|
"--max-num-seqs", "256",
|
|
"--gpu-memory-utilization", "0.75",
|
|
"--tokenizer-mode", "deepseek_v4",
|
|
"--reasoning-parser", "deepseek_v4",
|
|
"--spec-method", "dspark",
|
|
"--spec-model", "/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"--spec-tokens", "3",
|
|
"--no-disable-hybrid-kv-cache-manager",
|
|
"--disable-uvicorn-access-log",
|
|
"--port", "30004",
|
|
],
|
|
},
|
|
"dspark-st5": {
|
|
"model": "/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"port": 30004,
|
|
"spec_tokens": 5,
|
|
"cmd": [
|
|
"/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
|
|
"/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"--trust-remote-code",
|
|
"--tensor-parallel-size", "8",
|
|
"--kv-cache-dtype", "fp8",
|
|
"--block-size", "256",
|
|
"--max-model-len", "auto",
|
|
"--max-num-seqs", "256",
|
|
"--gpu-memory-utilization", "0.75",
|
|
"--tokenizer-mode", "deepseek_v4",
|
|
"--reasoning-parser", "deepseek_v4",
|
|
"--spec-method", "dspark",
|
|
"--spec-model", "/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"--spec-tokens", "5",
|
|
"--no-disable-hybrid-kv-cache-manager",
|
|
"--disable-uvicorn-access-log",
|
|
"--port", "30004",
|
|
],
|
|
},
|
|
"nospec": {
|
|
"model": "/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"port": 30004,
|
|
"cmd": [
|
|
"/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
|
|
"/data/models/DeepSeek-V4-Flash-DSpark",
|
|
"--trust-remote-code",
|
|
"--tensor-parallel-size", "8",
|
|
"--kv-cache-dtype", "fp8",
|
|
"--block-size", "256",
|
|
"--max-model-len", "auto",
|
|
"--max-num-seqs", "256",
|
|
"--gpu-memory-utilization", "0.75",
|
|
"--tokenizer-mode", "deepseek_v4",
|
|
"--reasoning-parser", "deepseek_v4",
|
|
"--no-disable-hybrid-kv-cache-manager",
|
|
"--disable-uvicorn-access-log",
|
|
"--port", "30004",
|
|
],
|
|
},
|
|
}
|
|
|
|
|
|
def log(msg):
|
|
print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {msg}", flush=True)
|
|
|
|
|
|
def wait_for_health(port, timeout=300):
|
|
url = f"http://{HOST}:{port}/health"
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=5) as resp:
|
|
if resp.status == 200:
|
|
return True
|
|
except Exception:
|
|
pass
|
|
time.sleep(2)
|
|
return False
|
|
|
|
|
|
def run_warmup(port, model, n=20):
|
|
"""Run a few warmup requests to stabilize CUDA graphs."""
|
|
log(f"Running {n} warmup requests...")
|
|
for i in range(n):
|
|
try:
|
|
req = urllib.request.Request(
|
|
f"http://{HOST}:{port}/v1/completions",
|
|
data=json.dumps({
|
|
"model": model,
|
|
"prompt": "Hello, how are you?",
|
|
"max_tokens": 32,
|
|
"temperature": 0,
|
|
"seed": SEED,
|
|
}).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=60) as resp:
|
|
resp.read()
|
|
except Exception as e:
|
|
log(f"Warmup request {i} failed: {e}")
|
|
|
|
|
|
def run_benchmark(port, model, concurrency, result_dir):
|
|
result_file = result_dir / f"profile_c{concurrency}.json"
|
|
log_file = result_dir / f"profile_c{concurrency}.log"
|
|
bench_cmd = [
|
|
"/data/user1/yy/envs/vllm-dspark/bin/vllm", "bench", "serve",
|
|
"--host", HOST,
|
|
"--port", str(port),
|
|
"--backend", "openai",
|
|
"--dataset-name", "sharegpt",
|
|
"--dataset-path", DATASET,
|
|
"--sharegpt-output-len", str(OUTPUT_LEN),
|
|
"--num-prompts", str(NUM_PROMPTS),
|
|
"--max-concurrency", str(concurrency),
|
|
"--endpoint", "/v1/completions",
|
|
"--model", model,
|
|
"--seed", str(SEED),
|
|
"--save-result",
|
|
"--result-dir", str(result_dir),
|
|
"--result-filename", result_file.name,
|
|
]
|
|
log(f"Running benchmark c={concurrency}...")
|
|
start = time.time()
|
|
with open(log_file, "w") as f:
|
|
proc = subprocess.Popen(bench_cmd, stdout=f, stderr=subprocess.STDOUT)
|
|
proc.wait()
|
|
duration = time.time() - start
|
|
log(f"Benchmark c={concurrency} finished in {duration:.1f}s, exit={proc.returncode}")
|
|
return result_file
|
|
|
|
|
|
def start_service(config, tool, result_dir):
|
|
name = config
|
|
cfg = CONFIGS[config]
|
|
port = cfg["port"]
|
|
base_cmd = cfg["cmd"]
|
|
|
|
if tool == "nsys":
|
|
nsys_output = result_dir / f"{name}_nsys"
|
|
cmd = [
|
|
"nsys", "profile",
|
|
"--sample=none",
|
|
"--cpuctxsw=none",
|
|
"--trace=cuda,nvtx,osrt",
|
|
"--output", str(nsys_output),
|
|
"--force-overwrite", "true",
|
|
"--delay", "30",
|
|
"--duration", "15",
|
|
] + base_cmd
|
|
else:
|
|
cmd = base_cmd
|
|
|
|
log(f"Starting service for {name} (tool={tool}) on port {port}...")
|
|
log_file = result_dir / f"{name}_service.log"
|
|
env = os.environ.copy()
|
|
env["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
|
|
env["TMPDIR"] = str(ROOT / "tmp")
|
|
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=open(log_file, "w"),
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
)
|
|
if not wait_for_health(port):
|
|
log(f"ERROR: {name} failed to start")
|
|
proc.terminate()
|
|
return None
|
|
log(f"{name} is ready")
|
|
return proc
|
|
|
|
|
|
def stop_service(proc, name):
|
|
if proc is None:
|
|
return
|
|
log(f"Stopping {name} (pid {proc.pid})...")
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=60)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
proc.wait()
|
|
log(f"{name} stopped")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", required=True, choices=list(CONFIGS.keys()))
|
|
parser.add_argument("--concurrency", type=int, default=64)
|
|
parser.add_argument("--tool", default="nsys", choices=["nsys"])
|
|
parser.add_argument("--skip-bench", action="store_true", help="Only capture service startup/warmup, skip benchmark")
|
|
args = parser.parse_args()
|
|
|
|
result_dir = RESULT_DIR / f"{args.config}_c{args.concurrency}_{args.tool}"
|
|
result_dir.mkdir(parents=True, exist_ok=True)
|
|
log(f"Results will be saved to {result_dir}")
|
|
|
|
proc = start_service(args.config, args.tool, result_dir)
|
|
if proc is None:
|
|
sys.exit(1)
|
|
|
|
try:
|
|
run_warmup(CONFIGS[args.config]["port"], CONFIGS[args.config]["model"], n=30)
|
|
if not args.skip_bench:
|
|
run_benchmark(
|
|
CONFIGS[args.config]["port"],
|
|
CONFIGS[args.config]["model"],
|
|
args.concurrency,
|
|
result_dir,
|
|
)
|
|
else:
|
|
log("--skip-bench set; sleeping 60s to capture warmup/steady state...")
|
|
time.sleep(60)
|
|
finally:
|
|
stop_service(proc, args.config)
|
|
|
|
log(f"All done. Results in {result_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|