#!/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()