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

[SKILL] Upgrade sglang profile and auto_benchmark skills (#24250)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-05-02 10:12:47 +08:00
committed by GitHub
parent 4c2ed9a254
commit 321298da75
105 changed files with 10996 additions and 3813 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py
+56 -4
View File
@@ -11,6 +11,12 @@ from typing import Dict, List, Optional, Sequence, Tuple
import triage_kernel_helpers as kernel_helpers
import triage_overlap_helpers as overlap_helpers
from profile_common import (
DEFAULT_DECODE_INPUT_LEN,
DEFAULT_DECODE_OUTPUT_LEN,
DEFAULT_PREFILL_INPUT_LEN,
DEFAULT_PREFILL_OUTPUT_LEN,
DEFAULT_WARMUP_STEPS,
PROFILE_WORKLOAD_CHOICES,
discover_trace_targets,
framework_display_name,
load_server_args,
@@ -57,8 +63,8 @@ def build_triage_parser() -> argparse.ArgumentParser:
default=None,
help=(
"Running server URL for single-trace triage. SGLang supports direct "
"capture through its profiler HTTP API. vLLM and TensorRT-LLM require "
"a server-side torch-profiler output path exposed via --output-dir."
"capture via sglang.profiler. vLLM and TensorRT-LLM require a server-side "
"torch-profiler output path exposed via --output-dir."
),
)
parser.add_argument(
@@ -132,7 +138,13 @@ def build_triage_parser() -> argparse.ArgumentParser:
"--num-steps",
type=int,
default=5,
help="Profiler steps when generating traces from URLs.",
help="Active profiler steps when generating traces from URLs.",
)
parser.add_argument(
"--warmup-steps",
type=int,
default=DEFAULT_WARMUP_STEPS,
help="Warmup steps to run before arming the profiler for URL capture.",
)
parser.add_argument(
"--profile-by-stage", action=argparse.BooleanOptionalAction, default=True
@@ -151,11 +163,45 @@ def build_triage_parser() -> argparse.ArgumentParser:
)
parser.add_argument("--probe-max-new-tokens", type=int, default=None)
parser.add_argument("--probe-delay", type=float, default=0.5)
parser.add_argument(
"--profile-workload",
choices=PROFILE_WORKLOAD_CHOICES,
default="both",
help=(
"Live-capture workload shape. Default 'both' captures separate "
"prefill and decode profiles instead of one mixed request. Use "
"'legacy' to keep the old --probe-prompt behavior."
),
)
parser.add_argument(
"--prefill-input-len",
type=int,
default=DEFAULT_PREFILL_INPUT_LEN,
help="Synthetic input length for the prefill profile workload.",
)
parser.add_argument(
"--prefill-output-len",
type=int,
default=DEFAULT_PREFILL_OUTPUT_LEN,
help="Output length for the prefill profile workload.",
)
parser.add_argument(
"--decode-input-len",
type=int,
default=DEFAULT_DECODE_INPUT_LEN,
help="Synthetic input length for the decode profile workload.",
)
parser.add_argument(
"--decode-output-len",
type=int,
default=DEFAULT_DECODE_OUTPUT_LEN,
help="Output length for the decode profile workload.",
)
parser.add_argument(
"--start-step",
type=int,
default=None,
help="SGLang-only profiler start step when generating traces from URLs.",
help="Pass through to sglang.profiler when generating traces from URLs.",
)
parser.add_argument(
"--pid-substring",
@@ -239,9 +285,15 @@ def resolve_profile_targets(
probe_prompt=args.probe_prompt,
probe_max_new_tokens=args.probe_max_new_tokens,
probe_delay=args.probe_delay,
warmup_steps=args.warmup_steps,
start_step=args.start_step,
framework=framework,
framework_hint_path=output_dir,
profile_workload=args.profile_workload,
prefill_input_len=args.prefill_input_len,
prefill_output_len=args.prefill_output_len,
decode_input_len=args.decode_input_len,
decode_output_len=args.decode_output_len,
)
traces, server_args = discover_trace_targets(target_dir, all_traces=False)
resolved_framework = resolve_framework(
.claude/skills/llm-torch-profiler-analysis/scripts/probe_llm_server.py
+1 -1
View File
@@ -15,7 +15,7 @@ from urllib import request
from profile_common import extract_openai_chat_text
DEFAULT_PROMPTS = [
"用一句中文介绍上海。",
"Introduce Shanghai in one short sentence.",
"What is 2+2? Answer briefly.",
"Write one short haiku about GPUs.",
]
.claude/skills/llm-torch-profiler-analysis/scripts/profile_common.py
+315 -50
View File
@@ -5,10 +5,12 @@ from __future__ import annotations
import gzip
import json
import re
import shutil
import sys
import tempfile
import time
from collections import Counter, defaultdict
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
@@ -42,6 +44,21 @@ TRACE_METADATA_NAMES = {
}
NON_KERNEL_TRACE_CATEGORIES = ("python_function", "cpu_op", "trace")
PYTHON_SCOPE_NAME_PREFIXES = ("python/", "nn.module:")
PROFILE_WORKLOAD_CHOICES = ("legacy", "prefill", "decode", "both")
DEFAULT_PREFILL_INPUT_LEN = 4090
DEFAULT_PREFILL_OUTPUT_LEN = 1
DEFAULT_DECODE_INPUT_LEN = 1
DEFAULT_DECODE_OUTPUT_LEN = 2048
DEFAULT_WARMUP_STEPS = 10
@dataclass(frozen=True)
class ProbePlan:
prompt: str
capture_max_new_tokens: int
capture_requests: int
warmup_max_new_tokens: int
warmup_requests: int
@lru_cache(maxsize=65536)
@@ -416,10 +433,16 @@ def resolve_framework(
def parse_stage(path: Path) -> str:
name = path.name.lower()
if "-extend" in name or "-prefill" in name:
parts = [part.lower() for part in path.parts[-6:]]
name = " ".join(parts)
segment_path = "/" + "/".join(parts) + "/"
if any(marker in name for marker in ("-extend", "-prefill", "_extend", "_prefill")):
return "extend"
if "-decode" in name:
if any(f"/{segment}/" in segment_path for segment in ("extend", "prefill")):
return "extend"
if any(marker in name for marker in ("-decode", "_decode")):
return "decode"
if "/decode/" in segment_path:
return "decode"
return "all"
@@ -514,8 +537,19 @@ def discover_trace_targets(
if path.is_file():
return [path], load_server_args(path)
trace_dir = newest_trace_dir(path)
traces = discover_trace_files(trace_dir, recursive=False)
direct_traces = discover_trace_files(path, recursive=False)
recursive_traces = discover_trace_files(path, recursive=True)
recursive_stages = {parse_stage(trace) for trace in recursive_traces}
if (
not direct_traces
and recursive_traces
and any(stage != "all" for stage in recursive_stages)
):
traces = recursive_traces
trace_dir = path
else:
trace_dir = newest_trace_dir(path)
traces = discover_trace_files(trace_dir, recursive=False)
if not traces:
raise FileNotFoundError(f"No trace files found under {trace_dir}")
@@ -606,6 +640,109 @@ def send_probe_request(
)
def unique_probe_prompt(prompt: str, probe_index: int) -> str:
marker = f"profile_probe_{max(0, int(probe_index))}"
parts = prompt.split(maxsplit=1)
suffix = parts[1] if len(parts) == 2 else prompt
return f"{marker} {suffix}".strip()
def send_probe_requests(
*,
url: str,
prompt: str,
max_new_tokens: int,
request_count: int,
framework: str,
model: Optional[str] = None,
sampling_seed_offset: int = 0,
) -> None:
request_count = max(0, int(request_count))
seed_offset = max(0, int(sampling_seed_offset))
for request_idx in range(request_count):
probe_index = seed_offset + request_idx
send_probe_request(
url=url,
prompt=unique_probe_prompt(prompt, probe_index),
max_new_tokens=max_new_tokens,
sampling_seed=probe_index,
framework=framework,
model=model,
)
def synthetic_prompt(input_len: int) -> str:
token_count = max(1, int(input_len))
return " ".join(["profile"] * token_count)
def workload_probe(
stage: str,
*,
prefill_input_len: int,
prefill_output_len: int,
decode_input_len: int,
decode_output_len: int,
) -> Tuple[str, int]:
if stage == "prefill":
return synthetic_prompt(prefill_input_len), max(1, int(prefill_output_len))
if stage == "decode":
return synthetic_prompt(decode_input_len), max(1, int(decode_output_len))
raise ValueError(f"unknown profile workload stage: {stage}")
def build_probe_plan(
stage: str,
*,
prompt: str,
max_new_tokens: int,
num_steps: int,
probe_requests: int,
warmup_steps: int,
) -> ProbePlan:
active_steps = max(1, int(num_steps))
requested_probes = max(1, int(probe_requests))
warmup_steps = max(0, int(warmup_steps))
max_new_tokens = max(1, int(max_new_tokens))
if stage == "prefill":
return ProbePlan(
prompt=prompt,
capture_max_new_tokens=max_new_tokens,
capture_requests=max(requested_probes, active_steps),
warmup_max_new_tokens=max_new_tokens,
warmup_requests=warmup_steps,
)
if stage == "decode":
return ProbePlan(
prompt=prompt,
capture_max_new_tokens=max_new_tokens,
capture_requests=requested_probes,
warmup_max_new_tokens=max(1, warmup_steps),
warmup_requests=1 if warmup_steps else 0,
)
return ProbePlan(
prompt=prompt,
capture_max_new_tokens=max_new_tokens,
capture_requests=requested_probes,
warmup_max_new_tokens=max_new_tokens,
warmup_requests=warmup_steps,
)
def expand_profile_workload(profile_workload: str) -> List[str]:
workload = normalize_text(profile_workload).lower()
if workload not in PROFILE_WORKLOAD_CHOICES:
raise ValueError(
f"--profile-workload must be one of {', '.join(PROFILE_WORKLOAD_CHOICES)}"
)
if workload == "both":
return ["prefill", "decode"]
if workload == "legacy":
return ["legacy"]
return [workload]
def discover_openai_model(url: str) -> str:
payload = try_get_json(url.rstrip("/") + "/v1/models", timeout=60.0)
if not isinstance(payload, dict):
@@ -690,35 +827,50 @@ def run_remote_profiler(
url: str,
output_dir: Optional[str],
framework: str,
probe_requests: int,
probe_prompt: str,
probe_max_new_tokens: Optional[int],
probe_plan: ProbePlan,
probe_delay: float,
num_steps: int,
stage: Optional[str] = None,
) -> Path:
framework = canonicalize_framework(framework)
output_path = ensure_remote_profiler_output_path(output_dir, framework)
if stage and output_path.is_file():
raise ValueError(
"--profile-workload both requires a directory output path for "
f"{framework_display_name(framework)} so each stage trace can be labeled."
)
before_traces = (
set(discover_trace_files(output_path, recursive=True))
if output_path.exists()
else set()
)
model = discover_openai_model(url) if framework in {"vllm", "trtllm"} else None
if probe_plan.warmup_requests > 0:
send_probe_requests(
url=url,
prompt=probe_plan.prompt,
max_new_tokens=probe_plan.warmup_max_new_tokens,
request_count=probe_plan.warmup_requests,
framework=framework,
model=model,
)
start_remote_profiler(url, framework)
stop_error: Optional[BaseException] = None
try:
if probe_requests > 0:
# Some profiler endpoints need a brief setup window after
if probe_plan.capture_requests > 0:
# `sglang.profiler` performs its own startup work before it reaches
# POST /start_profile. A very short delay can send probes too early
# and miss the profiling window entirely.
time.sleep(max(5.0, probe_delay))
effective_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8)
model = (
discover_openai_model(url) if framework in {"vllm", "trtllm"} else None
send_probe_requests(
url=url,
prompt=probe_plan.prompt,
max_new_tokens=probe_plan.capture_max_new_tokens,
request_count=probe_plan.capture_requests,
framework=framework,
model=model,
sampling_seed_offset=probe_plan.warmup_requests,
)
for request_idx in range(probe_requests):
send_probe_request(
url=url,
prompt=probe_prompt,
max_new_tokens=effective_max_new_tokens,
sampling_seed=request_idx,
framework=framework,
model=model,
)
finally:
try:
stop_remote_profiler(url, framework)
@@ -726,7 +878,22 @@ def run_remote_profiler(
stop_error = exc
if stop_error is not None:
raise stop_error
return wait_for_profiler_artifact(output_path)
artifact = wait_for_profiler_artifact(output_path)
if stage and output_path.is_dir():
after_traces = set(discover_trace_files(output_path, recursive=True))
new_traces = sorted(after_traces - before_traces, key=lambda item: item.name)
if new_traces:
stage_dir = output_path / stage
stage_dir.mkdir(parents=True, exist_ok=True)
for trace in new_traces:
if stage_dir in trace.parents:
continue
target = stage_dir / trace.name
if target.exists():
target = stage_dir / f"{time.time_ns()}-{trace.name}"
shutil.move(str(trace), str(target))
return stage_dir
return artifact
def run_sglang_profiler(
@@ -736,9 +903,7 @@ def run_sglang_profiler(
profile_by_stage: bool,
merge_profiles: bool,
profile_prefix: Optional[str],
probe_requests: int,
probe_prompt: str,
probe_max_new_tokens: Optional[int],
probe_plan: ProbePlan,
probe_delay: float,
start_step: Optional[int] = None,
) -> Path:
@@ -765,6 +930,15 @@ def run_sglang_profiler(
if start_step is not None:
payload["start_step"] = str(start_step)
if probe_plan.warmup_requests > 0:
send_probe_requests(
url=url,
prompt=probe_plan.prompt,
max_new_tokens=probe_plan.warmup_max_new_tokens,
request_count=probe_plan.warmup_requests,
framework="sglang",
)
req = request.Request(
url.rstrip("/") + "/start_profile",
data=json.dumps(payload).encode("utf-8"),
@@ -773,17 +947,20 @@ def run_sglang_profiler(
with request.urlopen(req, timeout=300.0):
pass
if probe_requests > 0:
if probe_plan.capture_requests > 0:
time.sleep(max(0.0, probe_delay))
effective_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8)
for request_idx in range(probe_requests):
send_probe_request(
url=url,
prompt=probe_prompt,
max_new_tokens=effective_max_new_tokens,
sampling_seed=request_idx,
framework="sglang",
)
send_probe_requests(
url=url,
prompt=probe_plan.prompt,
max_new_tokens=probe_plan.capture_max_new_tokens,
request_count=probe_plan.capture_requests,
framework="sglang",
sampling_seed_offset=probe_plan.warmup_requests,
)
try:
stop_remote_profiler(url, "sglang")
except RuntimeError:
pass
return wait_for_profiler_artifact(output_path, timeout_s=180.0)
@@ -799,9 +976,15 @@ def run_profiler(
probe_prompt: str,
probe_max_new_tokens: Optional[int],
probe_delay: float,
warmup_steps: int = DEFAULT_WARMUP_STEPS,
start_step: Optional[int] = None,
framework: str = "auto",
framework_hint_path: Optional[str] = None,
profile_workload: str = "both",
prefill_input_len: int = DEFAULT_PREFILL_INPUT_LEN,
prefill_output_len: int = DEFAULT_PREFILL_OUTPUT_LEN,
decode_input_len: int = DEFAULT_DECODE_INPUT_LEN,
decode_output_len: int = DEFAULT_DECODE_OUTPUT_LEN,
) -> Path:
resolved_framework = resolve_framework(
framework,
@@ -813,6 +996,58 @@ def run_profiler(
),
)
if resolved_framework == "sglang":
stages = expand_profile_workload(profile_workload)
if stages != ["legacy"]:
output_root = (
Path(output_dir).expanduser().resolve()
if output_dir
else Path(tempfile.mkdtemp(prefix="sglang-torch-profile-"))
)
output_root.mkdir(parents=True, exist_ok=True)
for stage in stages:
prompt, max_new_tokens = workload_probe(
stage,
prefill_input_len=prefill_input_len,
prefill_output_len=prefill_output_len,
decode_input_len=decode_input_len,
decode_output_len=decode_output_len,
)
probe_plan = build_probe_plan(
stage,
prompt=prompt,
max_new_tokens=max_new_tokens,
num_steps=num_steps,
probe_requests=probe_requests,
warmup_steps=warmup_steps,
)
# SGLang increments `forward_ct` before checking whether the
# profiler reached its target. Ask for one extra step so the
# requested stage forward is captured instead of stopping just
# before it runs.
stage_num_steps = max(1, int(num_steps)) + 1
run_sglang_profiler(
url=url,
output_dir=str(output_root / stage),
num_steps=stage_num_steps,
profile_by_stage=False,
merge_profiles=merge_profiles,
profile_prefix=(
f"{profile_prefix}-{stage}" if profile_prefix else stage
),
probe_plan=probe_plan,
probe_delay=probe_delay,
start_step=start_step,
)
return output_root
legacy_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8)
legacy_plan = build_probe_plan(
"legacy",
prompt=probe_prompt,
max_new_tokens=legacy_max_new_tokens,
num_steps=num_steps,
probe_requests=probe_requests,
warmup_steps=warmup_steps,
)
return run_sglang_profiler(
url=url,
output_dir=output_dir,
@@ -820,9 +1055,7 @@ def run_profiler(
profile_by_stage=profile_by_stage,
merge_profiles=merge_profiles,
profile_prefix=profile_prefix,
probe_requests=probe_requests,
probe_prompt=probe_prompt,
probe_max_new_tokens=probe_max_new_tokens,
probe_plan=legacy_plan,
probe_delay=probe_delay,
start_step=start_step,
)
@@ -844,16 +1077,48 @@ def run_profiler(
"--profile-prefix on the HTTP profiler control path.",
file=sys.stderr,
)
return run_remote_profiler(
url=url,
output_dir=output_dir,
framework=resolved_framework,
probe_requests=probe_requests,
probe_prompt=probe_prompt,
probe_max_new_tokens=probe_max_new_tokens,
probe_delay=probe_delay,
num_steps=num_steps,
)
stages = expand_profile_workload(profile_workload)
if stages == ["legacy"]:
legacy_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8)
return run_remote_profiler(
url=url,
output_dir=output_dir,
framework=resolved_framework,
probe_plan=build_probe_plan(
"legacy",
prompt=probe_prompt,
max_new_tokens=legacy_max_new_tokens,
num_steps=num_steps,
probe_requests=probe_requests,
warmup_steps=warmup_steps,
),
probe_delay=probe_delay,
)
output_root = ensure_remote_profiler_output_path(output_dir, resolved_framework)
for stage in stages:
prompt, max_new_tokens = workload_probe(
stage,
prefill_input_len=prefill_input_len,
prefill_output_len=prefill_output_len,
decode_input_len=decode_input_len,
decode_output_len=decode_output_len,
)
run_remote_profiler(
url=url,
output_dir=str(output_root),
framework=resolved_framework,
probe_plan=build_probe_plan(
stage,
prompt=prompt,
max_new_tokens=max_new_tokens,
num_steps=num_steps,
probe_requests=probe_requests,
warmup_steps=warmup_steps,
),
probe_delay=probe_delay,
stage=stage,
)
return output_root
def select_heaviest_pid(
.claude/skills/llm-torch-profiler-analysis/scripts/run_llm_single_model_matrix_host.sh
Executable
+274
View File
@@ -0,0 +1,274 @@
#!/usr/bin/env bash
set -euo pipefail
usage() {
cat <<'EOF'
Usage:
run_llm_single_model_matrix_host.sh \
--model-id gpt_oss_20b \
--model openai/gpt-oss-20b \
--root /data/bbuf/validate/unified_llm_profiler_skill/runs/20260423_h100_large_model_matrix \
--gpus 2,3,4,5 \
--sglang-port 30098 \
--vllm-formal-port 31098 \
--vllm-mapping-port 31099 \
--trt-formal-prefill-port 32098 \
--trt-formal-decode-port 32099 \
--trt-mapping-prefill-port 32198 \
--trt-mapping-decode-port 32199
This script is intended to run on the H100 host. It:
1. captures SGLang live profiling and writes `analysis_sglang.txt`
2. captures vLLM formal + eager mapping traces and writes `analysis_vllm.txt`
3. captures TensorRT-LLM formal + graph-off mapping traces and writes `analysis_trtllm.txt`
4. stores one benchmark JSON per framework under the model run directory
Default profiler workloads are stage-separated:
prefill: input 4090, output 1
decode: input 1, output 2048
Environment:
Export `HF_TOKEN` and `HUGGINGFACE_HUB_TOKEN` before running.
EOF
}
MODEL_ID=""
MODEL=""
ROOT=""
GPUS=""
TP_SIZE=""
SGLANG_PORT=""
VLLM_FORMAL_PORT=""
VLLM_MAPPING_PORT=""
TRT_FORMAL_PREFILL_PORT=""
TRT_FORMAL_DECODE_PORT=""
TRT_MAPPING_PREFILL_PORT=""
TRT_MAPPING_DECODE_PORT=""
SGLANG_MEM_FRACTION="0.85"
MAX_MODEL_LEN="4096"
KV_FRACTION="0.85"
SGLANG_SERVER_EXTRA=""
PROFILE_WORKLOAD="both"
PREFILL_INPUT_LEN=4090
PREFILL_OUTPUT_LEN=1
DECODE_INPUT_LEN=1
DECODE_OUTPUT_LEN=2048
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TRT_IMAGE="nvcr.io/nvidia/tensorrt-llm/release:latest"
TRT_OVERRIDE_ROOT="/data/bbuf/validate/unified_llm_profiler_skill/overrides/trtllm"
TRT_OVERRIDE_SOURCE="$TRT_OVERRIDE_ROOT/py_executor.original.py"
TRT_OVERRIDE_PATH="$TRT_OVERRIDE_ROOT/py_executor_with_stack.py"
while [[ $# -gt 0 ]]; do
case "$1" in
--model-id) MODEL_ID="$2"; shift 2 ;;
--model) MODEL="$2"; shift 2 ;;
--root) ROOT="$2"; shift 2 ;;
--gpus) GPUS="$2"; shift 2 ;;
--tp-size) TP_SIZE="$2"; shift 2 ;;
--sglang-port) SGLANG_PORT="$2"; shift 2 ;;
--vllm-formal-port) VLLM_FORMAL_PORT="$2"; shift 2 ;;
--vllm-mapping-port) VLLM_MAPPING_PORT="$2"; shift 2 ;;
--trt-formal-prefill-port) TRT_FORMAL_PREFILL_PORT="$2"; shift 2 ;;
--trt-formal-decode-port) TRT_FORMAL_DECODE_PORT="$2"; shift 2 ;;
--trt-mapping-prefill-port) TRT_MAPPING_PREFILL_PORT="$2"; shift 2 ;;
--trt-mapping-decode-port) TRT_MAPPING_DECODE_PORT="$2"; shift 2 ;;
--sglang-mem-fraction) SGLANG_MEM_FRACTION="$2"; shift 2 ;;
--sglang-server-extra) SGLANG_SERVER_EXTRA="$2"; shift 2 ;;
--max-model-len) MAX_MODEL_LEN="$2"; shift 2 ;;
--kv-fraction) KV_FRACTION="$2"; shift 2 ;;
--profile-workload) PROFILE_WORKLOAD="$2"; shift 2 ;;
--prefill-input-len) PREFILL_INPUT_LEN="$2"; shift 2 ;;
--prefill-output-len) PREFILL_OUTPUT_LEN="$2"; shift 2 ;;
--decode-input-len) DECODE_INPUT_LEN="$2"; shift 2 ;;
--decode-output-len) DECODE_OUTPUT_LEN="$2"; shift 2 ;;
--help|-h) usage; exit 0 ;;
*)
echo "Unknown argument: $1" >&2
usage >&2
exit 2
;;
esac
done
if [[ -z "${HF_TOKEN:-}" && -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
echo "Set HF_TOKEN or HUGGINGFACE_HUB_TOKEN before running." >&2
exit 2
fi
if [[ -z "${HF_TOKEN:-}" ]]; then
HF_TOKEN="$HUGGINGFACE_HUB_TOKEN"
fi
if [[ -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
HUGGINGFACE_HUB_TOKEN="$HF_TOKEN"
fi
for value in \
MODEL_ID MODEL ROOT GPUS \
SGLANG_PORT VLLM_FORMAL_PORT VLLM_MAPPING_PORT \
TRT_FORMAL_PREFILL_PORT TRT_FORMAL_DECODE_PORT \
TRT_MAPPING_PREFILL_PORT TRT_MAPPING_DECODE_PORT; do
if [[ -z "${!value}" ]]; then
echo "Missing required argument: $value" >&2
usage >&2
exit 2
fi
done
IFS=',' read -r -a GPU_LIST <<< "$GPUS"
GPU_COUNT="${#GPU_LIST[@]}"
if [[ "$GPU_COUNT" -lt 1 ]]; then
echo "Could not parse --gpus: $GPUS" >&2
exit 2
fi
if [[ -z "$TP_SIZE" ]]; then
TP_SIZE="$GPU_COUNT"
fi
if (( TP_SIZE < 1 || TP_SIZE > GPU_COUNT )); then
echo "--tp-size must be between 1 and the visible GPU count ($GPU_COUNT)." >&2
exit 2
fi
MODEL_ROOT="$ROOT/$MODEL_ID"
SGLANG_ANALYSIS="$MODEL_ROOT/analysis_sglang.txt"
VLLM_FORMAL_DIR="$MODEL_ROOT/vllm_formal"
VLLM_MAPPING_DIR="$MODEL_ROOT/vllm_mapping"
VLLM_ANALYSIS="$MODEL_ROOT/analysis_vllm.txt"
TRT_FORMAL_DIR="$MODEL_ROOT/trtllm_formal"
TRT_MAPPING_DIR="$MODEL_ROOT/trtllm_mapping"
TRT_ANALYSIS="$MODEL_ROOT/analysis_trtllm.txt"
docker exec sglang_bbuf bash -lc "mkdir -p '$MODEL_ROOT'"
if [[ ! -s "$TRT_OVERRIDE_SOURCE" ]]; then
echo "[bootstrap] TensorRT-LLM py_executor source snapshot"
docker exec sglang_bbuf bash -lc "mkdir -p '$TRT_OVERRIDE_ROOT'"
docker run --rm --entrypoint cat "$TRT_IMAGE" \
/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/py_executor.py \
| docker exec -i sglang_bbuf bash -lc "cat > '$TRT_OVERRIDE_SOURCE'"
fi
echo "[bootstrap] TensorRT-LLM py_executor override with with_stack=True and rank0-only trace export"
docker exec sglang_bbuf bash -lc "cd '$SCRIPT_DIR' && python3 make_trtllm_py_executor_override.py --source '$TRT_OVERRIDE_SOURCE' --output '$TRT_OVERRIDE_PATH'"
sglang_args=(
--model "$MODEL"
--run-dir "$MODEL_ROOT"
--port "$SGLANG_PORT"
--gpus "$GPUS"
--tp-size "$TP_SIZE"
--mem-fraction "$SGLANG_MEM_FRACTION"
--profile-workload "$PROFILE_WORKLOAD"
--prefill-input-len "$PREFILL_INPUT_LEN"
--prefill-output-len "$PREFILL_OUTPUT_LEN"
--decode-input-len "$DECODE_INPUT_LEN"
--decode-output-len "$DECODE_OUTPUT_LEN"
--trust-remote-code
)
if [[ -n "$SGLANG_SERVER_EXTRA" ]]; then
sglang_args+=(--server-extra "$SGLANG_SERVER_EXTRA")
fi
echo "[1/6] SGLang server + live triage"
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_sglang_torch_profile_host.sh" \
"${sglang_args[@]}"
echo "[2/6] vLLM formal"
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_vllm_torch_profile_host.sh" \
--model "$MODEL" \
--run-dir "$VLLM_FORMAL_DIR" \
--port "$VLLM_FORMAL_PORT" \
--gpus "$GPUS" \
--tensor-parallel-size "$TP_SIZE" \
--max-model-len "$MAX_MODEL_LEN" \
--profile-workload "$PROFILE_WORKLOAD" \
--prefill-input-len "$PREFILL_INPUT_LEN" \
--prefill-output-len "$PREFILL_OUTPUT_LEN" \
--decode-input-len "$DECODE_INPUT_LEN" \
--decode-output-len "$DECODE_OUTPUT_LEN" \
--trust-remote-code
echo "[3/6] vLLM mapping"
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_vllm_torch_profile_host.sh" \
--model "$MODEL" \
--run-dir "$VLLM_MAPPING_DIR" \
--port "$VLLM_MAPPING_PORT" \
--gpus "$GPUS" \
--tensor-parallel-size "$TP_SIZE" \
--profiler-active-iterations 2 \
--max-model-len "$MAX_MODEL_LEN" \
--profile-workload "$PROFILE_WORKLOAD" \
--prefill-input-len "$PREFILL_INPUT_LEN" \
--prefill-output-len "$PREFILL_OUTPUT_LEN" \
--decode-input-len "$DECODE_INPUT_LEN" \
--decode-output-len "$DECODE_OUTPUT_LEN" \
--trust-remote-code \
--enforce-eager
echo "[4/6] vLLM mapping-formal analysis"
docker exec sglang_bbuf bash -lc "cd '$SCRIPT_DIR' && python3 analyze_llm_torch_profile.py --framework vllm --mapping-input '$VLLM_MAPPING_DIR' --formal-input '$VLLM_FORMAL_DIR' > '$VLLM_ANALYSIS'"
echo "[5/6] TensorRT-LLM formal + mapping captures"
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_trtllm_pytorch_profile_host.sh" \
--model "$MODEL" \
--run-dir "$TRT_FORMAL_DIR" \
--stage prefill \
--port "$TRT_FORMAL_PREFILL_PORT" \
--gpus "$GPUS" \
--tp-size "$TP_SIZE" \
--kv-fraction "$KV_FRACTION" \
--input-len "$PREFILL_INPUT_LEN" \
--output-len "$PREFILL_OUTPUT_LEN" \
--override-py-executor "$TRT_OVERRIDE_PATH" \
--trust-remote-code
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_trtllm_pytorch_profile_host.sh" \
--model "$MODEL" \
--run-dir "$TRT_FORMAL_DIR" \
--stage decode \
--port "$TRT_FORMAL_DECODE_PORT" \
--gpus "$GPUS" \
--tp-size "$TP_SIZE" \
--kv-fraction "$KV_FRACTION" \
--input-len "$DECODE_INPUT_LEN" \
--output-len "$DECODE_OUTPUT_LEN" \
--override-py-executor "$TRT_OVERRIDE_PATH" \
--trust-remote-code
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_trtllm_pytorch_profile_host.sh" \
--model "$MODEL" \
--run-dir "$TRT_MAPPING_DIR" \
--stage prefill \
--port "$TRT_MAPPING_PREFILL_PORT" \
--gpus "$GPUS" \
--tp-size "$TP_SIZE" \
--kv-fraction "$KV_FRACTION" \
--input-len "$PREFILL_INPUT_LEN" \
--output-len "$PREFILL_OUTPUT_LEN" \
--override-py-executor "$TRT_OVERRIDE_PATH" \
--disable-cudagraph \
--trust-remote-code
HF_TOKEN="$HF_TOKEN" HUGGINGFACE_HUB_TOKEN="$HUGGINGFACE_HUB_TOKEN" \
"$SCRIPT_DIR/run_trtllm_pytorch_profile_host.sh" \
--model "$MODEL" \
--run-dir "$TRT_MAPPING_DIR" \
--stage decode \
--port "$TRT_MAPPING_DECODE_PORT" \
--gpus "$GPUS" \
--tp-size "$TP_SIZE" \
--kv-fraction "$KV_FRACTION" \
--input-len "$DECODE_INPUT_LEN" \
--output-len "$DECODE_OUTPUT_LEN" \
--override-py-executor "$TRT_OVERRIDE_PATH" \
--disable-cudagraph \
--trust-remote-code
echo "[6/6] TensorRT-LLM mapping-formal analysis"
docker exec sglang_bbuf bash -lc "cd '$SCRIPT_DIR' && python3 analyze_llm_torch_profile.py --framework trtllm --mapping-input '$TRT_MAPPING_DIR' --formal-input '$TRT_FORMAL_DIR' > '$TRT_ANALYSIS'"
echo "MODEL_ROOT=$MODEL_ROOT"
echo "ANALYSIS_SGLANG=$SGLANG_ANALYSIS"
echo "ANALYSIS_VLLM=$VLLM_ANALYSIS"
echo "ANALYSIS_TRTLLM=$TRT_ANALYSIS"
.claude/skills/llm-torch-profiler-analysis/scripts/run_sglang_torch_profile_host.sh
Executable
+241
View File
@@ -0,0 +1,241 @@
#!/usr/bin/env bash
set -euo pipefail
usage() {
cat <<'EOF'
Usage:
run_sglang_torch_profile_host.sh \
--model Qwen/Qwen3-8B \
--run-dir /data/bbuf/validate/unified_llm_profiler_skill/runs/example_sglang \
--port 30088 \
--gpus 0
run_sglang_torch_profile_host.sh \
--model openai/gpt-oss-20b \
--run-dir /data/bbuf/validate/unified_llm_profiler_skill/runs/example_sglang_4gpu \
--port 30088 \
--gpus 2,3,4,5 \
--tp-size 4
Options:
--model TEXT Model id or local path for SGLang.
--run-dir PATH Shared /data directory for logs and traces.
--port INT Server port.
--gpus TEXT CUDA_VISIBLE_DEVICES value, for example 0 or 2,3,4,5.
--gpu TEXT Alias for --gpus.
--tp-size INT Tensor parallel size. Defaults to the visible GPU count.
--trust-remote-code Pass --trust-remote-code.
--mem-fraction FLOAT SGLang static memory fraction.
--request-max-tokens INT Generation length for the probe request.
--prompt TEXT Probe prompt.
--warmup-steps INT Warmup steps before profiling. Defaults to 10.
--profile-workload TEXT legacy|prefill|decode|both. Defaults to both.
--prefill-input-len INT Synthetic prefill prompt length. Defaults to 4090.
--prefill-output-len INT Synthetic prefill output length. Defaults to 1.
--decode-input-len INT Synthetic decode prompt length. Defaults to 1.
--decode-output-len INT Synthetic decode output length. Defaults to 2048.
--repo-dir PATH SGLang repo path inside `sglang_bbuf`.
--server-extra TEXT Extra args appended to launch_server.
--help Show this message.
Notes:
- Run this on the H100 host. It uses `docker exec sglang_bbuf`.
- The server is launched first, then the profiler capture runs with
stage-separated prefill/decode workloads and `--profile-by-stage`.
- A small benchmark summary is written after profiling.
EOF
}
MODEL=""
RUN_DIR=""
PORT=""
GPUS=""
TP_SIZE=""
TRUST_REMOTE_CODE=0
MEM_FRACTION=0.85
REQUEST_MAX_TOKENS=12
PROMPT="Explain the difference between CUDA graph mode and eager mode in two sentences."
WARMUP_STEPS=10
PROFILE_WORKLOAD="both"
PREFILL_INPUT_LEN=4090
PREFILL_OUTPUT_LEN=1
DECODE_INPUT_LEN=1
DECODE_OUTPUT_LEN=2048
SGLANG_REPO_DIR="${SGLANG_REPO_DIR:-/data/bbuf/repos/sglang}"
SERVER_EXTRA=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
while [[ $# -gt 0 ]]; do
case "$1" in
--model)
MODEL="$2"
shift 2
;;
--run-dir)
RUN_DIR="$2"
shift 2
;;
--port)
PORT="$2"
shift 2
;;
--gpu)
GPUS="$2"
shift 2
;;
--gpus)
GPUS="$2"
shift 2
;;
--tp-size)
TP_SIZE="$2"
shift 2
;;
--trust-remote-code)
TRUST_REMOTE_CODE=1
shift
;;
--mem-fraction)
MEM_FRACTION="$2"
shift 2
;;
--request-max-tokens)
REQUEST_MAX_TOKENS="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--warmup-steps)
WARMUP_STEPS="$2"
shift 2
;;
--profile-workload)
PROFILE_WORKLOAD="$2"
shift 2
;;
--prefill-input-len)
PREFILL_INPUT_LEN="$2"
shift 2
;;
--prefill-output-len)
PREFILL_OUTPUT_LEN="$2"
shift 2
;;
--decode-input-len)
DECODE_INPUT_LEN="$2"
shift 2
;;
--decode-output-len)
DECODE_OUTPUT_LEN="$2"
shift 2
;;
--repo-dir)
SGLANG_REPO_DIR="$2"
shift 2
;;
--server-extra)
SERVER_EXTRA="$2"
shift 2
;;
--help|-h)
usage
exit 0
;;
*)
echo "Unknown argument: $1" >&2
usage >&2
exit 2
;;
esac
done
if [[ -z "$MODEL" || -z "$RUN_DIR" || -z "$PORT" || -z "$GPUS" ]]; then
usage >&2
exit 2
fi
IFS=',' read -r -a GPU_LIST <<< "$GPUS"
GPU_COUNT="${#GPU_LIST[@]}"
if [[ "$GPU_COUNT" -lt 1 ]]; then
echo "Could not parse --gpus: $GPUS" >&2
exit 2
fi
if [[ -z "$TP_SIZE" ]]; then
TP_SIZE="$GPU_COUNT"
fi
if (( TP_SIZE < 1 || TP_SIZE > GPU_COUNT )); then
echo "--tp-size must be between 1 and the visible GPU count ($GPU_COUNT)." >&2
exit 2
fi
LOG_PATH="$RUN_DIR/sglang_server.log"
ANALYSIS_PATH="$RUN_DIR/analysis_sglang.txt"
PROFILE_ROOT="$RUN_DIR/sglang_profile_live"
BENCHMARK_PATH="$RUN_DIR/benchmark_sglang.json"
PID_PATH="$RUN_DIR/sglang_server.pid"
LAUNCH_PATTERN="[s]glang.launch_server.*--port $PORT"
SERVER_ARGS="python3 -m sglang.launch_server --model-path \"$MODEL\" --port \"$PORT\" --tp-size \"$TP_SIZE\" --mem-fraction-static \"$MEM_FRACTION\""
if [[ "$TRUST_REMOTE_CODE" -eq 1 ]]; then
SERVER_ARGS="$SERVER_ARGS --trust-remote-code"
fi
if [[ -n "$SERVER_EXTRA" ]]; then
SERVER_ARGS="$SERVER_ARGS $SERVER_EXTRA"
fi
docker exec sglang_bbuf bash -lc "mkdir -p '$RUN_DIR' '$PROFILE_ROOT'"
docker exec sglang_bbuf bash -lc "pkill -f '$LAUNCH_PATTERN' >/dev/null 2>&1 || true"
docker exec sglang_bbuf bash -lc "mkdir -p '$RUN_DIR' '$PROFILE_ROOT' && cd '$SGLANG_REPO_DIR' && rm -f '$PID_PATH' && (CUDA_VISIBLE_DEVICES=$GPUS PYTHONPATH=python nohup $SERVER_ARGS > '$LOG_PATH' 2>&1 < /dev/null & echo \$! > '$PID_PATH')"
cleanup() {
docker exec sglang_bbuf bash -lc "pkill -f '$LAUNCH_PATTERN' >/dev/null 2>&1 || true" >/dev/null 2>&1 || true
}
trap cleanup EXIT
ready=0
for _ in $(seq 1 180); do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" >/dev/null; then
ready=1
break
fi
sleep 2
done
if [[ "$ready" -ne 1 ]]; then
echo "SGLang server did not become ready on port ${PORT}. Recent logs:" >&2
ssh_log=$(docker exec sglang_bbuf bash -lc "tail -n 120 '$LOG_PATH'" 2>/dev/null || true)
printf '%s\n' "$ssh_log" >&2
exit 1
fi
python3 - <<PY
import json
import urllib.request
payload = {
"text": ${PROMPT@Q},
"sampling_params": {
"temperature": 0.0,
"max_new_tokens": int(${REQUEST_MAX_TOKENS@Q}),
},
"stream": False,
}
req = urllib.request.Request(
"http://127.0.0.1:${PORT}/generate",
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=600) as resp:
body = json.loads(resp.read().decode())
text = body.get("text", "")
print(text[:400])
PY
docker exec sglang_bbuf bash -lc "cd '$SCRIPT_DIR' && python3 analyze_llm_torch_profile.py --framework sglang --url http://127.0.0.1:${PORT} --output-dir '$PROFILE_ROOT' --num-steps 5 --warmup-steps '$WARMUP_STEPS' --probe-requests 1 --profile-by-stage --profile-workload '$PROFILE_WORKLOAD' --prefill-input-len '$PREFILL_INPUT_LEN' --prefill-output-len '$PREFILL_OUTPUT_LEN' --decode-input-len '$DECODE_INPUT_LEN' --decode-output-len '$DECODE_OUTPUT_LEN' > '$ANALYSIS_PATH'"
python3 "$SCRIPT_DIR/probe_llm_server.py" \
--framework sglang \
--url "http://127.0.0.1:${PORT}" \
| docker exec -i sglang_bbuf bash -lc "cat > '$BENCHMARK_PATH'" >/dev/null
docker exec sglang_bbuf bash -lc "sed -n '1,240p' '$ANALYSIS_PATH'"
echo "BENCHMARK_PATH=$BENCHMARK_PATH"
.claude/skills/llm-torch-profiler-analysis/scripts/run_trtllm_pytorch_profile_host.sh
Executable
+407
View File
@@ -0,0 +1,407 @@
#!/usr/bin/env bash
set -euo pipefail
usage() {
cat <<'EOF'
Usage:
run_trtllm_pytorch_profile_host.sh \
--model Qwen/Qwen3-8B \
--run-dir /data/bbuf/validate/unified_llm_profiler_skill/runs/example \
--stage prefill \
--port 32188 \
--gpus 0
run_trtllm_pytorch_profile_host.sh \
--model openai/gpt-oss-20b \
--run-dir /data/bbuf/validate/unified_llm_profiler_skill/runs/example_4gpu \
--stage prefill \
--port 32188 \
--gpus 2,3,4,5 \
--tp-size 4
Options:
--model TEXT Hugging Face model id.
--run-dir PATH Shared /data run directory for logs and traces.
--stage prefill|decode Capture window. Prefill profiles 4090->1 by
default; decode profiles 1->2048 by default.
--port INT Host port for trtllm-serve.
--gpus TEXT CUDA_VISIBLE_DEVICES value, for example 0 or 2,3,4,5.
--gpu TEXT Alias for --gpus.
--tp-size INT Tensor parallel size. Defaults to the visible GPU count.
--image TEXT Container image.
--shared-root PATH Shared validation root mounted into the container.
--hf-cache PATH Host Hugging Face cache path.
--override-py-executor PATH Optional py_executor.py override path.
--disable-cudagraph Generate/use a YAML override with cuda_graph_config: null.
--input-len INT Synthetic prompt length for this stage.
Defaults: prefill 4090, decode 1.
--request-max-tokens INT Generation length for this stage.
Defaults: prefill 1, decode 2048.
--output-len INT Alias for --request-max-tokens.
--prompt TEXT Probe prompt. Defaults to a synthetic prompt
sized by --input-len.
--warmup-steps INT Warmup steps before the profiler window. Defaults to 10.
--active-steps INT Active profiler steps to capture. Defaults to 5.
--max-seq-len INT Serve max sequence length.
--kv-fraction FLOAT KV cache free GPU memory fraction.
--container-name TEXT Override container name.
--trust-remote-code Pass --trust_remote_code to trtllm-serve.
--help Show this message.
Environment:
HF_TOKEN or HUGGINGFACE_HUB_TOKEN must be set.
Notes:
- Run this on the H100 host, not inside `sglang_bbuf`.
- It always pins TensorRT-LLM to `--backend pytorch`.
- The default image tag is floating; record the resolved TensorRT-LLM version
in the run manifest and pass --image for reproducible validation.
- Profiling uses `TLLM_PROFILE_START_STOP` and `TLLM_TORCH_PROFILE_TRACE`.
- For Python-location recovery, prefer a `py_executor.py` override with `with_stack=True`.
- A small benchmark summary is written after the trace is emitted.
EOF
}
IMAGE="nvcr.io/nvidia/tensorrt-llm/release:latest"
SHARED_ROOT="/data/bbuf/validate/unified_llm_profiler_skill"
HF_CACHE="/data/.cache/huggingface"
OVERRIDE_PY_EXECUTOR=""
DISABLE_CUDAGRAPH=0
REQUEST_MAX_TOKENS=""
INPUT_LEN=""
PROMPT=""
WARMUP_STEPS=10
ACTIVE_STEPS=5
MAX_SEQ_LEN=4096
KV_FRACTION=0.85
CONTAINER_NAME=""
TRUST_REMOTE_CODE=0
TP_SIZE=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODEL=""
RUN_DIR=""
STAGE=""
PORT=""
GPUS=""
while [[ $# -gt 0 ]]; do
case "$1" in
--model)
MODEL="$2"
shift 2
;;
--run-dir)
RUN_DIR="$2"
shift 2
;;
--stage)
STAGE="$2"
shift 2
;;
--port)
PORT="$2"
shift 2
;;
--gpu)
GPUS="$2"
shift 2
;;
--gpus)
GPUS="$2"
shift 2
;;
--tp-size)
TP_SIZE="$2"
shift 2
;;
--image)
IMAGE="$2"
shift 2
;;
--shared-root)
SHARED_ROOT="$2"
shift 2
;;
--hf-cache)
HF_CACHE="$2"
shift 2
;;
--override-py-executor)
OVERRIDE_PY_EXECUTOR="$2"
shift 2
;;
--disable-cudagraph)
DISABLE_CUDAGRAPH=1
shift
;;
--input-len)
INPUT_LEN="$2"
shift 2
;;
--request-max-tokens)
REQUEST_MAX_TOKENS="$2"
shift 2
;;
--output-len)
REQUEST_MAX_TOKENS="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--warmup-steps)
WARMUP_STEPS="$2"
shift 2
;;
--active-steps)
ACTIVE_STEPS="$2"
shift 2
;;
--max-seq-len)
MAX_SEQ_LEN="$2"
shift 2
;;
--kv-fraction)
KV_FRACTION="$2"
shift 2
;;
--container-name)
CONTAINER_NAME="$2"
shift 2
;;
--trust-remote-code)
TRUST_REMOTE_CODE=1
shift
;;
--help|-h)
usage
exit 0
;;
*)
echo "Unknown argument: $1" >&2
usage >&2
exit 2
;;
esac
done
if [[ -z "${HF_TOKEN:-}" && -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
echo "Set HF_TOKEN or HUGGINGFACE_HUB_TOKEN before running." >&2
exit 2
fi
if [[ -z "${HF_TOKEN:-}" ]]; then
HF_TOKEN="$HUGGINGFACE_HUB_TOKEN"
fi
if [[ -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
HUGGINGFACE_HUB_TOKEN="$HF_TOKEN"
fi
if [[ -z "$MODEL" || -z "$RUN_DIR" || -z "$STAGE" || -z "$PORT" || -z "$GPUS" ]]; then
usage >&2
exit 2
fi
IFS=',' read -r -a GPU_LIST <<< "$GPUS"
GPU_COUNT="${#GPU_LIST[@]}"
if [[ "$GPU_COUNT" -lt 1 ]]; then
echo "Could not parse --gpus: $GPUS" >&2
exit 2
fi
if [[ -z "$TP_SIZE" ]]; then
TP_SIZE="$GPU_COUNT"
fi
if (( TP_SIZE < 1 || TP_SIZE > GPU_COUNT )); then
echo "--tp-size must be between 1 and the visible GPU count ($GPU_COUNT)." >&2
exit 2
fi
case "$STAGE" in
prefill)
TRACE_PATH="$RUN_DIR/trace-prefill.json"
LOG_PATH="$RUN_DIR/server-prefill.log"
BENCHMARK_PATH="$RUN_DIR/benchmark-prefill.json"
if [[ -z "$INPUT_LEN" ]]; then
INPUT_LEN=4090
fi
if [[ -z "$REQUEST_MAX_TOKENS" ]]; then
REQUEST_MAX_TOKENS=1
fi
;;
decode)
TRACE_PATH="$RUN_DIR/trace-decode.json"
LOG_PATH="$RUN_DIR/server-decode.log"
BENCHMARK_PATH="$RUN_DIR/benchmark-decode.json"
if [[ -z "$INPUT_LEN" ]]; then
INPUT_LEN=1
fi
if [[ -z "$REQUEST_MAX_TOKENS" ]]; then
REQUEST_MAX_TOKENS=2048
fi
;;
*)
echo "--stage must be prefill or decode." >&2
exit 2
;;
esac
if (( WARMUP_STEPS < 0 || ACTIVE_STEPS < 1 )); then
echo "--warmup-steps must be >= 0 and --active-steps must be >= 1." >&2
exit 2
fi
case "$STAGE" in
prefill)
profile_start=$((WARMUP_STEPS + 1))
;;
decode)
profile_start=$((WARMUP_STEPS + 2))
;;
esac
profile_stop=$((profile_start + ACTIVE_STEPS - 1))
PROFILE_START_STOP="${profile_start}-${profile_stop}"
if [[ -z "$CONTAINER_NAME" ]]; then
model_slug="${MODEL##*/}"
model_slug="${model_slug//\//-}"
model_slug="${model_slug//./-}"
model_slug="${model_slug//_/-}"
model_slug="${model_slug// /-}"
gpu_slug="${GPUS//,/-}"
CONTAINER_NAME="trtllm-${model_slug}-${STAGE}-g${gpu_slug}-p${PORT}"
fi
EXTRA_LLM_OPTIONS=""
if [[ "$DISABLE_CUDAGRAPH" -eq 1 ]]; then
EXTRA_CFG_PATH="$SHARED_ROOT/tmp/trt_no_cudagraph.yaml"
docker exec sglang_bbuf bash -lc "mkdir -p '$(dirname "$EXTRA_CFG_PATH")' && printf 'cuda_graph_config: null\n' > '$EXTRA_CFG_PATH'"
EXTRA_LLM_OPTIONS="--extra_llm_api_options $EXTRA_CFG_PATH"
fi
docker exec sglang_bbuf bash -lc "mkdir -p '$RUN_DIR'"
docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
docker_args=(
run -d --rm
--name "$CONTAINER_NAME"
--gpus all
--ipc=host
--network host
--entrypoint bash
-e "CUDA_VISIBLE_DEVICES=$GPUS"
-e "HF_TOKEN=$HF_TOKEN"
-e "HUGGINGFACE_HUB_TOKEN=$HUGGINGFACE_HUB_TOKEN"
-e "TLLM_PROFILE_START_STOP=$PROFILE_START_STOP"
-e "TLLM_LLMAPI_ENABLE_NVTX=1"
-e "TLLM_TORCH_PROFILE_TRACE=$TRACE_PATH"
-e "RUN_DIR=$RUN_DIR"
-e "LOG_PATH=$LOG_PATH"
-e "MODEL_ID=$MODEL"
-e "SERVE_PORT=$PORT"
-v "$HF_CACHE:/root/.cache/huggingface"
-v "$SHARED_ROOT:$SHARED_ROOT"
)
if [[ -n "$OVERRIDE_PY_EXECUTOR" ]]; then
docker_args+=(
-v "$OVERRIDE_PY_EXECUTOR:/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/py_executor.py:ro"
)
fi
trust_remote_code_arg=""
if [[ "$TRUST_REMOTE_CODE" -eq 1 ]]; then
trust_remote_code_arg="--trust_remote_code"
fi
container_cmd=$(
cat <<EOF
mkdir -p "$RUN_DIR" && trtllm-serve serve "$MODEL" \
--backend pytorch \
--tp_size "$TP_SIZE" \
--gpus_per_node "$GPU_COUNT" \
--host 0.0.0.0 \
--port "$PORT" \
--max_seq_len "$MAX_SEQ_LEN" \
--kv_cache_free_gpu_memory_fraction "$KV_FRACTION" \
$trust_remote_code_arg \
$EXTRA_LLM_OPTIONS \
> "$LOG_PATH" 2>&1
EOF
)
docker_args+=("$IMAGE" -lc "$container_cmd")
docker "${docker_args[@]}" >/dev/null
cleanup() {
docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
}
trap cleanup EXIT
ready=0
for _ in $(seq 1 180); do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" >/dev/null; then
ready=1
break
fi
sleep 2
done
if [[ "$ready" -ne 1 ]]; then
echo "Server did not become ready on port ${PORT}. Recent logs:" >&2
docker logs "$CONTAINER_NAME" 2>&1 | tail -n 120 >&2 || true
exit 1
fi
python3 - <<PY
import json
import sys
import urllib.request
sys.path.insert(0, ${SCRIPT_DIR@Q})
from profile_common import extract_openai_chat_text, synthetic_prompt
prompt = ${PROMPT@Q} or synthetic_prompt(int(${INPUT_LEN@Q}))
stage = ${STAGE@Q}
warmup_steps = int(${WARMUP_STEPS@Q})
active_steps = int(${ACTIVE_STEPS@Q})
request_count = warmup_steps + active_steps if stage == "prefill" else 1
payload = {
"model": ${MODEL@Q},
"messages": [{"role": "user", "content": prompt}],
"temperature": 0,
"max_tokens": int(${REQUEST_MAX_TOKENS@Q}),
}
for request_idx in range(request_count):
req = urllib.request.Request(
"http://127.0.0.1:${PORT}/v1/chat/completions",
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=600) as resp:
body = json.loads(resp.read().decode())
text, source = extract_openai_chat_text(body)
print(text[:400] if text else f"[empty completion; source={source}]")
PY
for _ in $(seq 1 120); do
if [[ -s "$TRACE_PATH" ]]; then
break
fi
sleep 2
done
if [[ ! -s "$TRACE_PATH" ]]; then
echo "Trace was not written: $TRACE_PATH" >&2
exit 1
fi
python3 "$SCRIPT_DIR/probe_llm_server.py" \
--framework trtllm \
--url "http://127.0.0.1:${PORT}" \
--model "$MODEL" \
| docker exec -i sglang_bbuf bash -lc "cat > '$BENCHMARK_PATH'" >/dev/null
echo "TRACE_PATH=$TRACE_PATH"
echo "LOG_PATH=$LOG_PATH"
echo "BENCHMARK_PATH=$BENCHMARK_PATH"
.claude/skills/llm-torch-profiler-analysis/scripts/run_vllm_torch_profile_host.sh
Executable
+343
View File
@@ -0,0 +1,343 @@
#!/usr/bin/env bash
set -euo pipefail
usage() {
cat <<'EOF'
Usage:
run_vllm_torch_profile_host.sh \
--model Qwen/Qwen3-8B \
--run-dir /data/bbuf/validate/unified_llm_profiler_skill/runs/example_vllm_formal \
--port 31088 \
--gpus 1
run_vllm_torch_profile_host.sh \
--model openai/gpt-oss-20b \
--run-dir /data/bbuf/validate/unified_llm_profiler_skill/runs/example_vllm_4gpu \
--port 31088 \
--gpus 2,3,4,5 \
--tensor-parallel-size 4
Options:
--model TEXT Hugging Face model id.
--run-dir PATH Shared /data directory for logs and traces.
--port INT Host port for vllm serve.
--gpus TEXT CUDA_VISIBLE_DEVICES value, for example 1 or 2,3,4,5.
--gpu TEXT Alias for --gpus.
--image TEXT Container image.
--hf-cache PATH Host Hugging Face cache path.
--gpu-memory-util FLOAT vLLM --gpu-memory-utilization.
--max-model-len INT vLLM --max-model-len.
--tensor-parallel-size INT vLLM --tensor-parallel-size. Defaults to the visible GPU count.
--profiler-active-iterations INT
Torch-profiler active iterations.
--enforce-eager Launch vLLM with --enforce-eager for mapping traces.
--trust-remote-code Pass --trust-remote-code.
--request-max-tokens INT Generation length for the probe request.
--prompt TEXT Probe prompt.
--warmup-steps INT Warmup steps before profiling. Defaults to 10.
--profile-workload TEXT legacy|prefill|decode|both. Defaults to both.
--prefill-input-len INT Synthetic prefill prompt length. Defaults to 4090.
--prefill-output-len INT Synthetic prefill output length. Defaults to 1.
--decode-input-len INT Synthetic decode prompt length. Defaults to 1.
--decode-output-len INT Synthetic decode output length. Defaults to 2048.
--container-name TEXT Override container name.
--help Show this message.
Environment:
HF_TOKEN or HUGGINGFACE_HUB_TOKEN must be set.
Notes:
- Run this on the H100 host, not inside `sglang_bbuf`.
- This uses the vLLM torch-profiler flow: `--profiler-config`, then POST
`/start_profile` and `/stop_profile`.
- Default capture is two labeled profiles: prefill 4090->1 and decode 1->2048.
- Current vLLM profiler config already defaults `torch_profiler_with_stack=true`.
- A small benchmark summary is written after profiling.
EOF
}
IMAGE="vllm/vllm-openai:latest"
HF_CACHE="/data/.cache/huggingface"
GPU_MEMORY_UTIL=0.90
MAX_MODEL_LEN=4096
TP_SIZE=""
ENFORCE_EAGER=0
TRUST_REMOTE_CODE=0
REQUEST_MAX_TOKENS=12
PROFILER_ACTIVE_ITERATIONS=5
PROMPT="Explain the difference between CUDA graph mode and eager mode in two sentences."
WARMUP_STEPS=10
PROFILE_WORKLOAD="both"
PREFILL_INPUT_LEN=4090
PREFILL_OUTPUT_LEN=1
DECODE_INPUT_LEN=1
DECODE_OUTPUT_LEN=2048
CONTAINER_NAME=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODEL=""
RUN_DIR=""
PORT=""
GPUS=""
while [[ $# -gt 0 ]]; do
case "$1" in
--model)
MODEL="$2"
shift 2
;;
--run-dir)
RUN_DIR="$2"
shift 2
;;
--port)
PORT="$2"
shift 2
;;
--gpu)
GPUS="$2"
shift 2
;;
--gpus)
GPUS="$2"
shift 2
;;
--image)
IMAGE="$2"
shift 2
;;
--hf-cache)
HF_CACHE="$2"
shift 2
;;
--gpu-memory-util)
GPU_MEMORY_UTIL="$2"
shift 2
;;
--max-model-len)
MAX_MODEL_LEN="$2"
shift 2
;;
--tensor-parallel-size)
TP_SIZE="$2"
shift 2
;;
--profiler-active-iterations)
PROFILER_ACTIVE_ITERATIONS="$2"
shift 2
;;
--enforce-eager)
ENFORCE_EAGER=1
shift
;;
--trust-remote-code)
TRUST_REMOTE_CODE=1
shift
;;
--request-max-tokens)
REQUEST_MAX_TOKENS="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--warmup-steps)
WARMUP_STEPS="$2"
shift 2
;;
--profile-workload)
PROFILE_WORKLOAD="$2"
shift 2
;;
--prefill-input-len)
PREFILL_INPUT_LEN="$2"
shift 2
;;
--prefill-output-len)
PREFILL_OUTPUT_LEN="$2"
shift 2
;;
--decode-input-len)
DECODE_INPUT_LEN="$2"
shift 2
;;
--decode-output-len)
DECODE_OUTPUT_LEN="$2"
shift 2
;;
--container-name)
CONTAINER_NAME="$2"
shift 2
;;
--help|-h)
usage
exit 0
;;
*)
echo "Unknown argument: $1" >&2
usage >&2
exit 2
;;
esac
done
if [[ -z "${HF_TOKEN:-}" && -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
echo "Set HF_TOKEN or HUGGINGFACE_HUB_TOKEN before running." >&2
exit 2
fi
if [[ -z "${HF_TOKEN:-}" ]]; then
HF_TOKEN="$HUGGINGFACE_HUB_TOKEN"
fi
if [[ -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
HUGGINGFACE_HUB_TOKEN="$HF_TOKEN"
fi
if [[ -z "$MODEL" || -z "$RUN_DIR" || -z "$PORT" || -z "$GPUS" ]]; then
usage >&2
exit 2
fi
IFS=',' read -r -a GPU_LIST <<< "$GPUS"
GPU_COUNT="${#GPU_LIST[@]}"
if [[ "$GPU_COUNT" -lt 1 ]]; then
echo "Could not parse --gpus: $GPUS" >&2
exit 2
fi
if [[ -z "$TP_SIZE" ]]; then
TP_SIZE="$GPU_COUNT"
fi
if (( TP_SIZE < 1 || TP_SIZE > GPU_COUNT )); then
echo "--tensor-parallel-size must be between 1 and the visible GPU count ($GPU_COUNT)." >&2
exit 2
fi
if (( PROFILER_ACTIVE_ITERATIONS < 1 )); then
echo "--profiler-active-iterations must be >= 1." >&2
exit 2
fi
PROFILE_DIR="$RUN_DIR/vllm_profile"
LOG_PATH="$RUN_DIR/server.log"
ANALYSIS_PATH="$RUN_DIR/analysis_vllm_live.txt"
BENCHMARK_PATH="$RUN_DIR/benchmark_vllm.json"
if [[ -z "$CONTAINER_NAME" ]]; then
model_slug="${MODEL##*/}"
model_slug="${model_slug//\//-}"
model_slug="${model_slug//./-}"
model_slug="${model_slug//_/-}"
gpu_slug="${GPUS//,/-}"
CONTAINER_NAME="vllm-${model_slug}-g${gpu_slug}-p${PORT}"
if [[ "$ENFORCE_EAGER" -eq 1 ]]; then
CONTAINER_NAME="${CONTAINER_NAME}-eager"
fi
fi
docker exec sglang_bbuf bash -lc "mkdir -p '$PROFILE_DIR'"
docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
profiler_config=$(python3 - <<PY
import json
print(json.dumps({
"profiler": "torch",
"torch_profiler_dir": ${PROFILE_DIR@Q},
"active_iterations": int(${PROFILER_ACTIVE_ITERATIONS@Q}),
}))
PY
)
docker_args=(
run -d --rm
--name "$CONTAINER_NAME"
--gpus all
--ipc=host
--network host
-e "CUDA_VISIBLE_DEVICES=$GPUS"
-e "HF_TOKEN=$HF_TOKEN"
-e "HUGGINGFACE_HUB_TOKEN=$HUGGINGFACE_HUB_TOKEN"
-e "VLLM_RPC_TIMEOUT=1800000"
-v "$HF_CACHE:/root/.cache/huggingface"
-v "$RUN_DIR:$RUN_DIR"
)
docker_cmd=(
"$IMAGE"
"$MODEL"
--host 0.0.0.0
--port "$PORT"
--tensor-parallel-size "$TP_SIZE"
--max-model-len "$MAX_MODEL_LEN"
--gpu-memory-utilization "$GPU_MEMORY_UTIL"
--profiler-config "$profiler_config"
)
if [[ "$ENFORCE_EAGER" -eq 1 ]]; then
docker_cmd+=(--enforce-eager)
fi
if [[ "$TRUST_REMOTE_CODE" -eq 1 ]]; then
docker_cmd+=(--trust-remote-code)
fi
docker "${docker_args[@]}" "${docker_cmd[@]}" >/dev/null
cleanup() {
docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
}
trap cleanup EXIT
ready=0
for _ in $(seq 1 180); do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" >/dev/null; then
ready=1
break
fi
sleep 2
done
if [[ "$ready" -ne 1 ]]; then
echo "Server did not become ready on port ${PORT}. Recent logs:" >&2
docker logs "$CONTAINER_NAME" 2>&1 | tail -n 120 >&2 || true
exit 1
fi
python3 "$SCRIPT_DIR/analyze_llm_torch_profile.py" \
--framework vllm \
--url "http://127.0.0.1:${PORT}" \
--output-dir "$PROFILE_DIR" \
--num-steps "$PROFILER_ACTIVE_ITERATIONS" \
--warmup-steps "$WARMUP_STEPS" \
--probe-requests 1 \
--no-profile-by-stage \
--profile-workload "$PROFILE_WORKLOAD" \
--probe-prompt "$PROMPT" \
--probe-max-new-tokens "$REQUEST_MAX_TOKENS" \
--prefill-input-len "$PREFILL_INPUT_LEN" \
--prefill-output-len "$PREFILL_OUTPUT_LEN" \
--decode-input-len "$DECODE_INPUT_LEN" \
--decode-output-len "$DECODE_OUTPUT_LEN" \
> "$ANALYSIS_PATH"
profile_found=0
for _ in $(seq 1 240); do
if find "$PROFILE_DIR" -type f \( -name '*.pt.trace.json' -o -name '*.pt.trace.json.gz' -o -name '*.trace.json' -o -name '*.trace.json.gz' \) | grep -q .; then
profile_found=1
break
fi
sleep 2
done
if [[ "$profile_found" -ne 1 ]]; then
echo "No vLLM profiler traces appeared under $PROFILE_DIR" >&2
docker logs "$CONTAINER_NAME" 2>&1 | tail -n 120 >&2 || true
exit 1
fi
python3 "$SCRIPT_DIR/probe_llm_server.py" \
--framework vllm \
--url "http://127.0.0.1:${PORT}" \
--model "$MODEL" \
| docker exec -i sglang_bbuf bash -lc "cat > '$BENCHMARK_PATH'" >/dev/null
docker logs "$CONTAINER_NAME" 2>&1 | docker exec -i sglang_bbuf bash -lc "cat > '$LOG_PATH'" || true
sed -n '1,240p' "$ANALYSIS_PATH"
echo "PROFILE_DIR=$PROFILE_DIR"
echo "LOG_PATH=$LOG_PATH"
echo "ANALYSIS_PATH=$ANALYSIS_PATH"
echo "BENCHMARK_PATH=$BENCHMARK_PATH"