From 7824903417b7398ffaf9befe8a221080627e152f Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Tue, 28 Apr 2026 17:05:36 +0800 Subject: [PATCH] [SKILL] Sync SGLang skill docs (#23921) --- .claude/skills/add-jit-kernel/SKILL.md | 8 +- .claude/skills/add-sgl-kernel/SKILL.md | 4 + .claude/skills/ci-workflow-guide/SKILL.md | 13 +- .claude/skills/generate-profile/SKILL.md | 6 +- .../llm-torch-profiler-analysis/SKILL.md | 330 +++++ .../references/fuse-overlap-catalog.md | 65 +- .../references/heuristics.md | 0 .../references/overlap-catalog.md | 40 +- .../references/source-map.md | 0 .../scripts/analyze_llm_torch_profile.py | 806 ++++++++++++ .../scripts/analyze_sglang_torch_profile.py | 16 + .../make_trtllm_py_executor_override.py | 132 ++ .../scripts/probe_llm_server.py | 230 ++++ .../scripts/profile_common.py | 880 +++++++++++++ .../scripts/render_triage_markdown_bundle.py | 259 ++++ .../scripts/triage_kernel_helpers.py | 1085 +++++++++++----- .../scripts/triage_overlap_helpers.py | 1097 ++++++++++------- .../sglang-bisect-ci-regression/SKILL.md | 5 + .../sglang-torch-profiler-analysis/SKILL.md | 189 --- .../scripts/analyze_sglang_torch_profile.py | 601 --------- .../scripts/profile_common.py | 333 ----- .claude/skills/write-sglang-test/SKILL.md | 8 +- .../sglang-diffusion-add-model/SKILL.md | 26 +- .../references/testing-and-accuracy.md | 10 +- .../sglang-diffusion-ako4all-kernel/SKILL.md | 5 + 25 files changed, 4199 insertions(+), 1949 deletions(-) create mode 100644 .claude/skills/llm-torch-profiler-analysis/SKILL.md rename .claude/skills/{sglang-torch-profiler-analysis => llm-torch-profiler-analysis}/references/fuse-overlap-catalog.md (88%) rename .claude/skills/{sglang-torch-profiler-analysis => llm-torch-profiler-analysis}/references/heuristics.md (100%) rename .claude/skills/{sglang-torch-profiler-analysis => llm-torch-profiler-analysis}/references/overlap-catalog.md (86%) rename .claude/skills/{sglang-torch-profiler-analysis => llm-torch-profiler-analysis}/references/source-map.md (100%) create mode 100644 .claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py create mode 100644 .claude/skills/llm-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py create mode 100644 .claude/skills/llm-torch-profiler-analysis/scripts/make_trtllm_py_executor_override.py create mode 100755 .claude/skills/llm-torch-profiler-analysis/scripts/probe_llm_server.py create mode 100644 .claude/skills/llm-torch-profiler-analysis/scripts/profile_common.py create mode 100644 .claude/skills/llm-torch-profiler-analysis/scripts/render_triage_markdown_bundle.py rename .claude/skills/{sglang-torch-profiler-analysis => llm-torch-profiler-analysis}/scripts/triage_kernel_helpers.py (68%) rename .claude/skills/{sglang-torch-profiler-analysis => llm-torch-profiler-analysis}/scripts/triage_overlap_helpers.py (59%) delete mode 100644 .claude/skills/sglang-torch-profiler-analysis/SKILL.md delete mode 100644 .claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py delete mode 100644 .claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index 739651de2..e63a7d77b 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -435,7 +435,7 @@ if torch.cuda.get_device_capability()[0] < 9: JIT kernel tests live under `python/sglang/jit_kernel/tests/`. **CI does not run `pytest` in that directory directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` there (and every `bench_*.py` under `benchmark/`), collects `register_*_ci(...)` calls by **statically parsing each file's AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check. -- **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `stage-b-kernel-unit-1-gpu-large` (see `.github/workflows/pr-test-jit-kernel.yml`: `python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large`). +- **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `stage-b-kernel-unit-1-gpu-large` on H100 and `stage-b-kernel-unit-1-gpu-b200` on B200/SM100 paths (see `.github/workflows/pr-test-jit-kernel.yml`). Multi-GPU JIT tests use `stage-b-kernel-unit-8-gpu-h200`. - **Nightly kernel suite**: `nightly-kernel-1-gpu` with `--nightly` — typically used with `SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1` in CI for expanded parameter grids (see `python/sglang/jit_kernel/utils.py` → `should_run_full_tests` / `get_ci_test_range`). Wired in `.github/workflows/nightly-test-nvidia.yml` (e.g. `python3 run_suite.py --hw cuda --suite nightly-kernel-1-gpu --nightly --continue-on-error`). Registration pattern (module level, **literal** `est_time` and `suite` strings — required for AST parsing): @@ -444,6 +444,8 @@ Registration pattern (module level, **literal** `est_time` and `suite` strings from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large") +# Optional B200/SM100 registration for tests that cover Blackwell-specific code paths +# register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-b200") # Optional second registration: same file also listed under the nightly kernel suite # register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) ``` @@ -455,7 +457,9 @@ Use `register_cuda_ci(..., disabled="reason")` if the file must stay in-tree but **Run like CI** (from repo root): ```bash -cd test && python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large +(cd test && python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large) +# For B200/SM100-specific coverage: +(cd test && python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-b200) ``` For fast iteration you can still run `pytest` on a single file locally; CI coverage is via `run_suite.py`. diff --git a/.claude/skills/add-sgl-kernel/SKILL.md b/.claude/skills/add-sgl-kernel/SKILL.md index 8f6c4639b..559b8751f 100644 --- a/.claude/skills/add-sgl-kernel/SKILL.md +++ b/.claude/skills/add-sgl-kernel/SKILL.md @@ -328,6 +328,10 @@ pytest sgl-kernel/tests/test_scale.py -q python sgl-kernel/benchmark/bench_scale.py ``` +PR CI also runs `pr-test-sgl-kernel.yml`, including the B200 job +`sgl-kernel-b200-test` when kernel changes are detected. Use that job as the +Blackwell coverage signal for AOT `sgl-kernel` changes. + --- ## Troubleshooting diff --git a/.claude/skills/ci-workflow-guide/SKILL.md b/.claude/skills/ci-workflow-guide/SKILL.md index 8877ba300..f315d17a4 100644 --- a/.claude/skills/ci-workflow-guide/SKILL.md +++ b/.claude/skills/ci-workflow-guide/SKILL.md @@ -276,6 +276,7 @@ Large suites are split across matrix jobs using the **LPT (Longest Processing Ti | `stage-b-test-2-gpu-large` | 4 | `2-gpu-h100` | — | | `stage-b-test-4-gpu-b200` | 1 (no matrix) | `4-gpu-b200` | — | | `stage-b-kernel-unit-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — | +| `stage-b-kernel-unit-1-gpu-b200` | 1 (no matrix) | `4-gpu-b200` | — | | `stage-b-kernel-unit-8-gpu-h200` | 1 (no matrix) | `8-gpu-h200` | — | | `stage-b-kernel-benchmark-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — | | `stage-c-test-4-gpu-h100` | 3 | `4-gpu-h100` | — | @@ -283,10 +284,12 @@ Large suites are split across matrix jobs using the **LPT (Longest Processing Ti | `stage-c-test-8-gpu-h20` | 2 | `8-gpu-h20` | — | | `stage-c-test-deepep-4-gpu-h100` | 1 (no matrix) | `4-gpu-h100` | — | | `stage-c-test-deepep-8-gpu-h200` | 1 (no matrix) | `8-gpu-h200` | — | -| `stage-c-test-4-gpu-b200` | 4 | `4-gpu-b200` | — | -| `stage-c-test-4-gpu-gb200` | 1 (no matrix) | `4-gpu-gb200` | — | +| `stage-c-test-4-gpu-b200` | 3 | `4-gpu-b200` | — | +| `stage-c-test-4-gpu-b200-small` | 3 | `4-gpu-b200-low-disk` | — | +| `stage-c-test-8-gpu-b200` | registered only | `8-gpu-b200` | — | +| `stage-c-test-4-gpu-gb200` | registered only | `4-gpu-gb200` | — | -> **Note**: Kernel suites (`stage-b-kernel-*`) run via `pr-test-jit-kernel.yml` and `pr-test-sgl-kernel.yml`, not the main `pr-test.yml`. Multimodal diffusion uses `python/sglang/multimodal_gen/test/run_suite.py`, not `test/run_suite.py`. +> **Note**: Kernel suites (`stage-b-kernel-*`) run via `pr-test-jit-kernel.yml` and `pr-test-sgl-kernel.yml`, not the main `pr-test.yml`. `stage-c-test-8-gpu-b200` is registered in `test/run_suite.py` but not wired to PR CI. The GB200 job is currently commented out in `pr-test.yml` until a company-owned runner is provisioned. Multimodal diffusion uses `python/sglang/multimodal_gen/test/run_suite.py`, not `test/run_suite.py`. **Workflow usage:** ```yaml @@ -317,11 +320,11 @@ Determines which test suites to run based on file changes. | Output | Triggers | |--------|----------| | `main_package` | Stage A/B/C test suites | -| `sgl_kernel` | Kernel wheel builds + kernel test suites | +| `sgl_kernel` | Kernel wheel builds + kernel test suites; also switches B200 jobs to kernel-build runner labels outside `target_stage` mode | | `jit_kernel` | JIT kernel test workflow | | `multimodal_gen` | Multimodal-gen test workflow | -> **Note**: `sgl_kernel` is forced to `false` when `target_stage` is set, because `sgl-kernel-build-wheels` won't run and wheel artifacts won't be available. +> **Note**: In `target_stage` mode, `sgl_kernel` is only active when `include_wheel_build=true`. Without that opt-in, kernel-change reruns fail validation instead of running a target stage without freshly built wheels. Outside `target_stage`, `sgl_kernel=true` switches B200 jobs from `4-gpu-b200` / `4-gpu-b200-low-disk` to `4-gpu-b200-kernel` / `4-gpu-b200-kernel-low-disk`. --- diff --git a/.claude/skills/generate-profile/SKILL.md b/.claude/skills/generate-profile/SKILL.md index 2b6201f6b..dae475cfa 100644 --- a/.claude/skills/generate-profile/SKILL.md +++ b/.claude/skills/generate-profile/SKILL.md @@ -46,11 +46,13 @@ Typical startup time: 30-90 seconds depending on model size and whether CUDA gra ### Step 3: Validate accuracy (sanity check) ```bash -python3 -m sglang.test.few_shot_gsm8k --num-q 20 +python3 -m sglang.test.run_eval --host 127.0.0.1 --port --eval-name gsm8k --num-examples 20 ``` - Expected accuracy: **> 0.8** for capable models (Qwen3-8B, Llama-3.1-8B-Instruct, etc.) - This is a quick sanity check, not a rigorous benchmark. +- `sglang.test.few_shot_gsm8k` is deprecated; use the unified `run_eval` entrypoint. +- If you intentionally need the old completion-style GSM8K path, add `--api completion`. - If accuracy is unexpectedly low, something is wrong — do not proceed to profiling. ### Step 4: Generate the profile @@ -108,7 +110,7 @@ for i in $(seq 1 120); do done # 3. Accuracy check -python3 -m sglang.test.few_shot_gsm8k --num-q 20 +python3 -m sglang.test.run_eval --host 127.0.0.1 --port 30000 --eval-name gsm8k --num-examples 20 # Expected: Accuracy > 0.8 # 4. Profile diff --git a/.claude/skills/llm-torch-profiler-analysis/SKILL.md b/.claude/skills/llm-torch-profiler-analysis/SKILL.md new file mode 100644 index 000000000..60f99e6bc --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/SKILL.md @@ -0,0 +1,330 @@ +--- +name: llm-torch-profiler-analysis +description: "Unified LLM torch-profiler triage skill for `sglang`, `vllm`, and `TensorRT-LLM`. Use it to inspect an existing `trace.json(.gz)` or profile directory, or to drive live profiling against a running server and return one three-table report with kernel, overlap-opportunity, and fuse-pattern tables." +--- + +# Unified LLM Torch Profiler Analysis + +## Overview + +Use this skill for `torch.profiler` analysis across: + +- `sglang` +- `vllm` +- `TensorRT-LLM` + +There is only one public workflow: + +- `triage` + +Preferred unified entrypoint: + +- [scripts/analyze_llm_torch_profile.py](scripts/analyze_llm_torch_profile.py) + +Backwards-compatibility shim (kept so older `docker exec ... analyze_sglang_torch_profile.py ...` calls keep working; it just forwards to the unified entrypoint): + +- [scripts/analyze_sglang_torch_profile.py](scripts/analyze_sglang_torch_profile.py) + +Markdown bundling helper: + +- [scripts/render_triage_markdown_bundle.py](scripts/render_triage_markdown_bundle.py) + +`triage` always prints the same three tables: + +- kernel table +- overlap-opportunity table +- fuse-pattern table + +By default, all three tables only render rows at or above `1.0%` cumulative GPU-time share. +Rows below that are hidden by default unless the user asks for a lower cutoff. + +Keep the fuse-pattern table source-backed and deterministic. +Do not turn it into a fuzzy matcher. + +If exact source-backed matching is weak but a kernel cluster is still close to a known family, +add one short note after the tables with exactly one of: + +- `high` +- `medium` +- `low` + +## Capability Matrix + +| Capability | SGLang | vLLM | TensorRT-LLM | +| --- | --- | --- | --- | +| Existing trace triage | yes | yes | yes | +| Single-trace live capture | yes | yes, if torch profiler is enabled on server | requires profiler control endpoints | +| Two-trace mapping+formal triage | yes | yes | yes | +| Stage-aware live capture | yes | no | no | +| `--profile-prefix` control | yes | usually ignored on HTTP profiler route | usually ignored on HTTP profiler route | + +For TensorRT-LLM, live capture only works when the server exposes `/start_profile` and +`/stop_profile`, and when the deployment already provides a shared trace path plus the +required env vars. + +## Validation Notes + +This unified workflow has been validated with a `4x H100` matrix across SGLang, +vLLM, and TensorRT-LLM. Use these model shapes as representative coverage when +refreshing or extending the skill: + +| Model | SGLang | vLLM | TensorRT-LLM | Result | +| --- | --- | --- | --- | --- | +| `mistralai/Mixtral-8x7B-Instruct-v0.1` | `4x H100` | `4x H100` | `4x H100` | three tables rendered correctly on all three frameworks; benchmark probes returned direct, non-empty text | +| `Qwen/Qwen2.5-32B-Instruct` | `4x H100` | `4x H100` | `4x H100` | three tables rendered correctly on all three frameworks; benchmark probes returned direct, non-empty text | +| `Qwen/Qwen3-32B` | `4x H100` | `4x H100` | `4x H100` | three tables rendered correctly on all three frameworks; vLLM and TensorRT-LLM chat probes often emitted `` prefixes | + +To render a validated run into one markdown document: + +```bash +python3 scripts/render_triage_markdown_bundle.py \ + --analysis-root /path/to/analysis_root \ + --output /path/to/analysis_bundle.md +``` + +The bundle groups by model and keeps the three tables for each framework. + +Validation notes: + +- all three frameworks now render kernel, overlap, and fuse tables with separate `extend/prefill` and `decode` sections when the trace contains a clean stage split +- SGLang live capture is validated and calls the server profiler API directly instead of shelling out to `sglang.profiler` +- SGLang trace flush can lag well beyond a few seconds, so the runner waits longer for artifacts than the earlier implementation +- SGLang kernel-site reconstruction keeps sampling disabled in the mapping path so the optimized parser does not perturb SGLang table output; equality rechecks matched for `Mixtral-8x7B-Instruct-v0.1`, `Qwen3-32B`, and `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8` +- vLLM live capture requires `--output-dir` to match the server `torch_profiler_dir`; the validated H100 flow uses `--profiler-config {"profiler":"torch","torch_profiler_dir":"..."}` and then drives `/start_profile` and `/stop_profile` +- TensorRT-LLM validation stays on `--backend pytorch`; the H100 flow writes the trace with `TLLM_TORCH_PROFILE_TRACE` and then analyzes the saved trace +- the 2026-04-22 TensorRT-LLM 1.0.0 `py_executor.py` profiler setup still needed a `with_stack=True` override for table-quality Python locations; re-check this on TensorRT-LLM 1.2.1 or any 1.3.x release-candidate image before assuming the override is still required + +## When To Use It + +- inspect a `torch.profiler` trace or profile directory from `sglang`, `vllm`, or `TensorRT-LLM` +- profile a live serving endpoint and analyze the result +- summarize which kernel families dominate prefill or decode +- map kernels back to Python code paths +- judge whether a code path still leaves overlap opportunity +- check whether an already-known fusion or overlap path should have applied + +## Diffusion Backend Gate + +For diffusion benchmark or profiling work, only analyze traces produced by the native +SGLang diffusion backend. + +If the run that generated the trace logs any of: + +- `Falling back to diffusers backend` +- `Using diffusers backend` +- `Loaded diffusers pipeline` + +stop the workflow instead of analyzing the trace. +Handle it as a backend-selection issue, not as native-kernel profiler evidence. + +## Main Flows + +### 1. Single-trace triage from an existing profile dir or trace + +```bash +python3 scripts/analyze_llm_torch_profile.py \ + --input /path/to/profile_dir_or_trace.json.gz +``` + +Use this when one trace is enough. +The overlap table stays conservative in single-trace mode and will tell you when a +mapping/formal pair is needed. + +### 2. Single-trace live capture from SGLang + +```bash +python3 scripts/analyze_llm_torch_profile.py \ + --framework sglang \ + --url http://127.0.0.1:30000 \ + --output-dir /tmp/llm-profiler/sglang_profile_live \ + --num-steps 5 \ + --profile-by-stage +``` + +The script sends `POST /start_profile` to the SGLang server directly. +The script writes `server_args.json`, sends the probe requests after profiling is armed, +and waits longer for trace flush than the earlier implementation. + +### 3. Single-trace live capture from vLLM + +Launch vLLM with torch profiler enabled, for example: + +```bash +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --profiler-config '{"profiler":"torch","torch_profiler_dir":"/tmp/llm-profiler/vllm_profile"}' +``` + +Then run: + +```bash +python3 scripts/analyze_llm_torch_profile.py \ + --framework vllm \ + --url http://127.0.0.1:8000 \ + --output-dir /tmp/llm-profiler/vllm_profile \ + --num-steps 5 \ + --no-profile-by-stage +``` + +For vLLM, `--output-dir` must point to the same `torch_profiler_dir` the server uses. +The current vLLM profiler config already defaults `torch_profiler_with_stack=true`, +so the runner only needs to set `torch_profiler_dir`. + +### 4. Single-trace live capture from TensorRT-LLM + +Use this only when the server exposes `POST /start_profile` and `POST /stop_profile`, +and the trace path is shared with the current machine. + +Typical env expectations are: + +- `TLLM_PROFILE_START_STOP=1` +- `TLLM_TORCH_PROFILE_TRACE=/shared/path/trace.json` or `.json.gz` + +Then run: + +```bash +python3 scripts/analyze_llm_torch_profile.py \ + --framework trtllm \ + --url http://127.0.0.1:8000 \ + --output-dir /shared/path \ + --num-steps 5 \ + --no-profile-by-stage +``` + +If the deployment does not expose the profiler control endpoints, fall back to analyzing +an existing trace instead of trying live capture. + +On the current TensorRT-LLM mainline path, `py_executor.py` creates the torch profiler +with `record_shapes=True` and `with_modules=True` but not `with_stack=True`. +For table-quality validation, use the override generator: + +```bash +python3 scripts/make_trtllm_py_executor_override.py \ + --source /path/to/original/py_executor.py \ + --output /tmp/llm-profiler/py_executor_with_stack.py +``` + +The validated TensorRT-LLM flow is: + +1. launch `trtllm-serve` with `TLLM_TORCH_PROFILE_TRACE=/shared/path/trace.json` +2. run a few benchmark requests +3. analyze the emitted trace with `--input /shared/path/trace.json` + +### 5. Two-trace triage from existing profile dirs or traces + +```bash +python3 scripts/analyze_llm_torch_profile.py triage \ + --mapping-input /path/to/graph_off_profile_dir \ + --formal-input /path/to/graph_on_profile_dir +``` + +Use this when you need stronger overlap attribution and kernel-to-source mapping. + +### 6. Two-trace triage from running servers + +```bash +python3 scripts/analyze_llm_torch_profile.py triage \ + --framework sglang \ + --mapping-url http://127.0.0.1:31025 \ + --formal-url http://127.0.0.1:31026 \ + --num-steps 5 \ + --profile-by-stage +``` + +For `vllm` or `TensorRT-LLM`, use the same shape but pass: + +- `--framework vllm` or `--framework trtllm` +- `--mapping-output-dir ...` +- `--formal-output-dir ...` +- `--no-profile-by-stage` + +## `profile_by_stage` + +`--profile-by-stage` is only meaningful on the SGLang live-capture path. + +- On ordinary non-PD SGLang serving, it is still useful because prefill and decode usually have very different bottlenecks. +- On the current profile-v2 path inside SGLang, stage-based profiling is effectively the normal path. +- PD-disaggregated serving adds one extra rule: prefill workers and decode workers must be profiled separately. That is stricter than ordinary `profile_by_stage`. +- For `vllm` and `TensorRT-LLM`, disable it with `--no-profile-by-stage`. + +## How To Choose The Triage Shape + +### Single-trace triage + +Use when you want the lowest-friction report: + +- one trace is already available +- you mainly want kernel share and fusion clues +- you are comparing two runs side by side by running triage once per trace + +Prefer this by default. + +### Two-trace triage + +Use when you need: + +- a stronger overlap answer +- graph-off source mapping plus graph-on final behavior +- more trustworthy overlap recommendations in the middle table + +1. mapping trace with graph disabled or with the lower-fusion / more-readable config +2. formal trace with the real serving optimizations enabled + +Do not call the mapping pass a "fast profile". +It exists to recover `kernel -> cpu_op -> python scope`. + +## Workflow + +### Single-trace workflow + +1. If the user only wants a diagnosis, one trace is enough. +2. Prefer one-rank traces over merged traces whenever the profiler emitted both. +3. For a live server, let the script drive the profiler only when the framework-specific prerequisites are already met. +4. Prefer SGLang `--profile-by-stage` unless the user explicitly wants an all-stage mixed trace. +5. Create or clean the target trace directory before live capture so the profiler can write artifacts without permission surprises. + +### Two-trace workflow + +1. Produce a mapping trace first with graph disabled or the lower-fusion configuration. +2. Produce a formal trace second with the real serving optimizations enabled. +3. Run `triage` for the three-table report. +4. Read the results in this order: + - kernel table + - overlap-opportunity table + - fuse-pattern table +5. Before calling something a "new" optimization idea, compare the top rows against both [references/fuse-overlap-catalog.md](references/fuse-overlap-catalog.md) and [references/overlap-catalog.md](references/overlap-catalog.md). Check mainline rows first, then the `PR-backed / in-flight` sections. Prefer reporting: + - an existing fused or overlap path that should already apply here + - an existing path that appears disabled, unsupported, or regressed in this trace + - an upstream pattern that is mainline elsewhere but missing locally, or still open upstream + - a truly new opportunity only when no catalog entry fits +6. If no exact pattern fully matches but the trace is still close to a known family, add one flat similarity note after the tables. + Use `high`, `medium`, or `low` only. + Base that note on the full pattern shape, not on one kernel name alone. + Prefer semantic cues such as producer-consumer chain, source locations, CPU op names, TP context, and model-specific structure. + Do not rewrite the script table itself to include these heuristic judgments. + +## References + +Load these only when needed: + +- [references/source-map.md](references/source-map.md) + - upstream SGLang profiler entrypoints and trace-writing paths; still most useful for SGLang-specific source follow-up +- [references/heuristics.md](references/heuristics.md) + - overlap labels, dependency-risk interpretation, and limits +- [references/fuse-overlap-catalog.md](references/fuse-overlap-catalog.md) + - mixed source-backed catalog of existing fuse and overlap patterns, including mainline rows plus PR-backed / in-flight rows +- [references/overlap-catalog.md](references/overlap-catalog.md) + - overlap-only lookup table across LLM, VLM, diffusion, disaggregation, HiSparse, and speculative scheduling + +## Output Contract + +Return: + +- trace path or generated profile path +- framework +- model/server args when available +- kernel table +- overlap-opportunity table +- fuse-pattern table +- optional similarity note with `high` / `medium` / `low` when exact matching is inconclusive +- one short summary of what dominates the run +- whether the overlap read came from single-trace triage or mapping/formal two-trace triage diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/fuse-overlap-catalog.md b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md similarity index 88% rename from .claude/skills/sglang-torch-profiler-analysis/references/fuse-overlap-catalog.md rename to .claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md index 1b2e0c5fe..f4e67a926 100644 --- a/.claude/skills/sglang-torch-profiler-analysis/references/fuse-overlap-catalog.md +++ b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md @@ -18,7 +18,7 @@ Use it like this: 3. If a finding matches an existing row, report it as: - an existing optimization path that is missing, disabled, regressed, or unsupported for the current backend, or - an already-known family that should be re-applied to the current model shape. -4. Check the `PR-backed / in-flight` sections too. If a match exists there, do not call it novel; call it an upstream or in-flight pattern instead. +4. Check the mainline comparison sections and the `PR-backed / in-flight` sections too. If a match exists there, do not call it novel; call it an upstream or in-flight pattern instead. 5. Only call a finding "new" when it does not fit any mainline or PR-backed row in this catalog. The `vLLM-origin` sections below are comparative references. They are not @@ -28,16 +28,33 @@ overlap opportunity as novel. The catalog is grouped by reusable optimization family, not by one specific model. +Refresh note `2026-04-22`: rescanned current `sglang`, `flashinfer`, +`TensorRT-LLM`, and `vllm` mainline plus rechecked referenced PR state via the +GitHub API on `2026-04-22`. Stable current-code families such as Qwen-style +shared-expert top-k append, TensorRT-LLM Triton fused add+RMSNorm+FP8 quant, +and vLLM `merge_attn_states` attention-output quant are folded into the +mainline rows below. Closed-unmerged SGLang +[#22410](https://github.com/sgl-project/sglang/pull/22410) and FlashInfer +[#2840](https://github.com/flashinfer-ai/flashinfer/pull/2840) were removed +from the PR-backed sections. Keep FlashInfer +[#3058](https://github.com/flashinfer-ai/flashinfer/pull/3058) / +[#3079](https://github.com/flashinfer-ai/flashinfer/pull/3079) in mind because +that branch was reverted, and keep vLLM +[#40057](https://github.com/vllm-project/vllm/pull/40057) in mind when using +B200 FP4 MoE test coverage as a signal: it disables some B200 FP4 MoE layer +tests rather than proving the kernel family is absent. + ## 1. LLM / SRT fused-kernel families | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| Fused residual add + RMSNorm | `fused_add_rmsnorm*`
`npu_add_rms_norm`
`add_rmsnorm_bias`
`gemma_fused_add_rmsnorm`
residual add right before norm | `python/sglang/srt/layers/layernorm.py`
`python/sglang/srt/layers/quantization/modelslim/modelslim.py` | Shared CUDA / ROCm / CPU / NPU fused add-RMSNorm implementations, including Gemma and NPU-bias variants | Treat split residual add + RMSNorm as an existing cross-backend fusion first, not a new idea. | +| Fused residual add + RMSNorm | `fused_add_rmsnorm*`
`npu_add_rms_norm`
`add_rmsnorm_bias`
`gemma_fused_add_rmsnorm`
`gemma_rmsnorm_residual_scalar`
`_gemma_rmsnorm_residual_kernel`
residual add right before norm | `python/sglang/srt/layers/layernorm.py`
`python/sglang/srt/layers/gemma4_fused_ops.py`
`python/sglang/srt/layers/quantization/modelslim/modelslim.py` | Shared CUDA / ROCm / CPU / NPU fused add-RMSNorm implementations, including Gemma, Gemma4 scalar-residual, and NPU-bias variants | Treat split residual add + RMSNorm as an existing cross-backend fusion first, not a new idea. | | FlashInfer unified `allreduce_fusion` | `cross_device_reduce_1stage*`
`all_reduce`
`FusedAddRMSNormKernel`
`rmsnorm*` | `python/sglang/srt/layers/flashinfer_comm_fusion.py`
`python/sglang/srt/layers/layernorm.py::forward_with_allreduce_fusion`
`python/sglang/srt/layers/communicator.py::apply_flashinfer_allreduce_fusion` | FlashInfer workspace creation plus `allreduce_fusion(..., pattern=AllReduceFusionPattern.kARResidualRMSNorm, ...)` | First suspect missing / disabled / unsupported FlashInfer allreduce fusion, not a brand new TP fusion idea. | | AITER allreduce fusion | ROCm all-reduce plus RMSNorm still split | `python/sglang/srt/layers/layernorm.py::forward_with_allreduce_fusion`
`python/sglang/srt/distributed/communication_op.py::tensor_model_parallel_fused_allreduce_rmsnorm`
`python/sglang/srt/layers/communicator.py::apply_aiter_all_reduce_fusion` | ROCm-side fused TP all-reduce + RMSNorm with fallback to plain all-reduce plus norm | On AMD, rule out existing AITER fusion before proposing a new communication fusion. | | Fused activation-and-mul (`SwiGLU` / `GeGLU`) | `silu_and_mul`
`gelu_and_mul`
`npu_swiglu` | `python/sglang/srt/layers/activation.py` | Single op covers activation plus elementwise multiply across CUDA / CPU / NPU / XPU backends | Treat separate activation + mul on packed MLP outputs as missing existing fusion. | | Fused dual residual RMSNorm | residual add plus two RMSNorm-like kernels around Grok blocks | `python/sglang/srt/layers/elementwise.py::fused_dual_residual_rmsnorm`
`python/sglang/srt/models/grok.py` | One Triton kernel computes intermediate residual update and next RMSNorm output together | On Grok-like residual layouts, treat split residual + norm as missing existing fusion. | | In-place QK RMSNorm | split `q_norm` / `k_norm` kernels | `python/sglang/srt/models/utils.py::apply_qk_norm`
`python/sglang/jit_kernel/norm.py::fused_inplace_qknorm` | In-place JIT QK norm plus optional `alt_stream` overlap for K | Check shape, dtype, deterministic mode, and in-place legality before proposing a new QK fuse. | +| TorchInductor horizontal Q/K norm combo-kernels | `combo_kernels`
`benchmark_combo_kernel`
`q_norm`
`k_norm`
`split_with_sizes` | `torch._inductor.config.combo_kernels` | TorchInductor can horizontally fuse sibling Q-norm and K-norm kernels in compiled traces, often deleting `split_with_sizes` / `clone` ladders | Treat separate Q/K norm ladders in compile-heavy traces as an existing compiler-fusion family first. | | MiniMax TP fused QK RMSNorm | `MiniMaxM2RMSNormTP`
`rms_sumsq_serial`
`rms_apply_serial`
`forward_qk` | `python/sglang/srt/models/minimax_m2.py` | Triton kernels compute Q / K sumsq together, TP all-reduces shared stats, then apply both RMSNorms together | On MiniMax traces, separate Q norm and K norm are usually a missed model-specific Triton fusion. | | Fused QK RMSNorm + RoPE | `qknorm*` + `rope*` + `rotary*` as separate steps | `python/sglang/jit_kernel/fused_qknorm_rope.py`
`python/sglang/srt/models/qwen3_moe.py` | One JIT kernel applies QK RMSNorm and RoPE in-place on packed QKV | For compatible LLMs, classify split QK norm + RoPE as a missing existing fusion. | | Fused QK RoPE reshape + KV cache write | `fused_qk_rope_reshape_and_cache*`
RoPE followed by reshape / cache DtoD | `python/sglang/srt/layers/attention/utils.py::fused_qk_rope_reshape_and_cache` | One Triton kernel applies RoPE to Q / K, reshapes cache layout, and writes K / V directly to paged cache | Treat separate RoPE + reshape + cache-write ladders as an existing attention-prep fusion family. | @@ -52,6 +69,7 @@ The catalog is grouped by reusable optimization family, not by one specific mode | Fused MLA KV cache write + FP8 quant | `set_mla_kv_buffer_fp8_quant*`
`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`
`python/sglang/srt/mem_cache/memory_pool.py` | MLA / NSA KV pool path can quantize K and write directly into KV storage without a separate concat-and-quant chain | Treat standalone quant + KV-buffer write on MLA paths as missing existing fusion first. | | Fused MoE router / top-k / softcapping | `FusedMoeRouter`
`fused_moe_router*`
router GEMM + `topk` + `tanh` | `python/sglang/srt/layers/moe/router.py` | Single fused router kernel covers router matmul, softcapping, and top-k selection | Treat exposed router matmul + softcap + top-k chains as an existing MoE fusion family. | | Fused MoE grouped-topk / gate kernels | `fused_topk_deepseek`
`moe_fused_gate`
`aiter_fused_topk`
`kimi_k2_moe_fused_gate` | `python/sglang/srt/layers/moe/topk.py` | CUDA / ROCm / FlashInfer kernels fuse bias, grouped-topk, renorm, and routed scaling into one gate op | Check backend / model eligibility before proposing a novel router-gate fusion. | +| Qwen-style shared-expert append into routed top-k output | `_append_shared_to_topk_output`
`fused_append_shared_experts_with_weights`
`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`
`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py` | Qwen-style MoE paths can append shared-expert ids and sigmoid gate weights to routed top-k output in one Triton kernel so the shared experts execute inside the fused MoE path | Treat routed top-k plus shared-expert pad / concat ladders as an existing MoE-prep fusion family first. | | Fused MoE dispatch / permute / combine | token permutation
dispatch / combine
grouped top-k
many small MoE support kernels | `python/sglang/srt/layers/moe/fused_moe_triton/layer.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | `FusedMoE` plus DeepEP / FlashInfer / FuseEP / standard dispatch backends and `permute_fusion=True` | First ask whether the model is missing an existing `FusedMoE`-style path or backend-specific dispatcher path. | | Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. | | Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/jit_kernel/nvfp4.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | @@ -102,7 +120,7 @@ The catalog is grouped by reusable optimization family, not by one specific mode | Fused add-RMSNorm and one-pass RMSNorm | residual add plus RMSNorm still split on short hidden sizes | `python/sglang/multimodal_gen/runtime/layers/layernorm.py`
`python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` | `fused_add_rmsnorm(...)` and `triton_one_pass_rms_norm(...)` | For short hidden-size diffusion blocks, this is already an established fusion family. | | Fused diffusion QK norm + RoPE | split QK norm and RoPE in diffusion attention blocks | `python/sglang/jit_kernel/diffusion/qknorm_rope.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py::apply_qk_norm_rope` | `fused_inplace_qknorm_rope(...)`, with fallback to QK norm plus `apply_flashinfer_rope_qk_inplace(...)` | Distinguish between missing fused qknorm + rope and the existing FlashInfer RoPE fallback. | | Z-Image fused `norm(x) * tanh(scale) + shift` | `fused_norm_tanh_mul_add`
`tanh(gate) * rmsnorm(x)` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | CuTeDSL kernel plus runtime helper for Z-Image residual-form modulation | Treat split Z-Image residual-form modulation as a missing existing diffusion fusion, not a novel idea. | -| Z-Image fused residual modulation + next norm-scale | `fused_norm_tanh_mul_add_norm_scale`
`residual + tanh(gate) * rmsnorm(x)`
`ffn_norm1(x) * scale_mlp` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/models/dits/zimage.py` | One CuTeDSL kernel fuses the first residual-form modulation and the next normalization / scale stage | If you see this chain split in Z-Image traces, report it as a missing existing merged fusion family. | +| Z-Image fused residual modulation + next norm-scale | `fused_norm_tanh_mul_add_norm_scale`
`residual + tanh(gate) * rmsnorm(x)`
`ffn_norm1(x) * scale_mlp` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/models/dits/zimage.py` | One CuTeDSL kernel fuses the first residual-form modulation and the next normalization / scale stage | If you see this chain split in Z-Image traces, report it as a missing existing mainline fusion family. | | Nunchaku fused GELU MLP | `_fused_gelu_mlp`
`fused_gelu_mlp` | `python/sglang/multimodal_gen/runtime/models/dits/flux.py` | Nunchaku path fuses `fc1 GEMM + GELU + shift + re-quant + fc2.lora_down` before the second GEMM | Treat split GELU-MLP on Nunchaku checkpoints as an existing fused family, not a new discovery. | ## 5. Diffusion kernel-overlap and async-communication families @@ -117,9 +135,8 @@ The catalog is grouped by reusable optimization family, not by one specific mode ## 6. PR-backed / in-flight fused-kernel families -These rows are intentionally not restricted to merged code. If the trace or -user request is about upstream work, use these rows to avoid calling an -already-known PR family "new". +These rows track still-open upstream work or status-sensitive PR families. +Stable entries should be folded into the mainline family rows above. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | @@ -128,9 +145,7 @@ already-known PR family "new". | PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`
`WRITE_PT`
`dequant_fp4_paged_decode` | `PR #21889`
`python/sglang/srt/layers/attention/nsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. | | PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`
`trtllm_fp8_block_scale_moe` | `PR #21491`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. | | PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`
`per_token_quant_fp8` | `PR #22005`
`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`
`python/sglang/jit_kernel/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. | -| PR `#21952` Gemma4 fused RMSNorm + residual + scalar | `gemma_rmsnorm_residual_scalar`
`_gemma_rmsnorm_residual_kernel`
`Gemma4` | `PR #21952`
`python/sglang/srt/layers/gemma4_fused_ops.py`
`python/sglang/srt/models/gemma4_causal.py` | Triton kernel fuses decoder post-FF RMSNorm, residual add, and per-layer scalar multiply into one pass | If Gemma4-style post-FF norm + residual + scalar steps appear split, treat them as an in-flight upstream Triton fuse family. | | PR `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`
`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`
`rotary_dim` | `PR #20667`
`python/sglang/srt/models/qwen3_5.py`
`python/sglang/srt/models/utils.py` | ROCm / AITER path fuses Q / K RMSNorm, partial or 3D RoPE, and direct KV cache write for Qwen3.5 attention | Treat split QK-norm + RoPE + cache-store on Qwen3.5 as a concrete in-flight upstream family, not a novel idea. | -| PR `#21977` TorchInductor combo-kernels horizontal Q/K norm fusion | `combo_kernels`
`benchmark_combo_kernel`
`q_norm`
`k_norm`
`split_with_sizes` | `PR #21977`
`torch._inductor.config.combo_kernels` | TorchInductor horizontally fuses sibling Q-norm and K-norm kernels, often deleting `split_with_sizes` / `clone` ladders in compiled traces | Treat separate Q/K norm ladders in compile-heavy traces as an in-flight compiler-fusion family first. | | PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`
`fp8_scaled_mm`
`nvjet`
`cudaMemsetAsync` | `PR #22392`
`sgl-kernel/python/sgl_kernel/gemm.py`
`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. | ## 7. PR-backed / in-flight kernel-overlap families @@ -138,7 +153,6 @@ already-known PR family "new". | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | | PR `#21877` fused down-GEMM + combine superseding SBO | `enable_fused_grouped_gemm_combine`
`combine`
`down_gemm` | `PR #21877`
`python/sglang/srt/server_args.py`
`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | Fused combine eliminates the standalone combine window, so SBO is intentionally disabled when this path is on | If the trace discussion is about combine overlap, first classify it as this upstream fused-overlap family. | -| PR `#22410` hiSparse H2D transfer overlap with hit-attention | `transfer_stream`
`execute_h2d_async`
`hit-attention`
`merge_state` | `PR #22410`
`python/sglang/srt/layers/attention/nsa_backend.py`
`python/sglang/srt/hisparse/hisparse_coordinator.py` | hiSparse decode overlaps host-to-device KV transfer on a transfer stream with hit-attention on the compute stream before merging miss-attention work | Treat hit-attention vs H2D KV transfer windows as a concrete in-flight SGLang overlap family first. | ## 8. FlashInfer mainline fused-kernel families @@ -151,10 +165,12 @@ only consumes a subset of that implementation. | FlashInfer activation / gate epilogues | `silu_and_mul`
`gelu_tanh_and_mul`
`gelu_and_mul`
`silu_and_mul_scaled_nvfp4_experts_quantize` | `flashinfer/activation.py`
`flashinfer/quantization/fp4_quantization.py` | FlashInfer covers both the plain activation-plus-mul epilogues and the NVFP4 expert-quantized extension used on MoE expert paths | Treat standalone activation, multiply, and expert-side quant ladders as one existing FlashInfer epilogue family first. | | FlashInfer norm / residual / quant epilogues | `rmsnorm_quant`
`fused_add_rmsnorm`
`fused_add_rmsnorm_quant`
`gemma_rmsnorm`
`gemma_fused_add_rmsnorm`
`fused_rmsnorm_silu`
`rmsnorm_fp4quant`
`add_rmsnorm_fp4quant` | `flashinfer/norm/__init__.py`
`flashinfer/cute_dsl/rmsnorm_fp4quant.py`
`flashinfer/cute_dsl/add_rmsnorm_fp4quant.py` | The norm family spans plain RMSNorm derivatives, residual-add epilogues, norm+activation, and direct FP8 / NVFP4 output variants instead of materializing each intermediate | Treat split residual add, norm, activation, and quant chains as one existing FlashInfer epilogue family first. | | FlashInfer allreduce + post-op fusion family | `allreduce_fusion`
`AllReduceFusionPattern`
`kARResidualRMSNorm`
`kARResidualRMSNormFP8Quant`
`kARResidualRMSNormFP4Quant`
`trtllm_mnnvl_allreduce_fusion` | `flashinfer/comm/allreduce.py`
`flashinfer/comm/trtllm_ar.py`
`flashinfer/comm/trtllm_mnnvl_ar.py` | TRTLLM and MNNVL backends fuse all-reduce with residual add, RMSNorm, and backend-appropriate quant / norm-output variants | Treat TP collective + norm (+ quant) ladders as an existing FlashInfer fused-collective family first. | -| FlashInfer RoPE + FP8 quant / cache-update family | `rope_quantize_fp8`
`mla_rope_quantize_fp8`
`rope_quantize_fp8_append_paged_kv_cache` | `flashinfer/rope.py` | The RoPE family covers both RoPE+FP8 output and the larger decode / prefill-prep path that also writes K / V directly into paged KV cache | Treat split RoPE, quant, and cache-write ladders as one existing FlashInfer attention-prep family first. | +| FlashInfer RoPE + FP8 quant / cache-update family | `rope_quantize_fp8`
`mla_rope_quantize_fp8`
`rope_quantize_fp8_append_paged_kv_cache`
`seqlen=0`
`batch_indices < 0` | `flashinfer/rope.py` | The RoPE family covers both RoPE+FP8 output and the larger decode / prefill-prep path that writes K / V directly into paged KV cache, including padding-token / zero-length sequence handling | Treat split RoPE, quant, cache-write, and padding-token ladders as one existing FlashInfer attention-prep family first. | | FlashInfer fused DeepSeek grouped-topk routing | `fused_topk_deepseek`
`NoAuxTc` | `flashinfer/fused_moe/fused_routing_dsv3.py` | One kernel performs sigmoid+bias, grouped score reduction, group top-k, expert top-k, and routed renorm for DeepSeek-V3-style routing | Treat router score activation -> grouped top-k -> renorm ladders as an existing FlashInfer router family first. | -| FlashInfer fused MoE expert execution | `cutlass_fused_moe`
`trtllm_bf16_moe`
`trtllm_fp8_per_tensor_scale_moe`
`trtllm_fp8_block_scale_moe`
`trtllm_fp4_block_scale_moe`
`trtllm_mxint4_block_scale_moe` | `flashinfer/fused_moe/core.py` | CUTLASS and TRTLLM backends collapse expert execution, routed combine, and quantized expert variants into fused MoE runners | Treat exposed expert-side tiny GEMM ladders as matching an existing FlashInfer fused-MoE family. | +| FlashInfer fused MoE expert execution | `cutlass_fused_moe`
`trtllm_bf16_moe`
`trtllm_fp8_per_tensor_scale_moe`
`trtllm_fp8_block_scale_moe`
`trtllm_fp4_block_scale_moe`
`trtllm_mxint4_block_scale_moe`
`non-gated` | `flashinfer/fused_moe/core.py` | CUTLASS and TRTLLM backends collapse expert execution, routed combine, and quantized expert variants into fused MoE runners, including gated and non-gated FP8 per-tensor cases | Treat exposed expert-side tiny GEMM or non-gated FP8 ladders as matching an existing FlashInfer fused-MoE family. | | FlashInfer CuTeDSL two-stage MoE fusion | `blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion_nvfp4`
`blockscaled_contiguous_grouped_gemm_finalize_fusion_nvfp4`
`moe_permute`
`moe_unpermute` | `flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py`
`flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_grouped_gemm_finalize_fusion.py` | The CuTeDSL path fuses gather+GEMM1+SwiGLU in the first stage and finalize+unpermute+scatter-reduce in the second stage, removing standalone `moe_permute` and `moe_unpermute` kernels | Treat multi-kernel MoE ladders around permute / finalize as one existing FlashInfer CuTeDSL family first. | +| FlashInfer SM120 FP4 / groupwise GEMM heuristics | `cutlass_fp4_gemm_sm120`
`CutlassTileConfigSM120`
`group_gemm_nvfp4_nt_groupwise`
`group_gemm_mxfp4_nt_groupwise` | `flashinfer/gemm/gemm_base.py`
`include/flashinfer/gemm/fp4_gemm_cutlass_template_sm120.h`
`include/flashinfer/gemm/group_gemm_nvfp4_groupwise_sm120.cuh`
`csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp` | FlashInfer mainline adds SM120-oriented FP4 GEMM selection and b12x CuTeDSL fused-MoE kernels | Treat SM120 FP4 MoE/GEMM tile selection and Blackwell-lite shape restrictions as an upstream FlashInfer kernel family before inventing a local heuristic. | +| FlashInfer MoE `routing_replay_out` support | `routing_replay_out`
`mPtrRoutingReplayOut`
`trtllm_fp8_block_scale_moe` | `flashinfer/fused_moe/core.py`
`csrc/trtllm_fused_moe_kernel_launcher.cu`
`csrc/fused_moe/noAuxTcKernels.cu` | TRTLLM-gen MoE kernels can optionally emit compact routing replay metadata without a separate routing-side reconstruction pass | Treat routing-replay writes in MoE traces as part of the upstream FlashInfer TRTLLM MoE family, not a separate postprocess opportunity. | ## 9. FlashInfer mainline kernel-overlap families @@ -168,10 +184,7 @@ only consumes a subset of that implementation. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| PR `#2792` RoPE + FP8 quant + paged KV append with padding-token support | `rope_quantize_fp8_append_paged_kv_cache`
`seqlen=0`
`batch_indices < 0` | `PR #2792`
`flashinfer/rope.py`
`include/flashinfer/pos_enc.cuh` | Extends the existing RoPE+quant+cache-write family to CUDA-graph padding tokens / zero-length sequences instead of introducing a separate kernel ladder | Treat split padding-token handling around RoPE+cache write as an in-flight upstream FlashInfer family first. | -| PR `#2840` CuTeDSL MoE aux-stream overlap race fix | `aux_stream`
`use_prealloc`
`use_cuda_graph` | `PR #2840`
`flashinfer/fused_moe/cute_dsl/fused_moe.py` | Clarifies that async memset overlap is only safe for the preallocated / CUDA-graph case; non-graph mode falls back to main-stream zeroing to avoid races | Treat missing aux-stream overlap in non-graph traces as an intentional safety rule, not a novel opportunity. | | PR `#2720` PDL runtime-API migration | `cudaGridDependencySynchronize`
`cudaTriggerProgrammaticLaunchCompletion`
`inline PTX` | `PR #2720`
`include/flashinfer/comm/trtllm_allreduce_fusion.cuh`
`include/flashinfer/pos_enc.cuh` | Repo-wide migration preserves the existing PDL overlap family while replacing inline PTX with CUDA runtime APIs across norm, RoPE, attention, and MoE codepaths | Treat PDL-looking launch groups as an upstream FlashInfer overlap family even when implementation details differ across revisions. | -| PR `#2882` FP8 per-tensor TRTLLM MoE non-gated activation | `trtllm_fp8_per_tensor_scale_moe`
`non-gated` | `PR #2882`
`csrc/trtllm_fused_moe_kernel_launcher.cu` | Extends the existing TRTLLM FP8 fused-MoE family to non-gated activations instead of requiring a separate expert path | Treat non-gated FP8 expert ladders as an in-flight upstream FlashInfer extension first. | ## 11. TensorRT-LLM-origin fused-kernel families @@ -184,9 +197,10 @@ the current `sglang` checkout only carries an analogous implementation. | TensorRT-LLM FlashInfer activation / gate epilogues | `flashinfer_silu_and_mul`
`flashinfer_gelu_tanh_and_mul`
`auto_deploy::silu_and_mul`
post-GEMM `silu` + `mul` | `tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py`
`tensorrt_llm/_torch/auto_deploy/transform/library/fuse_silu_mul.py`
`tensorrt_llm/_torch/models/modeling_gemma3.py` | Runtime custom ops and AutoDeploy rewrite `split/getitem + activation + mul` MLP epilogues into one FlashInfer op, including Gemma3 `gelu_tanh_and_mul` | Treat split gate activation + multiply as an existing TensorRT-LLM/FlashInfer epilogue family first. | | TensorRT-LLM FlashInfer RMSNorm family | `flashinfer_rmsnorm`
`flashinfer_gemma_rmsnorm`
`auto_deploy::flashinfer_rms_norm` | `tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py`
`tensorrt_llm/_torch/modules/rms_norm.py`
`tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py` | Runtime modules and AutoDeploy can lower plain RMSNorm and Gemma RMSNorm directly to FlashInfer kernels | Treat split RMSNorm ladders as an existing TensorRT-LLM norm family before calling them novel. | | TensorRT-LLM FlashInfer residual add + RMSNorm | `flashinfer_fused_add_rmsnorm`
`flashinfer_gemma_fused_add_rmsnorm`
`auto_deploy::flashinfer_fused_add_rms_norm_inplace` | `tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py`
`tensorrt_llm/_torch/modules/rms_norm.py`
`tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py` | Residual add immediately before RMSNorm can collapse to one in-place FlashInfer op, with Gemma variant support | Treat residual add + RMSNorm chains as an existing TensorRT-LLM fused epilogue family first. | +| TensorRT-LLM Triton fused residual add + RMSNorm + FP8 quant | `triton_fused_add_rms_norm_quant_fp8`
`fuse_rmsnorm_quant_fp8`
`fp8 static quant` | `tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/triton_fused_add_rms_norm_quant_fp8.py`
`tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_fp8.py` | Mainline AutoDeploy can rewrite residual-add plus RMSNorm plus FP8 static quant into one Triton op that emits BF16 norm output, FP8 quant output, and residual-add output together | Treat split add + norm + FP8 quant ladders as an existing TensorRT-LLM mainline family first. | | TensorRT-LLM FlashInfer RoPE with shared cos/sin cache | `flashinfer_apply_rope_with_cos_sin_cache_inplace`
`flashinfer_rope`
`cos_sin_cache` | `tensorrt_llm/_torch/modules/rotary_embedding.py`
`tensorrt_llm/_torch/auto_deploy/custom_ops/rope/flashinfer_rope.py`
`tensorrt_llm/_torch/auto_deploy/transform/library/rope.py` | Runtime path applies in-place RoPE from a shared cos/sin cache, while AutoDeploy can prebuild the full cache and lower diverse RoPE graphs to `flashinfer_rope` | Treat separate cos/sin gather + RoPE application ladders as an existing TensorRT-LLM attention-prep family. | | TensorRT-LLM FlashInfer cached paged attention | `append_paged_kv_cache`
`BatchPrefillWithPagedKVCacheWrapper`
`BatchDecodeWithPagedKVCacheWrapper`
`auto_deploy::flashinfer_attention_mha_with_cache`
`read_cache_only` | `tensorrt_llm/_torch/attention_backend/flashinfer.py`
`tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py`
`docs/source/features/attention.md` | FlashInfer attention backend fuses metadata setup, optional paged-KV append, and prefill/decode wrapper execution, including shared-KV and read-cache-only variants in AutoDeploy | Treat metadata + KV-append + cached-attention ladders as one existing TensorRT-LLM cached-attention family first. | -| TensorRT-LLM FlashInfer MLA regular prefill | `append_paged_mla_kv_cache`
`BatchPrefillWithRaggedKVCacheWrapper`
`flashinfer_mla` | `tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` | Regular MLA prefill writes compressed KV pages and runs FlashInfer ragged prefill instead of a split append-plus-prefill ladder | Treat MLA regular-prefill prep as an existing TensorRT-LLM FlashInfer family first. | +| TensorRT-LLM FlashInfer MLA regular prefill | `append_paged_mla_kv_cache`
`BatchPrefillWithRaggedKVCacheWrapper`
`flashinfer_mla`
`rank 256`
`gpu append kernel` | `tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` | Regular MLA prefill writes compressed KV pages and runs FlashInfer ragged prefill instead of a split append-plus-prefill ladder, with rank-256 paged-KV setups using the GPU append path | Treat MLA regular-prefill prep as an existing TensorRT-LLM FlashInfer family first. | | TensorRT-LLM FlashInfer MLA chunked prefill with absorbed `W_kn` | `BatchMLAPagedAttentionWrapper`
`chunked prefill`
`W_kn`
`W_v` | `tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` | Chunked prefill absorbs `W_kn` into the query-side projection, runs paged MLA attention in compressed space, then projects back with `W_v` | Treat split absorbed-proj + MLA + output-proj ladders as an existing TensorRT-LLM MLA family first. | | TensorRT-LLM FlashInfer MLA decode with absorbed `W_kn` + `W_v` | `plan_decode`
`BatchMLAPagedAttentionWrapper`
`decode`
`W_kn`
`W_v` | `tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` | Decode path reuses the absorbed-query MLA family and projects the compressed attention output back with `W_v` | Treat similar decode-time absorbed MLA ladders as an existing TensorRT-LLM family, not a new idea. | | TensorRT-LLM FlashInfer fused MoE backend | `flashinfer.fused_moe`
`trtllm_bf16_moe`
`trtllm_fp8_block_scale_moe`
`trtllm_fp4_block_scale_moe`
`TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER` | `tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py`
`tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py` | TRTLLM-gen MoE can route expert execution and quant helpers through FlashInfer instead of exposing per-expert eager ladders | Treat expert-side tiny GEMM ladders as matching an existing TensorRT-LLM FlashInfer MoE family first. | @@ -197,19 +211,16 @@ the current `sglang` checkout only carries an analogous implementation. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | | TensorRT-LLM multi-stream MLA attention | `multi_stream_mla_attn`
`record_event_passthrough`
`_aux`
`wait_event` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | AutoDeploy rewrites MLA Q/KV forks so the KV projection runs on an auxiliary stream while the Q path stays on the caller stream | Treat exposed Q-branch vs KV-branch overlap as an existing TensorRT-LLM multi-stream family first. | -| TensorRT-LLM multi-stream MoE shared-vs-routed overlap | `multi_stream_moe`
`begin_aux_stream_passthrough`
`end_aux_stream_passthrough`
`wait_aux_stream_passthrough` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Shared-expert work is moved to an auxiliary stream while routed-expert MoE work remains on the main stream and rejoins at the merge node | Treat shared-expert vs routed-expert windows as an existing TensorRT-LLM branch-overlap family. | +| TensorRT-LLM multi-stream MoE shared-vs-routed overlap | `multi_stream_moe`
`begin_aux_stream_passthrough`
`end_aux_stream_passthrough`
`wait_aux_stream_passthrough`
`mlir_elementwise_fusion`
`piecewise cudagraph`
`caller_stream.synchronize()` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Shared-expert work is moved to an auxiliary stream while routed-expert MoE work remains on the main stream and rejoins at the merge node; the same family includes synchronization rules for MLIR-fused kernels and piecewise cudagraph replay | Treat shared-expert vs routed-expert windows, including altered `multi_stream_moe` behavior under MLIR / piecewise graph modes, as an existing TensorRT-LLM branch-overlap family. | | TensorRT-LLM multi-stream FP8 GEMM fork parallelism | `multi_stream_gemm`
`trtllm_finegrained_fp8_linear`
`record_event_passthrough`
`_aux` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_gemm.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Compiler pass identifies fork points with multiple FP8 linears and moves the largest GEMM to the auxiliary stream so sibling GEMMs overlap | Treat sibling FP8 linear branches as an existing TensorRT-LLM overlap family before designing a new stream split. | ## 13. TensorRT-LLM-origin PR-backed / in-flight fused-kernel and kernel-overlap families | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| PR `#12674` fused residual add + RMSNorm + FP8 quant | `triton_fused_add_rms_norm_quant_fp8`
`residual_add`
`rms_norm`
`fp8 static quant` | `PR #12674`
`tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_fp8.py` | Open PR adds a pattern-matcher pass that replaces residual-add plus RMSNorm plus FP8 static quant with a fused FlashInfer 0.6.7-backed path | Treat split add + norm + FP8 quant ladders as an in-flight TensorRT-LLM family first. | -| PR `#12519` rank-256 `flashinfer_mla` extension | `flashinfer_mla`
`rank 256`
`paged KV-cache`
`gpu append kernel` | `PR #12519`
`tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` | Open PR extends the existing FlashInfer MLA family with a TRTLLM MLA operator, paged-KV support, and a GPU append kernel for rank-256 setups | Treat rank-256 MLA prep / decode ladders as an in-flight TensorRT-LLM MLA family, not a novel direction. | | PR `#12525` FlashInfer TRTLLM-gen FMHA paged-index / buffer rework | `shared paged index`
`trtllm-gen attention`
`flashinfer`
`kv cache buffer` | `PR #12525`
`tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py` | Open PR refines the existing FlashInfer TRTLLM-gen cached-attention family by disabling shared paged index and unifying KV-buffer construction | Treat these attention-prep changes as an in-flight implementation evolution of an existing family first. | | PR `#12544` NVFP4 KV cache support in TRTLLM-gen attention | `NVFP4 KV cache`
`trtllm-gen attention`
`flashinfer` | `PR #12544`
`tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py` | Open PR extends the cached-attention family so the FlashInfer-backed TRTLLM-gen path can build and consume NVFP4 KV buffers directly | Treat split KV-cache quant + buffer-build ladders as an in-flight TensorRT-LLM attention family first. | | PR `#12738` / `#12557` BF16 TRTLLM-gen MoE through FlashInfer | `bf16 trtllm-gen moe`
`flashinfer`
`trtllm_bf16_moe` | `PR #12738`
`PR #12557`
`tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py` | Open PRs extend the TRTLLM-gen MoE family so BF16 expert execution can route through FlashInfer instead of only CUTLASS-like paths | Treat BF16 expert ladders as an in-flight TensorRT-LLM FlashInfer MoE family. | -| PR `#12847` `multi_stream_moe` sync fix for MLIR and piecewise cudagraphs | `multi_stream_moe`
`mlir_elementwise_fusion`
`piecewise cudagraph`
`caller_stream.synchronize()` | `PR #12847`
`tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Open PR preserves the existing multi-stream MoE overlap family while tightening synchronization when MLIR-fused kernels or piecewise cudagraph replay are present | Treat missing or altered `multi_stream_moe` overlap under MLIR / piecewise graph modes as an in-flight TensorRT-LLM rule first. | ## 14. vLLM-origin fused-kernel families @@ -224,7 +235,7 @@ contain the same implementation. | vLLM-origin RMSNorm (+ residual add) + quant | `RMSNormQuantFusionPass`
`fused_add_rms_norm_static_fp8_quant`
`per_token_quant`
`per_group_quant` | `vllm/compilation/passes/fusion/rms_quant_fusion.py`
`vllm/compilation/passes/fusion/rocm_aiter_fusion.py` | Compile-time and ROCm AITER paths fuse RMSNorm or fused-add-RMSNorm with FP8 / FP4 quant output | Treat split norm/add + quant as an upstream fused family, not an unexplored direction. | | vLLM-origin SiLU+Mul + quant | `ActivationQuantFusionPass`
`SiluMulFp8*`
`Nvfp4`
`rocm_aiter` | `vllm/compilation/passes/fusion/act_quant_fusion.py`
`vllm/compilation/passes/fusion/rocm_aiter_fusion.py` | Activation epilogues fuse `SiLU+Mul` with FP8 / NVFP4 / AITER group quant instead of materializing the BF16 activation first | Treat standalone activation then quant kernels as matching a vLLM-origin precedent. | | vLLM-origin add + RMSNorm + pad | `fuse_act_padding`
`RocmAiterTritonAddRMSNormPadFusionPass`
`add_rmsnorm_pad` | `vllm/compilation/passes/fusion/rocm_aiter_fusion.py`
`docs/design/fusions.md` | ROCm / AITER path fuses residual add + RMSNorm directly into the padded layout expected by the next kernel | Treat norm-plus-padding ladders as an existing backend-specific fuse family first. | -| vLLM-origin attention + output quant | `fuse_attn_quant`
`AttnQuantFusionPass`
`output_scale`
`output_block_scale` | `vllm/compilation/passes/fusion/attn_quant_fusion.py`
`vllm/v1/attention/backends/`
`docs/design/fusions.md` | Compile-time fusion pushes FP8 / NVFP4 quantization into the attention epilogue on supported Triton / FlashInfer / ROCm / AITER backends | Treat attention-output quant kernels as a known upstream epilogue fusion family before calling them novel. | +| vLLM-origin attention + output quant | `fuse_attn_quant`
`AttnQuantFusionPass`
`merge_attn_states`
`output_scale`
`output_group_scale`
`output_block_scale` | `vllm/compilation/passes/fusion/attn_quant_fusion.py`
`vllm/v1/attention/ops/merge_attn_states.py`
`vllm/csrc/attention/merge_attn_states.cu`
`docs/design/fusions.md` | Compile-time fusion pushes FP8 / NVFP4 quantization into the attention epilogue on supported Triton / FlashInfer / ROCm / AITER backends, and mainline `merge_attn_states` kernels already support FP8 output when `output_scale` is provided | Treat attention-output quant and merged-attention quant epilogues as a known upstream family before calling them novel. | | vLLM-origin fused QK RMSNorm + RoPE | `fused_qk_norm_rope`
`QKNormRoPEFusionPass`
`qk norm + rope` | `vllm/compilation/passes/fusion/qk_norm_rope_fusion.py`
`vllm/_custom_ops.py`
`csrc/fused_qknorm_rope_kernel.cu` | Compile-time and direct custom-op paths fuse per-head Q / K RMSNorm with RoPE | Treat split QK norm + RoPE as a clear vLLM-origin precedent. | | vLLM-origin fused reshape + KV cache write | `reshape_and_cache`
`triton_reshape_and_cache_flash`
`kv cache write` | `vllm/v1/attention/ops/triton_reshape_and_cache_flash.py`
`vllm/v1/attention/backends/triton_attn.py` | Triton cache-update kernels reshape K / V into paged-cache layout and can include FP8 KV-cache scale/write logic | Treat reshape / transpose / cache-write ladders as an existing cache-store fusion family. | | vLLM-origin fused RoPE + KV cache update | `fuse_rope_kvcache`
`RopeKVCacheFusionPass`
`triton_rope_and_cache` | `vllm/compilation/passes/fusion/rope_kvcache_fusion.py`
`vllm/_aiter_ops.py`
`docs/design/fusions.md` | ROCm / AITER compile-time fusion combines RoPE with paged KV cache update instead of launching them separately | Treat split RoPE + cache-store as a known upstream family, especially on ROCm-like paths. | @@ -234,6 +245,8 @@ contain the same implementation. | vLLM-origin DSV3 router GEMM | `dsv3_router_gemm`
`allow_dsv3_router_gemm`
`router logits` | `vllm/_custom_ops.py`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`csrc/moe/dsv3_router_gemm_entry.cu`
`csrc/moe/dsv3_router_gemm_float_out.cu` | Hopper-class CUDA kernel specializes the DeepSeek router linear for small decode batches and can emit FP32 logits directly without a generic GEMM chain | Treat DeepSeek-style router linear paths as an existing upstream specialized fuse, distinct from grouped-topk itself. | | vLLM-origin GPT-OSS router GEMM | `gpt_oss_router_gemm`
`router gemm` | `vllm/_custom_ops.py`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`csrc/moe/gpt_oss_router_gemm.cu` | Model-specific CUDA kernel replaces the router linear plus bias path with one specialized GEMM op | Treat GPT-OSS-style router linear chains as an existing upstream specialized fuse. | | vLLM-origin DeepSeek min-latency fused QKV-A projection | `dsv3_fused_a_gemm`
`fused_qkv_a_proj`
`q_a_proj` | `vllm/model_executor/models/deepseek_v2.py`
`vllm/_custom_ops.py`
`csrc/dsv3_fused_a_gemm.cu` | Hopper-class CUDA kernel replaces the tiny-batch DeepSeek QKV-A projection path with one specialized min-latency GEMM instead of a generic linear launch | Treat small-batch DeepSeek QKV-A projection ladders as a known upstream fused kernel family first. | +| vLLM-origin DSV3.2 fused indexer projections | `wk_weights_proj`
`MergedColumnParallelLinear`
`weights_proj` | `vllm/model_executor/models/deepseek_v2.py`
`vllm/model_executor/models/deepseek_mtp.py` | DSV3.2 indexer paths can fuse the `wk` and `weights_proj` projections into one GEMM and carry the matching MTP weight-loading path | Treat paired indexer projection chains as a known upstream fused linear family before calling the opportunity novel. | +| vLLM-origin MiniMax allreduce_rms kernels | `minimax_allreduce_rms`
`minimax_allreduce_rmsnorm`
`MiniMax-M2.5`
`allreduce_rms` | `vllm/model_executor/models/minimax_m2.py` | TensorRT-LLM-derived MiniMax allreduce-plus-RMSNorm kernels are a concrete upstream TP decode family | Treat MiniMax TP norm + collective ladders as an upstream specialized fusion family. | | vLLM-origin CUTLASS scaled MM with scale / bias epilogue | `cutlass_scaled_mm`
`cutlass_scaled_mm_azp`
`scaled mm` | `vllm/_custom_ops.py`
`vllm/model_executor/kernels/linear/scaled_mm/cutlass.py`
`csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu` | CUTLASS kernels fuse activation scales, weight scales, matmul, and optional bias / AZP epilogues | Treat separate scale-mul + GEMM + bias ladders as a vLLM-origin fused linear family first. | | vLLM-origin fused MoE expert execution | `cpu_fused_moe`
`rocm_aiter_fused_moe`
`FusedMoE` | `vllm/model_executor/layers/fused_moe/layer.py`
`vllm/model_executor/layers/fused_moe/cpu_fused_moe.py`
`vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py`
`vllm/_aiter_ops.py` | MoE backends on CUDA / ROCm / CPU already collapse packed expert execution into fused expert kernels rather than per-expert eager GEMMs | Treat exposed expert-side tiny GEMM ladders as matching an upstream fused-MoE family. | | vLLM-origin fused MoE LoRA | `fused_moe_lora`
`fused_moe_lora_fp8`
`w13_shrink`
`w2_expand` | `vllm/lora/ops/triton_ops/fused_moe_lora_op.py`
`vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py`
`vllm/lora/layers/fused_moe.py` | Triton kernels fuse LoRA shrink / expand work into MoE expert execution, including FP8 variants | Treat MoE-LoRA adapter work as an upstream fused family before proposing a brand new kernel. | @@ -245,21 +258,19 @@ contain the same implementation. | --- | --- | --- | --- | --- | | vLLM-origin AsyncTP GEMM + collective overlap | `fuse_gemm_comms`
`fused_matmul_reduce_scatter`
`fused_all_gather_matmul` | `vllm/compilation/passes/fusion/collective_fusion.py`
`docs/design/fusions.md` | AsyncTP overlaps GEMM with reduce-scatter / all-gather via symmetric-memory collectives | Treat GEMM+comm windows as a clear vLLM-origin overlap precedent first. | | vLLM-origin Sequence Parallelism staging | `enable_sp`
`ReduceScatter`
`AllGather`
`SequenceParallelismPass` | `vllm/compilation/passes/fusion/sequence_parallelism.py`
`docs/design/fusions.md` | Sequence-parallel rewrites all-reduce into RS -> local norm -> AG so later passes can overlap comm and compute | Treat RS / AG staging around norm blocks as an upstream overlap-enabling family. | -| vLLM-origin shared-expert aux-stream overlap | `aux_stream`
`shared_experts_stream`
shared expert near router | `vllm/utils/torch_utils.py`
`vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py` | MoE shared experts can run on a dedicated aux stream and overlap with router-side work | Treat shared-expert vs router overlap as an existing upstream sparse-model family. | +| vLLM-origin shared-expert aux-stream overlap | `aux_stream`
`shared_experts_stream`
shared expert near router | `vllm/model_executor/layers/fused_moe/runner/shared_experts.py`
`vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py` | MoE shared experts can record the cloned input on `shared_experts_stream`, wait on the caller stream, run in parallel with router-side work, and rejoin before merge | Treat shared-expert vs router overlap as an existing upstream sparse-model family. | | vLLM-origin DCP async all-to-all overlap | `dcp_alltoall`
`all_to_all_single`
`async_op=True` | `vllm/v1/attention/ops/dcp_alltoall.py` | Output / LSE exchange uses async all-to-all handles instead of serializing collective completion on the main path | Treat DCP all-to-all windows as an upstream async-collective family. | ## 16. vLLM-origin PR-backed / in-flight fused-kernel and kernel-overlap families | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| PR `#35968` DSV3.2 multi-stream indexer overlap | `weights_proj`
`wk`
`k_norm`
`aux_stream` | `PR #35968`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/utils/torch_utils.py` | Open PR overlaps the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. | +| PR `#35968` DSV3.2 multi-stream indexer overlap | `weights_proj`
`wk`
`k_norm`
`aux_stream` | `PR #35968`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/utils/torch_utils.py` | Closed PR explored overlapping the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. | | PR `#37110` Triton attention + per-group FP8 dynamic quant | `group_size=128`
`group_size=64`
`output_group_scale`
`per-group FP8` | `PR #37110`
`vllm/compilation/passes/fusion/attn_quant_fusion.py`
`vllm/v1/attention/ops/triton_unified_attention.py` | In-flight Triton attention epilogue computes per-group FP8 scales and quantizes output directly instead of launching a separate group-quant kernel | Treat attention + per-group FP8 quant as a concrete upstream vLLM family, not a novel idea. | | PR `#38445` MiniMax-M2 FP32 gate kernel | `fp32_router_gemm`
`MiniMax-M2`
`gate kernel` | `PR #38445`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`vllm/model_executor/models/minimax_m2.py` | Draft CUDA kernel fuses BF16->FP32 conversion and low-batch router GEMM for MiniMax-M2, replacing up to three kernels on the gate path | Treat MiniMax-M2 gate ladders as an in-flight upstream fused router family first. | | PR `#38621` fused QK norm + RoPE + cache + quant | `fused_qk_norm_rope_cache_quant`
`QK Norm + RoPE + Cache + Quant` | `PR #38621`
`csrc/fused_qk_norm_rope_cache_quant.cu`
`vllm/compilation/passes/fusion/qk_norm_rope_cache_quant_fusion.py` | Draft CUDA kernel and compile-time pass try to fuse QK RMSNorm, RoPE, KV cache write, and optional FP8 quant for small-batch decode | Treat this as an in-flight upstream fusion family before calling a similar idea novel. | -| PR `#38684` DSV3.2 fused `wk + weights_proj` | `wk_weights_proj`
`MergedColumnParallelLinear`
`weights_proj` | `PR #38684`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/model_executor/models/deepseek_mtp.py` | Merged PR fuses the DSV3.2 indexer `wk` and `weights_proj` projections into one GEMM; FP8 weight-loading caveats are being handled in follow-up `PR #38870` | Treat paired indexer projections as a concrete upstream fused linear family before calling the opportunity novel. | | PR `#37646` ROCm AITER fused allreduce + RMSNorm | `rocm_aiter_fused_allreduce_rmsnorm`
`custom_fused_ar_rms`
`RocmAiterAllReduceFusionPass` | `PR #37646`
`vllm/_aiter_ops.py`
`vllm/compilation/passes/pass_manager.py` | ROCm-specific compile-time path swaps the generic all-reduce fusion pass for an AITER fused allreduce-plus-RMSNorm kernel family | Treat ROCm TP all-reduce + RMSNorm ladders as an in-flight upstream fused-collective family first. | | PR `#36413` FlashInfer RMSNorm + FP4 quant fusion | `fuse_norm_quant`
`flashinfer`
`NVFP4`
`rmsnorm + fp4 quant` | `PR #36413`
`vllm/compilation/passes/fusion/rms_quant_fusion.py`
`vllm/docs/design/fusions.md` | FlashInfer-backed norm-plus-FP4 quant fusion extends the existing RMSNorm+quant family to NVFP4 flows | Treat split RMSNorm + FP4 quant ladders as an upstream in-flight family, not a fresh idea. | -| PR `#37045` MiniMax TRTLLM `minimax_allreduce_rms` kernels | `minimax_allreduce_rms`
`MiniMax-M2.5`
`allreduce_rms` | `PR #37045`
`vllm/model_executor/models/minimax_m2.py` | Draft kernel ports TensorRT-LLM MiniMax allreduce-plus-RMSNorm kernels into vLLM for TP MiniMax decode | Treat MiniMax TP norm + collective ladders as an in-flight upstream specialized fusion family. | | PR `#39301` GLM5 router GEMM with PDL overlap | `TRTLLM_ENABLE_PDL`
`router_gemm`
`GLM5`
`FI AR RMS fusion` | `PR #39301`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`vllm/csrc/moe/dsv3_router_gemm_utils.h` | Extends the specialized router GEMM family to GLM5 hidden size and uses PDL to overlap the router launch with the preceding fused allreduce-plus-RMS block | Treat this as an in-flight upstream router-kernel plus launch-overlap family before calling it novel. | ## 17. Important toggles and caveats @@ -312,7 +323,7 @@ FLASHINFER_REPO=${FLASHINFER_REPO:-../flashinfer} TRTLLM_REPO=${TRTLLM_REPO:-../TensorRT-LLM} VLLM_REPO=${VLLM_REPO:-../vllm} -rg -n "fused_add_rmsnorm|gemma_fused_add_rmsnorm|silu_and_mul|gelu_and_mul|fused_qk_rope_reshape_and_cache|fused_set_kv_buffer|fused_metadata_copy|normal_decode_set_metadata" python/sglang +rg -n "fused_add_rmsnorm|gemma_fused_add_rmsnorm|silu_and_mul|gelu_and_mul|fused_qk_rope_reshape_and_cache|fused_set_kv_buffer|fused_metadata_copy|normal_decode_set_metadata|_append_shared_to_topk_output|fused_append_shared_experts_with_weights" python/sglang rg -n "MiniMaxM2RMSNormTP|fused_qknorm_rope|fused_qk_rope_cat_and_cache_mla|fused_qk_norm_mrope_3d_cache_pts_quant_shuffle|split_qkv_rmsnorm_rope|trtllm_fp8_kv_kernel|set_mla_kv_buffer_fp8_quant" python/sglang rg -n "FusedMoeRouter|fused_topk_deepseek|moe_fused_gate|aiter_fused_topk|fused_rms_fp8_group_quant|fast_topk_transform_fused|fused_store_index_k_cache|fused_temperature_softmax|fused_softcap" python/sglang rg -n "fused_qkvzba_split_reshape_cat|fused_gdn_gating|rms_norm_gated|layer_norm_gated|chunk_gated_delta_rule_fwd_kkt_solve_kernel|fused_recurrent_gated_delta_rule_update|fused_mamba_state_scatter_with_mask|_fused_gather_to_staging_kernel|_fused_scatter_from_staging_kernel" python/sglang @@ -322,11 +333,11 @@ rg -n "silu_and_mul|gelu_tanh_and_mul|gelu_and_mul|silu_and_mul_scaled_nvfp4_exp rg -n "AllReduceFusionPattern|allreduce_fusion|trigger_completion_at_end|rope_quantize_fp8|rope_quantize_fp8_append_paged_kv_cache|fused_topk_deepseek|cutlass_fused_moe|trtllm_.*_moe" "$FLASHINFER_REPO/flashinfer" rg -n "aux_stream|use_async_memset|split_device_green_ctx|split_device_green_ctx_by_sm_count|enable_pdl|launch_with_pdl" "$FLASHINFER_REPO/flashinfer" "$FLASHINFER_REPO/include" git -C "$FLASHINFER_REPO" log --all --format='%h %s' | rg -i 'fused|fusion|overlap|pdl|stream|rope|kv|quant|topk|moe' -rg -n "flashinfer_silu_and_mul|flashinfer_gelu_tanh_and_mul|flashinfer_rmsnorm|flashinfer_gemma_rmsnorm|flashinfer_fused_add_rmsnorm|flashinfer_apply_rope_with_cos_sin_cache_inplace" "$TRTLLM_REPO/tensorrt_llm/_torch" +rg -n "flashinfer_silu_and_mul|flashinfer_gelu_tanh_and_mul|flashinfer_rmsnorm|flashinfer_gemma_rmsnorm|flashinfer_fused_add_rmsnorm|flashinfer_apply_rope_with_cos_sin_cache_inplace|triton_fused_add_rms_norm_quant_fp8|fuse_rmsnorm_quant_fp8" "$TRTLLM_REPO/tensorrt_llm/_torch" rg -n "flashinfer_attention_mha_with_cache|append_paged_kv_cache|flashinfer_mla|append_paged_mla_kv_cache|flashinfer_cached_ssm|selective_state_update|flashinfer.fused_moe" "$TRTLLM_REPO/tensorrt_llm/_torch" "$TRTLLM_REPO/docs/source" rg -n "multi_stream_moe|multi_stream_mla_attn|multi_stream_gemm|record_event_passthrough|begin_aux_stream_passthrough|end_aux_stream_passthrough|wait_aux_stream_passthrough" "$TRTLLM_REPO/tensorrt_llm/_torch" git -C "$TRTLLM_REPO" log --all --format='%h %s' | rg -i 'fused|fusion|overlap|flashinfer|mla|kv cache|multi-stream|stream|rope|rmsnorm|moe' -rg -n "fused_add_rms_norm|fused_qk_norm_rope|grouped_topk|topk_softmax|topk_sigmoid|dsv3_router_gemm|dsv3_fused_a_gemm|concat_and_cache_mla_rope_fused|gpt_oss_router_gemm|cutlass_scaled_mm|cpu_fused_moe|fused_moe_lora|triton_pos_embed_interpolate" "$VLLM_REPO/vllm" "$VLLM_REPO/csrc" +rg -n "fused_add_rms_norm|merge_attn_states|fused_qk_norm_rope|grouped_topk|topk_softmax|topk_sigmoid|dsv3_router_gemm|dsv3_fused_a_gemm|concat_and_cache_mla_rope_fused|gpt_oss_router_gemm|cutlass_scaled_mm|cpu_fused_moe|fused_moe_lora|triton_pos_embed_interpolate" "$VLLM_REPO/vllm" "$VLLM_REPO/csrc" rg -n "fuse_allreduce_rms|fuse_norm_quant|fuse_act_quant|fuse_attn_quant|enable_qk_norm_rope_fusion|fuse_rope_kvcache|enable_sp|fuse_gemm_comms|RocmAiter|dcp_alltoall|shared_experts_stream|TRTLLM_ENABLE_PDL|wk_weights_proj" "$VLLM_REPO/vllm" "$VLLM_REPO/docs/design/fusions.md" "$VLLM_REPO/csrc" git -C "$VLLM_REPO" log --all --format='%h %s' | rg -i 'fused|fusion|overlap|triton|cuda|rope|kv cache|topk|router|allreduce|reduce-scatter|all-gather|all_to_all|quant' # GitHub PR scan terms for the connector or web UI: diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/heuristics.md b/.claude/skills/llm-torch-profiler-analysis/references/heuristics.md similarity index 100% rename from .claude/skills/sglang-torch-profiler-analysis/references/heuristics.md rename to .claude/skills/llm-torch-profiler-analysis/references/heuristics.md diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/overlap-catalog.md b/.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md similarity index 86% rename from .claude/skills/sglang-torch-profiler-analysis/references/overlap-catalog.md rename to .claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md index 6ba1084ab..5a38cb204 100644 --- a/.claude/skills/sglang-torch-profiler-analysis/references/overlap-catalog.md +++ b/.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md @@ -16,8 +16,8 @@ Use it like this: 3. If a match exists in the mainline sections, report it as an existing overlap family that is missing, disabled, regressed, or unsupported on the current backend. -4. If a match exists only in the `PR-backed / in-flight` section, report it as - an upstream overlap pattern, not a novel idea. +4. If a match exists only in the `PR-backed / in-flight` + section, report it as an upstream overlap pattern, not a novel idea. 5. Only call an overlap opportunity "new" when no row in this file or `fuse-overlap-catalog.md` fits. @@ -26,6 +26,18 @@ necessarily present in the checked-out `sglang` tree, but they should still be treated as upstream or analogous kernel-overlap families before labeling an overlap opportunity as novel. +Refresh note `2026-04-22`: rescanned current `sglang`, `flashinfer`, +`TensorRT-LLM`, and `vllm` mainline overlap paths plus rechecked referenced PR +state via the GitHub API on `2026-04-22`. Closed-unmerged SGLang +[#22410](https://github.com/sgl-project/sglang/pull/22410) and FlashInfer +[#2840](https://github.com/flashinfer-ai/flashinfer/pull/2840) were removed +from the PR-backed sections. SGLang +[#21877](https://github.com/sgl-project/sglang/pull/21877), FlashInfer +[#2720](https://github.com/flashinfer-ai/flashinfer/pull/2720), and vLLM +[#35968](https://github.com/vllm-project/vllm/pull/35968) / +[#39301](https://github.com/vllm-project/vllm/pull/39301) remain useful +upstream overlap references as of this refresh. + ## 1. LLM / SRT kernel-overlap families | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | @@ -64,7 +76,6 @@ overlap opportunity as novel. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | | PR `#21877` fused down-GEMM + combine superseding SBO | `enable_fused_grouped_gemm_combine`
`combine`
`down_gemm` | `PR #21877`
`python/sglang/srt/server_args.py`
`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | Fused combine eliminates the standalone combine window, so SBO is intentionally disabled when this path is on | If the trace discussion is about combine overlap, first classify it as this upstream fused-overlap family. | -| PR `#22410` hiSparse H2D transfer overlap with hit-attention | `transfer_stream`
`execute_h2d_async`
`hit-attention`
`merge_state` | `PR #22410`
`python/sglang/srt/layers/attention/nsa_backend.py`
`python/sglang/srt/hisparse/hisparse_coordinator.py` | hiSparse decode overlaps host-to-device KV transfer on a transfer stream with hit-attention on the compute stream before running miss-attention and merge | Treat hit-attention vs H2D KV transfer windows as an in-flight SGLang overlap family first. | ## 5. FlashInfer kernel-overlap families @@ -82,7 +93,6 @@ checkout only calls part of that implementation. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| PR `#2840` CuTeDSL MoE aux-stream overlap race fix | `aux_stream`
`use_prealloc`
`use_cuda_graph` | `PR #2840`
`flashinfer/fused_moe/cute_dsl/fused_moe.py` | Clarifies that async memset overlap is only safe for the preallocated / CUDA-graph case; non-graph mode falls back to main-stream zeroing to avoid races | Treat missing aux-stream overlap in non-graph traces as an intentional safety rule, not a novel opportunity. | | PR `#2720` PDL runtime-API migration | `cudaGridDependencySynchronize`
`cudaTriggerProgrammaticLaunchCompletion`
`inline PTX` | `PR #2720`
`include/flashinfer/comm/trtllm_allreduce_fusion.cuh`
`include/flashinfer/pos_enc.cuh` | Repo-wide migration preserves the existing PDL overlap family while replacing inline PTX with CUDA runtime APIs across norm, RoPE, attention, and MoE codepaths | Treat PDL-looking launch groups as an upstream FlashInfer overlap family even when implementation details differ across revisions. | ## 7. TensorRT-LLM-origin kernel-overlap families @@ -94,32 +104,26 @@ AutoDeploy rather than same-stream PDL windows. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | | TensorRT-LLM multi-stream MLA attention | `multi_stream_mla_attn`
`record_event_passthrough`
`_aux`
`wait_event` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | AutoDeploy rewrites MLA Q/KV forks so the KV projection runs on an auxiliary stream while the Q path stays on the caller stream | Treat exposed Q-branch vs KV-branch overlap as an existing TensorRT-LLM multi-stream family first. | -| TensorRT-LLM multi-stream MoE shared-vs-routed overlap | `multi_stream_moe`
`begin_aux_stream_passthrough`
`end_aux_stream_passthrough`
`wait_aux_stream_passthrough` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Shared-expert work is moved to an auxiliary stream while routed-expert MoE work remains on the main stream and rejoins at the merge node | Treat shared-expert vs routed-expert windows as an existing TensorRT-LLM branch-overlap family. | +| TensorRT-LLM multi-stream MoE shared-vs-routed overlap | `multi_stream_moe`
`begin_aux_stream_passthrough`
`end_aux_stream_passthrough`
`wait_aux_stream_passthrough`
`mlir_elementwise_fusion`
`piecewise cudagraph`
`caller_stream.synchronize()` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Shared-expert work is moved to an auxiliary stream while routed-expert MoE work remains on the main stream and rejoins at the merge node; the same family includes synchronization rules for MLIR-fused kernels and piecewise cudagraph replay | Treat shared-expert vs routed-expert windows, including altered behavior under MLIR / piecewise graph modes, as an existing TensorRT-LLM branch-overlap family. | | TensorRT-LLM multi-stream FP8 GEMM fork parallelism | `multi_stream_gemm`
`trtllm_finegrained_fp8_linear`
`record_event_passthrough`
`_aux` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_gemm.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Compiler pass identifies fork points with multiple FP8 linears and moves the largest GEMM to the auxiliary stream so sibling GEMMs overlap | Treat sibling FP8 linear branches as an existing TensorRT-LLM overlap family before designing a new stream split. | -## 8. TensorRT-LLM-origin PR-backed / in-flight kernel-overlap families - -| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | -| --- | --- | --- | --- | --- | -| PR `#12847` `multi_stream_moe` sync fix for MLIR and piecewise cudagraphs | `multi_stream_moe`
`mlir_elementwise_fusion`
`piecewise cudagraph`
`caller_stream.synchronize()` | `PR #12847`
`tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`
`tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py` | Open PR preserves the existing multi-stream MoE overlap family while tightening synchronization when MLIR-fused kernels or piecewise cudagraph replay are present | Treat missing or altered `multi_stream_moe` overlap under MLIR / piecewise graph modes as an in-flight TensorRT-LLM rule first. | - -## 9. vLLM-origin kernel-overlap families +## 8. vLLM-origin kernel-overlap families | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | | vLLM-origin AsyncTP GEMM + collective overlap | `fuse_gemm_comms`
`fused_matmul_reduce_scatter`
`fused_all_gather_matmul` | `vllm/compilation/passes/fusion/collective_fusion.py`
`docs/design/fusions.md` | AsyncTP overlaps GEMM with reduce-scatter / all-gather via symmetric-memory collectives | Treat GEMM+comm windows as a clear vLLM-origin overlap precedent first. | | vLLM-origin Sequence Parallelism staging | `enable_sp`
`ReduceScatter`
`AllGather`
`SequenceParallelismPass` | `vllm/compilation/passes/fusion/sequence_parallelism.py`
`docs/design/fusions.md` | Sequence-parallel rewrites all-reduce into RS -> local norm -> AG so later passes can overlap comm and compute | Treat RS / AG staging around norm blocks as an upstream overlap-enabling family. | -| vLLM-origin shared-expert aux-stream overlap | `aux_stream`
`shared_experts_stream`
shared expert near router | `vllm/utils/torch_utils.py`
`vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py` | MoE shared experts can run on a dedicated aux stream and overlap with router-side work | Treat shared-expert vs router overlap as an existing upstream sparse-model family. | +| vLLM-origin shared-expert aux-stream overlap | `aux_stream`
`shared_experts_stream`
shared expert near router | `vllm/model_executor/layers/fused_moe/runner/shared_experts.py`
`vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py` | MoE shared experts can record the cloned input on `shared_experts_stream`, wait on the caller stream, run in parallel with router-side work, and rejoin before merge | Treat shared-expert vs router overlap as an existing upstream sparse-model family. | | vLLM-origin DCP async all-to-all overlap | `dcp_alltoall`
`all_to_all_single`
`async_op=True` | `vllm/v1/attention/ops/dcp_alltoall.py` | Output / LSE exchange uses async all-to-all handles instead of serializing collective completion on the main path | Treat DCP all-to-all windows as an upstream async-collective family. | -## 10. vLLM-origin PR-backed / in-flight kernel-overlap families +## 9. vLLM-origin PR-backed / in-flight kernel-overlap families | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| PR `#35968` DSV3.2 multi-stream indexer overlap | `weights_proj`
`wk`
`k_norm`
`aux_stream` | `PR #35968`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/utils/torch_utils.py` | Open PR overlaps the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. | +| PR `#35968` DSV3.2 multi-stream indexer overlap | `weights_proj`
`wk`
`k_norm`
`aux_stream` | `PR #35968`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/utils/torch_utils.py` | Closed PR explored overlapping the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. | | PR `#39301` GLM5 router GEMM with PDL overlap | `TRTLLM_ENABLE_PDL`
`router_gemm`
`GLM5`
`FI AR RMS fusion` | `PR #39301`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`vllm/csrc/moe/dsv3_router_gemm_utils.h` | The GLM5 router GEMM path explicitly uses PDL so the router kernel can overlap with the preceding fused allreduce-plus-RMS block on supported GPUs | Treat router-GEMM launch overlap on GLM5-like traces as an in-flight upstream family first. | -## 11. Important toggles and caveats +## 10. Important toggles and caveats | Toggle / env | Location | Effect on trace interpretation | | --- | --- | --- | @@ -141,7 +145,7 @@ AutoDeploy rather than same-stream PDL windows. | `PassConfig.enable_sp` | `vllm/config/compilation.py` | Enables vLLM's sequence-parallel staging family that creates RS / AG overlap opportunities. | | `PassConfig.fuse_gemm_comms` | `vllm/config/compilation.py` | Enables AsyncTP GEMM + collective overlap and auto-enables `enable_sp` when valid. | -## 12. Suggested refresh commands +## 11. Suggested refresh commands These commands are only for maintainers refreshing this catalog by rescanning the local source trees. They are not used by the triage scripts at runtime. @@ -160,7 +164,7 @@ git -C "$FLASHINFER_REPO" log --all --format='%h %s' | rg -i 'fused|fusion|overl rg -n "multi_stream_moe|multi_stream_mla_attn|multi_stream_gemm|record_event_passthrough|begin_aux_stream_passthrough|end_aux_stream_passthrough|wait_aux_stream_passthrough" "$TRTLLM_REPO/tensorrt_llm/_torch" rg -n "mlir_elementwise_fusion|piecewise|cudagraph|caller_stream.synchronize" "$TRTLLM_REPO/tensorrt_llm/_torch" git -C "$TRTLLM_REPO" log --all --format='%h %s' | rg -i 'overlap|multi-stream|aux stream|cudagraph|mlir|stream|flashinfer|moe|mla' -rg -n "fuse_gemm_comms|enable_sp|fused_matmul_reduce_scatter|fused_all_gather_matmul|shared_experts_stream|dcp_alltoall|async_op=True|aux_stream|maybe_execute_in_parallel" "$VLLM_REPO/vllm" "$VLLM_REPO/docs/design/fusions.md" +rg -n "fuse_gemm_comms|enable_sp|fused_matmul_reduce_scatter|fused_all_gather_matmul|shared_experts_stream|maybe_sync_shared_experts_stream|dcp_alltoall|async_op=True|aux_stream|maybe_execute_in_parallel" "$VLLM_REPO/vllm" "$VLLM_REPO/docs/design/fusions.md" git -C "$VLLM_REPO" log --all --format='%h %s' | rg -i 'fused|fusion|overlap|allreduce|reduce-scatter|all-gather|all_to_all|stream|multi-stream|triton|cuda|router' # GitHub PR scan terms for the connector or web UI: # "fused OR overlap repo:sgl-project/sglang" diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/source-map.md b/.claude/skills/llm-torch-profiler-analysis/references/source-map.md similarity index 100% rename from .claude/skills/sglang-torch-profiler-analysis/references/source-map.md rename to .claude/skills/llm-torch-profiler-analysis/references/source-map.md diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py b/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py new file mode 100644 index 000000000..6b6a78550 --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py @@ -0,0 +1,806 @@ +"""Compact triage entrypoint for unified LLM torch-profiler analysis.""" + +from __future__ import annotations + +import argparse +import sys +from collections import defaultdict +from pathlib import Path +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 ( + discover_trace_targets, + framework_display_name, + load_server_args, + load_trace_json, + parse_stage, + resolve_framework, + run_profiler, +) + +MIN_RENDER_SHARE_PCT = 1.0 +MAPPING_KERNEL_SAMPLE_LIMIT_PER_NAME = 16 + + +def build_triage_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="analyze_llm_torch_profile.py", + description=( + "Compact LLM torch-profiler triage entrypoint for SGLang, vLLM, and " + "TensorRT-LLM. " + "This prints three tables: kernel mapping, overlap opportunities, " + "and fuse opportunities. " + "Use either a single trace/profile input or a mapping+formal two-trace pair." + ), + ) + parser.add_argument( + "--framework", + type=str, + default="auto", + choices=["auto", "sglang", "vllm", "trtllm", "tllm", "tensorrt-llm"], + help=( + "Serving framework. Use auto to detect from trace contents, path hints, " + "or URL features." + ), + ) + parser.add_argument( + "--input", + type=str, + default=None, + help="Single trace file or profile directory to triage.", + ) + parser.add_argument( + "--url", + type=str, + 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." + ), + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help=( + "Trace output dir when using --url. For vLLM this should match the " + "server's torch_profiler_dir. For TensorRT-LLM it should match the " + "directory or file path configured by TLLM_TORCH_PROFILE_TRACE." + ), + ) + parser.add_argument( + "--profile-prefix", + type=str, + default="triage-trace", + help=( + "Profile prefix when generating a trace from --url. SGLang uses it " + "directly; vLLM and TensorRT-LLM may ignore it on the HTTP profiler path." + ), + ) + parser.add_argument( + "--mapping-input", + type=str, + default=None, + help="Graph-off mapping trace file or directory.", + ) + parser.add_argument( + "--mapping-url", + type=str, + default=None, + help="Running graph-off server URL for the mapping trace.", + ) + parser.add_argument( + "--formal-input", + type=str, + default=None, + help="Formal graph-on trace file or directory.", + ) + parser.add_argument( + "--formal-url", + type=str, + default=None, + help="Running graph-on server URL for the formal trace.", + ) + parser.add_argument( + "--mapping-output-dir", + type=str, + default=None, + help="Trace output dir when using --mapping-url.", + ) + parser.add_argument( + "--formal-output-dir", + type=str, + default=None, + help="Trace output dir when using --formal-url.", + ) + parser.add_argument( + "--mapping-profile-prefix", + type=str, + default="mapping-trace", + help="Profile prefix for the mapping trace.", + ) + parser.add_argument( + "--formal-profile-prefix", + type=str, + default="formal-trace", + help="Profile prefix for the formal trace.", + ) + parser.add_argument( + "--num-steps", + type=int, + default=5, + help="Profiler steps when generating traces from URLs.", + ) + parser.add_argument( + "--profile-by-stage", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--merge-profiles", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument("--probe-requests", type=int, default=1) + parser.add_argument( + "--probe-prompt", + type=str, + default=( + "Repeat the word profiler many times with spaces so the server performs several decode steps. " + "Do not add explanations." + ), + ) + parser.add_argument("--probe-max-new-tokens", type=int, default=None) + parser.add_argument("--probe-delay", type=float, default=0.5) + parser.add_argument( + "--start-step", + type=int, + default=None, + help="SGLang-only profiler start step when generating traces from URLs.", + ) + parser.add_argument( + "--pid-substring", + type=str, + default=None, + help="Restrict overlap analysis to PIDs containing this substring.", + ) + parser.add_argument( + "--kernel-table-limit", + type=int, + default=0, + help="How many kernel rows to print per stage. Use 0 for all kernels.", + ) + parser.add_argument( + "--overlap-table-limit", + type=int, + default=0, + help="How many overlap rows to print per stage. Use 0 for all kernels.", + ) + return parser + + +def parse_triage_args(argv: Sequence[str]) -> argparse.Namespace: + parser = build_triage_parser() + args = parser.parse_args(argv) + + single_trace_mode = bool(args.input) or bool(args.url) + dual_trace_mode = any( + [ + args.mapping_input, + args.mapping_url, + args.formal_input, + args.formal_url, + ] + ) + + if single_trace_mode and dual_trace_mode: + parser.error( + "Use either single-trace mode (--input/--url) or two-trace mode " + "(--mapping-* plus --formal-*), not both." + ) + + if single_trace_mode: + if bool(args.input) == bool(args.url): + parser.error("Provide exactly one of --input or --url.") + return args + + if bool(args.mapping_input) == bool(args.mapping_url): + parser.error("Provide exactly one of --mapping-input or --mapping-url.") + if bool(args.formal_input) == bool(args.formal_url): + parser.error("Provide exactly one of --formal-input or --formal-url.") + return args + + +def resolve_profile_targets( + *, + label: str, + input_path: Optional[str], + url: Optional[str], + output_dir: Optional[str], + profile_prefix: Optional[str], + args: argparse.Namespace, +) -> Tuple[List[Path], Optional[dict], str]: + if bool(input_path) == bool(url): + raise ValueError(f"{label} trace requires exactly one of input path or URL.") + + if url: + framework = resolve_framework( + args.framework, + input_path=Path(output_dir).resolve() if output_dir else None, + url=url, + ) + target_dir = run_profiler( + url=url, + output_dir=output_dir, + num_steps=args.num_steps, + profile_by_stage=args.profile_by_stage, + merge_profiles=args.merge_profiles, + profile_prefix=profile_prefix, + probe_requests=max(0, args.probe_requests), + probe_prompt=args.probe_prompt, + probe_max_new_tokens=args.probe_max_new_tokens, + probe_delay=args.probe_delay, + start_step=args.start_step, + framework=framework, + framework_hint_path=output_dir, + ) + traces, server_args = discover_trace_targets(target_dir, all_traces=False) + resolved_framework = resolve_framework( + args.framework, + input_path=target_dir, + url=url, + server_args=server_args, + ) + return traces, server_args, resolved_framework + + resolved = Path(input_path).resolve() + traces, server_args = discover_trace_targets(resolved, all_traces=False) + if server_args is None: + server_args = load_server_args(resolved) + framework = resolve_framework( + args.framework, input_path=resolved, server_args=server_args + ) + return traces, server_args, framework + + +def build_mapping_kernel_map(trace_paths: Sequence[Path], framework: str) -> dict: + stage_site_stats = defaultdict( + lambda: defaultdict(lambda: defaultdict(kernel_helpers.MappingSiteAggregate)) + ) + stage_kernel_categories: Dict[str, Dict[str, str]] = defaultdict(dict) + global_site_stats = defaultdict( + lambda: defaultdict(kernel_helpers.MappingSiteAggregate) + ) + global_kernel_categories: Dict[str, str] = {} + + for trace_path in trace_paths: + trace = load_trace_json(trace_path) + kernels, cpu_ops, python_frames, launch_events, _, _ = ( + kernel_helpers.extract_trace_data(trace) + ) + if not kernels: + continue + cpu_ops_by_external_id = kernel_helpers.build_cpu_op_index(cpu_ops) + launches_by_correlation = kernel_helpers.build_launch_index(launch_events) + site_context_cache = {} + default_stage = parse_stage(trace_path) + for stage, stage_kernels in kernel_helpers.group_kernels_by_stage( + kernels, default_stage + ).items(): + sampled_stage_kernels = ( + stage_kernels + if framework == "sglang" + else sample_kernels_for_mapping(stage_kernels) + ) + local_site_stats = kernel_helpers.aggregate_kernel_sites( + sampled_stage_kernels, + cpu_ops_by_external_id, + python_frames, + launches_by_correlation=launches_by_correlation, + site_context_cache=site_context_cache, + ) + kernel_categories = { + kernel.canonical_name: kernel.category for kernel in stage_kernels + } + kernel_helpers.merge_site_stats(stage_site_stats[stage], local_site_stats) + kernel_helpers.merge_site_stats(global_site_stats, local_site_stats) + stage_kernel_categories[stage].update(kernel_categories) + global_kernel_categories.update(kernel_categories) + + stage_payloads = { + stage: kernel_helpers.build_stage_payload( + dict(site_stats), stage_kernel_categories.get(stage, {}) + ) + for stage, site_stats in stage_site_stats.items() + } + global_payload = kernel_helpers.build_stage_payload( + dict(global_site_stats), global_kernel_categories + ) + return {"stages": stage_payloads, "global": global_payload} + + +def stage_index(stage: str) -> int: + return {"extend": 0, "prefill": 0, "decode": 1, "all": 2}.get(stage, 99) + + +def sample_kernels_for_mapping( + kernels: Sequence[kernel_helpers.KernelEvent], + per_name_limit: int = MAPPING_KERNEL_SAMPLE_LIMIT_PER_NAME, +) -> List[kernel_helpers.KernelEvent]: + if per_name_limit <= 0: + return list(kernels) + + grouped: Dict[str, List[kernel_helpers.KernelEvent]] = defaultdict(list) + for kernel in kernels: + grouped[kernel.canonical_name].append(kernel) + + sampled: List[kernel_helpers.KernelEvent] = [] + for kernel_name in sorted(grouped): + items = grouped[kernel_name] + if len(items) <= per_name_limit: + sampled.extend(items) + continue + for sample_idx in range(per_name_limit): + pos = round(sample_idx * (len(items) - 1) / (per_name_limit - 1)) + sampled.append(items[pos]) + sampled.sort(key=lambda kernel: (kernel.ts, kernel.name)) + return sampled + + +def stage_display(stage: str) -> str: + return kernel_helpers.stage_label(stage) + + +def pick_stage_value(stage_to_value: Dict[str, object], stage: str) -> Optional[object]: + if stage in stage_to_value: + return stage_to_value[stage] + if "all" in stage_to_value: + return stage_to_value["all"] + if len(stage_to_value) == 1: + return next(iter(stage_to_value.values())) + return None + + +def render_stages(stage_to_value: Dict[str, object]) -> List[str]: + stages = set(stage_to_value) + if any(stage != "all" for stage in stages): + stages.discard("all") + return sorted(stages, key=stage_index) + + +def build_overlap_stage_bundle_map( + trace_paths: Sequence[Path], + *, + label_prefix: str, + server_args: Optional[dict], + pid_substring: Optional[str], +) -> Dict[str, overlap_helpers.TraceBundle]: + stage_bundles: Dict[str, overlap_helpers.TraceBundle] = {} + for trace_path in sorted( + trace_paths, key=lambda item: (stage_index(parse_stage(item)), item.name) + ): + trace_json = load_trace_json(trace_path) + raw_events = trace_json.get( + "traceEvents", + trace_json if isinstance(trace_json, list) else [], + ) + events, pid = overlap_helpers.extract_kernel_events(trace_json, pid_substring) + if not events: + continue + default_stage = parse_stage(trace_path) + stage_groups = overlap_helpers.group_events_by_stage(events, default_stage) + for stage in render_stages(stage_groups): + if stage in stage_bundles: + continue + stage_bundles[stage] = overlap_helpers.TraceBundle( + label=f"{label_prefix}-{stage}", + trace_path=trace_path, + server_args=server_args, + raw_events=raw_events, + events=stage_groups[stage], + pid=pid, + ) + if "all" in stage_groups and not stage_bundles: + stage_bundles["all"] = overlap_helpers.TraceBundle( + label=f"{label_prefix}-all", + trace_path=trace_path, + server_args=server_args, + raw_events=raw_events, + events=stage_groups["all"], + pid=pid, + ) + return stage_bundles + + +def group_rows_by_stage(rows: Sequence[dict]) -> List[Tuple[str, List[dict]]]: + grouped: Dict[str, List[dict]] = defaultdict(list) + for row in rows: + grouped[str(row.get("stage") or "all")].append(row) + return [ + (stage, grouped[stage]) for stage in sorted(grouped.keys(), key=stage_index) + ] + + +def render_kernel_table_for_stage(rows: Sequence[dict]) -> List[str]: + lines = [ + "| Kernel | Category | GPU time | Share | Launches | Python location (site share) | CPU op |", + "| --- | --- | ---: | ---: | ---: | --- | --- |", + ] + if not rows: + lines.append( + "| No kernel rows at or above 1.0% share. | - | - | - | - | - | - |" + ) + return lines + for row in rows: + lines.append( + "| {kernel} | {category} | {gpu_time} | {share:.1f}% | {launches} | {location} | {cpu_op} |".format( + kernel=kernel_helpers.escape_md_cell(row["kernel"]), + category=kernel_helpers.escape_md_cell(row["category"]), + gpu_time=kernel_helpers.format_ms(row["total_us"]), + share=row["share_pct"], + launches=row["launches"], + location=kernel_helpers.escape_md_cell(row["location"]), + cpu_op=kernel_helpers.escape_md_cell(row["cpu_op"]), + ) + ) + return lines + + +def render_stage_section_tables( + rows: Sequence[dict], + *, + render_stage_fn, + stage_label_prefix: str = "#####", +) -> List[str]: + if not rows: + return render_stage_fn([]) + stage_groups = group_rows_by_stage(rows) + if len(stage_groups) == 1 and stage_groups[0][0] == "all": + return render_stage_fn(stage_groups[0][1]) + + lines: List[str] = [] + for index, (stage, stage_rows) in enumerate(stage_groups): + lines.append(f"{stage_label_prefix} {stage_display(stage)}") + lines.extend(render_stage_fn(stage_rows)) + if index != len(stage_groups) - 1: + lines.append("") + return lines + + +def render_kernel_tables(rows: Sequence[dict]) -> List[str]: + return render_stage_section_tables( + rows, render_stage_fn=render_kernel_table_for_stage + ) + + +def render_overlap_table_for_stage(rows: Sequence[dict]) -> List[str]: + lines = [ + "| Priority | Verdict | Kernel | Python scope | Formal signal | Dep risk | Recommendation |", + "| --- | --- | --- | --- | --- | --- | --- |", + ] + if not rows: + lines.append( + "| - | - | No rows cleared the 1.0% reporting bar. Use mapping/formal mode for overlap attribution. | - | - | - | - |" + ) + return lines + for row in rows: + formal_signal = ( + f"{row['total_us']:.1f} us, share {row['share_pct']:.1f}%, " + f"excl {row['exclusive_ratio'] * 100:.1f}% / hid {row['hidden_ratio'] * 100:.1f}%" + ) + lines.append( + "| " + + " | ".join( + [ + row["priority"], + row["verdict"], + kernel_helpers.escape_md_cell(row["kernel"]), + kernel_helpers.escape_md_cell(row["python_scope"]), + kernel_helpers.escape_md_cell(formal_signal), + overlap_helpers.dependency_risk_label(row["dependency_signal"]), + row["recommendation"], + ] + ) + + " |" + ) + return lines + + +def render_overlap_tables(rows: Sequence[dict]) -> List[str]: + return render_stage_section_tables( + rows, + render_stage_fn=render_overlap_table_for_stage, + ) + + +def render_fuse_table_for_stage(rows: Sequence[dict]) -> List[str]: + lines = [ + "| Pattern | Confidence | Related GPU time | Share | Evidence kernels | Current kernel Python location | Candidate fused Python path | Rationale |", + "| --- | --- | ---: | ---: | --- | --- | --- | --- |", + ] + if not rows: + lines.append( + "| No medium-confidence source-backed fusion opportunity matched this trace. | - | - | - | - | - | - | - |" + ) + return lines + for row in rows: + lines.append( + "| {pattern} | {confidence} | {gpu_time} | {share:.1f}% | {evidence} | {current_locations} | {candidate_path} | {rationale} |".format( + pattern=kernel_helpers.escape_md_cell(row["pattern"]), + confidence=kernel_helpers.escape_md_cell(row["confidence"]), + gpu_time=kernel_helpers.format_ms(row["related_us"]), + share=row["share_pct"], + evidence=kernel_helpers.escape_md_cell(row["evidence"]), + current_locations=kernel_helpers.escape_md_cell( + row["current_locations"] + ), + candidate_path=kernel_helpers.escape_md_cell(row["candidate_path"]), + rationale=kernel_helpers.escape_md_cell(row["rationale"]), + ) + ) + return lines + + +def render_fuse_tables(rows: Sequence[dict]) -> List[str]: + return render_stage_section_tables( + rows, + render_stage_fn=render_fuse_table_for_stage, + ) + + +def run_triage(args: argparse.Namespace) -> int: + single_trace_mode = bool(args.input) or bool(args.url) + if single_trace_mode: + formal_traces, formal_server_args, formal_framework = resolve_profile_targets( + label="input", + input_path=args.input, + url=args.url, + output_dir=args.output_dir, + profile_prefix=args.profile_prefix, + args=args, + ) + mapping_traces = formal_traces + mapping_server_args = formal_server_args + mapping_framework = formal_framework + else: + mapping_traces, mapping_server_args, mapping_framework = ( + resolve_profile_targets( + label="mapping", + input_path=args.mapping_input, + url=args.mapping_url, + output_dir=args.mapping_output_dir, + profile_prefix=args.mapping_profile_prefix, + args=args, + ) + ) + formal_traces, formal_server_args, formal_framework = resolve_profile_targets( + label="formal", + input_path=args.formal_input, + url=args.formal_url, + output_dir=args.formal_output_dir, + profile_prefix=args.formal_profile_prefix, + args=args, + ) + + mapping_kernel_map = build_mapping_kernel_map(mapping_traces, mapping_framework) + + kernel_rows_rendered: List[dict] = [] + fuse_rows_rendered: List[dict] = [] + formal_stage_payloads: Dict[str, dict] = {} + + for formal_trace in formal_traces: + trace = load_trace_json(formal_trace) + kernels, cpu_ops, python_frames, launch_events, _, _ = ( + kernel_helpers.extract_trace_data(trace) + ) + if not kernels: + continue + default_stage = parse_stage(formal_trace) + stage_groups = kernel_helpers.group_kernels_by_stage(kernels, default_stage) + formal_cpu_ops_by_external_id = kernel_helpers.build_cpu_op_index(cpu_ops) + formal_launches_by_correlation = kernel_helpers.build_launch_index( + launch_events + ) + formal_site_context_cache = {} + for stage_name, stage_kernels in stage_groups.items(): + local_site_stats = kernel_helpers.aggregate_kernel_sites( + stage_kernels, + formal_cpu_ops_by_external_id, + python_frames, + launches_by_correlation=formal_launches_by_correlation, + site_context_cache=formal_site_context_cache, + ) + formal_stage_payloads[stage_name] = kernel_helpers.build_stage_payload( + local_site_stats, + {kernel.canonical_name: kernel.category for kernel in stage_kernels}, + ) + trace_total_us = sum(kernel.dur for kernel in kernels) + for stage in sorted(stage_groups, key=stage_index): + stage_kernels = stage_groups[stage] + if not stage_kernels: + continue + total_us = sum(kernel.dur for kernel in stage_kernels) + if ( + stage == "all" + and default_stage == "all" + and kernel_helpers.pct(total_us, trace_total_us) < MIN_RENDER_SHARE_PCT + ): + continue + kernel_stats = kernel_helpers.aggregate( + stage_kernels, key_fn=lambda item: item.canonical_name + ) + kernel_categories = { + kernel.canonical_name: kernel.category for kernel in stage_kernels + } + full_kernel_rows = kernel_helpers.build_kernel_rows( + stage=stage, + kernel_stats=kernel_stats, + kernel_categories=kernel_categories, + local_stage_payload=formal_stage_payloads.get(stage, {"kernels": {}}), + external_kernel_map=mapping_kernel_map, + ) + visible_kernel_rows = kernel_helpers.limit_kernel_rows( + full_kernel_rows, args.kernel_table_limit + ) + for row in visible_kernel_rows: + share_pct = kernel_helpers.pct(row.total_us, total_us) + if share_pct < MIN_RENDER_SHARE_PCT: + continue + kernel_rows_rendered.append( + { + "stage": stage, + "kernel": row.name, + "category": row.category, + "total_us": row.total_us, + "share_pct": share_pct, + "launches": row.aggregate.count, + "location": row.location, + "cpu_op": row.cpu_op, + } + ) + for item in kernel_helpers.detect_fusion_opportunities( + kernel_rows=full_kernel_rows, + total_us=total_us, + server_args=formal_server_args or mapping_server_args, + framework=formal_framework, + ): + share_pct = kernel_helpers.pct(item.related_us, total_us) + if share_pct < MIN_RENDER_SHARE_PCT: + continue + fuse_rows_rendered.append( + { + "stage": stage, + "pattern": item.pattern, + "confidence": item.confidence, + "related_us": item.related_us, + "share_pct": share_pct, + "evidence": item.evidence, + "current_locations": item.current_locations, + "candidate_path": item.candidate_path, + "rationale": item.rationale, + } + ) + + overlap_rows_rendered: List[dict] = [] + if not single_trace_mode: + mapping_overlap_bundles = build_overlap_stage_bundle_map( + mapping_traces, + label_prefix="mapping", + server_args=mapping_server_args, + pid_substring=args.pid_substring, + ) + formal_overlap_bundles = build_overlap_stage_bundle_map( + formal_traces, + label_prefix="formal", + server_args=formal_server_args, + pid_substring=args.pid_substring, + ) + for stage in render_stages(formal_overlap_bundles): + formal_bundle = pick_stage_value(formal_overlap_bundles, stage) + mapping_bundle = pick_stage_value(mapping_overlap_bundles, stage) + if formal_bundle is None or mapping_bundle is None: + continue + formal_bundle.overlap_stats = overlap_helpers.analyze_overlap( + formal_bundle.events + ) + aggregates = overlap_helpers.aggregate_events(formal_bundle.events) + source_map = overlap_helpers.build_kernel_source_map( + mapping_bundle, + kernel_map_entry_lookup=lambda stage_name, kernel_name: ( + kernel_helpers.lookup_kernel_map_entry( + mapping_kernel_map, stage_name, kernel_name + ) + if mapping_kernel_map + else None + ), + stage=stage, + ) + source_map = overlap_helpers.merge_source_map_from_kernel_payload( + source_map, + pick_stage_value(formal_stage_payloads, stage), + ) + stage_rows = overlap_helpers.build_action_rows( + aggregates, + source_map, + formal_bundle.events, + formal_bundle.overlap_stats["total_busy_us"], + table_limit=max(0, args.overlap_table_limit), + ) + for row in stage_rows: + if row.share_pct < MIN_RENDER_SHARE_PCT: + continue + overlap_rows_rendered.append( + { + "stage": stage, + "priority": row.priority, + "verdict": row.verdict, + "kernel": row.kernel, + "python_scope": row.python_scope, + "total_us": row.total_us, + "share_pct": row.share_pct, + "exclusive_ratio": row.exclusive_ratio, + "hidden_ratio": row.hidden_ratio, + "dependency_signal": row.dependency_signal, + "recommendation": row.recommendation, + } + ) + + lines: List[str] = [] + lines.append("Triage View") + lines.append(f"Mode: {'single-trace' if single_trace_mode else 'mapping-formal'}") + if single_trace_mode: + lines.append(f"Framework: {framework_display_name(formal_framework)}") + lines.append(f"Input traces: {', '.join(str(path) for path in formal_traces)}") + else: + if mapping_framework == formal_framework: + lines.append(f"Framework: {framework_display_name(formal_framework)}") + else: + lines.append( + f"Mapping framework: {framework_display_name(mapping_framework)}" + ) + lines.append( + f"Formal framework: {framework_display_name(formal_framework)}" + ) + lines.append( + f"Mapping traces: {', '.join(str(path) for path in mapping_traces)}" + ) + lines.append(f"Formal traces: {', '.join(str(path) for path in formal_traces)}") + if formal_server_args or mapping_server_args: + server_args = formal_server_args or mapping_server_args + model = server_args.get("model_path") or server_args.get("model") + if model: + lines.append(f"Model: {model}") + lines.append("") + lines.append("Kernel Table") + lines.extend(render_kernel_tables(kernel_rows_rendered)) + lines.append("") + lines.append("Overlap Opportunity Table") + lines.extend(render_overlap_tables(overlap_rows_rendered)) + lines.append("") + lines.append("Fuse Opportunity Table") + lines.extend(render_fuse_tables(fuse_rows_rendered)) + print("\n".join(lines).rstrip()) + return 0 + + +def main(argv: Optional[Sequence[str]] = None) -> int: + argv = list(argv or sys.argv[1:]) + triage_parser = build_triage_parser() + + if not argv or argv[0] in {"-h", "--help"}: + triage_parser.print_help() + return 0 + + if argv[0] == "triage": + argv = argv[1:] + elif not argv[0].startswith("-"): + triage_parser.error( + "This skill exposes only the triage workflow. " + "Use single-trace mode (--input/--url) or mapping+formal two-trace mode." + ) + return 2 + + return run_triage(parse_triage_args(argv)) + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py b/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py new file mode 100644 index 000000000..35aabc4c5 --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py @@ -0,0 +1,16 @@ +"""Backwards-compatibility shim for the unified LLM torch-profiler entrypoint. + +The real implementation now lives in ``analyze_llm_torch_profile`` because this +skill covers SGLang, vLLM, and TensorRT-LLM. Older scripts and runbooks that +still invoke ``analyze_sglang_torch_profile.py`` keep working by forwarding to +that module. +""" + +from __future__ import annotations + +import sys + +from analyze_llm_torch_profile import main + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/make_trtllm_py_executor_override.py b/.claude/skills/llm-torch-profiler-analysis/scripts/make_trtllm_py_executor_override.py new file mode 100644 index 000000000..597665d67 --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/make_trtllm_py_executor_override.py @@ -0,0 +1,132 @@ +"""Generate a TensorRT-LLM py_executor override for stable torch-profiler capture.""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +START_MARKER = "torch_profiler = torch.profiler.profile(" + + +@dataclass +class ProfileCallSpan: + start: int + end: int + block: str + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Create a py_executor.py override that enables with_stack=True for " + "TensorRT-LLM torch-profiler traces." + ) + ) + parser.add_argument("--source", required=True, help="Original py_executor.py path.") + parser.add_argument("--output", required=True, help="Override file path to write.") + return parser.parse_args() + + +def find_profile_call_span(text: str) -> ProfileCallSpan: + start = text.find(START_MARKER) + if start == -1: + raise SystemExit("Could not find torch profiler setup in source file.") + + open_paren = text.find("(", start) + if open_paren == -1: + raise SystemExit("Malformed torch profiler setup in source file.") + + depth = 0 + for index in range(open_paren, len(text)): + char = text[index] + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + return ProfileCallSpan( + start=start, + end=index + 1, + block=text[start : index + 1], + ) + raise SystemExit("Could not find the end of the torch profiler call.") + + +def inject_with_stack(block: str) -> str: + if "with_stack=" in block: + return block + + lines = block.splitlines() + if not lines: + raise SystemExit("Unexpected torch profiler block format.") + + last_line = lines[-1] + if not last_line.strip(): + raise SystemExit("Unexpected torch profiler block terminator.") + + if last_line.strip() == ")": + if len(lines) < 2: + raise SystemExit("Could not find the last torch profiler argument line.") + last_arg_index = len(lines) - 2 + last_arg_line = lines[last_arg_index] + indent = last_arg_line[: len(last_arg_line) - len(last_arg_line.lstrip())] + if not last_arg_line.rstrip().endswith(","): + lines[last_arg_index] = last_arg_line.rstrip() + "," + lines.insert(len(lines) - 1, f"{indent}with_stack=True") + return "\n".join(lines) + + if not last_line.rstrip().endswith(")"): + raise SystemExit("Unexpected torch profiler block terminator.") + + indent = last_line[: len(last_line) - len(last_line.lstrip())] + last_arg_text = last_line.rstrip()[:-1].rstrip() + if not last_arg_text.endswith(","): + last_arg_text += "," + lines[-1] = last_arg_text + lines.append(f"{indent}with_stack=True)") + return "\n".join(lines) + + +def inject_rank0_trace_guard(text: str) -> str: + needle = ( + " enable_torch_trace = bool(torch_trace_path and profile_start_stop)\n" + ) + replacement = ( + " # Multi-rank PyTorch backend workers race on the same chrome-trace " + "path.\n" + " # Keep the full torch-profiler trace on rank 0 and let the other " + "ranks\n" + " # continue with CUDA-profiler gating only.\n" + " enable_torch_trace = bool(\n" + " torch_trace_path and profile_start_stop and self.dist.rank == 0\n" + " )\n" + ) + if replacement in text: + return text + if needle not in text: + raise SystemExit("Could not find enable_torch_trace assignment in source file.") + return text.replace(needle, replacement, 1) + + +def main() -> int: + args = parse_args() + source = Path(args.source).expanduser().resolve() + output = Path(args.output).expanduser().resolve() + text = source.read_text(encoding="utf-8") + span = find_profile_call_span(text) + patched_block = inject_with_stack(span.block) + patched = ( + text + if patched_block == span.block + else (text[: span.start] + patched_block + text[span.end :]) + ) + patched = inject_rank0_trace_guard(patched) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(patched, encoding="utf-8") + print(output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/probe_llm_server.py b/.claude/skills/llm-torch-profiler-analysis/scripts/probe_llm_server.py new file mode 100755 index 000000000..97008b789 --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/probe_llm_server.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""Run a small correctness and latency probe against an LLM server.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time +from pathlib import Path +from typing import Any, Dict, List, Optional +from urllib import request + +from profile_common import extract_openai_chat_text + +DEFAULT_PROMPTS = [ + "用一句中文介绍上海。", + "What is 2+2? Answer briefly.", + "Write one short haiku about GPUs.", +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Send a few short requests to an LLM server and record latency plus " + "sample outputs." + ) + ) + parser.add_argument( + "--framework", + required=True, + choices=("sglang", "vllm", "trtllm"), + help="Serving framework.", + ) + parser.add_argument( + "--url", + required=True, + help="Server base URL, for example http://127.0.0.1:30000.", + ) + parser.add_argument( + "--model", + default=None, + help="OpenAI model id. Auto-discovered for vLLM and TensorRT-LLM when omitted.", + ) + parser.add_argument( + "--requests", + type=int, + default=6, + help="How many probe requests to send.", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=48, + help="Generation length for each request.", + ) + parser.add_argument( + "--timeout", + type=float, + default=180.0, + help="Per-request timeout in seconds.", + ) + parser.add_argument( + "--prompt", + action="append", + default=[], + help="Optional prompt override. Repeat to add more prompts.", + ) + parser.add_argument( + "--output", + default=None, + help="Optional JSON output path.", + ) + return parser.parse_args() + + +def post_json(url: str, payload: Dict[str, Any], timeout: float) -> Dict[str, Any]: + req = request.Request( + url=url, + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with request.urlopen(req, timeout=timeout) as resp: + raw = resp.read() + return json.loads(raw.decode("utf-8")) if raw else {} + + +def get_json(url: str, timeout: float) -> Dict[str, Any]: + req = request.Request(url=url, method="GET") + with request.urlopen(req, timeout=timeout) as resp: + raw = resp.read() + return json.loads(raw.decode("utf-8")) if raw else {} + + +def discover_openai_model(base_url: str, timeout: float) -> str: + payload = get_json(base_url.rstrip("/") + "/v1/models", timeout=timeout) + data = payload.get("data") + if not isinstance(data, list) or not data: + raise RuntimeError(f"No models returned by {base_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 {base_url.rstrip('/')}") + + +def p95(values: List[float]) -> Optional[float]: + if not values: + return None + ordered = sorted(values) + index = max(0, math.ceil(len(ordered) * 0.95) - 1) + return ordered[index] + + +def sglang_request(base_url: str, prompt: str, max_tokens: int, timeout: float) -> str: + payload = { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": max_tokens, + }, + "stream": False, + } + body = post_json(base_url.rstrip("/") + "/generate", payload, timeout=timeout) + return str(body.get("text", "")) + + +def openai_request( + base_url: str, + model: str, + prompt: str, + max_tokens: int, + timeout: float, +) -> Dict[str, str]: + payload = { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "temperature": 0.0, + "max_tokens": max_tokens, + "stream": False, + } + body = post_json( + base_url.rstrip("/") + "/v1/chat/completions", + payload, + timeout=timeout, + ) + text, source = extract_openai_chat_text(body) + return {"text": text, "source": source} + + +def run_probe(args: argparse.Namespace) -> Dict[str, Any]: + prompts = args.prompt or list(DEFAULT_PROMPTS) + model = args.model + if args.framework in {"vllm", "trtllm"} and not model: + model = discover_openai_model(args.url, timeout=args.timeout) + + latencies: List[float] = [] + samples: List[Dict[str, Any]] = [] + errors: List[Dict[str, str]] = [] + + for request_idx in range(args.requests): + prompt = prompts[request_idx % len(prompts)] + start = time.time() + try: + if args.framework == "sglang": + text = sglang_request( + args.url, + prompt, + max_tokens=args.max_tokens, + timeout=args.timeout, + ) + source = "generate.text" + else: + assert model is not None + result = openai_request( + args.url, + model, + prompt, + max_tokens=args.max_tokens, + timeout=args.timeout, + ) + text = result["text"] + source = result["source"] + elapsed = time.time() - start + latencies.append(elapsed) + samples.append( + { + "prompt": prompt, + "latency_s": round(elapsed, 3), + "content": text[:240], + "source": source, + "non_empty": bool(text.strip()), + } + ) + except Exception as exc: # pragma: no cover - runtime probe path + errors.append({"prompt": prompt, "error": repr(exc)}) + + return { + "framework": args.framework, + "url": args.url, + "model": model, + "requests": args.requests, + "success": len(samples), + "errors": len(errors), + "all_non_empty": ( + all(sample["non_empty"] for sample in samples) if samples else False + ), + "avg_latency_s": round(statistics.mean(latencies), 3) if latencies else None, + "p95_latency_s": round(p95(latencies), 3) if latencies else None, + "samples": samples[:3], + "error_samples": errors[:3], + } + + +def main() -> int: + args = parse_args() + summary = run_probe(args) + rendered = json.dumps(summary, ensure_ascii=False, indent=2) + print(rendered) + if args.output: + output_path = Path(args.output).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(rendered + "\n", encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/profile_common.py b/.claude/skills/llm-torch-profiler-analysis/scripts/profile_common.py new file mode 100644 index 000000000..ee589e26c --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/profile_common.py @@ -0,0 +1,880 @@ +"""Shared helpers for unified LLM torch-profiler skill scripts.""" + +from __future__ import annotations + +import gzip +import json +import re +import sys +import tempfile +import time +from collections import Counter, defaultdict +from functools import lru_cache +from pathlib import Path +from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple +from urllib import request + +STAGE_ORDER = {"extend": 0, "prefill": 0, "decode": 1, "all": 2} +FRAMEWORK_LABELS = { + "auto": "auto", + "sglang": "SGLang", + "vllm": "vLLM", + "trtllm": "TensorRT-LLM", +} +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:") + + +@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", + } + 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/"), + ("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 ( + "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} + if lowered_keys & { + "attention_backend", + "sampling_backend", + "disable_cuda_graph", + "disable_piecewise_cuda_graph", + "chunked_prefill_size", + "schedule_policy", + }: + return "sglang" + return detect_framework_from_text(json.dumps(server_args, sort_keys=True)) + + +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" + 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: + name = path.name.lower() + if "-extend" in name or "-prefill" in name: + return "extend" + if "-decode" in name: + 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) + + 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 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) -> None: + try: + post_json(url.rstrip("/") + "/start_profile", 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=1 and TLLM_TORCH_PROFILE_TRACE=/shared/path." + ) from exc + raise + + +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_requests: int, + probe_prompt: str, + probe_max_new_tokens: Optional[int], + probe_delay: float, + num_steps: int, +) -> Path: + framework = canonicalize_framework(framework) + output_path = ensure_remote_profiler_output_path(output_dir, framework) + start_remote_profiler(url, framework) + stop_error: Optional[BaseException] = None + try: + if probe_requests > 0: + # Some profiler endpoints need a brief setup window after + # 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 + ) + 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) + except BaseException as exc: # pragma: no cover - preserve original failure + stop_error = exc + if stop_error is not None: + raise stop_error + return wait_for_profiler_artifact(output_path) + + +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_requests: int, + probe_prompt: str, + probe_max_new_tokens: Optional[int], + 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) + + 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_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", + ) + + 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, + start_step: Optional[int] = None, + framework: str = "auto", + framework_hint_path: Optional[str] = None, +) -> 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": + 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_requests=probe_requests, + probe_prompt=probe_prompt, + probe_max_new_tokens=probe_max_new_tokens, + 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 or TensorRT-LLM." + ) + if merge_profiles: + raise ValueError( + "--merge-profiles is only supported for SGLang live capture. " + "Disable it when profiling vLLM or TensorRT-LLM." + ) + if profile_prefix: + print( + f"Note: {framework_display_name(resolved_framework)} ignores " + "--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, + ) + + +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]) diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/render_triage_markdown_bundle.py b/.claude/skills/llm-torch-profiler-analysis/scripts/render_triage_markdown_bundle.py new file mode 100644 index 000000000..cd12429c8 --- /dev/null +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/render_triage_markdown_bundle.py @@ -0,0 +1,259 @@ +"""Bundle one or more triage text reports into a single markdown document.""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +FRAMEWORK_LABELS = { + "sglang": "SGLang", + "vllm": "vLLM", + "trtllm": "TensorRT-LLM", +} + +FRAMEWORK_ORDER = {"sglang": 0, "vllm": 1, "trtllm": 2} + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Render multiple profiler triage text outputs into one markdown file. " + "Input files are expected to be the existing analysis_*.txt outputs " + "already emitted by analyze_llm_torch_profile.py." + ) + ) + parser.add_argument( + "--analysis-root", + type=str, + default=None, + help=( + "Root directory to scan recursively for analysis_*.txt files. " + "Parent directory names are used as model section ids." + ), + ) + parser.add_argument( + "--analysis-file", + action="append", + default=[], + help=( + "Explicit analysis file entry. Use either PATH or LABEL=PATH. " + "When LABEL is omitted, the parent directory name is used." + ), + ) + parser.add_argument( + "--title", + type=str, + default="Unified LLM Torch Profiler Triage Bundle", + help="Top-level markdown title.", + ) + parser.add_argument( + "--output", + type=str, + default=None, + help="Write the bundled markdown to this file. Prints to stdout when omitted.", + ) + parser.add_argument( + "--include-toc", + action=argparse.BooleanOptionalAction, + default=True, + help="Include a simple table of contents.", + ) + args = parser.parse_args(argv) + if not args.analysis_root and not args.analysis_file: + parser.error("Provide at least one of --analysis-root or --analysis-file.") + return args + + +def framework_key_from_path(path: Path) -> str: + lowered = path.name.lower() + if "sglang" in lowered: + return "sglang" + if "vllm" in lowered: + return "vllm" + if "trtllm" in lowered or "tensorrt" in lowered: + return "trtllm" + return "other" + + +def framework_label(framework_key: str) -> str: + return FRAMEWORK_LABELS.get(framework_key, framework_key) + + +def discover_analysis_files(root: Path) -> List[Tuple[str, Path]]: + entries: List[Tuple[str, Path]] = [] + for path in sorted(root.rglob("analysis*.txt")): + entries.append((path.parent.name, path)) + return entries + + +def parse_explicit_entry(raw: str) -> Tuple[str, Path]: + if "=" in raw: + label, path_text = raw.split("=", 1) + path = Path(path_text).expanduser().resolve() + return label.strip(), path + path = Path(raw).expanduser().resolve() + return path.parent.name, path + + +def slugify(text: str) -> str: + chars = [] + last_dash = False + for char in text.lower(): + if char.isalnum(): + chars.append(char) + last_dash = False + elif not last_dash: + chars.append("-") + last_dash = True + return "".join(chars).strip("-") + + +def extract_model_name(report_text: str) -> Optional[str]: + for line in report_text.splitlines(): + if line.startswith("Model: "): + return line.split("Model: ", 1)[1].strip() + return None + + +def choose_model_display_name( + current: Optional[str], + candidate: Optional[str], + *, + label: str, +) -> str: + if candidate and candidate != label: + if not current or current == label: + return candidate + if len(candidate) > len(current): + return candidate + return current + if current: + return current + return label + + +def normalize_report_text(report_text: str) -> str: + text = report_text.replace("\r\n", "\n").strip() + if not text: + return "_Empty analysis output._" + heading_map = { + "Triage View": "#### Triage View", + "Kernel Table": "#### Kernel Table", + "Overlap Opportunity Table": "#### Overlap Opportunity Table", + "Fuse Opportunity Table": "#### Fuse Opportunity Table", + } + normalized_lines = [] + for line in text.splitlines(): + normalized_lines.append(heading_map.get(line, line)) + return "\n".join(normalized_lines) + + +def build_bundle_markdown( + *, + title: str, + labeled_paths: Sequence[Tuple[str, Path]], + include_toc: bool, +) -> str: + grouped: Dict[str, List[Tuple[str, Path, str]]] = defaultdict(list) + model_display: Dict[str, str] = {} + + for label, path in labeled_paths: + raw_text = path.read_text(encoding="utf-8") + report_text = normalize_report_text(raw_text) + model_name = extract_model_name(report_text) + grouped[label].append((framework_key_from_path(path), path, report_text)) + model_display[label] = choose_model_display_name( + model_display.get(label), + model_name, + label=label, + ) + + ordered_labels = sorted( + grouped, + key=lambda item: (model_display[item].lower(), item.lower()), + ) + + lines: List[str] = [f"# {title}", ""] + lines.append( + f"_Generated on {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}_" + ) + lines.append("") + + if include_toc: + lines.append("## Contents") + lines.append("") + for label in ordered_labels: + lines.append( + f"- [{model_display[label]}](#{slugify(model_display[label])})" + ) + lines.append("") + + for label in ordered_labels: + display_name = model_display[label] + lines.append(f"## {display_name}") + lines.append("") + lines.append(f"Model id: `{label}`") + lines.append("") + + records = sorted( + grouped[label], + key=lambda item: ( + FRAMEWORK_ORDER.get(item[0], 99), + item[1].name.lower(), + ), + ) + + for framework_key, path, report_text in records: + lines.append(f"### {framework_label(framework_key)}") + lines.append("") + lines.append(f"Source: `{path}`") + lines.append("") + lines.append(report_text) + lines.append("") + + return "\n".join(lines).rstrip() + "\n" + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + + labeled_paths: List[Tuple[str, Path]] = [] + if args.analysis_root: + labeled_paths.extend( + discover_analysis_files(Path(args.analysis_root).expanduser().resolve()) + ) + for raw_entry in args.analysis_file: + labeled_paths.append(parse_explicit_entry(raw_entry)) + + existing = [] + missing = [] + for label, path in labeled_paths: + if path.is_file(): + existing.append((label, path)) + else: + missing.append(str(path)) + if missing: + raise SystemExit("Missing analysis files:\n" + "\n".join(missing)) + if not existing: + raise SystemExit("No analysis files found.") + + markdown = build_bundle_markdown( + title=args.title, + labeled_paths=existing, + include_toc=args.include_toc, + ) + + if args.output: + output_path = Path(args.output).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(markdown, encoding="utf-8") + else: + print(markdown, end="") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/triage_kernel_helpers.py b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py similarity index 68% rename from .claude/skills/sglang-torch-profiler-analysis/scripts/triage_kernel_helpers.py rename to .claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py index ba43fb2de..ca1d47d18 100644 --- a/.claude/skills/sglang-torch-profiler-analysis/scripts/triage_kernel_helpers.py +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py @@ -4,8 +4,10 @@ from __future__ import annotations import json import re +from bisect import bisect_right from collections import Counter, defaultdict from dataclasses import dataclass, field +from functools import lru_cache from pathlib import Path from typing import DefaultDict, Dict, Iterable, List, Optional, Sequence, Tuple @@ -21,7 +23,6 @@ from profile_common import ( looks_like_python_scope_name, normalize_repo_relative_path, normalize_text, - parse_stage, select_heaviest_pid, ) @@ -200,12 +201,31 @@ LOW_LEVEL_FRAME_PREFIXES = ( "torch/nn/modules/module.py", ) +LOW_SIGNAL_FUNCTION_TOKENS = ( + "__torch_function__", + "__torch_dispatch__", + "__call__", + "_call_impl", + "_wrapped_call_impl", +) + +LOW_SIGNAL_PATH_TOKENS = ( + "model_executor/parameter.py:", + "model_executor/cuda_graph_runner.py:", + "compilation/cuda_graph.py:", + "pyexecutor/cuda_graph_runner.py:", + "pyexecutor/py_executor.py:", + "_torch/utils.py:", + "torch/fx/graph_module.py:", +) + @dataclass class KernelEvent: name: str canonical_name: str category: str + stage: str pid: str tid: str ts: float @@ -244,10 +264,36 @@ class PythonFrame: dur: float python_id: Optional[int] parent_id: Optional[int] + end_ts: float + priority: int - @property - def end_ts(self) -> float: - return self.ts + self.dur + +@dataclass +class TimedEventIndex: + events: List[object] + start_ts: List[float] + + +@dataclass +class FrameResolution: + location: str + stack: str + + +@dataclass(frozen=True) +class StageAnnotation: + stage: str + ts: float + end_ts: float + external_id: Optional[int] + is_gpu: bool + + +@dataclass(frozen=True) +class StageWindow: + stage: str + ts: float + end_ts: float @dataclass @@ -362,8 +408,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("rmsnorm", "layernorm", "fused_add_rmsnorm", "layernorm.py"), ), rationale_hint=( - "FlashInfer already exposes a TP all-reduce plus residual/RMSNorm" - " fusion path." + "FlashInfer has a TP all-reduce plus residual/RMSNorm fusion path." ), require_tp=True, min_tp_size=2, @@ -429,9 +474,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("apply_qk_norm", "q_norm", "k_norm", "qknorm"), ("apply_rope", "rotary", "rope", "mrope"), ), - rationale_hint=( - "SGLang already ships a fused QK-norm plus RoPE kernel family." - ), + rationale_hint=("SGLang has a fused QK-norm plus RoPE kernel family."), min_share=0.3, likely_share=2.0, priority=30, @@ -611,6 +654,30 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( priority=50, subsumes=("Fused MoE router / top-k / softcapping",), ), + FusionPatternSpec( + pattern="Qwen-style shared-expert append into routed top-k output", + candidate_path=( + "python/sglang/srt/models/qwen2_moe.py" + "
python/sglang/srt/layers/moe/moe_runner/triton_utils/" + "fused_moe_triton_kernels.py" + ), + active_keywords=( + "_append_shared_to_topk_output", + "fused_append_shared_experts_with_weights", + "_fused_append_shared_experts_with_weights_kernel", + ), + split_groups=( + ("_append_shared_to_topk_output", "topk", "grouped_topk"), + ("shared_expert", "shared_expert_gate", "sigmoid"), + ), + rationale_hint=( + "Qwen-style shared experts can already be appended into routed top-k" + " output in one Triton prep kernel before fused MoE execution." + ), + min_share=0.05, + likely_share=0.5, + priority=55, + ), FusionPatternSpec( pattern="Fused MoE sum + all-reduce", candidate_path=("python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py"), @@ -753,8 +820,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("cache", "kv_buffer", "cache write"), ), rationale_hint=( - "An open SGLang ROCm PR already wires a fused QK-norm plus RoPE" - " plus KV-cache family for Qwen3.5." + "Open SGLang ROCm PR wires a fused QK-norm plus RoPE plus KV-cache" + " family for Qwen3.5." ), origin="inflight", model_include=("qwen3.5", "qwen3_5"), @@ -780,8 +847,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("memset", "memcpy128"), ), rationale_hint=( - "An open SGLang PR already replaces nvjet FP8 GEMM with CUTLASS to" - " remove memset bubbles and extra copies." + "Open SGLang PR replaces nvjet FP8 GEMM with CUTLASS to remove" + " memset bubbles and extra copies." ), origin="inflight", min_share=0.2, @@ -792,11 +859,14 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( pattern="vLLM-origin Attention + Quantization", candidate_path=( "vllm/compilation/passes/fusion/attn_quant_fusion.py" + "
vllm/v1/attention/ops/merge_attn_states.py" + "
vllm/csrc/attention/merge_attn_states.cu" "
vllm/docs/design/fusions.md" ), active_keywords=( "merge_attn_states", "attn_quant_fusion", + "output_scale", "output_group_scale", ), split_groups=( @@ -804,13 +874,31 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("quant", "fp8", "nvfp4", "group_scale"), ), rationale_hint=( - "vLLM already treats attention-epilogue quantization as a reusable" - " fused family." + "vLLM combines attention merge with attention-epilogue quantization." ), origin="upstream", min_share=0.3, likely_share=1.5, ), + FusionPatternSpec( + pattern="vLLM-origin DSV3.2 fused indexer projections", + candidate_path=( + "vllm/model_executor/models/deepseek_v2.py" + "
vllm/model_executor/models/deepseek_mtp.py" + ), + active_keywords=("wk_weights_proj",), + split_groups=( + ("wk_weights_proj", "wk", "weights_proj"), + ("mergedcolumnparallellinear", "gemm", "matmul"), + ), + rationale_hint=( + "vLLM already fuses the paired `wk` and `weights_proj` indexer" + " projections into one DSV3.2 linear family." + ), + origin="upstream", + min_share=0.2, + likely_share=1.0, + ), FusionPatternSpec( pattern="vLLM-origin RMSNorm + Quantization", candidate_path=( @@ -848,9 +936,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("silu", "gelu", "act_and_mul"), ("quant", "fp8", "fp4", "block_quant"), ), - rationale_hint=( - "vLLM already treats activation-plus-quant as a reusable fusion" " family." - ), + rationale_hint=("vLLM has an activation-plus-quant fusion family."), origin="upstream", min_share=0.3, likely_share=1.5, @@ -867,13 +953,31 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("gemm", "matmul", "cublas", "cutlass"), ), rationale_hint=( - "vLLM already has a specialized DeepSeek router GEMM family for" - " small decode batches." + "vLLM has a specialized DeepSeek router GEMM family for small" + " decode batches." ), origin="upstream", min_share=0.3, likely_share=1.5, ), + FusionPatternSpec( + pattern="vLLM-origin GPT-OSS router GEMM", + candidate_path=( + "vllm/_custom_ops.py" + "
vllm/model_executor/layers/fused_moe/router/gate_linear.py" + "
vllm/csrc/moe/gpt_oss_router_gemm.cu" + ), + active_keywords=("gpt_oss_router_gemm",), + split_groups=( + ("router", "gate", "router logits", "gpt_oss"), + ("gemm", "matmul", "cublas", "cutlass"), + ), + rationale_hint=("vLLM has a GPT-OSS-specific router GEMM path."), + origin="upstream", + model_include=("gpt-oss", "gpt_oss"), + min_share=0.3, + likely_share=1.5, + ), FusionPatternSpec( pattern="vLLM-origin DeepSeek min-latency fused QKV-A projection", candidate_path=( @@ -886,8 +990,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("gemm", "matmul", "cutlass", "cublas"), ), rationale_hint=( - "vLLM already has a fused DeepSeek QKV-A projection family for" - " decode-latency reduction." + "vLLM has a fused DeepSeek QKV-A projection family for decode" + " latency reduction." ), origin="upstream", model_include=("deepseek", "glm"), @@ -909,8 +1013,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("quant", "fp8", "nvfp4"), ), rationale_hint=( - "An open vLLM PR already treats QK-norm plus RoPE plus cache plus" - " quant as a concrete in-flight fusion family." + "Open vLLM PR covers QK-norm plus RoPE plus cache plus quant as" + " one fusion family." ), origin="inflight", min_share=0.4, @@ -919,22 +1023,143 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( subsumes=("vLLM-origin Attention + Quantization",), ), FusionPatternSpec( - pattern="PR #37045 MiniMax allreduce_rms kernels", - candidate_path=("PR #37045" "
vllm/model_executor/models/minimax_m2.py"), + pattern="vLLM-origin MiniMax allreduce_rms kernels", + candidate_path="vllm/model_executor/models/minimax_m2.py", active_keywords=("minimax_allreduce_rms", "minimax_allreduce_rmsnorm"), split_groups=( ("q_norm", "k_norm", "rmsnorm", "minimax"), ("allreduce", "all_reduce", "cross_device_reduce"), ), rationale_hint=( - "An open vLLM PR already ports TRTLLM MiniMax allreduce-plus-RMSNorm" - " kernels." + "vLLM includes the TRTLLM-derived MiniMax allreduce-plus-RMSNorm" + " kernel family." ), - origin="inflight", + origin="upstream", model_include=("minimax",), min_share=0.3, likely_share=1.5, ), + FusionPatternSpec( + pattern="vLLM fused residual add + RMSNorm", + candidate_path=( + "vllm/_custom_ops.py" + "
vllm/compilation/passes/fusion/rms_quant_fusion.py" + ), + active_keywords=( + "fused_add_rms_norm", + "fused_add_rms_norm_static_fp8_quant", + ), + rationale_hint=( + "vLLM exposes fused residual-add-plus-RMSNorm kernels and matching" + " compile-time hooks." + ), + origin="upstream", + min_share=0.1, + likely_share=1.0, + ), + FusionPatternSpec( + pattern="vLLM fused activation-and-mul", + candidate_path=( + "vllm/_custom_ops.py" + "
vllm/compilation/passes/fusion/act_quant_fusion.py" + ), + active_keywords=( + "silu_and_mul", + "silu_and_mul_quant", + "silu_and_mul_per_block_quant", + "act_and_mul", + ), + rationale_hint=( + "vLLM ships fused activation-and-multiply kernels plus quantized" + " variants for the MLP epilogue." + ), + origin="upstream", + min_share=0.1, + likely_share=1.0, + ), + FusionPatternSpec( + pattern="TensorRT-LLM FlashInfer residual add + RMSNorm", + candidate_path=( + "tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py" + "
tensorrt_llm/_torch/modules/rms_norm.py" + "
tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py" + ), + active_keywords=( + "flashinfer_fused_add_rmsnorm", + "flashinfer_gemma_fused_add_rmsnorm", + "flashinfer::norm::FusedAddRMSNormKernel", + "FusedAddRMSNormKernel", + "auto_deploy::flashinfer_fused_add_rms_norm_inplace", + ), + rationale_hint=( + "TensorRT-LLM exposes a FlashInfer fused residual-add plus RMSNorm" + " family, including AutoDeploy rewrites." + ), + origin="upstream", + min_share=0.1, + likely_share=1.0, + ), + FusionPatternSpec( + pattern="TensorRT-LLM Triton fused residual add + RMSNorm + FP8 quant", + candidate_path=( + "tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/" + "triton_fused_add_rms_norm_quant_fp8.py" + "
tensorrt_llm/_torch/auto_deploy/transform/library/" + "fuse_rmsnorm_quant_fp8.py" + ), + active_keywords=( + "triton_fused_add_rms_norm_quant_fp8", + "fuse_rmsnorm_quant_fp8", + ), + rationale_hint=( + "TensorRT-LLM mainline has a Triton residual-add plus RMSNorm plus" + " FP8-quant family in AutoDeploy." + ), + origin="upstream", + min_share=0.2, + likely_share=1.0, + priority=20, + ), + FusionPatternSpec( + pattern="TensorRT-LLM FlashInfer RMSNorm family", + candidate_path=( + "tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py" + "
tensorrt_llm/_torch/modules/rms_norm.py" + "
tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py" + ), + active_keywords=( + "flashinfer_rmsnorm", + "flashinfer_gemma_rmsnorm", + "auto_deploy::flashinfer_rms_norm", + ), + rationale_hint=( + "TensorRT-LLM lowers RMSNorm-style ladders to FlashInfer kernels" + " and AutoDeploy custom ops." + ), + origin="upstream", + min_share=0.1, + likely_share=1.0, + ), + FusionPatternSpec( + pattern="TensorRT-LLM FlashInfer activation / gate epilogues", + candidate_path=( + "tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py" + "
tensorrt_llm/_torch/auto_deploy/transform/library/fuse_silu_mul.py" + "
tensorrt_llm/_torch/models/modeling_gemma3.py" + ), + active_keywords=( + "flashinfer_silu_and_mul", + "flashinfer_gelu_tanh_and_mul", + "auto_deploy::silu_and_mul", + ), + rationale_hint=( + "TensorRT-LLM already rewrites gate activation plus multiply" + " ladders into FlashInfer epilogue kernels." + ), + origin="upstream", + min_share=0.1, + likely_share=1.0, + ), ) @@ -945,6 +1170,7 @@ def short_name(name: str, max_len: int = 96) -> str: return text[: max_len - 3] + "..." +@lru_cache(maxsize=65536) def canonicalize_name(name: str) -> str: text = normalize_text(name) text = re.sub(r"0x[0-9a-fA-F]+", "0xADDR", text) @@ -965,6 +1191,7 @@ def canonicalize_name(name: str) -> str: return text +@lru_cache(maxsize=65536) def classify_kernel(name: str) -> str: # Keep the matching order explicit: strong communication/memory signals win # first, then we fall back to weaker category hints. @@ -987,6 +1214,7 @@ def classify_kernel(name: str) -> str: return "other" +@lru_cache(maxsize=65536) def normalize_source_location(name: str) -> str: text = normalize_text(name) match = re.match(r"(?P.+?)\((?P\d+)\): (?P.+)$", text) @@ -1000,18 +1228,23 @@ def source_location_priority(location: str) -> int: text = str(location).strip() if not text or text == "unresolved": return -100 + penalty = 80 if is_low_signal_source_location(text) else 0 if text.startswith("python/sglang/"): - return 300 + return 300 - penalty if text.startswith("sglang/"): - return 290 + return 290 - penalty + if text.startswith("vllm/"): + return 285 - penalty + if text.startswith("tensorrt_llm/"): + return 280 - penalty if text.startswith("sgl_kernel/"): - return 260 + return 260 - penalty if text.startswith("python/"): - return 180 + return 180 - penalty if text.startswith("torch/") or "/torch/" in text: return 20 if ".py:" in text: - return 120 + return 120 - penalty return 0 @@ -1020,6 +1253,8 @@ def is_preferred_source_location(location: str) -> bool: return ( text.startswith("python/sglang/") or text.startswith("sglang/") + or text.startswith("vllm/") + or text.startswith("tensorrt_llm/") or text.startswith("sgl_kernel/") ) @@ -1043,7 +1278,9 @@ def extract_preferred_stack_location(stack: Optional[str]) -> Optional[str]: def site_display_location(site: dict) -> str: location = str(site.get("location") or "unresolved").strip() - if is_preferred_source_location(location): + if is_preferred_source_location(location) and not is_low_signal_source_location( + location + ): return location stack_location = extract_preferred_stack_location(site.get("stack")) if stack_location: @@ -1066,27 +1303,43 @@ def choose_best_location(locations: Dict[str, MappingSiteAggregate]) -> str: return ranked[0][0] +@lru_cache(maxsize=65536) def frame_priority(frame_name: str) -> int: raw_text = str(frame_name).strip() normalized_text = normalize_source_location(raw_text) + penalty = 80 if is_low_signal_source_location(normalized_text) else 0 if raw_text.startswith(NOISE_FRAME_PREFIXES): return -20 if normalized_text.startswith("python/sglang/"): - return 300 + return 300 - penalty if normalized_text.startswith("sglang/"): - return 290 + return 290 - penalty + if normalized_text.startswith("vllm/"): + return 285 - penalty + if normalized_text.startswith("tensorrt_llm/"): + return 280 - penalty if normalized_text.startswith("sgl_kernel/"): - return 260 + return 260 - penalty if normalized_text.startswith("triton_kernels/"): - return 220 + return 220 - penalty if normalized_text.startswith(LOW_LEVEL_FRAME_PREFIXES): return 0 if raw_text.startswith("/data/") or raw_text.startswith("/Users/"): if "/sglang/" in raw_text: return 120 + if "/vllm/" in raw_text: + return 118 + if "/TensorRT-LLM/" in raw_text or "/tensorrt_llm/" in raw_text: + return 116 return 100 if ".py(" in raw_text and "/sglang/" in raw_text: return 110 + if ".py(" in raw_text and "/vllm/" in raw_text: + return 108 + if ".py(" in raw_text and ( + "/TensorRT-LLM/" in raw_text or "/tensorrt_llm/" in raw_text + ): + return 106 if ".py:" in normalized_text and ( "site-packages" in raw_text or normalized_text.startswith("torch/") ): @@ -1098,6 +1351,16 @@ def frame_priority(frame_name: str) -> int: return 0 +@lru_cache(maxsize=65536) +def is_low_signal_source_location(location: str) -> bool: + lowered = str(location).strip().lower() + if not lowered: + return False + return any(token in lowered for token in LOW_SIGNAL_FUNCTION_TOKENS) or any( + token in lowered for token in LOW_SIGNAL_PATH_TOKENS + ) + + def stage_label(stage: str) -> str: if stage == "extend": return "extend/prefill" @@ -1126,7 +1389,8 @@ def format_ms(value_us: float) -> str: return f"{value_us / 1000.0:.2f} ms" -def is_cuda_launch_event(name: str, cat: str) -> bool: +@lru_cache(maxsize=16384) +def _is_cuda_launch_event_cached(name: str, cat: str) -> bool: lowered_name = normalize_text(name).lower() lowered_cat = normalize_text(cat).lower() if lowered_cat not in {"cuda_runtime", "cuda_driver"}: @@ -1134,6 +1398,10 @@ def is_cuda_launch_event(name: str, cat: str) -> bool: return "launch" in lowered_name +def is_cuda_launch_event(name: str, cat: str) -> bool: + return _is_cuda_launch_event_cached(str(name), str(cat)) + + def is_gpu_kernel_event(event: dict) -> bool: # Be conservative here: first drop trace metadata / Python scopes / # annotations, then only accept entries with clear GPU-kernel markers. @@ -1155,6 +1423,142 @@ def is_gpu_kernel_event(event: dict) -> bool: return has_stream_marker(args) +def infer_stage_from_annotation_name(name: str) -> Optional[str]: + lowered = normalize_text(name).lower() + if not lowered: + return None + if "generation_1" in lowered or "decode" in lowered: + return "decode" + if "generation_0" in lowered or "prefill" in lowered: + return "extend" + return None + + +def build_stage_annotations( + raw_events: Sequence[dict], +) -> Tuple[ + Dict[int, StageAnnotation], + List[StageWindow], + List[StageWindow], +]: + by_external_id: Dict[int, StageAnnotation] = {} + gpu_annotations: List[StageAnnotation] = [] + cpu_annotations: List[StageAnnotation] = [] + + def should_replace(current: StageAnnotation, candidate: StageAnnotation) -> bool: + if candidate.is_gpu != current.is_gpu: + return candidate.is_gpu + return (candidate.end_ts - candidate.ts) > (current.end_ts - current.ts) + + for event in raw_events: + if not is_complete_duration_event(event): + continue + category = normalize_text(event.get("cat", "")).lower() + if category not in {"user_annotation", "gpu_user_annotation"}: + continue + stage = infer_stage_from_annotation_name(str(event.get("name", ""))) + if not stage: + continue + annotation = StageAnnotation( + stage=stage, + ts=float(event.get("ts", 0.0)), + end_ts=float(event.get("ts", 0.0)) + float(event.get("dur", 0.0)), + external_id=coerce_optional_int( + (event.get("args") or {}).get("External id") + ), + is_gpu=(category == "gpu_user_annotation"), + ) + if annotation.external_id is not None: + existing = by_external_id.get(annotation.external_id) + if existing is None or should_replace(existing, annotation): + by_external_id[annotation.external_id] = annotation + if annotation.is_gpu: + gpu_annotations.append(annotation) + else: + cpu_annotations.append(annotation) + + gpu_annotations.sort(key=lambda item: (item.ts, item.end_ts)) + cpu_annotations.sort(key=lambda item: (item.ts, item.end_ts)) + return ( + by_external_id, + merge_stage_windows(gpu_annotations), + merge_stage_windows(cpu_annotations), + ) + + +def merge_stage_windows(annotations: Sequence[StageAnnotation]) -> List[StageWindow]: + merged: List[StageWindow] = [] + for annotation in annotations: + if ( + merged + and merged[-1].stage == annotation.stage + and annotation.ts <= merged[-1].end_ts + 1e-3 + ): + merged[-1] = StageWindow( + stage=merged[-1].stage, + ts=merged[-1].ts, + end_ts=max(merged[-1].end_ts, annotation.end_ts), + ) + continue + merged.append( + StageWindow( + stage=annotation.stage, + ts=annotation.ts, + end_ts=annotation.end_ts, + ) + ) + return merged + + +def resolve_stage_from_windows( + probe_ts: float, + windows: Sequence[StageWindow], +) -> Tuple[Optional[str], Optional[float]]: + nearest_stage: Optional[str] = None + nearest_gap: Optional[float] = None + for window in windows: + if window.ts <= probe_ts <= window.end_ts + 1e-3: + return window.stage, 0.0 + gap = min(abs(probe_ts - window.ts), abs(probe_ts - window.end_ts)) + if nearest_gap is None or gap < nearest_gap: + nearest_gap = gap + nearest_stage = window.stage + return nearest_stage, nearest_gap + + +def resolve_kernel_stage( + *, + kernel_ts: float, + external_id: Optional[int], + annotations_by_external_id: Dict[int, StageAnnotation], + gpu_annotations: Sequence[StageWindow], + cpu_annotations: Sequence[StageWindow], +) -> str: + if external_id is not None: + annotation = annotations_by_external_id.get(external_id) + if annotation is not None: + return annotation.stage + probe_ts = kernel_ts + 1e-3 + nearest_stage: Optional[str] = None + nearest_gap: Optional[float] = None + for windows in (gpu_annotations, cpu_annotations): + stage, gap = resolve_stage_from_windows(probe_ts, windows) + if gap == 0.0 and stage is not None: + return stage + if stage is not None and ( + nearest_gap is None or (gap is not None and gap < nearest_gap) + ): + nearest_stage = stage + nearest_gap = gap + if ( + nearest_stage is not None + and nearest_gap is not None + and nearest_gap <= 20_000.0 + ): + return nearest_stage + return "all" + + def extract_trace_data( trace: dict, ) -> Tuple[ @@ -1170,6 +1574,11 @@ def extract_trace_data( # source attribution, and CUDA launch calls for correlation-based fallback. raw_events = extract_trace_events(trace) correlation_external = build_correlation_external_lookup(raw_events) + ( + annotations_by_external_id, + gpu_stage_annotations, + cpu_stage_annotations, + ) = build_stage_annotations(raw_events) chosen_pid = select_heaviest_pid( raw_events, is_gpu_kernel_event, @@ -1206,6 +1615,8 @@ def extract_trace_data( dur=dur, python_id=coerce_optional_int(args.get("Python id")), parent_id=coerce_optional_int(args.get("Python parent id")), + end_ts=ts + dur, + priority=frame_priority(name), ) ) @@ -1246,6 +1657,13 @@ def extract_trace_data( name=name, canonical_name=canonicalize_name(name), category=classify_kernel(name), + stage=resolve_kernel_stage( + kernel_ts=ts, + external_id=external_id, + annotations_by_external_id=annotations_by_external_id, + gpu_annotations=gpu_stage_annotations, + cpu_annotations=cpu_stage_annotations, + ), pid=pid, tid=tid, ts=ts, @@ -1273,17 +1691,27 @@ def build_correlation_external_lookup(raw_events: Sequence[dict]) -> Dict[int, i return lookup -def build_cpu_op_index(cpu_ops: Sequence[CpuOpEvent]) -> Dict[int, List[CpuOpEvent]]: +def build_timed_event_index(events: Sequence[object]) -> TimedEventIndex: + ordered = list(events) + ordered.sort(key=lambda item: item.ts) + return TimedEventIndex( + events=ordered, + start_ts=[float(item.ts) for item in ordered], + ) + + +def build_cpu_op_index(cpu_ops: Sequence[CpuOpEvent]) -> Dict[int, TimedEventIndex]: output: DefaultDict[int, List[CpuOpEvent]] = defaultdict(list) for cpu_op in cpu_ops: output[cpu_op.external_id].append(cpu_op) - for items in output.values(): - items.sort(key=lambda item: item.ts) - return dict(output) + return { + external_id: build_timed_event_index(items) + for external_id, items in output.items() + } def match_cpu_op( - kernel: KernelEvent, cpu_ops_by_external_id: Dict[int, List[CpuOpEvent]] + kernel: KernelEvent, cpu_ops_by_external_id: Dict[int, TimedEventIndex] ) -> Optional[CpuOpEvent]: if kernel.external_id is None: return None @@ -1294,17 +1722,18 @@ def match_cpu_op( def build_launch_index( launch_events: Sequence[LaunchEvent], -) -> Dict[int, List[LaunchEvent]]: +) -> Dict[int, TimedEventIndex]: output: DefaultDict[int, List[LaunchEvent]] = defaultdict(list) for launch in launch_events: output[launch.correlation].append(launch) - for items in output.values(): - items.sort(key=lambda item: item.ts) - return dict(output) + return { + correlation: build_timed_event_index(items) + for correlation, items in output.items() + } def match_launch_event( - kernel: KernelEvent, launches_by_correlation: Dict[int, List[LaunchEvent]] + kernel: KernelEvent, launches_by_correlation: Dict[int, TimedEventIndex] ) -> Optional[LaunchEvent]: if kernel.correlation is None: return None @@ -1313,7 +1742,26 @@ def match_launch_event( ) -def match_timed_event(events: Sequence, probe_ts: float): +def match_timed_event(index: object, probe_ts: float): + if not index: + return None + if isinstance(index, TimedEventIndex): + events = index.events + if not events: + return None + right = bisect_right(index.start_ts, probe_ts + 1e-3) + candidates: List[object] = [] + if right > 0: + candidates.extend(events[max(0, right - 4) : right]) + if right < len(events): + candidates.extend(events[right : min(len(events), right + 2)]) + if not candidates: + return None + earlier = [item for item in candidates if item.ts <= probe_ts + 1e-3] + if earlier: + return min(earlier, key=lambda item: abs((item.ts + item.dur) - probe_ts)) + return min(candidates, key=lambda item: abs(item.ts - probe_ts)) + events = list(index) if not events: return None earlier = [item for item in events if item.ts <= probe_ts + 1e-3] @@ -1322,6 +1770,75 @@ def match_timed_event(events: Sequence, probe_ts: float): return min(events, key=lambda item: abs(item.ts - probe_ts)) +def resolve_active_frames_linear( + frames: Sequence[PythonFrame], probe_ts: float +) -> List[PythonFrame]: + active = [item for item in frames if item.ts <= probe_ts <= item.end_ts] + active.sort(key=lambda item: (item.ts, item.end_ts)) + return active + + +def thread_has_crossing_frames(frames: Sequence[PythonFrame]) -> bool: + ordered_frames = sorted(frames, key=lambda item: (item.ts, -item.end_ts)) + stack: List[PythonFrame] = [] + for frame in ordered_frames: + while stack and stack[-1].end_ts < frame.ts: + stack.pop() + if stack and frame.end_ts > stack[-1].end_ts + 1e-3: + return True + stack.append(frame) + return False + + +def render_frame_resolution( + active_frames: Sequence[PythonFrame], +) -> Optional[FrameResolution]: + if not active_frames: + return None + chosen_frame = choose_mapping_frame(active_frames) + if chosen_frame is None: + return None + return FrameResolution( + location=chosen_frame.normalized_name, + stack=build_stack_display(active_frames), + ) + + +def resolve_thread_query_times( + frames: Sequence[PythonFrame], query_times: Sequence[float] +) -> Dict[float, Optional[FrameResolution]]: + if not frames or not query_times: + return {} + ordered_frames = sorted(frames, key=lambda item: (item.ts, -item.end_ts)) + ordered_queries = sorted(set(float(ts) for ts in query_times)) + results: Dict[float, Optional[FrameResolution]] = {} + active_frames: List[PythonFrame] = [] + frame_idx = 0 + total_frames = len(ordered_frames) + + for ts in ordered_queries: + while frame_idx < total_frames and ordered_frames[frame_idx].ts <= ts: + active_frames.append(ordered_frames[frame_idx]) + frame_idx += 1 + if active_frames: + active_frames = [ + frame for frame in active_frames if frame.end_ts >= ts - 1e-3 + ] + results[ts] = render_frame_resolution(active_frames) + return results + + +def build_frame_resolution_index( + python_frames: Dict[Tuple[str, str], List[PythonFrame]], + query_times_by_thread: Dict[Tuple[str, str], Sequence[float]], +) -> Dict[Tuple[str, str], Dict[float, Optional[FrameResolution]]]: + output: Dict[Tuple[str, str], Dict[float, Optional[FrameResolution]]] = {} + for thread_key, query_times in query_times_by_thread.items(): + frames = python_frames.get(thread_key, []) + output[thread_key] = resolve_thread_query_times(frames, query_times) + return output + + def find_active_python_frames( cpu_op: CpuOpEvent, python_frames: Dict[Tuple[str, str], List[PythonFrame]], @@ -1330,9 +1847,7 @@ def find_active_python_frames( if not frames: return [] probe_ts = cpu_op.ts + min(cpu_op.dur * 0.5, 1.0) - active = [item for item in frames if item.ts <= probe_ts <= item.end_ts] - active.sort(key=lambda item: (item.ts, item.end_ts)) - return active + return resolve_active_frames_linear(frames, probe_ts) def find_active_python_frames_at_ts( @@ -1345,9 +1860,7 @@ def find_active_python_frames_at_ts( frames = python_frames.get((pid, tid), []) if not frames: return [] - active = [item for item in frames if item.ts <= ts <= item.end_ts] - active.sort(key=lambda item: (item.ts, item.end_ts)) - return active + return resolve_active_frames_linear(frames, ts) def render_kernel_site( @@ -1361,21 +1874,38 @@ def render_kernel_site( def resolve_kernel_site_context( kernel: KernelEvent, - cpu_ops_by_external_id: Dict[int, List[CpuOpEvent]], + cpu_ops_by_external_id: Dict[int, TimedEventIndex], python_frames: Dict[Tuple[str, str], List[PythonFrame]], - launches_by_correlation: Dict[int, List[LaunchEvent]], + launches_by_correlation: Dict[int, TimedEventIndex], + frame_resolution_index: Optional[ + Dict[Tuple[str, str], Dict[float, Optional[FrameResolution]]] + ] = None, ) -> Tuple[str, str, str]: # Prefer the normal External-id path first. If the kernel dropped that link, # fall back to the correlated CUDA launch and reuse the Python frames that # were active when the launch happened. cpu_op = match_cpu_op(kernel, cpu_ops_by_external_id) if cpu_op is not None: + probe_ts = cpu_op.ts + min(cpu_op.dur * 0.5, 1.0) + if frame_resolution_index is not None: + resolved = frame_resolution_index.get((cpu_op.pid, cpu_op.tid), {}).get( + probe_ts + ) + if resolved is not None: + return resolved.location, resolved.stack, cpu_op.name active_frames = find_active_python_frames(cpu_op, python_frames) if active_frames: return render_kernel_site(active_frames, cpu_op.name) launch_event = match_launch_event(kernel, launches_by_correlation) if launch_event is not None: + if frame_resolution_index is not None: + resolved = frame_resolution_index.get( + (launch_event.pid, launch_event.tid), {} + ).get(launch_event.ts) + if resolved is not None: + cpu_op_name = cpu_op.name if cpu_op is not None else launch_event.name + return resolved.location, resolved.stack, cpu_op_name active_frames = find_active_python_frames_at_ts( pid=launch_event.pid, tid=launch_event.tid, @@ -1394,19 +1924,20 @@ def resolve_kernel_site_context( def choose_mapping_frame(active_frames: Sequence[PythonFrame]) -> Optional[PythonFrame]: if not active_frames: return None - ranked = sorted( - active_frames, - key=lambda item: (frame_priority(item.name), item.ts, -item.dur), - ) - return ranked[-1] + best = active_frames[0] + best_key = (best.priority, best.ts, -best.dur) + for item in active_frames[1:]: + key = (item.priority, item.ts, -item.dur) + if key > best_key: + best = item + best_key = key + return best def build_stack_display(active_frames: Sequence[PythonFrame]) -> str: if not active_frames: return "" - filtered = [ - item.normalized_name for item in active_frames if frame_priority(item.name) > 0 - ] + filtered = [item.normalized_name for item in active_frames if item.priority > 0] if not filtered: filtered = [active_frames[-1].normalized_name] return " -> ".join(filtered[-4:]) @@ -1423,11 +1954,24 @@ def aggregate(events: Iterable[KernelEvent], key_fn) -> Dict[str, Aggregate]: return output +def group_kernels_by_stage( + kernels: Sequence[KernelEvent], default_stage: str +) -> Dict[str, List[KernelEvent]]: + grouped: DefaultDict[str, List[KernelEvent]] = defaultdict(list) + for kernel in kernels: + stage = default_stage if default_stage != "all" else (kernel.stage or "all") + grouped[stage].append(kernel) + return dict(grouped) + + def aggregate_kernel_sites( kernels: Sequence[KernelEvent], - cpu_ops_by_external_id: Dict[int, List[CpuOpEvent]], + cpu_ops_by_external_id: Dict[int, TimedEventIndex], python_frames: Dict[Tuple[str, str], List[PythonFrame]], - launches_by_correlation: Optional[Dict[int, List[LaunchEvent]]] = None, + launches_by_correlation: Optional[Dict[int, TimedEventIndex]] = None, + site_context_cache: Optional[ + Dict[Tuple[str, str, float, Optional[int], Optional[int]], Tuple[str, str, str]] + ] = None, ) -> Dict[str, Dict[str, MappingSiteAggregate]]: # Each kernel is mapped independently so the fallback behavior stays easy to # reason about and easy to regression-test. @@ -1435,13 +1979,41 @@ def aggregate_kernel_sites( lambda: defaultdict(MappingSiteAggregate) ) launch_index = launches_by_correlation or {} + query_times_by_thread: DefaultDict[Tuple[str, str], List[float]] = defaultdict(list) for kernel in kernels: - location, stack, cpu_op_name = resolve_kernel_site_context( - kernel, - cpu_ops_by_external_id, - python_frames, - launch_index, + cpu_op = match_cpu_op(kernel, cpu_ops_by_external_id) + if cpu_op is not None: + query_times_by_thread[(cpu_op.pid, cpu_op.tid)].append( + cpu_op.ts + min(cpu_op.dur * 0.5, 1.0) + ) + launch_event = match_launch_event(kernel, launch_index) + if launch_event is not None: + query_times_by_thread[(launch_event.pid, launch_event.tid)].append( + launch_event.ts + ) + frame_resolution_index = build_frame_resolution_index( + python_frames, query_times_by_thread + ) + resolved_cache = site_context_cache if site_context_cache is not None else {} + for kernel in kernels: + cache_key = ( + kernel.pid, + kernel.tid, + kernel.ts, + kernel.external_id, + kernel.correlation, ) + cached = resolved_cache.get(cache_key) + if cached is None: + cached = resolve_kernel_site_context( + kernel, + cpu_ops_by_external_id, + python_frames, + launch_index, + frame_resolution_index=frame_resolution_index, + ) + resolved_cache[cache_key] = cached + location, stack, cpu_op_name = cached item = output[kernel.canonical_name][location] item.total_us += kernel.dur @@ -1560,7 +2132,7 @@ def relaxed_kernel_entry_lookup( # recover the higher-level Python callsite from the mapping trace. lowered_compact = normalize_match_text(kernel_name) if len(lowered_compact) < 96: - return None + return alias_kernel_entry_lookup(kernels, kernel_name) def common_prefix_len(left: str, right: str) -> int: count = 0 @@ -1588,7 +2160,9 @@ def relaxed_kernel_entry_lookup( if score > best_score: best_key = candidate_key best_score = score - return kernels.get(best_key) if best_key else None + if best_key: + return kernels.get(best_key) + return alias_kernel_entry_lookup(kernels, kernel_name) def lookup_kernel_map_entry( @@ -1646,7 +2220,9 @@ def resolve_kernel_entry( kernel_entry = lookup_kernel_map_entry(external_kernel_map, stage, kernel_name) if kernel_entry: return kernel_entry - return local_stage_payload.get("kernels", {}).get(kernel_name) + return relaxed_kernel_entry_lookup( + local_stage_payload.get("kernels", {}), kernel_name + ) def build_kernel_rows( @@ -1725,6 +2301,89 @@ def normalize_match_text(text: object) -> str: return re.sub(r"[^0-9A-Za-z]+", "", normalize_text(text)).lower() +def kernel_entry_total_us(entry: Optional[dict]) -> float: + if not entry: + return 0.0 + return sum(float(site.get("total_us", 0.0)) for site in entry.get("sites", [])) + + +def kernel_entry_lookup_text(kernel_name: str, entry: Optional[dict]) -> str: + parts = [kernel_name] + if entry: + parts.append(str(entry.get("best_location") or "")) + for site in entry.get("sites", [])[:4]: + parts.append(str(site.get("location") or "")) + parts.append(str(site.get("display_location") or "")) + parts.append(str(site.get("top_cpu_op") or "")) + parts.append(str(site.get("stack") or "")) + return normalize_match_text(" ".join(parts)) + + +def kernel_alias_token_groups(kernel_name: str) -> List[Tuple[str, ...]]: + lowered = normalize_match_text(kernel_name) + groups: List[Tuple[str, ...]] = [] + if "flashattnfwdcombine" in lowered: + groups.append( + ( + "flashattnfwdsm90", + "flashattnvarlenfunc", + "vllmflashattnflashattninterface", + "vllmfa3cfwd", + ) + ) + if "kernelmha" in lowered: + groups.append( + ( + "maskedmultiheadattentionkernel", + "attentioninplace", + "attentionbackendtrtllm", + ) + ) + if "applybiasropeupdatekvcachev2" in lowered: + groups.append( + ( + "fusedqknormropekernel", + "applyqknormrope", + "modelingqwen3py98applyqknormrope", + ) + ) + if lowered.startswith("memset"): + groups.append(("memset",)) + return groups + + +def alias_kernel_entry_lookup( + kernels: Dict[str, dict], kernel_name: str +) -> Optional[dict]: + alias_groups = kernel_alias_token_groups(kernel_name) + if not alias_groups: + return None + + best_key = None + best_score = -1 + for candidate_key, entry in kernels.items(): + candidate_text = kernel_entry_lookup_text(candidate_key, entry) + score = 0 + for group_index, group in enumerate(alias_groups): + group_score = max( + (len(token) for token in group if token in candidate_text), + default=0, + ) + if group_score: + score += 1000 * (group_index + 1) + group_score + if score <= 0: + continue + score += max( + source_location_priority(str(entry.get("best_location") or "")), + source_location_priority(best_site_summary(entry)[0]), + ) + score += int(kernel_entry_total_us(entry) // 10) + if score > best_score: + best_key = candidate_key + best_score = score + return kernels.get(best_key) if best_key else None + + def row_matches(row: KernelRow, *needles: str) -> bool: lowered = " ".join([row.name, row.location, row.cpu_op]).lower() lowered_compact = normalize_match_text(lowered) @@ -1774,6 +2433,30 @@ def model_path_from_server_args(server_args: Optional[dict]) -> str: return str(server_args.get("model_path") or server_args.get("model") or "") +def fusion_framework_hints(spec: FusionPatternSpec) -> set[str]: + text = normalize_text(spec.candidate_path).lower() + hints: set[str] = set() + if "vllm/" in text: + hints.add("vllm") + if "tensorrt_llm/" in text: + hints.add("trtllm") + if any(token in text for token in ("python/sglang/", "sgl-kernel/", "sgl_kernel/")): + hints.add("sglang") + return hints + + +def pattern_supports_framework( + spec: FusionPatternSpec, framework: Optional[str] +) -> bool: + normalized = normalize_text(framework).lower() + if not normalized or normalized == "auto": + return True + hints = fusion_framework_hints(spec) + if not hints: + return True + return normalized in hints + + def matching_rows_for_keywords( kernel_rows: Sequence[KernelRow], keywords: Sequence[str], @@ -1812,10 +2495,10 @@ def pattern_model_matches(spec: FusionPatternSpec, model_path: str) -> bool: def pattern_status(spec: FusionPatternSpec, has_active_match: bool) -> str: if spec.origin == "mainline": - return "active fused path" if has_active_match else "split candidate" + return "mainline direct" if has_active_match else "mainline split" if spec.origin == "upstream": - return "upstream precedent" if has_active_match else "upstream split precedent" - return "in-flight precedent" if has_active_match else "in-flight split precedent" + return "upstream direct" if has_active_match else "upstream split" + return "pending direct" if has_active_match else "pending split" def build_pattern_rationale( @@ -1828,20 +2511,20 @@ def build_pattern_rationale( if spec.origin == "mainline": if has_active_match: return ( - f"This trace already hits the `{spec.pattern}` family directly at {share:.1f}% related GPU time. " + f"`{spec.pattern}` is present in this trace ({share:.1f}% related GPU time). " f"{spec.rationale_hint}" ) return ( - f"Related split kernels occupy {share:.1f}% of cumulative GPU time, and the checked-out SGLang tree " - f"already exposes this fusion family. {spec.rationale_hint}" + f"Split kernels in this family take {share:.1f}% of GPU time. " + f"This tree already has a matching path. {spec.rationale_hint}" ) if spec.origin == "upstream": return ( - f"This trace matches a reusable upstream vLLM precedent at {share:.1f}% related GPU time. " + f"Matches an upstream path ({share:.1f}% related GPU time). " f"{spec.rationale_hint}" ) return ( - f"This trace matches a PR-backed / in-flight pattern at {share:.1f}% related GPU time. " + f"Matches an open upstream path ({share:.1f}% related GPU time). " f"{spec.rationale_hint}" ) @@ -1865,9 +2548,12 @@ def detect_pattern_match( total_us: float, model_path: str, tp_size: int, + framework: Optional[str], ) -> Optional[FusionOpportunity]: if total_us <= 0: return None + if not pattern_supports_framework(spec, framework): + return None if spec.require_tp and tp_size < spec.min_tp_size: return None if not pattern_model_matches(spec, model_path): @@ -1894,9 +2580,9 @@ def detect_pattern_match( pattern=spec.pattern, status=pattern_status(spec, has_active_match), confidence=( - "Likely" + "Confirmed" if has_active_match or pct(related_us, total_us) >= spec.likely_share - else "Conditional" + else "Candidate" ), related_us=related_us, evidence=summarize_evidence(related_rows, total_us), @@ -1919,10 +2605,10 @@ def detect_pattern_match( def detect_fusion_opportunities( - stage: str, kernel_rows: Sequence[KernelRow], total_us: float, server_args: Optional[dict], + framework: Optional[str] = None, ) -> List[FusionOpportunity]: opportunities: List[FusionOpportunity] = [] if total_us <= 0: @@ -1941,6 +2627,7 @@ def detect_fusion_opportunities( total_us=total_us, model_path=model_path, tp_size=tp_size, + framework=framework, ) if opportunity is not None: raw_matches.append(opportunity) @@ -1959,225 +2646,3 @@ def detect_fusion_opportunities( consumed_row_keys.update(opportunity.covered_row_keys) blocked_patterns.update(opportunity.subsumes) return opportunities - - -def generate_takeaways( - stage: str, - total_us: float, - window_us: float, - category_stats: Dict[str, Aggregate], - resolved_us: float, - server_args: Optional[dict], - fusion_opportunities: Sequence[FusionOpportunity], -) -> List[str]: - items = sorted( - category_stats.items(), key=lambda pair: pair[1].total_us, reverse=True - ) - if not items: - return ["No GPU kernel events were found in the selected trace."] - - takeaways: List[str] = [] - top_name, top_agg = items[0] - takeaways.append( - f"{stage_label(stage)} is dominated by `{top_name}` at {pct(top_agg.total_us, total_us):.1f}% of cumulative GPU kernel time." - ) - if len(items) > 1: - second_name, second_agg = items[1] - combined = pct(top_agg.total_us + second_agg.total_us, total_us) - takeaways.append( - f"The top two categories are `{top_name}` + `{second_name}` at {combined:.1f}% combined." - ) - - comm_share = pct( - category_stats.get("communication", Aggregate()).total_us, total_us - ) - if comm_share >= 10.0: - tp = server_args.get("tp_size") if isinstance(server_args, dict) else None - if tp and tp > 1: - takeaways.append( - f"`communication` already accounts for {comm_share:.1f}% of cumulative GPU time in this TP={tp} run." - ) - else: - takeaways.append( - f"`communication` shows up at {comm_share:.1f}% even without an obvious large-TP context." - ) - - if pct(resolved_us, total_us) >= 70.0: - takeaways.append( - f"Kernel-to-Python mapping covers {pct(resolved_us, total_us):.1f}% of cumulative GPU time, so the table is representative enough for code triage." - ) - - if window_us > 0: - parallelism = total_us / window_us - if parallelism >= 1.15: - takeaways.append( - f"Summed kernel time is {parallelism:.2f}x the GPU time window, so these percentages are cumulative launch share rather than wall time share." - ) - if fusion_opportunities: - top_pattern = fusion_opportunities[0] - takeaways.append( - f"The strongest source-backed fuse-pattern match is `{top_pattern.pattern}` ({top_pattern.status}) with {pct(top_pattern.related_us, total_us):.1f}% related GPU time in this stage." - ) - return takeaways - - -def print_mapping_table( - kernel_rows: Sequence[KernelRow], - total_us: float, - table_limit: int, -) -> float: - resolved_us = 0.0 - rendered_rows = limit_kernel_rows(kernel_rows, table_limit) - label = "all kernels" if table_limit <= 0 else f"first {len(rendered_rows)} kernels" - print(f"\nKernel-to-Python mapping (Markdown, {label}):") - print( - "| Kernel | Category | GPU time | Share | Launches | Python location (site share) | CPU op |" - ) - print("| --- | --- | ---: | ---: | ---: | --- | --- |") - for row in rendered_rows: - if row.location != "unresolved": - resolved_us += row.total_us - print( - "| {kernel} | {category} | {gpu_time} | {share:.1f}% | {launches} | {location} | {cpu_op} |".format( - kernel=escape_md_cell(row.name), - category=escape_md_cell(row.category), - gpu_time=format_ms(row.total_us), - share=pct(row.total_us, total_us), - launches=row.aggregate.count, - location=escape_md_cell(row.location), - cpu_op=escape_md_cell(row.cpu_op), - ) - ) - return resolved_us - - -def print_fusion_opportunity_table( - opportunities: Sequence[FusionOpportunity], - total_us: float, -) -> None: - print("\nKernel fuse pattern matches (Markdown):") - print( - "| Pattern | Status | Confidence | Related GPU time | Share | Evidence kernels | Current kernel Python location | Reference path | Why it matters |" - ) - print("| --- | --- | --- | ---: | ---: | --- | --- | --- | --- |") - if not opportunities: - print( - "| No source-backed fuse pattern matched this trace. | - | - | - | - | - | - | - | - |" - ) - return - for item in opportunities: - print( - "| {pattern} | {status} | {confidence} | {gpu_time} | {share:.1f}% | {evidence} | {current_locations} | {candidate_path} | {rationale} |".format( - pattern=escape_md_cell(item.pattern), - status=escape_md_cell(item.status), - confidence=escape_md_cell(item.confidence), - gpu_time=format_ms(item.related_us), - share=pct(item.related_us, total_us), - evidence=escape_md_cell(item.evidence), - current_locations=escape_md_cell(item.current_locations), - candidate_path=escape_md_cell(item.candidate_path), - rationale=escape_md_cell(item.rationale), - ) - ) - - -def print_report( - trace_path: Path, - server_args: Optional[dict], - kernels: List[KernelEvent], - chosen_pid: Optional[str], - window_us: float, - local_stage_payload: dict, - external_kernel_map: Optional[dict], - top_k: int, - kernel_table_limit: int, - table_only: bool, -) -> None: - stage = parse_stage(trace_path) - total_us = sum(kernel.dur for kernel in kernels) - print(f"Trace: {trace_path}") - print(f"Stage: {stage_label(stage)}") - if chosen_pid: - print(f"Selected PID: {chosen_pid}") - - if server_args: - model_path = server_args.get("model_path") or server_args.get("model") - tp_size = server_args.get("tp_size") - dp_size = server_args.get("dp_size") - print(f"Model: {model_path}") - if tp_size or dp_size: - print(f"Parallelism: tp={tp_size or 1} dp={dp_size or 1}") - - if not kernels: - print("No GPU kernel events found.\n") - return - - print( - f"GPU kernels: {len(kernels)} | cumulative kernel time: {format_ms(total_us)} | " - f"GPU window: {format_ms(window_us)} | avg parallelism: {total_us / window_us:.2f}x" - if window_us - else f"GPU kernels: {len(kernels)} | cumulative kernel time: {format_ms(total_us)}" - ) - - category_stats = aggregate(kernels, key_fn=lambda item: item.category) - kernel_stats = aggregate(kernels, key_fn=lambda item: item.canonical_name) - kernel_categories = {kernel.canonical_name: kernel.category for kernel in kernels} - kernel_rows = build_kernel_rows( - stage=stage, - kernel_stats=kernel_stats, - kernel_categories=kernel_categories, - local_stage_payload=local_stage_payload, - external_kernel_map=external_kernel_map, - ) - fusion_opportunities = detect_fusion_opportunities( - stage=stage, - kernel_rows=kernel_rows, - total_us=total_us, - server_args=server_args, - ) - - if not table_only: - print("\nTop categories by cumulative GPU kernel time:") - for idx, (name, aggregate_item) in enumerate( - sorted( - category_stats.items(), key=lambda pair: pair[1].total_us, reverse=True - )[:8], - start=1, - ): - print( - f" {idx}. {name:<16} {format_ms(aggregate_item.total_us):>10} " - f"{pct(aggregate_item.total_us, total_us):>5.1f}% launches={aggregate_item.count}" - ) - - print("\nTop kernels by cumulative GPU kernel time:") - for idx, (name, aggregate_item) in enumerate( - sorted( - kernel_stats.items(), key=lambda pair: pair[1].total_us, reverse=True - )[:top_k], - start=1, - ): - print( - f" {idx}. {short_name(name, 76):<76} {format_ms(aggregate_item.total_us):>10} " - f"{pct(aggregate_item.total_us, total_us):>5.1f}% launches={aggregate_item.count} avg={format_ms(aggregate_item.avg_us)}" - ) - - resolved_us = print_mapping_table( - kernel_rows=kernel_rows, - total_us=total_us, - table_limit=kernel_table_limit, - ) - print_fusion_opportunity_table(fusion_opportunities, total_us) - - if not table_only: - print("\nTakeaways:") - for takeaway in generate_takeaways( - stage, - total_us, - window_us, - category_stats, - resolved_us, - server_args, - fusion_opportunities, - ): - print(f" - {takeaway}") - print() diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/triage_overlap_helpers.py b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_overlap_helpers.py similarity index 59% rename from .claude/skills/sglang-torch-profiler-analysis/scripts/triage_overlap_helpers.py rename to .claude/skills/llm-torch-profiler-analysis/scripts/triage_overlap_helpers.py index 98dbd9d3c..0b38e588c 100644 --- a/.claude/skills/sglang-torch-profiler-analysis/scripts/triage_overlap_helpers.py +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_overlap_helpers.py @@ -2,14 +2,14 @@ from __future__ import annotations -import argparse -import math import re +from bisect import bisect_left, bisect_right from collections import Counter, defaultdict from dataclasses import dataclass, field from pathlib import Path from typing import Dict, Iterable, List, Optional, Sequence, Tuple +import triage_kernel_helpers as kernel_helpers from profile_common import ( coerce_optional_int, contains_any_keyword, @@ -19,17 +19,14 @@ from profile_common import ( is_complete_duration_event, is_non_kernel_trace_category, is_trace_metadata_name, - load_server_args, - load_trace_json, looks_like_python_scope_name, normalize_repo_relative_path, normalize_text, -) -from profile_common import run_profiler as shared_run_profiler -from profile_common import ( select_heaviest_pid, ) +SOURCE_MAP_SAMPLE_LIMIT_PER_NAME = 16 + COMMUNICATION_STRONG_KEYWORDS = ( "allreduce", "all_reduce", @@ -109,13 +106,30 @@ COMPUTE_KEYWORDS = ( "mm_kernel", ) -CATEGORY_CHARS = { - "compute": "#", - "communication": "=", - "elementwise": "~", - "memory": "+", - "other": "*", -} +LOW_SIGNAL_FUNCTION_TOKENS = ( + "__torch_function__", + "__torch_dispatch__", + "__call__", + "_call_impl", + "_wrapped_call_impl", +) + +LOW_SIGNAL_PATH_TOKENS = ( + "model_executor/parameter.py(", + "model_executor/parameter.py:", + "model_executor/cuda_graph_runner.py(", + "model_executor/cuda_graph_runner.py:", + "compilation/cuda_graph.py(", + "compilation/cuda_graph.py:", + "pyexecutor/cuda_graph_runner.py(", + "pyexecutor/cuda_graph_runner.py:", + "pyexecutor/py_executor.py(", + "pyexecutor/py_executor.py:", + "_torch/utils.py(", + "_torch/utils.py:", + "torch/fx/graph_module.py(", + "torch/fx/graph_module.py:", +) CATEGORY_PRIORITY = { "compute": 4, @@ -170,6 +184,7 @@ class KernelEvent: ts: float dur: float end: float + stage: str = "all" external_id: Optional[int] = None correlation: Optional[int] = None hidden_us: float = 0.0 @@ -209,6 +224,8 @@ class PythonScope: ts: float dur: float end: float + is_meaningful: bool = False + is_fallback: bool = False @dataclass @@ -231,6 +248,7 @@ class KernelSourceStats: scope_counter: Counter = field(default_factory=Counter) chain_counter: Counter = field(default_factory=Counter) launch_op_counter: Counter = field(default_factory=Counter) + site_share_counter: Counter = field(default_factory=Counter) @property def mapping_ratio(self) -> float: @@ -396,6 +414,10 @@ def is_meaningful_python_scope(name: str) -> bool: return True if normalized.startswith("sglang/"): return True + if normalized.startswith("vllm/"): + return True + if normalized.startswith("tensorrt_llm/"): + return True if normalized.startswith("sgl_kernel/"): return True return ".py(" in normalized @@ -446,6 +468,11 @@ def extract_kernel_events( raw_events = extract_trace_events(trace) thread_names = extract_thread_names(raw_events) correlation_external = build_correlation_external_lookup(raw_events) + ( + annotations_by_external_id, + gpu_stage_annotations, + cpu_stage_annotations, + ) = kernel_helpers.build_stage_annotations(raw_events) chosen_pid = select_heaviest_pid( raw_events, is_kernel_event, @@ -484,6 +511,13 @@ def extract_kernel_events( name=name, canonical_name=canonicalize_name(name), category=classify_kernel(name), + stage=kernel_helpers.resolve_kernel_stage( + kernel_ts=ts, + external_id=external_id, + annotations_by_external_id=annotations_by_external_id, + gpu_annotations=gpu_stage_annotations, + cpu_annotations=cpu_stage_annotations, + ), pid=pid, tid=tid, stream=str(stream), @@ -498,6 +532,16 @@ def extract_kernel_events( return kernel_events, chosen_pid +def group_events_by_stage( + events: Sequence[KernelEvent], default_stage: str +) -> Dict[str, List[KernelEvent]]: + grouped: Dict[str, List[KernelEvent]] = defaultdict(list) + for event in events: + stage = default_stage if default_stage != "all" else (event.stage or "all") + grouped[stage].append(event) + return dict(grouped) + + def dominant_overlap_name( event: KernelEvent, active_events: Iterable[KernelEvent] ) -> Optional[str]: @@ -643,91 +687,6 @@ def top_overlap_opportunities( return (primary + fallback)[:5] -def choose_window_events( - events: Sequence[KernelEvent], - representative_idx: int, - window_us: Optional[float], -) -> Tuple[float, float, List[KernelEvent]]: - center = next(event for event in events if event.idx == representative_idx) - span = window_us if window_us is not None else max(40.0, center.dur * 6.0) - start = max(0.0, center.ts - span * 0.35) - end = center.end + span * 0.65 - window_events = [ - event for event in events if event.end >= start and event.ts <= end - ] - return start, end, window_events - - -def render_ascii_timeline( - events: Sequence[KernelEvent], - representative_idx: int, - window_us: Optional[float], - width: int, -) -> str: - start, end, window_events = choose_window_events( - events, representative_idx, window_us - ) - if not window_events: - return "No events found in the selected window." - - streams = sorted( - {event.stream for event in window_events}, key=lambda item: (len(item), item) - ) - symbol_map: Dict[int, str] = {} - legend_events = sorted(window_events, key=lambda event: event.dur, reverse=True)[:8] - symbol_alphabet = list( - "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" - ) - for index, event in enumerate(legend_events): - symbol_map[event.idx] = symbol_alphabet[index] - - label_width = max(len(stream) for stream in streams) - lines = [] - marker_positions = [0, width // 4, width // 2, (3 * width) // 4, width - 1] - header = [" "] * width - for position in marker_positions: - header[position] = "|" - lines.append("time(us) " + "".join(header)) - - time_line = [" "] * width - markers = [ - start, - start + (end - start) * 0.25, - start + (end - start) * 0.5, - start + (end - start) * 0.75, - end, - ] - for position, value in zip(marker_positions, markers): - text = f"{value:.1f}" - begin = min(max(position - len(text) // 2, 0), max(0, width - len(text))) - for offset, char in enumerate(text): - time_line[begin + offset] = char - lines.append(" " + "".join(time_line)) - - for stream in streams: - row = ["."] * width - row_events = [event for event in window_events if event.stream == stream] - for event in row_events: - char = symbol_map.get(event.idx, CATEGORY_CHARS[event.category]) - left = int((event.ts - start) / max(end - start, 1.0) * (width - 1)) - right = int( - math.ceil((event.end - start) / max(end - start, 1.0) * (width - 1)) - ) - right = max(left + 1, min(right, width - 1)) - for pos in range(max(0, left), min(width, right + 1)): - row[pos] = char - lines.append(f"{stream:<{label_width}} " + "".join(row)) - - if legend_events: - lines.append("legend:") - for event in legend_events: - symbol = symbol_map[event.idx] - lines.append( - f" {symbol} [{event.category[:4]}] {short_name(event.canonical_name, 72)} ({event.dur:.1f} us)" - ) - return "\n".join(lines) - - def choose_best_scope(scope_chain: Sequence[str]) -> Optional[str]: ranked: List[Tuple[float, str]] = [] for index, scope in enumerate(scope_chain): @@ -736,6 +695,10 @@ def choose_best_scope(scope_chain: Sequence[str]) -> Optional[str]: score += 50.0 elif scope.startswith("sglang/"): score += 48.0 + elif scope.startswith("vllm/"): + score += 46.0 + elif scope.startswith("tensorrt_llm/"): + score += 44.0 elif scope.startswith("sgl_kernel/"): score += 30.0 elif ".py(" in scope: @@ -744,10 +707,21 @@ def choose_best_scope(scope_chain: Sequence[str]) -> Optional[str]: score -= 15.0 if "scheduler_profiler_mixin.py" in scope: score -= 20.0 + if is_low_signal_scope(scope): + score -= 25.0 ranked.append((score, scope)) return max(ranked, key=lambda item: item[0])[1] if ranked else None +def is_low_signal_scope(scope: str) -> bool: + lowered = canonicalize_python_scope_name(scope).lower() + if not lowered: + return False + return any(token in lowered for token in LOW_SIGNAL_FUNCTION_TOKENS) or any( + token in lowered for token in LOW_SIGNAL_PATH_TOKENS + ) + + def scope_chain_key(scope_chain: Sequence[str]) -> Optional[str]: if not scope_chain: return None @@ -755,101 +729,346 @@ def scope_chain_key(scope_chain: Sequence[str]) -> Optional[str]: return " -> ".join(trimmed) +def normalize_match_text(text: object) -> str: + return re.sub(r"[^0-9A-Za-z]+", "", normalize_text(text)).lower() + + +def source_scope_priority(scope: Optional[str]) -> int: + normalized = canonicalize_python_scope_name(scope or "") + if not normalized or normalized == "unmapped": + return 0 + penalty = 80 if is_low_signal_scope(normalized) else 0 + if normalized.startswith("python/sglang/"): + return 300 - penalty + if normalized.startswith("sglang/"): + return 290 - penalty + if normalized.startswith("vllm/"): + return 285 - penalty + if normalized.startswith("tensorrt_llm/"): + return 280 - penalty + if normalized.startswith("sgl_kernel/"): + return 260 - penalty + if ".py(" in normalized: + return 120 - penalty + return 0 + + +def kernel_alias_token_groups(kernel_name: str) -> List[Tuple[str, ...]]: + lowered = normalize_match_text(kernel_name) + groups: List[Tuple[str, ...]] = [] + if "flashattnfwdcombine" in lowered: + groups.append( + ( + "flashattnfwdsm90", + "flashattnvarlenfunc", + "vllmflashattnflashattninterface", + "vllmfa3cfwd", + ) + ) + if "kernelmha" in lowered: + groups.append( + ( + "maskedmultiheadattentionkernel", + "attentioninplace", + "attentionbackendtrtllm", + ) + ) + if "applybiasropeupdatekvcachev2" in lowered: + groups.append( + ( + "fusedqknormropekernel", + "applyqknormrope", + "modelingqwen3py98applyqknormrope", + ) + ) + if lowered.startswith("memset"): + groups.append(("memset",)) + return groups + + +def source_stats_lookup_text( + kernel_name: str, stats: Optional[KernelSourceStats] +) -> str: + parts = [kernel_name] + if stats: + parts.append(str(stats.best_scope or "")) + parts.append(str(stats.best_chain or "")) + parts.append(str(stats.best_launch_op or "")) + return normalize_match_text(" ".join(parts)) + + +def relaxed_source_stats_lookup( + source_map: Dict[str, KernelSourceStats], kernel_name: str +) -> Optional[KernelSourceStats]: + if kernel_name in source_map: + return source_map[kernel_name] + + lowered = kernel_name.lower() + best_key = None + best_score = -1 + for candidate_key in source_map: + candidate_lowered = candidate_key.lower() + if candidate_lowered.startswith(lowered) or lowered.startswith( + candidate_lowered + ): + score = min(len(candidate_lowered), len(lowered)) + elif candidate_lowered in lowered or lowered in candidate_lowered: + score = min(len(candidate_lowered), len(lowered)) // 2 + else: + continue + if score > best_score: + best_key = candidate_key + best_score = score + if best_key: + return source_map.get(best_key) + + lowered_compact = normalize_match_text(kernel_name) + if len(lowered_compact) >= 96: + + def common_prefix_len(left: str, right: str) -> int: + count = 0 + for left_ch, right_ch in zip(left, right): + if left_ch != right_ch: + break + count += 1 + return count + + best_key = None + best_score = -1 + for candidate_key in source_map: + candidate_compact = normalize_match_text(candidate_key) + if len(candidate_compact) < 96: + continue + prefix_len = common_prefix_len(lowered_compact, candidate_compact) + shorter_len = min(len(lowered_compact), len(candidate_compact)) + if prefix_len < 64 or prefix_len < int(shorter_len * 0.4): + continue + score = prefix_len + if lowered_compact.startswith( + "voidcutlassdevicekernelflash" + ) and candidate_compact.startswith("voidcutlassdevicekernelflash"): + score += 32 + if score > best_score: + best_key = candidate_key + best_score = score + if best_key: + return source_map.get(best_key) + + alias_groups = kernel_alias_token_groups(kernel_name) + if not alias_groups: + return None + best_key = None + best_score = -1 + for candidate_key, stats in source_map.items(): + candidate_text = source_stats_lookup_text(candidate_key, stats) + score = 0 + for group_index, group in enumerate(alias_groups): + group_score = max( + (len(token) for token in group if token in candidate_text), + default=0, + ) + if group_score: + score += 1000 * (group_index + 1) + group_score + if score <= 0: + continue + score += source_scope_priority(stats.best_scope) + score += int(stats.mapping_ratio * 100) + if score > best_score: + best_key = candidate_key + best_score = score + return source_map.get(best_key) if best_key else None + + def extract_cpu_launch_contexts( raw_events: Sequence[dict], + target_external_ids: Optional[set[int]] = None, ) -> Dict[int, List[CPUOpContext]]: - # Rebuild `External id -> CPU op -> active Python scopes` so formal-trace - # kernels can be mapped back to readable Python locations from the mapping - # trace even when launches are interleaved on the same thread. - scopes_by_thread: Dict[Tuple[str, str], List[PythonScope]] = defaultdict(list) + # Rebuild `External id -> CPU op -> active Python scopes` only for the + # small set of launch ids that the source-map step will actually consume. + # vLLM eager traces can have millions of Python frames on one thread, so + # avoid global timeline reconstruction across unrelated threads and ids. cpu_ops_by_thread: Dict[Tuple[str, str], List[CPUOpContext]] = defaultdict(list) for event in raw_events: if not is_complete_duration_event(event): continue - cat = str(event.get("cat", "")) + if str(event.get("cat", "")) != "cpu_op": + continue + args = event.get("args", {}) or {} + external_id = coerce_optional_int(args.get("External id")) + if external_id is None: + continue + if target_external_ids is not None and external_id not in target_external_ids: + continue pid = str(event.get("pid")) tid = str(event.get("tid")) ts = float(event.get("ts", 0.0)) dur = float(event.get("dur", 0.0)) - args = event.get("args", {}) or {} - if cat == "python_function": - name = canonicalize_python_scope_name(event.get("name", "")) - scopes_by_thread[(pid, tid)].append( - PythonScope( - name=str(event.get("name", "")), - normalized_name=name, - pid=pid, - tid=tid, - ts=ts, - dur=dur, - end=ts + dur, - ) + cpu_ops_by_thread[(pid, tid)].append( + CPUOpContext( + external_id=external_id, + cpu_op_name=str(event.get("name", "")), + pid=pid, + tid=tid, + ts=ts, + dur=dur, + end=ts + dur, + scope_chain=(), ) - elif cat == "cpu_op": - external_id = coerce_optional_int(args.get("External id")) - if external_id is None: - continue - cpu_ops_by_thread[(pid, tid)].append( - CPUOpContext( - external_id=external_id, - cpu_op_name=str(event.get("name", "")), - pid=pid, - tid=tid, - ts=ts, - dur=dur, - end=ts + dur, - scope_chain=(), - ) + ) + + if not cpu_ops_by_thread: + return {} + + scopes_by_thread: Dict[Tuple[str, str], List[PythonScope]] = defaultdict(list) + relevant_threads = set(cpu_ops_by_thread) + for event in raw_events: + if not is_complete_duration_event(event): + continue + if str(event.get("cat", "")) != "python_function": + continue + pid = str(event.get("pid")) + tid = str(event.get("tid")) + thread_key = (pid, tid) + if thread_key not in relevant_threads: + continue + normalized_name = canonicalize_python_scope_name(event.get("name", "")) + is_meaningful = is_meaningful_python_scope(normalized_name) + is_fallback = is_fallback_python_scope(normalized_name) + if not is_meaningful and not is_fallback: + continue + ts = float(event.get("ts", 0.0)) + dur = float(event.get("dur", 0.0)) + scopes_by_thread[thread_key].append( + PythonScope( + name=str(event.get("name", "")), + normalized_name=normalized_name, + pid=pid, + tid=tid, + ts=ts, + dur=dur, + end=ts + dur, + is_meaningful=is_meaningful, + is_fallback=is_fallback, ) + ) contexts_by_external_id: Dict[int, List[CPUOpContext]] = defaultdict(list) - for thread_key in set(scopes_by_thread) | set(cpu_ops_by_thread): + for thread_key in relevant_threads: scopes = scopes_by_thread.get(thread_key, []) cpu_ops = cpu_ops_by_thread.get(thread_key, []) timeline = [] - for scope in scopes: - timeline.append((scope.ts, 0, scope)) - timeline.append((scope.end, 2, scope)) - for cpu_op in cpu_ops: - timeline.append((cpu_op.ts, 1, cpu_op)) + for scope_idx, scope in enumerate(scopes): + timeline.append((scope.ts, 0, scope_idx)) + timeline.append((scope.end, 2, scope_idx)) + for cpu_op_idx, cpu_op in enumerate(cpu_ops): + timeline.append((cpu_op.ts, 1, cpu_op_idx)) timeline.sort(key=lambda item: (item[0], item[1])) - active_scopes: List[PythonScope] = [] + active_scopes: Dict[int, PythonScope] = {} for _, kind, payload in timeline: if kind == 0: - active_scopes.append(payload) + active_scopes[payload] = scopes[payload] elif kind == 1: - normalized_chain = [scope.normalized_name for scope in active_scopes] meaningful = [ - scope - for scope in normalized_chain - if is_meaningful_python_scope(scope) - ] - fallback = [ - scope - for scope in normalized_chain - if is_fallback_python_scope(scope) + scope.normalized_name + for scope in active_scopes.values() + if scope.is_meaningful ] + fallback = ( + [] + if meaningful + else [ + scope.normalized_name + for scope in active_scopes.values() + if scope.is_fallback + ] + ) chosen_chain = tuple((meaningful or fallback)[-6:]) - contexts_by_external_id[payload.external_id].append( + cpu_op = cpu_ops[payload] + contexts_by_external_id[cpu_op.external_id].append( CPUOpContext( - external_id=payload.external_id, - cpu_op_name=payload.cpu_op_name, - pid=payload.pid, - tid=payload.tid, - ts=payload.ts, - dur=payload.dur, - end=payload.end, + external_id=cpu_op.external_id, + cpu_op_name=cpu_op.cpu_op_name, + pid=cpu_op.pid, + tid=cpu_op.tid, + ts=cpu_op.ts, + dur=cpu_op.dur, + end=cpu_op.end, scope_chain=chosen_chain, ) ) else: - if payload in active_scopes: - active_scopes.remove(payload) + active_scopes.pop(payload, None) return contexts_by_external_id +def is_cuda_launch_event(name: str, cat: str) -> bool: + lowered_name = normalize_text(name).lower() + lowered_cat = normalize_text(cat).lower() + if lowered_cat == "cuda_runtime": + return lowered_name in { + "cudaLaunchKernel", + "cudaLaunchKernelExC", + } + return lowered_name in { + "cuLaunchKernel", + "cuLaunchKernelEx", + "cudaLaunchKernel", + "cudaLaunchKernelExC", + } + + +@dataclass +class LaunchContext: + correlation: int + pid: str + tid: str + ts: float + dur: float + end: float + launch_name: str + + +def build_launch_contexts( + raw_events: Sequence[dict], +) -> Dict[int, List[LaunchContext]]: + output: Dict[int, List[LaunchContext]] = defaultdict(list) + for event in raw_events: + if not is_complete_duration_event(event): + continue + cat = str(event.get("cat", "")) + name = str(event.get("name", "")) + args = event.get("args", {}) or {} + correlation = coerce_optional_int(args.get("correlation")) + if correlation is None or not is_cuda_launch_event(name, cat): + continue + ts = float(event.get("ts", 0.0)) + dur = float(event.get("dur", 0.0)) + output[correlation].append( + LaunchContext( + correlation=correlation, + pid=str(event.get("pid")), + tid=str(event.get("tid")), + ts=ts, + dur=dur, + end=ts + dur, + launch_name=name, + ) + ) + for items in output.values(): + items.sort(key=lambda item: item.ts) + return output + + +def choose_launch_context( + contexts: Sequence[LaunchContext], kernel_ts: float +) -> Optional[LaunchContext]: + if not contexts: + return None + return min(contexts, key=lambda context: (abs(context.ts - kernel_ts), context.dur)) + + def choose_cpu_context( contexts: Sequence[CPUOpContext], kernel_ts: float ) -> Optional[CPUOpContext]: @@ -901,21 +1120,158 @@ def choose_temporal_scope_chain( return tuple(chain[-6:]) +def build_temporal_scope_lookup( + scopes: Sequence[PythonScope], + query_points: Sequence[Tuple[int, float]], +) -> Dict[int, Tuple[str, ...]]: + if not scopes or not query_points: + return {} + + timeline: List[Tuple[float, int, object]] = [] + for scope in scopes: + timeline.append((scope.ts, 0, scope)) + timeline.append((scope.end, 2, scope)) + for event_idx, probe_ts in query_points: + timeline.append((probe_ts, 1, event_idx)) + timeline.sort(key=lambda item: (item[0], item[1])) + + active_scopes: List[PythonScope] = [] + resolved: Dict[int, Tuple[str, ...]] = {} + for _, kind, payload in timeline: + if kind == 0: + active_scopes.append(payload) + continue + if kind == 2: + if payload in active_scopes: + active_scopes.remove(payload) + continue + + chain: List[str] = [] + seen: set[str] = set() + for scope in sorted( + active_scopes, + key=lambda scope: (scope.ts, -scope.dur, scope.normalized_name), + ): + name = scope.normalized_name + if name in seen: + continue + seen.add(name) + chain.append(name) + resolved[payload] = tuple(chain[-6:]) + return resolved + + +def build_temporal_scope_lookup_from_raw_events( + raw_events: Sequence[dict], + query_points: Sequence[Tuple[int, float]], +) -> Dict[int, Tuple[str, ...]]: + if not query_points: + return {} + + ordered_queries = sorted( + ((float(query_ts), int(query_id)) for query_id, query_ts in query_points), + key=lambda item: item[0], + ) + query_times = [query_ts for query_ts, _ in ordered_queries] + query_ids = [query_id for _, query_id in ordered_queries] + first_query_ts = query_times[0] + last_query_ts = query_times[-1] + + matches_by_query: Dict[int, List[PythonScope]] = defaultdict(list) + for event in raw_events: + if not is_complete_duration_event(event): + continue + if str(event.get("cat", "")) != "python_function": + continue + ts = float(event.get("ts", 0.0)) + dur = float(event.get("dur", 0.0)) + end = ts + dur + if end < first_query_ts or ts > last_query_ts: + continue + + normalized_name = canonicalize_python_scope_name(event.get("name", "")) + if not is_meaningful_python_scope(normalized_name): + continue + + left = bisect_left(query_times, ts - 1e-3) + right = bisect_right(query_times, end + 1e-3) + if left >= right: + continue + + scope = PythonScope( + name=str(event.get("name", "")), + normalized_name=normalized_name, + pid=str(event.get("pid")), + tid=str(event.get("tid")), + ts=ts, + dur=dur, + end=end, + is_meaningful=True, + is_fallback=False, + ) + for pos in range(left, right): + matches_by_query[query_ids[pos]].append(scope) + + resolved: Dict[int, Tuple[str, ...]] = {} + for query_id, scopes in matches_by_query.items(): + chain: List[str] = [] + seen: set[str] = set() + for scope in sorted( + scopes, + key=lambda scope: (scope.ts, -scope.dur, scope.normalized_name), + ): + name = scope.normalized_name + if name in seen: + continue + seen.add(name) + chain.append(name) + resolved[query_id] = tuple(chain[-6:]) + return resolved + + def build_kernel_source_map( mapping_bundle: TraceBundle, + kernel_map_entry_lookup=None, + stage: str = "all", ) -> Dict[str, KernelSourceStats]: - contexts_by_external_id = extract_cpu_launch_contexts(mapping_bundle.raw_events) - temporal_scopes = extract_meaningful_python_scopes(mapping_bundle.raw_events) + sampled_events = sample_source_map_events(mapping_bundle.events) + target_external_ids = { + event.external_id for event in sampled_events if event.external_id is not None + } + contexts_by_external_id = extract_cpu_launch_contexts( + mapping_bundle.raw_events, + target_external_ids=target_external_ids or None, + ) + correlation_external = build_correlation_external_lookup(mapping_bundle.raw_events) + launch_contexts_by_correlation = build_launch_contexts(mapping_bundle.raw_events) + fallback_queries = [ + (event.idx, event.ts) + for event in sampled_events + if event.external_id is None + or not contexts_by_external_id.get(event.external_id) + ] + temporal_scope_lookup = build_temporal_scope_lookup_from_raw_events( + mapping_bundle.raw_events, + fallback_queries, + ) source_map: Dict[str, KernelSourceStats] = {} - for event in mapping_bundle.events: + for event in sampled_events: stats = source_map.setdefault( event.canonical_name, KernelSourceStats(name=event.canonical_name) ) stats.total_count += 1 + kernel_entry = ( + kernel_map_entry_lookup(stage, event.canonical_name) + if kernel_map_entry_lookup is not None + else None + ) cpu_context = None - if event.external_id is not None: + effective_external_id = event.external_id + if effective_external_id is None and event.correlation is not None: + effective_external_id = correlation_external.get(event.correlation) + if effective_external_id is not None: cpu_context = choose_cpu_context( - contexts_by_external_id.get(event.external_id, []), event.ts + contexts_by_external_id.get(effective_external_id, []), event.ts ) launch_op = None @@ -924,17 +1280,54 @@ def build_kernel_source_map( launch_op = canonicalize_cpu_op_name(cpu_context.cpu_op_name) scope_chain = cpu_context.scope_chain else: - scope_chain = choose_temporal_scope_chain(temporal_scopes, event.ts) + launch_context = ( + choose_launch_context( + launch_contexts_by_correlation.get(event.correlation, []), event.ts + ) + if event.correlation is not None + else None + ) + if launch_context is not None: + scope_chain = build_temporal_scope_lookup_from_raw_events( + mapping_bundle.raw_events, + [(event.idx, launch_context.ts)], + ).get(event.idx, ()) + if scope_chain: + launch_op = canonicalize_cpu_op_name(launch_context.launch_name) + if not scope_chain: + scope_chain = temporal_scope_lookup.get(event.idx, ()) if scope_chain: launch_op = "time-window fallback" if not scope_chain: + if kernel_entry: + best_location = str(kernel_entry.get("best_location") or "").strip() + if best_location and best_location != "unresolved": + stats.mapped_count += 1 + stats.scope_counter[best_location] += 1 + stats.site_share_counter[best_location] += 1 + for site in kernel_entry.get("sites") or []: + display_location = str( + site.get("display_location") or site.get("location") or "" + ).strip() + if display_location and display_location != "unresolved": + launches = int(site.get("launches") or 0) + stats.site_share_counter[display_location] += max( + 1, launches + ) + if launches > 0: + stats.scope_counter[display_location] += launches + top_cpu_op = site.get("top_cpu_op") + if top_cpu_op: + launches = int(site.get("launches") or 0) + stats.launch_op_counter[str(top_cpu_op)] += max(1, launches) continue stats.mapped_count += 1 best_scope = choose_best_scope(scope_chain) if best_scope: stats.scope_counter[best_scope] += 1 + stats.site_share_counter[best_scope] += 1 chain = scope_chain_key(scope_chain) if chain: stats.chain_counter[chain] += 1 @@ -943,6 +1336,68 @@ def build_kernel_source_map( return source_map +def merge_source_map_from_kernel_payload( + source_map: Dict[str, KernelSourceStats], + stage_payload: Optional[dict], +) -> Dict[str, KernelSourceStats]: + if not stage_payload: + return source_map + + for kernel_name, entry in (stage_payload.get("kernels") or {}).items(): + sites = entry.get("sites") or [] + best_location = str(entry.get("best_location") or "").strip() + if not sites and (not best_location or best_location == "unresolved"): + continue + + stats = source_map.setdefault(kernel_name, KernelSourceStats(name=kernel_name)) + if sites: + for site in sites: + location = str(site.get("location") or best_location or "").strip() + launches = max(1, int(site.get("launches") or 0)) + stats.total_count += launches + if location and location != "unresolved": + stats.mapped_count += launches + stats.scope_counter[location] += launches + stats.site_share_counter[location] += launches + top_cpu_op = str(site.get("top_cpu_op") or "").strip() + if top_cpu_op: + stats.launch_op_counter[top_cpu_op] += launches + stack = str(site.get("stack") or "").strip() + if stack: + stats.chain_counter[stack] += launches + continue + + stats.total_count += 1 + stats.mapped_count += 1 + stats.scope_counter[best_location] += 1 + stats.site_share_counter[best_location] += 1 + return source_map + + +def sample_source_map_events( + events: Sequence[KernelEvent], + per_name_limit: int = SOURCE_MAP_SAMPLE_LIMIT_PER_NAME, +) -> List[KernelEvent]: + if per_name_limit <= 0: + return list(events) + + grouped: Dict[str, List[KernelEvent]] = defaultdict(list) + for event in events: + grouped[event.canonical_name].append(event) + + sampled: List[KernelEvent] = [] + for kernel_name in sorted(grouped): + items = grouped[kernel_name] + if len(items) <= per_name_limit: + sampled.extend(items) + continue + for sample_idx in range(per_name_limit): + pos = round(sample_idx * (len(items) - 1) / (per_name_limit - 1)) + sampled.append(items[pos]) + sampled.sort(key=lambda event: (event.ts, event.idx)) + return sampled + + def format_overlap_counter(counter: Counter, limit: int = 2) -> str: if not counter: return "n/a" @@ -954,27 +1409,27 @@ def format_overlap_counter(counter: Counter, limit: int = 2) -> str: def build_headroom_suggestion(stats: AggregateStats) -> str: if stats.category == "communication": - return "Exposed comm path. Check whether this code path can overlap with nearby compute." + return "Communication is still exposed. Check overlap with nearby compute." if stats.category in {"elementwise", "memory"}: - return "Still exposed. Try to fuse it or move it under a nearby compute-heavy window." - return "Meaningful exposed time remains. Inspect stream placement and surrounding dependencies." + return "This work is still exposed. Check fusion or nearby compute coverage." + return ( + "This work is still exposed. Check stream placement and immediate dependencies." + ) def build_hidden_suggestion(stats: AggregateStats) -> str: overlap = format_overlap_counter(stats.overlap_with, limit=1) if overlap != "n/a": - return f"Mostly hidden under {overlap}. Standalone tuning is probably low ROI." - return ( - "Mostly hidden already. Optimize it only if you also change fusion or schedule." - ) + return f"Mostly hidden under {overlap}. Revisit only if schedule or fusion changes." + return "Mostly hidden already. Revisit only if schedule or fusion changes." def build_other_suggestion(stats: AggregateStats) -> str: if stats.exclusive_ratio >= 0.6: - return "Some exposed time remains, but it did not rank among the strongest headroom rows." + return "Still exposed, but not one of the leading overlap targets." if stats.hidden_ratio >= 0.6: - return "Often hidden already. Usually secondary unless it also drives launch count." - return "Mixed exposure and overlap. Inspect after the stronger rows above." + return "Often hidden already. Revisit it if launch count or schedule changes." + return "Mixed exposure and overlap. Inspect it after the higher-share rows above." def parse_scope_signature(scope: str) -> Tuple[str, str]: @@ -1039,7 +1494,7 @@ def describe_neighbor( ) -> str: if neighbor is None: return "none" - source = source_map.get(neighbor.canonical_name) + source = relaxed_source_stats_lookup(source_map, neighbor.canonical_name) scope = source.best_scope if source and source.best_scope else "unmapped" if gap_us is not None: gap_us = max(gap_us, 0.0) @@ -1068,10 +1523,14 @@ def classify_dependency_signal( prev_gap = current.ts - prev_event.end if prev_event is not None else None next_gap = next_event.ts - current.end if next_event is not None else None prev_source = ( - source_map.get(prev_event.canonical_name) if prev_event is not None else None + relaxed_source_stats_lookup(source_map, prev_event.canonical_name) + if prev_event is not None + else None ) next_source = ( - source_map.get(next_event.canonical_name) if next_event is not None else None + relaxed_source_stats_lookup(source_map, next_event.canonical_name) + if next_event is not None + else None ) prev_scope = ( prev_source.best_scope if prev_source and prev_source.best_scope else "unmapped" @@ -1157,7 +1616,7 @@ def build_priority_and_recommendation( ) -> Tuple[str, str]: dep_label = dependency_risk_label(dependency_signal) if share_pct < 1.0: - return "P5", "skip overlap" + return "P5", "skip" if verdict == "headroom": if dep_label == "low": @@ -1167,17 +1626,17 @@ def build_priority_and_recommendation( return "P2", "check deps" if verdict == "low-roi-hidden": - return "P4", "skip overlap" + return "P4", "skip" if stats.exclusive_ratio >= 0.85 and dep_label == "low": - return "P3", "observe later" + return "P3", "defer" if stats.hidden_ratio >= 0.7: - return "P5", "skip overlap" + return "P5", "skip" if dep_label == "high": return "P4", "check deps" if dep_label == "unclear": - return "P4", "manual check" - return "P4", "observe later" + return "P4", "inspect" + return "P4", "defer" def make_action_row( @@ -1189,7 +1648,7 @@ def make_action_row( neighbor_index: Dict[int, Tuple[Optional[KernelEvent], Optional[KernelEvent]]], total_busy_us: float, ) -> ActionRow: - source = source_map.get(stats.name) + source = relaxed_source_stats_lookup(source_map, stats.name) representative_idx = stats.representative_idx dependency_signal = "adjacency unclear" prev_neighbor = "none" @@ -1286,271 +1745,3 @@ def build_action_rows( if table_limit > 0: return rows[:table_limit] return rows - - -def render_action_table(rows: Sequence[ActionRow]) -> List[str]: - lines = [ - "| Priority | Verdict | Kernel | Python scope | Formal signal | Dep risk | Recommendation |", - "| --- | --- | --- | --- | --- | --- | --- |", - ] - if not rows: - lines.append( - "| - | No actionable overlap rows stood out from the formal trace. | - | - | - | - | - |" - ) - return lines - for row in rows: - formal_signal = ( - f"share {row.share_pct:.1f}%, " - f"excl {row.exclusive_ratio * 100:.1f}% / " - f"hid {row.hidden_ratio * 100:.1f}%" - ) - lines.append( - "| " - + " | ".join( - [ - row.priority, - row.verdict, - row.kernel, - row.python_scope, - f"{row.total_us:.1f} us, {formal_signal}", - dependency_risk_label(row.dependency_signal), - row.recommendation, - ] - ) - + " |" - ) - return lines - - -def trace_summary_line(bundle: TraceBundle) -> str: - events = bundle.events - streams = sorted({event.stream for event in events}) - if bundle.overlap_stats is None: - return f"{bundle.label}: {len(events)} kernel events, {len(streams)} streams" - overlap_ratio = ( - bundle.overlap_stats["total_overlap_us"] / bundle.overlap_stats["total_busy_us"] - if bundle.overlap_stats["total_busy_us"] - else 0.0 - ) - return ( - f"{bundle.label}: {len(events)} kernel events, {len(streams)} streams, " - f"busy={bundle.overlap_stats['total_busy_us']:.1f} us, " - f"2+ stream overlap={bundle.overlap_stats['total_overlap_us']:.1f} us " - f"({overlap_ratio * 100:.1f}%), " - f"peak_concurrency={int(bundle.overlap_stats['max_concurrent_streams'])}" - ) - - -def launch_summary(server_args: Optional[dict]) -> Optional[str]: - if not server_args: - return None - model_path = server_args.get("model_path") or server_args.get("model") - shape_bits = [] - if model_path: - shape_bits.append(f"model={model_path}") - for key in ("tp_size", "dp_size", "pp_size", "ep_size", "enable_dp_attention"): - if key in server_args: - shape_bits.append(f"{key}={server_args[key]}") - return ", ".join(shape_bits) if shape_bits else None - - -def build_report( - mapping_bundle: TraceBundle, - formal_bundle: TraceBundle, - source_map: Dict[str, KernelSourceStats], - aggregates: Dict[Tuple[str, str], AggregateStats], - rows: Sequence[ActionRow], - window_us: Optional[float], - timeline_count: int, - width: int, - table_only: bool, -) -> str: - lines: List[str] = [] - lines.append(f"Mapping Trace: {mapping_bundle.trace_path}") - mapping_launch = launch_summary(mapping_bundle.server_args) - if mapping_launch: - lines.append(f"Mapping Launch: {mapping_launch}") - if mapping_bundle.pid: - lines.append(f"Mapping PID slice: {mapping_bundle.pid}") - lines.append(trace_summary_line(mapping_bundle)) - - lines.append("") - lines.append(f"Formal Trace: {formal_bundle.trace_path}") - formal_launch = launch_summary(formal_bundle.server_args) - if formal_launch: - lines.append(f"Formal Launch: {formal_launch}") - if formal_bundle.pid: - lines.append(f"Formal PID slice: {formal_bundle.pid}") - lines.append(trace_summary_line(formal_bundle)) - - mapped_kernels = sum(1 for stats in source_map.values() if stats.mapped_count > 0) - table_mapped = sum(1 for row in rows if row.python_scope != "unmapped") - lines.append("") - lines.append( - "Source Map Coverage: " - f"{mapped_kernels}/{len(source_map)} mapping-trace kernels found a Python scope, " - f"{table_mapped}/{len(rows)} table rows were mapped back to code." - ) - - lines.append("") - lines.append("Action Table") - lines.extend(render_action_table(rows)) - - if not table_only: - detail_lookup = {row.kernel: row for row in rows} - focus_rows = list(rows) - lines.append("") - lines.append("Source Context") - if not focus_rows: - lines.append(" No source-mapped rows to expand.") - else: - for index, row in enumerate(focus_rows, start=1): - stats = source_map.get(row.kernel) - lines.append( - f" {index}. {short_name(row.kernel, 88)} [{row.priority}, {row.verdict}, {row.category}] " - f"mapping={row.mapping_ratio * 100:.1f}%" - ) - lines.append(f" time share: {row.share_pct:.1f}%") - lines.append(f" python scope: {row.python_scope}") - lines.append(f" launch op: {row.launch_op}") - lines.append( - f" dependency signal: {dependency_risk_label(row.dependency_signal)}" - ) - lines.append(f" prev neighbor: {row.prev_neighbor}") - lines.append(f" next neighbor: {row.next_neighbor}") - lines.append(f" recommendation: {row.recommendation}") - if stats and stats.best_chain: - lines.append( - f" call chain: {short_name(stats.best_chain, 132)}" - ) - lines.append(f" conclusion: {row.suggestion}") - - timeline_targets: List[int] = [] - for row in rows: - if ( - row.representative_idx is not None - and row.representative_idx not in timeline_targets - ): - timeline_targets.append(row.representative_idx) - timeline_targets = timeline_targets[:timeline_count] - - if timeline_targets: - lines.append("") - lines.append("ASCII Timelines") - for index, representative_idx in enumerate(timeline_targets, start=1): - event = next( - event - for event in formal_bundle.events - if event.idx == representative_idx - ) - mapped_scope = ( - detail_lookup.get(event.canonical_name).python_scope - if event.canonical_name in detail_lookup - else "unmapped" - ) - lines.append( - f" Window {index}: {short_name(event.canonical_name, 90)} " - f"[{event.category}] ts={event.ts:.1f} us dur={event.dur:.1f} us" - ) - lines.append(f" mapped scope: {short_name(mapped_scope, 120)}") - lines.append( - render_ascii_timeline( - formal_bundle.events, representative_idx, window_us, width - ) - ) - lines.append("") - - lines.append("Notes") - lines.append( - " - The mapping trace should be graph-off so kernel-to-code attribution stays readable." - ) - lines.append( - " - The formal trace should keep the real serving optimizations enabled; overlap conclusions come from this trace." - ) - lines.append( - " - A mapped Python scope is a launch-site clue, not proof that the code is dependency-free to reorder." - ) - return "\n".join(lines).rstrip() - - -def discover_trace_file(path: Path) -> Tuple[Path, Optional[dict]]: - if path.is_file(): - return path, load_server_args(path) - - traces = sorted( - path.glob("*.trace.json.gz"), key=lambda candidate: candidate.stat().st_mtime - ) - traces.extend( - sorted( - [ - candidate - for candidate in path.glob("*.trace.json") - if candidate.name not in {trace.name[:-3] for trace in traces} - ], - key=lambda candidate: candidate.stat().st_mtime, - ) - ) - if not traces: - child_dirs = sorted( - [candidate for candidate in path.iterdir() if candidate.is_dir()], - key=lambda candidate: candidate.stat().st_mtime, - ) - for child_dir in reversed(child_dirs): - child_traces = list(child_dir.glob("*.trace.json.gz")) + list( - child_dir.glob("*.trace.json") - ) - if child_traces: - return discover_trace_file(child_dir) - if not traces: - raise FileNotFoundError(f"No trace files found under {path}") - - non_merged = [trace for trace in traces if not trace.name.startswith("merged-")] - tp0 = [ - trace for trace in non_merged if "-TP-0" in trace.name or "TP-0" in trace.name - ] - chosen = tp0[-1] if tp0 else (non_merged[-1] if non_merged else traces[-1]) - return chosen, load_server_args(path) - - -def resolve_trace_source( - label: str, - input_path: Optional[str], - url: Optional[str], - output_dir: Optional[str], - profile_prefix: Optional[str], - args: argparse.Namespace, -) -> TraceBundle: - if bool(input_path) == bool(url): - raise ValueError(f"{label} trace requires exactly one of input path or URL.") - - if url: - target_dir = shared_run_profiler( - url=url, - output_dir=output_dir, - num_steps=args.num_steps, - profile_by_stage=args.profile_by_stage, - merge_profiles=args.merge_profiles, - profile_prefix=profile_prefix, - probe_requests=max(0, args.probe_requests), - probe_prompt=args.probe_prompt, - probe_max_new_tokens=args.probe_max_new_tokens, - probe_delay=args.probe_delay, - start_step=args.start_step, - ) - trace_path, server_args = discover_trace_file(target_dir) - else: - trace_path, server_args = discover_trace_file(Path(input_path).resolve()) - - trace = load_trace_json(trace_path) - raw_events = trace.get("traceEvents", trace if isinstance(trace, list) else []) - events, pid = extract_kernel_events(trace, args.pid_substring) - if not events: - raise RuntimeError(f"No GPU kernel events found in {trace_path}") - return TraceBundle( - label=label, - trace_path=trace_path, - server_args=server_args, - raw_events=raw_events, - events=events, - pid=pid, - ) diff --git a/.claude/skills/sglang-bisect-ci-regression/SKILL.md b/.claude/skills/sglang-bisect-ci-regression/SKILL.md index 4eb39227c..0afd15972 100644 --- a/.claude/skills/sglang-bisect-ci-regression/SKILL.md +++ b/.claude/skills/sglang-bisect-ci-regression/SKILL.md @@ -1,3 +1,8 @@ +--- +name: sglang-bisect-ci-regression +description: Investigate consistently failing SGLang CI tests by extracting the failure signature from scheduled or rerun workflows, bisecting the passing/failing commit window, checking runner or hardware specificity, and optionally reproducing on a remote GPU host. +--- + # SGLang Bisect CI Regression Investigate a consistently failing CI test to find the root cause - whether it's a code regression from a specific PR, a hardware/runner-specific issue, or an environment change. Optionally reproduce the failure on a remote GPU server. diff --git a/.claude/skills/sglang-torch-profiler-analysis/SKILL.md b/.claude/skills/sglang-torch-profiler-analysis/SKILL.md deleted file mode 100644 index 1d7d5d5e9..000000000 --- a/.claude/skills/sglang-torch-profiler-analysis/SKILL.md +++ /dev/null @@ -1,189 +0,0 @@ ---- -name: sglang-torch-profiler-analysis -description: "Compact SGLang torch-profiler triage skill. Use when Codex should inspect an existing `trace.json(.gz)` or profile directory, trigger `sglang.profiler` against a live server, and return one compact report with kernel, overlap-opportunity, and fuse-pattern tables. Single-trace triage is enough for quick diagnosis; mapping+formal two-trace triage gives stronger overlap conclusions." ---- - -# SGLang Torch Profiler Analysis - -## Overview - -Use this skill for SGLang `torch.profiler` analysis. - -There is only one public workflow: - -- `triage` - -Use the unified entrypoint: - -- [scripts/analyze_sglang_torch_profile.py](scripts/analyze_sglang_torch_profile.py) - -`triage` always prints the same three tables: - -- kernel table -- overlap-opportunity table -- fuse-pattern table - -By default, all three tables only render rows at or above `1.0%` cumulative GPU-time share. -Treat anything below that as noise unless the user explicitly asks for a lower cutoff. - -The script-level fuse-pattern table should stay source-backed and deterministic. -Do not build a fuzzy string-matching engine into the script for typo-tolerance. - -If exact/source-backed matching is weak but the agent judges that a cluster of kernels -still looks semantically close to a known pattern, add a short AI note after the table -with one of these labels: - -- `high`: very likely the same pattern family; naming drift or minor implementation reshaping is the main uncertainty -- `medium`: several signals line up, but one important piece is still ambiguous -- `low`: weak resemblance only; mention it only if it is still worth a human follow-up - -## When To Use It - -- inspect an SGLang torch profiler trace or profile directory -- profile a live SGLang server and immediately analyze the output -- summarize which kernel families dominate prefill or decode -- map kernels back to Python code paths -- judge whether a code path still has overlap headroom -- check whether an already-known fusion or overlap path should have applied - -## Diffusion Backend Gate - -For diffusion benchmark or profiling work, only analyze traces produced by the native -SGLang diffusion backend. - -If the run that generated the trace logs any of: -- `Falling back to diffusers backend` -- `Using diffusers backend` -- `Loaded diffusers pipeline` - -stop the workflow instead of analyzing the trace. Treat it as a backend-selection issue, -not as valid SGLang diffusion profiler evidence. - -## Main Flows - -### 1. Single-trace triage from an existing profile dir or trace - -```bash -python3 scripts/analyze_sglang_torch_profile.py \ - --input /path/to/profile_dir_or_trace.json.gz -``` - -Use this when you want the fastest read on kernel share and likely fused-kernel pattern matches. -The overlap table stays conservative in single-trace mode and will tell you when a mapping/formal pair is needed. - -### 2. Single-trace triage from a running server - -```bash -python3 scripts/analyze_sglang_torch_profile.py \ - --url http://127.0.0.1:30000 \ - --num-steps 5 \ - --profile-by-stage -``` - -### 3. Two-trace triage from existing profile dirs or traces - -```bash -python3 scripts/analyze_sglang_torch_profile.py triage \ - --mapping-input /path/to/graph_off_profile_dir \ - --formal-input /path/to/graph_on_profile_dir -``` - -Use this when you need stronger overlap conclusions and cleaner kernel-to-source attribution. - -### 4. Two-trace triage from running servers - -```bash -python3 scripts/analyze_sglang_torch_profile.py triage \ - --mapping-url http://127.0.0.1:31025 \ - --formal-url http://127.0.0.1:31026 \ - --num-steps 5 \ - --profile-by-stage -``` - -## `profile_by_stage` - -`profile_by_stage` is not only for PD disaggregation. - -- On ordinary non-PD serving, it is still useful because prefill and decode usually have very different bottlenecks. -- On the current profile-v2 path inside SGLang, stage-based profiling is effectively the normal path. -- PD-disaggregated serving adds one extra rule: prefill workers and decode workers must be profiled separately. That is stricter than ordinary `profile_by_stage`. - -## How To Choose The Triage Shape - -### Single-trace triage - -Use when you want the lowest-friction report: - -- one trace is already available -- you mainly want kernel share and fusion clues -- you are comparing two runs side by side by running triage once per trace - -This is the recommended default. - -### Two-trace triage - -Use when you need: - -- a stronger answer about overlap headroom -- graph-off source mapping plus graph-on final behavior -- more trustworthy overlap recommendations in the middle table - -1. mapping trace with `--disable-cuda-graph --disable-piecewise-cuda-graph` -2. formal trace with the real serving optimizations enabled - -Do not call the mapping pass a "fast profile". It exists to recover `kernel -> cpu_op -> python scope`. - -## Workflow - -### Single-trace workflow - -1. If the user only wants a quick diagnosis, one trace is enough. -2. Prefer rank-local `TP-0` traces over merged traces. -3. For a live server, this skill can call `sglang.profiler` and automatically send a small probe request. -4. Prefer `--profile-by-stage` even on standard serving unless the user explicitly wants an all-stage mixed trace. - -### Two-trace workflow - -1. Produce a mapping trace first with graph disabled. -2. Produce a formal trace second with graph enabled and the real serving flags kept on. -3. Run `triage` for the compact three-table report. -4. Read the results in this order: - - kernel table - - overlap-opportunity table - - fuse-pattern table -5. Before calling something a "new" optimization idea, compare the top rows against both [references/fuse-overlap-catalog.md](references/fuse-overlap-catalog.md) and [references/overlap-catalog.md](references/overlap-catalog.md). Always check the `PR-backed / in-flight` sections too. Prefer reporting: - - an existing fused or overlap path that should already apply here - - an existing path that appears disabled, unsupported, or regressed in this trace - - an upstream PR-backed pattern that already exists but is not merged into the checked-out tree - - a truly new opportunity only when no catalog entry fits -6. If no exact pattern fully matches but the trace still looks semantically close to a known family, add one flat `AI similarity judgment` note after the tables. - Use `high`, `medium`, or `low` only. - Base that note on the full pattern shape, not on one kernel name alone. - Prefer semantic cues such as producer-consumer chain, source locations, CPU op names, TP context, and model-specific structure. - Do not rewrite the script table itself to include these heuristic judgments. - -## References - -Load these only when needed: - -- [references/source-map.md](references/source-map.md) - - upstream SGLang profiler entrypoints and trace-writing source paths -- [references/heuristics.md](references/heuristics.md) - - overlap labels, dependency-risk interpretation, and limits -- [references/fuse-overlap-catalog.md](references/fuse-overlap-catalog.md) - - mixed source-backed catalog of existing fuse and overlap patterns, including PR-backed / in-flight rows -- [references/overlap-catalog.md](references/overlap-catalog.md) - - overlap-only lookup table across LLM, VLM, diffusion, disaggregation, HiSparse, and speculative scheduling - -## Output Contract - -Return: - -- trace path or generated profile path -- model/server args when available -- kernel table -- overlap-opportunity table -- fuse-pattern table -- optional `AI similarity judgment` note with `high` / `medium` / `low` when exact matching is inconclusive -- one short conclusion about what dominates the run -- whether the overlap conclusion came from single-trace triage or mapping/formal two-trace triage diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py deleted file mode 100644 index 83b584312..000000000 --- a/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py +++ /dev/null @@ -1,601 +0,0 @@ -"""Compact triage entrypoint for SGLang torch-profiler analysis.""" - -from __future__ import annotations - -import argparse -import sys -from collections import defaultdict -from pathlib import Path -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 ( - discover_trace_targets, - load_server_args, - load_trace_json, - parse_stage, - run_profiler, -) - -MIN_RENDER_SHARE_PCT = 1.0 - - -def build_triage_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - prog="analyze_sglang_torch_profile.py", - description=( - "Compact SGLang torch-profiler triage entrypoint. " - "This prints three tables: kernel mapping, overlap opportunities, " - "and fuse opportunities. " - "Use either a single trace/profile input or a mapping+formal two-trace pair." - ), - ) - parser.add_argument( - "--input", - type=str, - default=None, - help="Single trace file or profile directory to triage.", - ) - parser.add_argument( - "--url", - type=str, - default=None, - help="Running SGLang server URL for single-trace triage.", - ) - parser.add_argument( - "--output-dir", - type=str, - default=None, - help="Trace output dir when using --url.", - ) - parser.add_argument( - "--profile-prefix", - type=str, - default="triage-trace", - help="Profile prefix when generating a single trace from --url.", - ) - parser.add_argument( - "--mapping-input", - type=str, - default=None, - help="Graph-off mapping trace file or directory.", - ) - parser.add_argument( - "--mapping-url", - type=str, - default=None, - help="Running graph-off SGLang server URL for the mapping trace.", - ) - parser.add_argument( - "--formal-input", - type=str, - default=None, - help="Formal graph-on trace file or directory.", - ) - parser.add_argument( - "--formal-url", - type=str, - default=None, - help="Running graph-on SGLang server URL for the formal trace.", - ) - parser.add_argument( - "--mapping-output-dir", - type=str, - default=None, - help="Trace output dir when using --mapping-url.", - ) - parser.add_argument( - "--formal-output-dir", - type=str, - default=None, - help="Trace output dir when using --formal-url.", - ) - parser.add_argument( - "--mapping-profile-prefix", - type=str, - default="mapping-trace", - help="Profile prefix for the mapping trace.", - ) - parser.add_argument( - "--formal-profile-prefix", - type=str, - default="formal-trace", - help="Profile prefix for the formal trace.", - ) - parser.add_argument( - "--num-steps", - type=int, - default=5, - help="Profiler steps when generating traces from URLs.", - ) - parser.add_argument( - "--profile-by-stage", action=argparse.BooleanOptionalAction, default=True - ) - parser.add_argument( - "--merge-profiles", action=argparse.BooleanOptionalAction, default=False - ) - parser.add_argument("--probe-requests", type=int, default=1) - parser.add_argument( - "--probe-prompt", - type=str, - default=( - "Repeat the word profiler many times with spaces so the server performs several decode steps. " - "Do not add explanations." - ), - ) - parser.add_argument("--probe-max-new-tokens", type=int, default=None) - parser.add_argument("--probe-delay", type=float, default=0.5) - parser.add_argument( - "--start-step", - type=int, - default=None, - help="Pass through to sglang.profiler when generating traces from URLs.", - ) - parser.add_argument( - "--pid-substring", - type=str, - default=None, - help="Restrict overlap analysis to PIDs containing this substring.", - ) - parser.add_argument( - "--kernel-table-limit", - type=int, - default=0, - help="How many kernel rows to print per stage. Use 0 for all kernels.", - ) - parser.add_argument( - "--overlap-table-limit", - type=int, - default=0, - help="How many overlap rows to print per stage. Use 0 for all kernels.", - ) - return parser - - -def parse_triage_args(argv: Sequence[str]) -> argparse.Namespace: - parser = build_triage_parser() - args = parser.parse_args(argv) - - single_trace_mode = bool(args.input) or bool(args.url) - dual_trace_mode = any( - [ - args.mapping_input, - args.mapping_url, - args.formal_input, - args.formal_url, - ] - ) - - if single_trace_mode and dual_trace_mode: - parser.error( - "Use either single-trace mode (--input/--url) or two-trace mode " - "(--mapping-* plus --formal-*), not both." - ) - - if single_trace_mode: - if bool(args.input) == bool(args.url): - parser.error("Provide exactly one of --input or --url.") - return args - - if bool(args.mapping_input) == bool(args.mapping_url): - parser.error("Provide exactly one of --mapping-input or --mapping-url.") - if bool(args.formal_input) == bool(args.formal_url): - parser.error("Provide exactly one of --formal-input or --formal-url.") - return args - - -def resolve_profile_targets( - *, - label: str, - input_path: Optional[str], - url: Optional[str], - output_dir: Optional[str], - profile_prefix: Optional[str], - args: argparse.Namespace, -) -> Tuple[List[Path], Optional[dict]]: - if bool(input_path) == bool(url): - raise ValueError(f"{label} trace requires exactly one of input path or URL.") - - if url: - target_dir = run_profiler( - url=url, - output_dir=output_dir, - num_steps=args.num_steps, - profile_by_stage=args.profile_by_stage, - merge_profiles=args.merge_profiles, - profile_prefix=profile_prefix, - probe_requests=max(0, args.probe_requests), - probe_prompt=args.probe_prompt, - probe_max_new_tokens=args.probe_max_new_tokens, - probe_delay=args.probe_delay, - start_step=args.start_step, - ) - traces, server_args = discover_trace_targets(target_dir, all_traces=False) - return traces, server_args - - resolved = Path(input_path).resolve() - traces, server_args = discover_trace_targets(resolved, all_traces=False) - if server_args is None: - server_args = load_server_args(resolved) - return traces, server_args - - -def build_mapping_kernel_map(trace_paths: Sequence[Path]) -> dict: - stage_site_stats = defaultdict( - lambda: defaultdict(lambda: defaultdict(kernel_helpers.MappingSiteAggregate)) - ) - stage_kernel_categories: Dict[str, Dict[str, str]] = defaultdict(dict) - global_site_stats = defaultdict( - lambda: defaultdict(kernel_helpers.MappingSiteAggregate) - ) - global_kernel_categories: Dict[str, str] = {} - - for trace_path in trace_paths: - trace = load_trace_json(trace_path) - kernels, cpu_ops, python_frames, launch_events, _, _ = ( - kernel_helpers.extract_trace_data(trace) - ) - cpu_ops_by_external_id = kernel_helpers.build_cpu_op_index(cpu_ops) - launches_by_correlation = kernel_helpers.build_launch_index(launch_events) - local_site_stats = kernel_helpers.aggregate_kernel_sites( - kernels, - cpu_ops_by_external_id, - python_frames, - launches_by_correlation=launches_by_correlation, - ) - stage = parse_stage(trace_path) - kernel_categories = { - kernel.canonical_name: kernel.category for kernel in kernels - } - kernel_helpers.merge_site_stats(stage_site_stats[stage], local_site_stats) - kernel_helpers.merge_site_stats(global_site_stats, local_site_stats) - stage_kernel_categories[stage].update(kernel_categories) - global_kernel_categories.update(kernel_categories) - - stage_payloads = { - stage: kernel_helpers.build_stage_payload( - dict(site_stats), stage_kernel_categories.get(stage, {}) - ) - for stage, site_stats in stage_site_stats.items() - } - global_payload = kernel_helpers.build_stage_payload( - dict(global_site_stats), global_kernel_categories - ) - return {"stages": stage_payloads, "global": global_payload} - - -def stage_index(stage: str) -> int: - return {"extend": 0, "prefill": 0, "decode": 1, "all": 2}.get(stage, 99) - - -def stage_display(stage: str) -> str: - return kernel_helpers.stage_label(stage) - - -def pick_trace_for_stage(stage_to_trace: Dict[str, Path], stage: str) -> Optional[Path]: - if stage in stage_to_trace: - return stage_to_trace[stage] - if "all" in stage_to_trace: - return stage_to_trace["all"] - if len(stage_to_trace) == 1: - return next(iter(stage_to_trace.values())) - return None - - -def build_stage_trace_map(trace_paths: Sequence[Path]) -> Dict[str, Path]: - stage_map: Dict[str, Path] = {} - for trace_path in sorted( - trace_paths, key=lambda item: (stage_index(parse_stage(item)), item.name) - ): - stage_map[parse_stage(trace_path)] = trace_path - return stage_map - - -def render_kernel_table(rows: Sequence[dict]) -> List[str]: - lines = [ - "| Stage | Kernel | Category | GPU time | Share | Launches | Python location (site share) | CPU op |", - "| --- | --- | --- | ---: | ---: | ---: | --- | --- |", - ] - for row in rows: - lines.append( - "| {stage} | {kernel} | {category} | {gpu_time} | {share:.1f}% | {launches} | {location} | {cpu_op} |".format( - stage=kernel_helpers.escape_md_cell(stage_display(row["stage"])), - kernel=kernel_helpers.escape_md_cell(row["kernel"]), - category=kernel_helpers.escape_md_cell(row["category"]), - gpu_time=kernel_helpers.format_ms(row["total_us"]), - share=row["share_pct"], - launches=row["launches"], - location=kernel_helpers.escape_md_cell(row["location"]), - cpu_op=kernel_helpers.escape_md_cell(row["cpu_op"]), - ) - ) - return lines - - -def render_overlap_table(rows: Sequence[dict]) -> List[str]: - lines = [ - "| Stage | Priority | Verdict | Kernel | Python scope | Formal signal | Dep risk | Recommendation |", - "| --- | --- | --- | --- | --- | --- | --- | --- |", - ] - if not rows: - lines.append( - "| - | - | - | No actionable overlap rows. Use mapping/formal two-trace triage for stronger overlap conclusions. | - | - | - | - |" - ) - return lines - for row in rows: - formal_signal = ( - f"{row['total_us']:.1f} us, share {row['share_pct']:.1f}%, " - f"excl {row['exclusive_ratio'] * 100:.1f}% / hid {row['hidden_ratio'] * 100:.1f}%" - ) - lines.append( - "| " - + " | ".join( - [ - kernel_helpers.escape_md_cell(stage_display(row["stage"])), - row["priority"], - row["verdict"], - kernel_helpers.escape_md_cell(row["kernel"]), - kernel_helpers.escape_md_cell(row["python_scope"]), - kernel_helpers.escape_md_cell(formal_signal), - overlap_helpers.dependency_risk_label(row["dependency_signal"]), - row["recommendation"], - ] - ) - + " |" - ) - return lines - - -def render_fuse_table(rows: Sequence[dict]) -> List[str]: - lines = [ - "| Stage | Pattern | Confidence | Related GPU time | Share | Evidence kernels | Current kernel Python location | Candidate fused Python path | Rationale |", - "| --- | --- | --- | ---: | ---: | --- | --- | --- | --- |", - ] - if not rows: - lines.append( - "| - | No medium-confidence source-backed fusion opportunity matched this trace. | - | - | - | - | - | - | - |" - ) - return lines - for row in rows: - lines.append( - "| {stage} | {pattern} | {confidence} | {gpu_time} | {share:.1f}% | {evidence} | {current_locations} | {candidate_path} | {rationale} |".format( - stage=kernel_helpers.escape_md_cell(stage_display(row["stage"])), - pattern=kernel_helpers.escape_md_cell(row["pattern"]), - confidence=kernel_helpers.escape_md_cell(row["confidence"]), - gpu_time=kernel_helpers.format_ms(row["related_us"]), - share=row["share_pct"], - evidence=kernel_helpers.escape_md_cell(row["evidence"]), - current_locations=kernel_helpers.escape_md_cell( - row["current_locations"] - ), - candidate_path=kernel_helpers.escape_md_cell(row["candidate_path"]), - rationale=kernel_helpers.escape_md_cell(row["rationale"]), - ) - ) - return lines - - -def run_triage(args: argparse.Namespace) -> int: - single_trace_mode = bool(args.input) or bool(args.url) - if single_trace_mode: - formal_traces, formal_server_args = resolve_profile_targets( - label="input", - input_path=args.input, - url=args.url, - output_dir=args.output_dir, - profile_prefix=args.profile_prefix, - args=args, - ) - mapping_traces = formal_traces - mapping_server_args = formal_server_args - else: - mapping_traces, mapping_server_args = resolve_profile_targets( - label="mapping", - input_path=args.mapping_input, - url=args.mapping_url, - output_dir=args.mapping_output_dir, - profile_prefix=args.mapping_profile_prefix, - args=args, - ) - formal_traces, formal_server_args = resolve_profile_targets( - label="formal", - input_path=args.formal_input, - url=args.formal_url, - output_dir=args.formal_output_dir, - profile_prefix=args.formal_profile_prefix, - args=args, - ) - - mapping_kernel_map = build_mapping_kernel_map(mapping_traces) - - kernel_rows_rendered: List[dict] = [] - fuse_rows_rendered: List[dict] = [] - - for formal_trace in formal_traces: - trace = load_trace_json(formal_trace) - kernels, _, _, _, _, _ = kernel_helpers.extract_trace_data(trace) - if not kernels: - continue - stage = parse_stage(formal_trace) - total_us = sum(kernel.dur for kernel in kernels) - kernel_stats = kernel_helpers.aggregate( - kernels, key_fn=lambda item: item.canonical_name - ) - kernel_categories = { - kernel.canonical_name: kernel.category for kernel in kernels - } - full_kernel_rows = kernel_helpers.build_kernel_rows( - stage=stage, - kernel_stats=kernel_stats, - kernel_categories=kernel_categories, - local_stage_payload=mapping_kernel_map.get("stages", {}).get( - stage, {"kernels": {}} - ), - external_kernel_map=mapping_kernel_map, - ) - visible_kernel_rows = kernel_helpers.limit_kernel_rows( - full_kernel_rows, args.kernel_table_limit - ) - for row in visible_kernel_rows: - share_pct = kernel_helpers.pct(row.total_us, total_us) - if share_pct < MIN_RENDER_SHARE_PCT: - continue - kernel_rows_rendered.append( - { - "stage": stage, - "kernel": row.name, - "category": row.category, - "total_us": row.total_us, - "share_pct": share_pct, - "launches": row.aggregate.count, - "location": row.location, - "cpu_op": row.cpu_op, - } - ) - for item in kernel_helpers.detect_fusion_opportunities( - stage=stage, - kernel_rows=full_kernel_rows, - total_us=total_us, - server_args=formal_server_args or mapping_server_args, - ): - share_pct = kernel_helpers.pct(item.related_us, total_us) - if share_pct < MIN_RENDER_SHARE_PCT: - continue - fuse_rows_rendered.append( - { - "stage": stage, - "pattern": item.pattern, - "confidence": item.confidence, - "related_us": item.related_us, - "share_pct": share_pct, - "evidence": item.evidence, - "current_locations": item.current_locations, - "candidate_path": item.candidate_path, - "rationale": item.rationale, - } - ) - - mapping_stage_map = build_stage_trace_map(mapping_traces) - formal_stage_map = build_stage_trace_map(formal_traces) - overlap_rows_rendered: List[dict] = [] - if not single_trace_mode: - for stage in sorted(formal_stage_map, key=stage_index): - formal_trace = formal_stage_map[stage] - mapping_trace = pick_trace_for_stage(mapping_stage_map, stage) - if mapping_trace is None: - continue - mapping_trace_json = load_trace_json(mapping_trace) - mapping_events, mapping_pid = overlap_helpers.extract_kernel_events( - mapping_trace_json, args.pid_substring - ) - if not mapping_events: - continue - formal_trace_json = load_trace_json(formal_trace) - formal_events, formal_pid = overlap_helpers.extract_kernel_events( - formal_trace_json, args.pid_substring - ) - if not formal_events: - continue - mapping_bundle = overlap_helpers.TraceBundle( - label=f"mapping-{stage}", - trace_path=mapping_trace, - server_args=mapping_server_args, - raw_events=mapping_trace_json.get( - "traceEvents", - mapping_trace_json if isinstance(mapping_trace_json, list) else [], - ), - events=mapping_events, - pid=mapping_pid, - ) - formal_bundle = overlap_helpers.TraceBundle( - label=f"formal-{stage}", - trace_path=formal_trace, - server_args=formal_server_args, - raw_events=formal_trace_json.get( - "traceEvents", - formal_trace_json if isinstance(formal_trace_json, list) else [], - ), - events=formal_events, - pid=formal_pid, - ) - formal_bundle.overlap_stats = overlap_helpers.analyze_overlap( - formal_bundle.events - ) - aggregates = overlap_helpers.aggregate_events(formal_bundle.events) - source_map = overlap_helpers.build_kernel_source_map(mapping_bundle) - stage_rows = overlap_helpers.build_action_rows( - aggregates, - source_map, - formal_bundle.events, - formal_bundle.overlap_stats["total_busy_us"], - table_limit=max(0, args.overlap_table_limit), - ) - for row in stage_rows: - if row.share_pct < MIN_RENDER_SHARE_PCT: - continue - overlap_rows_rendered.append( - { - "stage": stage, - "priority": row.priority, - "verdict": row.verdict, - "kernel": row.kernel, - "python_scope": row.python_scope, - "total_us": row.total_us, - "share_pct": row.share_pct, - "exclusive_ratio": row.exclusive_ratio, - "hidden_ratio": row.hidden_ratio, - "dependency_signal": row.dependency_signal, - "recommendation": row.recommendation, - } - ) - - lines: List[str] = [] - lines.append("Triage View") - if single_trace_mode: - lines.append(f"Input traces: {', '.join(str(path) for path in formal_traces)}") - else: - lines.append( - f"Mapping traces: {', '.join(str(path) for path in mapping_traces)}" - ) - lines.append(f"Formal traces: {', '.join(str(path) for path in formal_traces)}") - if formal_server_args or mapping_server_args: - server_args = formal_server_args or mapping_server_args - model = server_args.get("model_path") or server_args.get("model") - if model: - lines.append(f"Model: {model}") - lines.append("") - lines.append("Kernel Table") - lines.extend(render_kernel_table(kernel_rows_rendered)) - lines.append("") - lines.append("Overlap Opportunity Table") - lines.extend(render_overlap_table(overlap_rows_rendered)) - lines.append("") - lines.append("Fuse Opportunity Table") - lines.extend(render_fuse_table(fuse_rows_rendered)) - print("\n".join(lines).rstrip()) - return 0 - - -def main(argv: Optional[Sequence[str]] = None) -> int: - argv = list(argv or sys.argv[1:]) - triage_parser = build_triage_parser() - - if not argv or argv[0] in {"-h", "--help"}: - triage_parser.print_help() - return 0 - - if argv[0] == "triage": - argv = argv[1:] - elif not argv[0].startswith("-"): - triage_parser.error( - "This skill now exposes only the compact triage workflow. " - "Use single-trace mode (--input/--url) or mapping+formal two-trace mode." - ) - return 2 - - return run_triage(parse_triage_args(argv)) - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py b/.claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py deleted file mode 100644 index 602064d6f..000000000 --- a/.claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py +++ /dev/null @@ -1,333 +0,0 @@ -"""Shared helpers for SGLang torch-profiler skill scripts.""" - -from __future__ import annotations - -import gzip -import json -import re -import subprocess -import sys -import tempfile -import time -from collections import Counter, defaultdict -from pathlib import Path -from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple -from urllib import request - -STAGE_ORDER = {"extend": 0, "prefill": 0, "decode": 1, "all": 2} -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:") - - -def normalize_text(value: object) -> str: - return re.sub(r"\s+", " ", str(value)).strip() - - -def normalize_repo_relative_path(path: object) -> str: - text = normalize_text(path).replace("\\", "/") - for marker in ("python/sglang/", "sgl_kernel/"): - idx = text.find(marker) - if idx != -1: - return text[idx:].lstrip("/") - idx = text.find("sglang/") - if idx != -1: - return ("python/" + text[idx:]).lstrip("/") - return text.lstrip("/") - - -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 parse_stage(path: Path) -> str: - name = path.name.lower() - if "-extend" in name or "-prefill" in name: - return "extend" - if "-decode" in name: - return "decode" - return "all" - - -def parse_tp_rank(path: Path) -> Optional[int]: - match = re.search(r"TP-(\d+)", path.name) - return int(match.group(1)) if match else None - - -def newest_trace_dir(path: Path) -> Path: - if path.is_file(): - return path.parent - direct = list(path.glob("*.trace.json")) + list(path.glob("*.trace.json.gz")) - if direct: - return path - child_candidates = [item for item in path.rglob("*") if item.is_dir()] - trace_dirs = [ - candidate - for candidate in child_candidates - if list(candidate.glob("*.trace.json")) - or list(candidate.glob("*.trace.json.gz")) - ] - if not trace_dirs: - raise FileNotFoundError(f"No trace files found under {path}") - trace_dirs.sort(key=lambda item: item.stat().st_mtime) - 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) - - trace_dir = newest_trace_dir(path) - traces = sorted( - list(trace_dir.glob("*.trace.json")) + list(trace_dir.glob("*.trace.json.gz")), - key=lambda item: item.stat().st_mtime, - ) - 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: dict, timeout: float = 60.0) -> Optional[dict]: - req = request.Request( - url=url, - data=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 -) -> None: - 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) - - -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, - start_step: Optional[int] = None, -) -> Path: - if output_dir is None: - output_dir = tempfile.mkdtemp(prefix="sglang-torch-profile-") - output_path = Path(output_dir).resolve() - output_path.mkdir(parents=True, exist_ok=True) - - cmd = [ - sys.executable, - "-m", - "sglang.profiler", - "--url", - url, - "--output-dir", - str(output_path), - "--num-steps", - str(num_steps), - "--cpu", - "--gpu", - "--merge-profiles" if merge_profiles else "--no-merge-profiles", - "--profile-by-stage" if profile_by_stage else "--no-profile-by-stage", - ] - if profile_prefix: - cmd.extend(["--profile-prefix", profile_prefix]) - if start_step is not None: - cmd.extend(["--start-step", str(start_step)]) - - profiler_proc = subprocess.Popen(cmd) - try: - if probe_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, - ) - if profiler_proc.poll() is not None: - break - return_code = profiler_proc.wait() - finally: - if profiler_proc.poll() is None: - profiler_proc.kill() - - if return_code != 0: - raise subprocess.CalledProcessError(return_code, cmd) - - deadline = time.time() + 15.0 - while time.time() < deadline: - child_dirs = [path for path in output_path.iterdir() if path.is_dir()] - if child_dirs: - child_dirs.sort(key=lambda path: path.stat().st_mtime) - newest_child = child_dirs[-1] - if any(newest_child.glob("*.trace.json*")): - return newest_child - time.sleep(0.5) - - child_dirs = [path for path in output_path.iterdir() if path.is_dir()] - if child_dirs: - child_dirs.sort(key=lambda path: path.stat().st_mtime) - return child_dirs[-1] - return output_path - - -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]) diff --git a/.claude/skills/write-sglang-test/SKILL.md b/.claude/skills/write-sglang-test/SKILL.md index 51ab8919c..8bd49a8b6 100644 --- a/.claude/skills/write-sglang-test/SKILL.md +++ b/.claude/skills/write-sglang-test/SKILL.md @@ -66,6 +66,7 @@ Defined in `python/sglang/test/test_utils.py`: | `stage-b-test-2-gpu-large` | `2-gpu-h100` | Two-GPU correctness and parallelism (TP/PP) on H100 | | `stage-b-test-4-gpu-b200` | `4-gpu-b200` | Early Blackwell coverage (SM100+ paths) on four GPUs | | `stage-b-kernel-unit-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `python/sglang/jit_kernel/tests/` | +| `stage-b-kernel-unit-1-gpu-b200` | `4-gpu-b200` | JIT kernel correctness tests for Blackwell / SM100-specific paths | | `stage-b-kernel-unit-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `python/sglang/jit_kernel/tests/` | | `stage-b-kernel-benchmark-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `python/sglang/jit_kernel/benchmark/` | | `stage-c-test-4-gpu-h100` | `4-gpu-h100` | Large 4-GPU H100 integration and scaling tests | @@ -75,7 +76,8 @@ Defined in `python/sglang/test/test_utils.py`: | `stage-c-test-deepep-8-gpu-h200` | `8-gpu-h200` | DeepEP at 8-GPU H200 scale | | `stage-c-test-8-gpu-b200` | `8-gpu-b200` | 8-GPU B200 suite (registered but not yet wired to a workflow) | | `stage-c-test-4-gpu-b200` | `4-gpu-b200` | 4-GPU B200 suite for large models on Blackwell | -| `stage-c-test-4-gpu-gb200` | `4-gpu-gb200` | 4-GPU GB200 suite for large models on Grace Blackwell | +| `stage-c-test-4-gpu-b200-small` | `4-gpu-b200` | Smaller 4-GPU B200 suite split onto low-disk B200 runners | +| `stage-c-test-4-gpu-gb200` | `4-gpu-gb200` | 4-GPU GB200 suite for Grace Blackwell; registered in `run_suite.py`, but the PR workflow is currently disabled until a runner is provisioned | #### Per-commit (AMD) @@ -107,7 +109,7 @@ Defined in `python/sglang/test/test_utils.py`: #### Nightly -Nightly suites are listed in `NIGHTLY_SUITES` in [`test/run_suite.py`](../../../test/run_suite.py). They run via `nightly-test-nvidia.yml`, `nightly-test-amd.yml` amd `nightly-test-npu.yml`, not `pr-test.yml`. Examples: +Nightly suites are listed in `NIGHTLY_SUITES` in [`test/run_suite.py`](../../../test/run_suite.py). They run via `nightly-test-nvidia.yml`, `nightly-test-amd.yml`, and `nightly-test-npu.yml`, not `pr-test.yml`. Examples: - `nightly-1-gpu` (CUDA) - `nightly-kernel-1-gpu` (CUDA, JIT kernel full grids) @@ -132,6 +134,7 @@ Use the lightest suite that meets your test's needs: - **Most small GPU tests** → `stage-b-test-1-gpu-small` (default choice) - **Need H100 memory or Hopper features** → `stage-b-test-1-gpu-large` - **JIT kernel correctness** → `stage-b-kernel-unit-1-gpu-large` +- **JIT kernel correctness for B200 / SM100 paths** → `stage-b-kernel-unit-1-gpu-b200` - **JIT kernel benchmarks** → `stage-b-kernel-benchmark-1-gpu-large` - **Multi-GPU** → only when the test actually needs multiple GPUs @@ -352,6 +355,7 @@ from sglang.test.ci.ci_register import register_cuda_ci # Correctness tests in python/sglang/jit_kernel/tests/ register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-b200") register_cuda_ci(est_time=120, suite="stage-b-kernel-unit-8-gpu-h200") # Benchmarks in python/sglang/jit_kernel/benchmark/ diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/SKILL.md index 3e7e0075f..cb325ce25 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/SKILL.md @@ -246,10 +246,14 @@ The `PipelineConfig` holds static model configuration and defines callback metho from dataclasses import dataclass, field +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( - ImagePipelineConfig, # for image generation - # SpatialImagePipelineConfig, # alternative base - # VideoPipelineConfig, # for video generation + ImagePipelineConfig, + ModelTaskType, + # PipelineConfig, # common base for many video pipelines + # SpatialImagePipelineConfig, # alternative base for spatial image models ) from sglang.multimodal_gen.configs.models.dits.mymodel import MyModelDitConfig from sglang.multimodal_gen.configs.models.vaes.mymodel import MyModelVAEConfig @@ -314,6 +318,11 @@ class MyModelPipelineConfig(ImagePipelineConfig): return frames ``` +There is no separate `VideoPipelineConfig` base class. For video models, choose +`ModelTaskType.T2V`, `ModelTaskType.I2V`, or `ModelTaskType.TI2V`, and follow +existing video configs such as Wan, LTX, Hunyuan, Helios, or MOVA when deciding +whether to subclass `PipelineConfig` directly or use a model-specific base. + **Important**: The `prepare_pos_cond_kwargs` / `prepare_neg_cond_kwargs` methods define what the DiT receives at each denoising step. These must match the DiT's `forward()` signature. ### Step 6: Implement the BeforeDenoisingStage (Core Step) @@ -502,15 +511,21 @@ In `python/sglang/multimodal_gen/registry.py`, register your configs: ```python register_configs( - model_family="my_model", sampling_param_cls=MyModelSamplingParams, pipeline_config_cls=MyModelPipelineConfig, hf_model_paths=[ "org/my-model-name", # HuggingFace model ID(s) ], + model_detectors=[ + lambda path: "my-model" in path.lower(), + ], ) ``` +`register_configs()` does not take a `model_family` argument. It registers the +sampling and pipeline config classes, then resolves models by exact +`hf_model_paths` or optional detector predicates. + The `EntryClass` in your pipeline file is automatically discovered by the registry's `_discover_and_register_pipelines()` function -- no additional registration needed for the pipeline class itself. ### Step 9: Verify Output Quality @@ -590,4 +605,5 @@ After the model produces non-noise output, read [references/testing-and-accuracy.md](references/testing-and-accuracy.md) before adding GPU cases, component-accuracy skips/hooks, suite entries, or benchmark claims. That reference tracks the current `gpu_cases.py` / `testcase_configs.py` -/ `run_suite.py` split and the component-accuracy decision rules. +/ `accuracy_testcase_configs.py` / `run_suite.py` split and the component-accuracy +decision rules. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/references/testing-and-accuracy.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/references/testing-and-accuracy.md index 846d3b82e..3e5f8c9b2 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/references/testing-and-accuracy.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/references/testing-and-accuracy.md @@ -7,6 +7,9 @@ produce a non-noise image or video. - Add concrete GPU integration cases in `python/sglang/multimodal_gen/test/server/gpu_cases.py`. - Keep reusable dataclasses, constants, thresholds, and testcase factory helpers in `python/sglang/multimodal_gen/test/server/testcase_configs.py`. +- Add the case id to `python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py` + only when it should be part of component-accuracy coverage. Adding a GPU case + alone does not enroll it there. - Let `python/sglang/multimodal_gen/test/run_suite.py` own suite selection, runtime-based partitioning, and standalone test files. Do not hard-code CI shard lists elsewhere. - If a new standalone test file is added to a suite, update `STANDALONE_FILE_EST_TIMES` after the first measured CI/runtime value is known. @@ -22,8 +25,8 @@ PYTHONPATH=python python3 python/sglang/multimodal_gen/test/run_suite.py --suite If you add a new entry to `ONE_GPU_CASES`, `TWO_GPU_CASES`, or a B200-specific case group in `gpu_cases.py`, treat component accuracy as part of the -model-adding workflow. Do not assume the new testcase will automatically fit the -existing component-accuracy harness. +model-adding workflow. Do not assume the new testcase will automatically fit or +enter the existing component-accuracy harness. The component-accuracy harness compares SGLang components against Diffusers/HF reference components. This is stricter than pipeline-level inference. New GPU @@ -46,9 +49,12 @@ cases commonly fail here for one of three reasons: When adding a new GPU case, make this decision explicitly: +- if the case should have component-accuracy coverage, add its case id to + `accuracy_testcase_configs.py` - if the family needs minimal harness wiring, add the smallest possible change in `accuracy_hooks.py` - if the case is only a variant of an already covered source component and topology, add a skip in `accuracy_config.py` - if the HF/Diffusers reference component cannot be compared faithfully, add a skip in `accuracy_config.py` +- if the case is intentionally GPU-smoke-only, leave it out of `accuracy_testcase_configs.py` and keep that choice explicit in the PR notes Do not add a new GPU case and wait for CI to discover missing component-accuracy wiring. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md index e35497b33..aaecd641f 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md @@ -32,6 +32,7 @@ First use [../sglang-diffusion-benchmark-profile/SKILL.md](../sglang-diffusion-b - collect the perf dump baseline - capture one representative `torch.profiler` trace - rule out existing mainline fast paths +- prove the run stayed on the native SGLang diffusion backend, not a diffusers fallback If a future specialized optimization skill matches the kernel family better than AKO4ALL, hand off there instead. The diagnosis contract stays the same. @@ -127,5 +128,9 @@ See [references/ako-loop.md](references/ako-loop.md) for the checklist and commo - Treat AKO4ALL repo hygiene as a gate, not a suggestion. - Prefer exact local snapshot validation over hand-wavy “remote tree is close enough”. +- Do not start or justify kernel work from traces collected after + `Falling back to diffusers backend`, `Using diffusers backend`, or + `Loaded diffusers pipeline`; fix backend selection and rerun the + benchmark/profile workflow first. - Keep model-level validation honest: if microbench improves but denoise does not, do not keep the AKO-only variant in the main code path. - When writing conclusions, explain the win in terms of measurable causes such as lower registers per thread, higher occupancy, fewer executed instructions, or better scheduler eligibility.