ossrules.md

llm-torch-profiler-analysis

Unified LLM torch-profiler triage skill for `sglang`, `vllm`, `TensorRT-LLM`, and `TokenSpeed`. Use it to inspect an existing `trace.json(.gz)` or profile directory, or to drive live profiling against a running server when supported and return one three-table report with kernel, overlap-opportunity, and fuse-pattern tables.

Install
npx skills add 'https://github.com/sgl-project/sglang/tree/main/.claude/skills/llm-torch-profiler-analysis'
Download bundle ↓
main · a9fb1c3 ↗Scanned 2026-09-17

Contributors

GitHub-linked commit authors for this SKILL.md at the saved revision. Co-authors and history before file renames are not included.

File history ↗

scripts/profile_common.py

scripts/profile_common.pyBrowse 18 files
View on GitHub
9,073 tokens

40,177 bytes
Token encoding: o200k_base
Snapshot a9fb1c3

← Back to SKILL.md
"""Shared helpers for unified LLM torch-profiler skill scripts.""" from __future__ import annotations import gzipimport jsonimport reimport shutilimport sysimport tempfileimport timefrom collections import Counter, defaultdictfrom dataclasses import dataclassfrom functools import lru_cachefrom pathlib import Pathfrom typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuplefrom urllib import request STAGE_ORDER = {"extend": 0, "prefill": 0, "decode": 1, "all": 2}FRAMEWORK_LABELS = {    "auto": "auto",    "sglang": "SGLang",    "vllm": "vLLM",    "trtllm": "TensorRT-LLM",    "tokenspeed": "TokenSpeed",}TRACE_FILE_PATTERNS = (    "*.trace.json",    "*.trace.json.gz",    "*.pt.trace.json",    "*.pt.trace.json.gz",    "*.json",    "*.json.gz",)TRACE_FILE_IGNORE_NAMES = {    "server_args.json",    "metadata.json",    "config.json",}TRACE_METADATA_NAMES = {    "process_name",    "thread_name",    "process_sort_index",    "thread_sort_index",}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 = 4090DEFAULT_PREFILL_OUTPUT_LEN = 1DEFAULT_DECODE_INPUT_LEN = 1DEFAULT_DECODE_OUTPUT_LEN = 2048DEFAULT_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)def _normalize_text_cached(text: str) -> str:    text = text.strip()    if not text:        return ""    for token in (" ", "\t", "\n", "\r", "\v", "\f"):        if token in text:            return " ".join(text.split())    return text  def normalize_text(value: object) -> str:    return _normalize_text_cached(value if isinstance(value, str) else str(value))  def canonicalize_framework(value: object) -> str:    lowered = normalize_text(value).lower().replace("_", "-")    aliases = {        "": "auto",        "auto": "auto",        "sglang": "sglang",        "sgl": "sglang",        "vllm": "vllm",        "trt": "trtllm",        "tllm": "trtllm",        "trtllm": "trtllm",        "tensorrt-llm": "trtllm",        "tensorrtllm": "trtllm",        "tokenspeed": "tokenspeed",        "token-speed": "tokenspeed",        "ts": "tokenspeed",    }    return aliases.get(lowered, "auto")  def framework_display_name(value: object) -> str:    return FRAMEWORK_LABELS.get(canonicalize_framework(value), str(value))  @lru_cache(maxsize=65536)def _normalize_repo_relative_path_cached(text: str) -> str:    text = text.replace("\\", "/")    lowered = text.lower()    for marker, normalized_marker in (        ("python/sglang/", "python/sglang/"),        ("sgl_kernel/", "sgl_kernel/"),        ("vllm/", "vllm/"),        ("python/tokenspeed/", "python/tokenspeed/"),        ("tokenspeed/", "tokenspeed/"),        ("tensorrt_llm/", "tensorrt_llm/"),        ("tensorrt-llm/", "tensorrt_llm/"),    ):        idx = lowered.find(marker)        if idx != -1:            suffix = text[idx + len(marker) :].lstrip("/")            return f"{normalized_marker}{suffix}".lstrip("/")    idx = lowered.find("sglang/")    if idx != -1:        return ("python/" + text[idx:]).lstrip("/")    return text.lstrip("/")  def normalize_repo_relative_path(path: object) -> str:    return _normalize_repo_relative_path_cached(normalize_text(path))  def contains_any_keyword(text: str, keywords: Iterable[str]) -> bool:    return any(keyword in text for keyword in keywords)  def coerce_optional_int(value: object) -> Optional[int]:    if value in (None, "", "None"):        return None    if isinstance(value, int):        return value    if isinstance(value, float):        return int(value) if value.is_integer() else None    try:        return int(str(value))    except (TypeError, ValueError):        return None  def extract_trace_events(trace: object) -> Sequence[dict]:    if isinstance(trace, dict):        events = trace.get("traceEvents", [])        return events if isinstance(events, list) else []    if isinstance(trace, list):        return trace    return []  def is_trace_metadata_name(name: object) -> bool:    return str(name) in TRACE_METADATA_NAMES  def is_complete_duration_event(event: dict) -> bool:    if event.get("ph") != "X":        return False    dur = event.get("dur")    ts = event.get("ts")    if dur is None or ts is None:        return False    try:        return float(dur) > 0    except (TypeError, ValueError):        return False  def is_annotation_event(name: object, category: object) -> bool:    lowered_name = normalize_text(name).lower()    lowered_category = normalize_text(category).lower()    return "annotation" in lowered_category or lowered_name.startswith("## call ")  def is_non_kernel_trace_category(category: object) -> bool:    lowered_category = normalize_text(category).lower()    return any(token in lowered_category for token in NON_KERNEL_TRACE_CATEGORIES)  def looks_like_python_scope_name(name: object) -> bool:    lowered_name = normalize_text(name).lower()    return ".py(" in lowered_name or lowered_name.startswith(PYTHON_SCOPE_NAME_PREFIXES)  def has_stream_marker(args: Optional[dict]) -> bool:    trace_args = args or {}    return "stream" in trace_args or "cuda_stream" in trace_args  def load_trace_json(path: Path) -> dict:    if path.suffix == ".gz":        with gzip.open(path, "rt", encoding="utf-8") as handle:            return json.load(handle)    with open(path, "r", encoding="utf-8") as handle:        return json.load(handle)  def load_server_args(path: Path) -> Optional[dict]:    resolved = path.resolve()    candidate_dirs: List[Path] = []    if resolved.is_file():        candidate_dirs.extend([resolved.parent, resolved.parent.parent])    else:        candidate_dirs.extend([resolved, resolved.parent])     seen: set[Path] = set()    for candidate_dir in candidate_dirs:        if candidate_dir in seen:            continue        seen.add(candidate_dir)        candidate = candidate_dir / "server_args.json"        if candidate.exists():            with open(candidate, "r", encoding="utf-8") as handle:                return json.load(handle)    return None  def try_get_json(url: str, timeout: float = 60.0) -> Optional[object]:    try:        with request.urlopen(url, timeout=timeout) as response:            raw = response.read()    except Exception:        return None    if not raw:        return None    try:        return json.loads(raw.decode("utf-8"))    except json.JSONDecodeError:        return None  def _flatten_chat_text_parts(value: object) -> List[str]:    if value is None:        return []    if isinstance(value, str):        text = value.strip()        return [text] if text else []    if isinstance(value, list):        parts: List[str] = []        for item in value:            parts.extend(_flatten_chat_text_parts(item))        return parts    if isinstance(value, dict):        parts: List[str] = []        text_keys = (            "text",            "content",            "reasoning_content",            "reasoning",            "output_text",        )        if any(key in value for key in text_keys):            for key in text_keys:                parts.extend(_flatten_chat_text_parts(value.get(key)))            if parts:                return parts        item_type = normalize_text(value.get("type")).lower()        if item_type in {"text", "output_text", "input_text"}:            for key in ("text", "content", "value"):                parts.extend(_flatten_chat_text_parts(value.get(key)))        elif item_type in {"reasoning", "thinking"}:            for key in ("text", "content", "reasoning_content", "reasoning"):                parts.extend(_flatten_chat_text_parts(value.get(key)))        return parts    return []  def flatten_chat_text(value: object) -> str:    return "\n".join(_flatten_chat_text_parts(value)).strip()  def extract_openai_chat_text(body: object) -> Tuple[str, str]:    if not isinstance(body, dict):        return "", "invalid_body"     choices = body.get("choices")    if not isinstance(choices, list) or not choices:        fallback = flatten_chat_text(body.get("output_text"))        if fallback:            return fallback, "body.output_text"        return "", "missing_choices"     first_choice = choices[0]    if not isinstance(first_choice, dict):        return "", "invalid_choice"     message = first_choice.get("message")    if isinstance(message, dict):        for key in ("content", "reasoning_content", "reasoning"):            text = flatten_chat_text(message.get(key))            if text:                return text, f"message.{key}"     for key in ("text", "content", "reasoning_content", "reasoning"):        text = flatten_chat_text(first_choice.get(key))        if text:            return text, f"choice.{key}"     delta = first_choice.get("delta")    if isinstance(delta, dict):        for key in ("content", "reasoning_content", "reasoning"):            text = flatten_chat_text(delta.get(key))            if text:                return text, f"delta.{key}"     fallback = flatten_chat_text(body.get("output_text"))    if fallback:        return fallback, "body.output_text"    return "", "empty"  def detect_framework_from_text(text: object) -> Optional[str]:    lowered = normalize_text(text).lower()    if not lowered:        return None    if any(token in lowered for token in ("tokenspeed", "token-speed", "/ts/")):        return "tokenspeed"    if any(        token in lowered        for token in (            "tensorrt_llm",            "tensorrt-llm",            "trtllm",            "pyexecutor",        )    ):        return "trtllm"    if "vllm" in lowered:        return "vllm"    if any(token in lowered for token in ("python/sglang/", "sgl_kernel/", "sglang/")):        return "sglang"    return None  def detect_framework_from_server_args(server_args: Optional[dict]) -> Optional[str]:    if not isinstance(server_args, dict) or not server_args:        return None    lowered_keys = {normalize_text(key).lower() for key in server_args}    text = json.dumps(server_args, sort_keys=True)    if any(token in text.lower() for token in ("tokenspeed", "token-speed")):        return "tokenspeed"    if lowered_keys & {        "attn_tp_size",        "dense_tp_size",        "moe_tp_size",        "enable_mla_l1_5_cache",        "mla_chunk_multiplier",        "comm_fusion_max_num_tokens",        "enable_allreduce_fusion",    }:        return "tokenspeed"    text_hint = detect_framework_from_text(text)    if text_hint:        return text_hint    if lowered_keys & {        "attention_backend",        "sampling_backend",        "disable_cuda_graph",        "disable_piecewise_cuda_graph",        "chunked_prefill_size",        "schedule_policy",    }:        return "sglang"    return None  def detect_framework_from_trace(trace: object) -> Optional[str]:    text_samples: List[str] = []    for event in extract_trace_events(trace)[:256]:        text_samples.extend(            [                str(event.get("name", "")),                str(event.get("cat", "")),                str(event.get("pid", "")),            ]        )        trace_args = event.get("args")        if isinstance(trace_args, dict):            for key, value in list(trace_args.items())[:8]:                text_samples.append(str(key))                if isinstance(value, str):                    text_samples.append(value)    return detect_framework_from_text(" ".join(text_samples))  def detect_framework_from_path(path: Path) -> Optional[str]:    hint = detect_framework_from_text(str(path))    if hint:        return hint    server_args = load_server_args(path)    hint = detect_framework_from_server_args(server_args)    if hint:        return hint    if path.is_file():        try:            return detect_framework_from_trace(load_trace_json(path))        except Exception:            return None    trace_files = discover_trace_files(path, recursive=True, limit=3)    for trace_file in trace_files:        try:            hint = detect_framework_from_trace(load_trace_json(trace_file))        except Exception:            hint = None        if hint:            return hint    return None  def detect_framework_from_url(    url: str, output_dir: Optional[str] = None) -> Optional[str]:    hint = detect_framework_from_text(output_dir or "")    if hint:        return hint    server_info = try_get_json(url.rstrip("/") + "/server_info")    if isinstance(server_info, dict) and (        "internal_states" in server_info        or "tokenizer_path" in server_info        or "prefill" in server_info        or "decode" in server_info    ):        return "sglang"    readiness = try_get_json(url.rstrip("/") + "/readiness", timeout=5.0)    if readiness is not None:        return "tokenspeed"    models = try_get_json(url.rstrip("/") + "/v1/models")    if isinstance(models, dict) and isinstance(models.get("data"), list):        return "vllm"    return None  def resolve_framework(    requested: object,    *,    input_path: Optional[Path] = None,    url: Optional[str] = None,    server_args: Optional[dict] = None,) -> str:    explicit = canonicalize_framework(requested)    if explicit != "auto":        return explicit    for hint in (        detect_framework_from_server_args(server_args),        detect_framework_from_path(input_path) if input_path else None,        (            detect_framework_from_url(url, str(input_path) if input_path else None)            if url            else None        ),    ):        if hint:            return hint    return "sglang"  def parse_stage(path: Path) -> str:    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 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"  def parse_tp_rank(path: Path) -> Optional[int]:    for pattern in (        r"(?:^|[_-])tp(\d+)(?:[_.-]|$)",        r"TP-(\d+)",        r"(?:^|[_-])rank(\d+)(?:[_.-]|$)",        r"(?:^|[_-])worker(\d+)(?:[_.-]|$)",    ):        match = re.search(pattern, path.name, re.IGNORECASE)        if match:            return int(match.group(1))    return None  def file_looks_like_trace(path: Path) -> bool:    name = path.name.lower()    if name in TRACE_FILE_IGNORE_NAMES:        return False    if path.is_dir():        return False    if any(name.endswith(suffix) for suffix in (".trace.json", ".trace.json.gz")):        return True    if ".pt.trace.json" in name:        return True    if not any(name.endswith(suffix) for suffix in (".json", ".json.gz")):        return False    try:        trace = load_trace_json(path)    except Exception:        return False    if isinstance(trace, dict):        return isinstance(trace.get("traceEvents"), list)    if isinstance(trace, list):        return bool(trace) and all(isinstance(item, dict) for item in trace[:8])    return False  def discover_trace_files(    path: Path,    *,    recursive: bool,    limit: Optional[int] = None,) -> List[Path]:    if path.is_file():        return [path] if file_looks_like_trace(path) else []     candidates: List[Path] = []    seen: set[Path] = set()    for pattern in TRACE_FILE_PATTERNS:        iterator = path.rglob(pattern) if recursive else path.glob(pattern)        for candidate in iterator:            resolved = candidate.resolve()            if resolved in seen:                continue            seen.add(resolved)            candidates.append(resolved)    candidates = [        candidate        for candidate in candidates        if candidate.exists() and file_looks_like_trace(candidate)    ]    candidates.sort(key=lambda item: item.stat().st_mtime)    if limit is not None and limit >= 0:        return candidates[-limit:] if limit else []    return candidates  def newest_trace_dir(path: Path) -> Path:    if path.is_file():        return path.parent    direct = discover_trace_files(path, recursive=False)    if direct:        return path    traces = discover_trace_files(path, recursive=True)    trace_dirs = list({trace.parent for trace in traces})    if not trace_dirs:        raise FileNotFoundError(f"No trace files found under {path}")    trace_dirs.sort(        key=lambda item: max(            trace.stat().st_mtime for trace in traces if trace.parent == item        )    )    return trace_dirs[-1]  def discover_trace_targets(    path: Path, all_traces: bool) -> Tuple[List[Path], Optional[dict]]:    if path.is_file():        return [path], load_server_args(path)     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}")     non_merged = [trace for trace in traces if not trace.name.startswith("merged-")]    selected = non_merged or traces    if not all_traces:        ranks = sorted(            {                rank                for rank in (parse_tp_rank(trace) for trace in selected)                if rank is not None            }        )        if ranks:            rank = 0 if 0 in ranks else ranks[0]            selected = [trace for trace in selected if parse_tp_rank(trace) == rank]        grouped: Dict[str, List[Path]] = defaultdict(list)        for trace in selected:            grouped[parse_stage(trace)].append(trace)        selected = [            sorted(group, key=lambda item: item.stat().st_mtime)[-1]            for group in grouped.values()        ]     selected.sort(key=lambda item: (STAGE_ORDER.get(parse_stage(item), 99), item.name))    return selected, load_server_args(trace_dir)  def post_json(    url: str, payload: Optional[dict] = None, timeout: float = 60.0) -> Optional[dict]:    req = request.Request(        url=url,        data=(None if payload is None else json.dumps(payload).encode("utf-8")),        headers={"Content-Type": "application/json"},        method="POST",    )    with request.urlopen(req, timeout=timeout) as response:        raw = response.read()    return json.loads(raw.decode("utf-8")) if raw else None  def send_probe_request(    url: str,    prompt: str,    max_new_tokens: int,    sampling_seed: int,    framework: str,    model: Optional[str] = None,) -> None:    framework = canonicalize_framework(framework)    if framework == "sglang":        payload = {            "text": prompt,            "sampling_params": {                "sampling_seed": sampling_seed,                "temperature": 0.0,                "max_new_tokens": max_new_tokens,            },            "stream": False,        }        post_json(url.rstrip("/") + "/generate", payload, timeout=300.0)        return     resolved_model = model or discover_openai_model(url)    chat_payload = {        "model": resolved_model,        "messages": [{"role": "user", "content": prompt}],        "temperature": 0.0,        "max_tokens": max_new_tokens,        "stream": False,    }    try:        post_json(url.rstrip("/") + "/v1/chat/completions", chat_payload, timeout=300.0)        return    except Exception:        completion_payload = {            "model": resolved_model,            "prompt": prompt,            "temperature": 0.0,            "max_tokens": max_new_tokens,            "stream": False,        }        post_json(            url.rstrip("/") + "/v1/completions",            completion_payload,            timeout=300.0,        )  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):        raise RuntimeError(f"Could not read {url.rstrip('/')}/v1/models")    data = payload.get("data")    if not isinstance(data, list) or not data:        raise RuntimeError(f"No models returned by {url.rstrip('/')}/v1/models")    first = data[0]    if isinstance(first, dict) and first.get("id"):        return str(first["id"])    raise RuntimeError(f"Malformed /v1/models payload from {url.rstrip('/')}")  def ensure_remote_profiler_output_path(    output_dir: Optional[str], framework: str) -> Path:    if not output_dir:        raise ValueError(            f"{framework_display_name(framework)} live capture requires --output-dir "            "to point at the server-side torch profiler trace path that is visible "            "from this machine."        )    output_path = Path(output_dir).expanduser().resolve()    if output_path.suffix in {".json", ".gz"}:        output_path.parent.mkdir(parents=True, exist_ok=True)    else:        output_path.mkdir(parents=True, exist_ok=True)    return output_path  def wait_for_profiler_artifact(path: Path, timeout_s: float = 60.0) -> Path:    deadline = time.time() + timeout_s    while time.time() < deadline:        if path.is_file() and file_looks_like_trace(path):            return path        if path.exists():            trace_files = discover_trace_files(path, recursive=True)            if trace_files:                return newest_trace_dir(path)            if path.is_dir():                child_dirs = [item for item in path.iterdir() if item.is_dir()]                if child_dirs:                    child_dirs.sort(key=lambda item: item.stat().st_mtime)                    newest_child = child_dirs[-1]                    child_traces = discover_trace_files(newest_child, recursive=True)                    if child_traces:                        return newest_child        time.sleep(0.5)    return path  def start_remote_profiler(    url: str, framework: str, payload: Optional[dict] = None) -> None:    try:        post_json(url.rstrip("/") + "/start_profile", payload=payload, timeout=60.0)    except Exception as exc:        if framework == "vllm":            raise RuntimeError(                "vLLM live torch profiling requires the server to be launched with "                '--profiler-config \'{"profiler":"torch","torch_profiler_dir":"..."}\' '                "and to expose POST /start_profile."            ) from exc        if framework == "trtllm":            raise RuntimeError(                "TensorRT-LLM live torch profiling requires "                "a server build that exposes POST /start_profile plus the env vars "                "TLLM_PROFILE_START_STOP=<start>-<stop> and "                "TLLM_TORCH_PROFILE_TRACE=/shared/path."            ) from exc        if framework == "tokenspeed":            raise RuntimeError(                "TokenSpeed live torch profiling requires a server build that "                "exposes POST /start_profile and POST /stop_profile. The helper "                "passes output_dir, activities, and profile_id in the start payload."            ) from exc        raise  def build_remote_profiler_start_payload(    framework: str,    output_path: Path,    profile_prefix: Optional[str],    stage: Optional[str],) -> Optional[dict]:    if framework != "tokenspeed":        return None     profile_id = profile_prefix or "triage-trace"    if stage:        profile_id = f"{profile_id}-{stage}"     return {        "output_dir": str(output_path),        "activities": ["CPU", "GPU"],        "with_stack": True,        "record_shapes": False,        "profile_id": profile_id,    }  def stop_remote_profiler(url: str, framework: str) -> None:    try:        post_json(url.rstrip("/") + "/stop_profile", timeout=300.0)    except Exception as exc:        raise RuntimeError(            f"Failed to stop {framework_display_name(framework)} profiler via "            f"{url.rstrip('/')}/stop_profile"        ) from exc  def run_remote_profiler(    url: str,    output_dir: Optional[str],    framework: str,    probe_plan: ProbePlan,    probe_delay: float,    profile_prefix: Optional[str] = None,    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", "tokenspeed"}        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_payload = build_remote_profiler_start_payload(        framework=framework,        output_path=output_path,        profile_prefix=profile_prefix,        stage=stage,    )    start_remote_profiler(url, framework, payload=start_payload)    stop_error: Optional[BaseException] = None    try:        if probe_plan.capture_requests > 0:            # Server-side profilers may do setup work after POST /start_profile.            # A very short delay can send probes too early and miss the window.            time.sleep(max(5.0, probe_delay))            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,            )    finally:        try:            stop_remote_profiler(url, framework)        except BaseException as exc:  # pragma: no cover - preserve original failure            stop_error = exc    if stop_error is not None:        raise stop_error    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(    url: str,    output_dir: Optional[str],    num_steps: int,    profile_by_stage: bool,    merge_profiles: bool,    profile_prefix: Optional[str],    probe_plan: ProbePlan,    probe_delay: float,    start_step: Optional[int] = None,) -> Path:    if output_dir is None:        output_dir = tempfile.mkdtemp(prefix="sglang-torch-profile-")    output_root = Path(output_dir).resolve()    output_root.mkdir(parents=True, exist_ok=True)    output_path = output_root / str(time.time())    output_path.mkdir(parents=True, exist_ok=True)     server_args = try_get_json(url.rstrip("/") + "/server_info", timeout=60.0)    if server_args is not None:        with open(output_path / "server_args.json", "w", encoding="utf-8") as handle:            json.dump(server_args, handle)     payload = {        "output_dir": str(output_path),        "num_steps": str(num_steps),        "activities": ["CPU", "GPU"],        "profile_by_stage": profile_by_stage,        "merge_profiles": merge_profiles,        "profile_prefix": profile_prefix,    }    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"),        headers={"Content-Type": "application/json"},    )    with request.urlopen(req, timeout=300.0):        pass     if probe_plan.capture_requests > 0:        time.sleep(max(0.0, probe_delay))        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)  def run_profiler(    url: str,    output_dir: Optional[str],    num_steps: int,    profile_by_stage: bool,    merge_profiles: bool,    profile_prefix: Optional[str],    probe_requests: int,    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,        url=url,        input_path=(            Path(framework_hint_path).expanduser().resolve()            if framework_hint_path            else None        ),    )    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,            num_steps=num_steps,            profile_by_stage=profile_by_stage,            merge_profiles=merge_profiles,            profile_prefix=profile_prefix,            probe_plan=legacy_plan,            probe_delay=probe_delay,            start_step=start_step,        )    if start_step is not None:        raise ValueError("--start-step is only supported for SGLang live capture.")    if profile_by_stage:        raise ValueError(            "--profile-by-stage is only supported for SGLang live capture. "            "Disable it when profiling vLLM, TensorRT-LLM, or TokenSpeed."        )    if merge_profiles:        raise ValueError(            "--merge-profiles is only supported for SGLang live capture. "            "Disable it when profiling vLLM, TensorRT-LLM, or TokenSpeed."        )    if profile_prefix and resolved_framework in {"vllm", "trtllm"}:        print(            f"Note: {framework_display_name(resolved_framework)} ignores "            "--profile-prefix on the HTTP profiler control path.",            file=sys.stderr,        )    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,            profile_prefix=profile_prefix,        )    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,            profile_prefix=profile_prefix,            stage=stage,        )    return output_root  def select_heaviest_pid(    events: Sequence[dict],    event_filter: Callable[[dict], bool],    pid_substring: Optional[str] = None,    preferred_substrings: Iterable[str] = (),) -> Optional[str]:    durations: Counter = Counter()    for event in events:        if not event_filter(event):            continue        pid = str(event.get("pid"))        if pid_substring and pid_substring not in pid:            continue        durations[pid] += float(event["dur"])    if not durations:        return None     for substring in preferred_substrings:        preferred = [pid for pid in durations if substring in pid]        if preferred:            return max(preferred, key=lambda pid: durations[pid])    return max(durations, key=lambda pid: durations[pid])
← All SGLang skillsProject instructions →