[SKILL] Sync SGLang skill docs (#23921)
This commit is contained in:
@@ -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 `<think>` 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
|
||||
@@ -0,0 +1,353 @@
|
||||
# Fuse And Overlap Catalog
|
||||
|
||||
This catalog is the source-backed lookup table that the profiler skill should
|
||||
consult before labeling a fuse or overlap opportunity as novel.
|
||||
|
||||
For overlap-only triage, also load `references/overlap-catalog.md`.
|
||||
|
||||
This revision is intentionally kernel-scoped. Keep rows here only when they map
|
||||
to one fused GPU/NPU kernel family, one fused collective-plus-kernel family, or
|
||||
one profiler-visible stream overlap among GPU kernels / collective kernels.
|
||||
Host-only scheduler, event-loop, executor, offload, and load-path patterns are
|
||||
intentionally excluded.
|
||||
|
||||
Use it like this:
|
||||
|
||||
1. Start from the three `triage` tables.
|
||||
2. Match top rows against the `Trace keywords` and `Primary code` columns below.
|
||||
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 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
|
||||
necessarily present in the checked-out `sglang` tree, but they should still be
|
||||
treated as upstream or analogous kernel families before labeling a fuse or
|
||||
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*`<br>`npu_add_rms_norm`<br>`add_rmsnorm_bias`<br>`gemma_fused_add_rmsnorm`<br>`gemma_rmsnorm_residual_scalar`<br>`_gemma_rmsnorm_residual_kernel`<br>residual add right before norm | `python/sglang/srt/layers/layernorm.py`<br>`python/sglang/srt/layers/gemma4_fused_ops.py`<br>`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*`<br>`all_reduce`<br>`FusedAddRMSNormKernel`<br>`rmsnorm*` | `python/sglang/srt/layers/flashinfer_comm_fusion.py`<br>`python/sglang/srt/layers/layernorm.py::forward_with_allreduce_fusion`<br>`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`<br>`python/sglang/srt/distributed/communication_op.py::tensor_model_parallel_fused_allreduce_rmsnorm`<br>`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`<br>`gelu_and_mul`<br>`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`<br>`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`<br>`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`<br>`benchmark_combo_kernel`<br>`q_norm`<br>`k_norm`<br>`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`<br>`rms_sumsq_serial`<br>`rms_apply_serial`<br>`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`<br>`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*`<br>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. |
|
||||
| Fused RoPE + KV cache store | `fused_set_kv_buffer`<br>RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`<br>`python/sglang/srt/models/utils.py::enable_fused_set_kv_buffer` | Shared entrypoints can route to fused RoPE + KV-store or model-side `fused_set_kv_buffer` fast paths | Compare against the fused cache-store path before proposing a new KV rewrite. |
|
||||
| Fused decode metadata setup | `normal_decode_set_metadata`<br>`cache_seqlens_int32`<br>`cu_seqlens_k`<br>`page_table`<br>`swa_page_table` | `python/sglang/srt/layers/attention/flashattention_backend.py::normal_decode_set_metadata` | Triton decode path fuses seq-len cast/add, prefix-sum, req-to-token gather, page-table divide, and optional SWA metadata build into 1-2 kernels | If decode exposes multiple tiny metadata kernels before attention, first compare against this existing fused metadata-prep path. |
|
||||
| NSA fused metadata copy for graph replay | `fused_metadata_copy`<br>`fused_metadata_copy_multi`<br>`fused_nsa_cache_seqlens`<br>`fused_flashmla_metadata` | `python/sglang/jit_kernel/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around NSA replay as a missed existing replay fusion. |
|
||||
| DeepSeek MLA fused projection + norm + RoPE | `qkv_proj_with_rope_fused_weight`<br>`fused_qkv_a_proj_with_mqa`<br>`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`<br>`python/sglang/srt/models/deepseek_v2.py` | CPU / ROCm paths fuse DeepSeek MLA projection packing with q / k norm, RoPE, and cache-oriented MLA prep | For DeepSeek MLA, split proj / norm / rope prep is usually an existing backend-specific fuse that did not fire. |
|
||||
| Fused QK RoPE concat + MLA cache write | `fused_qk_rope_cat_and_cache_mla`<br>`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` | ROCm MLA path can fuse Q / K RoPE packing, concat, and MLA cache write in one backend-specific op | On DeepSeek / MLA traces, separate RoPE-cat-cache steps are not automatically novel. |
|
||||
| Qwen3 decode fused QK norm + 3D mRoPE + KV cache write | `fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`mrope`<br>decode cache write | `python/sglang/srt/models/qwen3.py` | ROCm / AITER decode path fuses QK norm, 3D mRoPE, and paged KV cache write | On Qwen3-style decode, separate norm + mRoPE + cache-store kernels are not a novel opportunity. |
|
||||
| NPU fused split-QKV + RMSNorm + RoPE | `split_qkv_rmsnorm_rope` | `python/sglang/srt/models/llama.py`<br>`python/sglang/srt/models/qwen3.py`<br>`python/sglang/srt/models/qwen3_moe.py`<br>`python/sglang/srt/models/glm4_moe.py` | Ascend path fuses QKV split, Q / K RMSNorm, and RoPE in one op | On NPU traces, separate split / norm / rope kernels usually mean the fused path is unavailable or bypassed. |
|
||||
| Fused FP8 quantize + paged KV cache write | `trtllm_fp8_kv_kernel`<br>`fp8 kv cache write`<br>`paged KV cache write` | `python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py` | TRTLLM MHA path fuses FP8 quantization, scale computation, and paged K / V cache write | If FP8 KV cache traces show standalone quant plus write kernels, first compare against this existing Triton fuse. |
|
||||
| Fused MLA KV cache write + FP8 quant | `set_mla_kv_buffer_fp8_quant*`<br>`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`<br>`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`<br>`fused_moe_router*`<br>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`<br>`moe_fused_gate`<br>`aiter_fused_topk`<br>`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`<br>`fused_append_shared_experts_with_weights`<br>`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`<br>`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<br>dispatch / combine<br>grouped top-k<br>many small MoE support kernels | `python/sglang/srt/layers/moe/fused_moe_triton/layer.py`<br>`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`<br>`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*`<br>`npu_dequant_swiglu_quant`<br>`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`<br>`python/sglang/jit_kernel/nvfp4.py`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`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. |
|
||||
| DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`<br>`fused_rms_mxfp4_quant`<br>`fused_flatten_fp8_group_quant`<br>`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py` | DeepSeek MLA / MHA ROCm paths fuse RMSNorm or flatten with FP8 / MXFP4 quantization for comm / attention prep | On DeepSeek quant traces, split norm + quant or flatten + quant is an existing family, not a new idea. |
|
||||
| NSA fused top-k transform / page-table build | `fast_topk_transform_fused`<br>`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/nsa_backend.py` | NSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If NSA top-k metadata work is split, check `SGLANG_NSA_FUSE_TOPK` and backend support first. |
|
||||
| NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`<br>`python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into NSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. |
|
||||
| Fused sampling temperature + softmax | `fused_temperature_softmax*` | `python/sglang/srt/layers/fused_sampling.py`<br>`python/sglang/srt/layers/sampler.py` | Triton single-pass / multi-pass kernels fuse temperature scaling and softmax during decode | Separate temp-divide + softmax at decode batch sizes is often a missed existing fusion. |
|
||||
| Fused logit softcap | `fused_softcap`<br>`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`<br>`python/sglang/srt/layers/logits_processor.py` | Triton kernels fuse cast-to-float and softcap / tanh math for logits or generic elementwise softcapping | Treat exposed cast + softcap ladders as an existing Triton fuse family. |
|
||||
| Linear-attention packed projection reshuffle | `fused_qkvzba_split_reshape_cat*`<br>`qkvz_proj`<br>`ba_proj`<br>`qkvabz_proj`<br>`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`python/sglang/srt/models/jet_nemotron.py` | GDN / Kimi / Jet-style linear-attn models pack multiple projections, then fuse split / reshape / cat into one kernel | Treat split reshape / transpose / cat ladders as an existing linear-attention fusion family. |
|
||||
| Fused GDN gating prep | `fused_gdn_gating`<br>`softplus`<br>`beta_output` | `python/sglang/srt/layers/attention/fla/fused_gdn_gating.py` | Triton kernel computes GDN gate preparation such as `-exp(A_log) * softplus(...)` and `sigmoid(b)` together | On GDN traces, treat split gate-prep elementwise kernels as missing existing fusion first. |
|
||||
| Fused RMSNorm-gated linear-attention output | `FusedRMSNormGated`<br>`layer_norm_gated_fwd` | `python/sglang/srt/layers/attention/fla/fused_norm_gate.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/kimi_linear.py` | One Triton op covers residual-aware (RMS)Norm plus sigmoid / swish gating | If norm and output gate appear as separate kernels in GDN / Kimi-like blocks, first suspect a missing existing fusion. |
|
||||
| Fused gated RMSNorm / LayerNorm | `rms_norm_gated`<br>`layer_norm_gated` | `python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py` | Mamba-derived kernels can fuse normalization with the gating branch `z * sigmoid(z)` | Treat split norm and gate post-processing on Mamba-style blocks as an existing fusion family. |
|
||||
| Fused linear-attention chunk KKT + solve_tril | `chunk_gated_delta_rule_fwd_kkt_solve_kernel`<br>`scaled_dot_kkt`<br>`solve_tril`<br>`recompute_w_u` | `python/sglang/srt/layers/attention/fla/chunk_fwd.py`<br>`python/sglang/srt/layers/attention/fla/kda.py` | GDN / KDA chunk forward fuses `scaled_dot_kkt + solve_tril` in the prefill / intra-chunk path, then finishes `recompute_w_u` as the next step | Treat split KKT + triangular-solve ladders as an existing linear-attention fusion family first. |
|
||||
| Fused linear-attention recurrent / KDA update | `fused_sigmoid_gating_delta_rule_update`<br>`fused_recurrent_gated_delta_rule_update`<br>`fused_kda_gate` | `python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py`<br>`python/sglang/srt/layers/attention/fla/fused_recurrent.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`python/sglang/srt/models/jet_nemotron.py` | Triton / CuTeDSL kernels fuse gating math, optional QK l2norm, recurrent state update, and output generation | Treat split gating + recurrent-update chains as existing linear-attention fusion, not a novel opportunity. |
|
||||
| Fused Mamba state gather/scatter with mask | `fused_mamba_state_scatter_with_mask`<br>`index_elementwise_kernel` | `python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py` | Triton kernel replaces multiple masked gather / scatter index kernels with one fused update | If Mamba verify/update shows many tiny index kernels, first compare against this existing fused path. |
|
||||
| Staging-buffer fused gather / scatter | `_fused_gather_to_staging_kernel`<br>`_fused_scatter_from_staging_kernel` | `python/sglang/srt/disaggregation/common/staging_buffer.py` | Triton kernels gather scattered KV slices into contiguous staging memory and scatter them back into KV cache on decode | Treat ladders of tiny gather/scatter/copy kernels in heterogeneous TP staging as missing an existing Triton fusion. |
|
||||
|
||||
## 2. LLM / SRT kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Single-batch overlap (SBO) | MoE combine, down-gemm, shared-expert work in nearby two-stream windows | `python/sglang/srt/batch_overlap/single_batch_overlap.py` | combine vs down-gemm overlap, combine vs shared-expert overlap, one-stream dispatch+shared overlap, explicit SM partitioning and events | If exposed MoE combine sits near neighboring compute, classify it against SBO before calling it new overlap. |
|
||||
| Q and K normalization on different streams | Q-side norm and K-side norm on different streams | `python/sglang/srt/models/utils.py::apply_qk_norm`<br>`python/sglang/srt/models/qwen3.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py` | Q stays on current stream, K can run on `alt_stream` in capture mode | Treat split Q / K norm as an existing overlap family when `alt_stream` is already wired. |
|
||||
| DeepSeek shared-expert / routed-expert overlap | shared-expert GEMMs near DeepEP dispatch / combine | `python/sglang/srt/models/deepseek_v2.py`<br>`python/sglang/srt/batch_overlap/single_batch_overlap.py` | shared experts on `alt_stream`, overlap with dispatch / combine and down-gemm, Blackwell-specific env gating | This is an established routed-vs-shared branch overlap pattern, not a novel idea. |
|
||||
| Llama4 shared branch vs routed branch overlap | shared expert branch plus routed MoE branch as adjacent windows | `python/sglang/srt/models/llama4.py` | shared expert on current stream, router + topk + routed experts on `alt_stream` | Use Llama4 as the first precedent for branch-level overlap in similar sparse models. |
|
||||
| ExaoneMoE shared experts vs router experts overlap | shared expert output and router-expert output form a two-branch window | `python/sglang/srt/models/exaone_moe.py::forward_normal_dual_stream` | shared experts on current stream, router + routed experts on `alt_stream`, explicit join before combine | This is an existing dual-stream MoE overlap family. |
|
||||
| Grok residual-MoE branch overlap | dense MLP and block-sparse MoE branches in parallel | `python/sglang/srt/models/grok.py::moe_with_rmoe` | dense MLP on current stream, MoE on `alt_stream`, fused dual residual RMSNorm around boundaries | Treat exposed Grok branch overlap as an existing pattern. |
|
||||
| NSA dual-stream overlap | Q-proj, K-proj, RoPE, cache-store, quantization in tight two-stream windows | `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Q / K projection split, RoPE split, cache-store vs quantization overlap | NSA already contains several dual-stream overlap precedents. |
|
||||
| MoriEP async dispatch / combine comm stream | `MoriEP`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. |
|
||||
| Heterogeneous-TP staging scatter overlap | `scatter_stream`<br>`_scatter_stream`<br>`staging` | `python/sglang/srt/disaggregation/common/staging_handler.py`<br>`python/sglang/srt/disaggregation/common/staging_buffer.py` | decode-side staging scatter kernels can run on a dedicated stream while forward continues on the main stream | If decode traces show staging scatter kernels adjacent to forward kernels, classify them against this existing overlap family first. |
|
||||
| Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. |
|
||||
|
||||
## 3. VLM-specific kernel families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Vision QK norm with aux stream | vision-side QK norm or norm-like kernels before attention | `python/sglang/srt/layers/attention/vision.py` | vision QK normalization can call shared `apply_qk_norm(...)`, with K-side work on `aux_stream` | If vision QK prep is split, first check this existing aux-stream path. |
|
||||
| ViT CUDA graph disables vision aux stream | expected vision overlap is absent under ViT graph | `python/sglang/srt/models/internvl.py`<br>`python/sglang/srt/layers/attention/vision.py`<br>`python/sglang/srt/environ.py::SGLANG_VIT_ENABLE_CUDA_GRAPH` | vision `aux_stream` is intentionally disabled when ViT CUDA graph is on | Missing vision overlap may be intentional, not a regression. |
|
||||
| Fused multimodal RoPE kernel | `triton_mrope_fused`<br>`multimodal_rotary_embedding_cpu`<br>`npu_mrope`<br>`MRotaryEmbedding` | `python/sglang/srt/layers/rotary_embedding/mrope.py`<br>`python/sglang/srt/layers/rotary_embedding/triton_kernels.py`<br>`python/sglang/srt/models/qwen3.py` | CUDA Triton, CPU `sgl_kernel`, and NPU paths already fuse multimodal t / h / w position lookup plus in-place Q / K rotary application | If VLM traces show separate mRoPE gather / shuffle / apply steps, first classify them as a missing existing mRoPE fusion. |
|
||||
|
||||
## 4. Diffusion fused-kernel families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Fused residual + norm + scale + shift | residual add, norm, scale, shift, gate around DiT blocks | `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py`<br>`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | `fused_scale_residual_norm_scale_shift(...)` | Treat split residual + norm + modulation as a missing existing diffusion fusion first. |
|
||||
| Fused norm + scale + shift | norm followed by scale / shift elementwise kernels | `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py`<br>`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | `fused_norm_scale_shift(...)` | Existing modulation fusion already covers this family. |
|
||||
| Triton scale / shift and gate-select kernels | tiny scale / shift or gate-select kernels dominate modulation blocks | `python/sglang/jit_kernel/diffusion/triton/scale_shift.py`<br>`python/sglang/multimodal_gen/runtime/layers/elementwise.py` | `fuse_scale_shift_kernel(...)` and `fuse_layernorm_scale_shift_gate_select01_kernel(...)` | Check whether the runtime is missing these existing Triton fusions. |
|
||||
| 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`<br>`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`<br>`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`<br>`tanh(gate) * rmsnorm(x)` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`<br>`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`<br>`residual + tanh(gate) * rmsnorm(x)`<br>`ffn_norm1(x) * scale_mlp` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`<br>`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`<br>`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
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Ulysses sequence-parallel attention | exposed `all_to_all` around attention blocks | `python/sglang/multimodal_gen/runtime/layers/attention/layer.py`<br>`python/sglang/multimodal_gen/runtime/distributed/communication_op.py` | head / sequence redistribution before and after attention | Treat sequence-parallel all-to-all as an existing distributed attention family. |
|
||||
| USP attention with all-to-all and ring attention | `all_to_all`, ring-attention comm, head / sequence reshards | `python/sglang/multimodal_gen/runtime/layers/attention/layer.py` | `_usp_input_all_to_all(...)`, `_usp_output_all_to_all(...)`, `ring_attn(...)` | This is the primary existing overlap / comm family for many diffusion models. |
|
||||
| Turbo-layer async all-to-all pipelining | pipelined A2A windows with explicit waits on a comm stream | `python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py` | looped `all_to_all_single(..., async_op=True)` plus staged postprocess on a comm stream | Treat exposed turbo A2A windows as an existing pipelined overlap pattern. |
|
||||
| TorchInductor compute / communication reorder | compiled traces with compute and comm partially interleaved | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py`<br>`python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py` | `torch._inductor.config.reorder_for_compute_comm_overlap = True` | Existing compile-time reordering may already explain partial overlap in diffusion traces. |
|
||||
| Dual-stream diffusion models | two nearby compute branches inside one DiT / UNet block | `python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py` | `use_dual_stream = True` | Treat dual-branch diffusion execution as an existing overlap family. |
|
||||
|
||||
## 6. PR-backed / in-flight fused-kernel families
|
||||
|
||||
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 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PR `#21877` fused grouped down-GEMM + combine | `grouped_gemm_nt_masked`<br>`combine`<br>`fused grouped gemm combine` | `PR #21877`<br>`python/sglang/srt/layers/moe/ep_moe/flashinfer_cutedsl_moe.py`<br>`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | FlashInfer CuTeDSL kernel fuses the second expert GEMM with DeepEP low-latency combine | Treat this as a concrete upstream MoE fuse / overlap family, not a new thought experiment. |
|
||||
| PR `#21889` fused BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`<br>`fp4 kv cache` | `PR #21889`<br>`python/sglang/srt/mem_cache/utils.py` | Triton kernel writes FP4 NSA KV pages directly while quantizing BF16 input | If NSA FP4 KV paths are split into quant plus store, classify them as an in-flight upstream fuse family. |
|
||||
| PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`<br>`WRITE_PT`<br>`dequant_fp4_paged_decode` | `PR #21889`<br>`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`<br>`trtllm_fp8_block_scale_moe` | `PR #21491`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`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`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`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 `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`<br>`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`rotary_dim` | `PR #20667`<br>`python/sglang/srt/models/qwen3_5.py`<br>`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 `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`sgl-kernel/python/sgl_kernel/gemm.py`<br>`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
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PR `#21877` fused down-GEMM + combine superseding SBO | `enable_fused_grouped_gemm_combine`<br>`combine`<br>`down_gemm` | `PR #21877`<br>`python/sglang/srt/server_args.py`<br>`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. |
|
||||
|
||||
## 8. FlashInfer mainline fused-kernel families
|
||||
|
||||
These rows are comparative references from `flashinfer`. Use them when a trace
|
||||
looks like an upstream FlashInfer family even if the current `sglang` checkout
|
||||
only consumes a subset of that implementation.
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| FlashInfer activation / gate epilogues | `silu_and_mul`<br>`gelu_tanh_and_mul`<br>`gelu_and_mul`<br>`silu_and_mul_scaled_nvfp4_experts_quantize` | `flashinfer/activation.py`<br>`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`<br>`fused_add_rmsnorm`<br>`fused_add_rmsnorm_quant`<br>`gemma_rmsnorm`<br>`gemma_fused_add_rmsnorm`<br>`fused_rmsnorm_silu`<br>`rmsnorm_fp4quant`<br>`add_rmsnorm_fp4quant` | `flashinfer/norm/__init__.py`<br>`flashinfer/cute_dsl/rmsnorm_fp4quant.py`<br>`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`<br>`AllReduceFusionPattern`<br>`kARResidualRMSNorm`<br>`kARResidualRMSNormFP8Quant`<br>`kARResidualRMSNormFP4Quant`<br>`trtllm_mnnvl_allreduce_fusion` | `flashinfer/comm/allreduce.py`<br>`flashinfer/comm/trtllm_ar.py`<br>`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`<br>`mla_rope_quantize_fp8`<br>`rope_quantize_fp8_append_paged_kv_cache`<br>`seqlen=0`<br>`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`<br>`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`<br>`trtllm_bf16_moe`<br>`trtllm_fp8_per_tensor_scale_moe`<br>`trtllm_fp8_block_scale_moe`<br>`trtllm_fp4_block_scale_moe`<br>`trtllm_mxint4_block_scale_moe`<br>`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`<br>`blockscaled_contiguous_grouped_gemm_finalize_fusion_nvfp4`<br>`moe_permute`<br>`moe_unpermute` | `flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py`<br>`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`<br>`CutlassTileConfigSM120`<br>`group_gemm_nvfp4_nt_groupwise`<br>`group_gemm_mxfp4_nt_groupwise` | `flashinfer/gemm/gemm_base.py`<br>`include/flashinfer/gemm/fp4_gemm_cutlass_template_sm120.h`<br>`include/flashinfer/gemm/group_gemm_nvfp4_groupwise_sm120.cuh`<br>`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`<br>`mPtrRoutingReplayOut`<br>`trtllm_fp8_block_scale_moe` | `flashinfer/fused_moe/core.py`<br>`csrc/trtllm_fused_moe_kernel_launcher.cu`<br>`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
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| FlashInfer PDL launch-overlap family | `enable_pdl`<br>`launch_with_pdl`<br>`cudaGridDependencySynchronize`<br>`cudaTriggerProgrammaticLaunchCompletion`<br>`trigger_completion_at_end=False`<br>`allreduce_fusion` | `flashinfer/norm/__init__.py`<br>`flashinfer/activation.py`<br>`flashinfer/rope.py`<br>`flashinfer/comm/allreduce.py`<br>`flashinfer/comm/trtllm_ar.py` | FlashInfer uses Programmatic Dependent Launch broadly, and the allreduce path can further advance completion so the next PDL-aware kernel overlaps on the same stream | Treat tight same-stream dependent windows and allreduce-followed-by-kernel windows as one existing FlashInfer launch-overlap family first. |
|
||||
| FlashInfer CuTeDSL MoE aux-stream async-memset overlap | `aux_stream`<br>`main_event`<br>`memset_event`<br>`use_async_memset` | `flashinfer/fused_moe/cute_dsl/fused_moe.py` | Preallocated MoE output is zeroed on an auxiliary CUDA stream while GEMM1 runs on the main stream, then both streams join before finalize | Treat GEMM1 vs output-zero windows as an existing FlashInfer multi-stream overlap family. |
|
||||
| FlashInfer green-context SM partition overlap | `split_device_green_ctx`<br>`split_device_green_ctx_by_sm_count`<br>`green_ctx` | `flashinfer/green_ctx.py` | CUDA green contexts partition SMs and create dedicated streams for concurrent kernel families on separate SM slices | Treat full-device two-stream traces and SM-partitioned traces as different manifestations of an existing FlashInfer overlap mechanism. |
|
||||
|
||||
## 10. FlashInfer PR-backed / in-flight fused-kernel and kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PR `#2720` PDL runtime-API migration | `cudaGridDependencySynchronize`<br>`cudaTriggerProgrammaticLaunchCompletion`<br>`inline PTX` | `PR #2720`<br>`include/flashinfer/comm/trtllm_allreduce_fusion.cuh`<br>`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. |
|
||||
|
||||
## 11. TensorRT-LLM-origin fused-kernel families
|
||||
|
||||
These rows are comparative references from `TensorRT-LLM`. Use them when a
|
||||
trace looks like a TensorRT-LLM or TensorRT-LLM-plus-FlashInfer family even if
|
||||
the current `sglang` checkout only carries an analogous implementation.
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| TensorRT-LLM FlashInfer activation / gate epilogues | `flashinfer_silu_and_mul`<br>`flashinfer_gelu_tanh_and_mul`<br>`auto_deploy::silu_and_mul`<br>post-GEMM `silu` + `mul` | `tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py`<br>`tensorrt_llm/_torch/auto_deploy/transform/library/fuse_silu_mul.py`<br>`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`<br>`flashinfer_gemma_rmsnorm`<br>`auto_deploy::flashinfer_rms_norm` | `tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py`<br>`tensorrt_llm/_torch/modules/rms_norm.py`<br>`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`<br>`flashinfer_gemma_fused_add_rmsnorm`<br>`auto_deploy::flashinfer_fused_add_rms_norm_inplace` | `tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py`<br>`tensorrt_llm/_torch/modules/rms_norm.py`<br>`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`<br>`fuse_rmsnorm_quant_fp8`<br>`fp8 static quant` | `tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/triton_fused_add_rms_norm_quant_fp8.py`<br>`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`<br>`flashinfer_rope`<br>`cos_sin_cache` | `tensorrt_llm/_torch/modules/rotary_embedding.py`<br>`tensorrt_llm/_torch/auto_deploy/custom_ops/rope/flashinfer_rope.py`<br>`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`<br>`BatchPrefillWithPagedKVCacheWrapper`<br>`BatchDecodeWithPagedKVCacheWrapper`<br>`auto_deploy::flashinfer_attention_mha_with_cache`<br>`read_cache_only` | `tensorrt_llm/_torch/attention_backend/flashinfer.py`<br>`tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py`<br>`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`<br>`BatchPrefillWithRaggedKVCacheWrapper`<br>`flashinfer_mla`<br>`rank 256`<br>`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`<br>`chunked prefill`<br>`W_kn`<br>`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`<br>`BatchMLAPagedAttentionWrapper`<br>`decode`<br>`W_kn`<br>`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`<br>`trtllm_bf16_moe`<br>`trtllm_fp8_block_scale_moe`<br>`trtllm_fp4_block_scale_moe`<br>`TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER` | `tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py`<br>`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. |
|
||||
| TensorRT-LLM FlashInfer cached SSM / Mamba update | `flashinfer_cached_ssm`<br>`selective_state_update`<br>`flashinfer_ssm` | `tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py`<br>`tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py` | Mamba2 paths can lower cached SSM state updates to FlashInfer selective-state-update kernels instead of many smaller state ops | Treat split cached-SSM state update ladders as an existing TensorRT-LLM FlashInfer family first. |
|
||||
|
||||
## 12. TensorRT-LLM-origin kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| TensorRT-LLM multi-stream MLA attention | `multi_stream_mla_attn`<br>`record_event_passthrough`<br>`_aux`<br>`wait_event` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`<br>`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`<br>`begin_aux_stream_passthrough`<br>`end_aux_stream_passthrough`<br>`wait_aux_stream_passthrough`<br>`mlir_elementwise_fusion`<br>`piecewise cudagraph`<br>`caller_stream.synchronize()` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`<br>`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`<br>`trtllm_finegrained_fp8_linear`<br>`record_event_passthrough`<br>`_aux` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_gemm.py`<br>`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 `#12525` FlashInfer TRTLLM-gen FMHA paged-index / buffer rework | `shared paged index`<br>`trtllm-gen attention`<br>`flashinfer`<br>`kv cache buffer` | `PR #12525`<br>`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`<br>`trtllm-gen attention`<br>`flashinfer` | `PR #12544`<br>`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`<br>`flashinfer`<br>`trtllm_bf16_moe` | `PR #12738`<br>`PR #12557`<br>`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. |
|
||||
|
||||
## 14. vLLM-origin fused-kernel families
|
||||
|
||||
These rows are comparative references from `vllm`. Use them when a trace looks
|
||||
similar to an upstream family even if the current `sglang` checkout does not
|
||||
contain the same implementation.
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| vLLM-origin fused residual add + RMSNorm | `fused_add_rms_norm*`<br>residual add right before RMSNorm | `vllm/model_executor/layers/layernorm.py`<br>`vllm/_custom_ops.py`<br>`csrc/layernorm_kernels.cu`<br>`csrc/cpu/layernorm.cpp` | Custom CUDA / CPU fused add-RMSNorm op reused directly and as a building block for later compile-time fusions | Treat split residual add + RMSNorm as a long-standing vLLM-origin precedent before calling the opportunity novel in sglang. |
|
||||
| vLLM-origin AllReduce + RMSNorm (+ residual / quant) | `fuse_allreduce_rms`<br>`AllReduceFusionPass`<br>`allreduce + rmsnorm` | `vllm/compilation/passes/fusion/allreduce_rms_fusion.py`<br>`docs/design/fusions.md` | Compile-time patterns cover `AllReduce -> RMSNorm(+residual_add)` and optional FP8 / NVFP4 quant suffixes | Treat TP collective + norm (+ quant) ladders as a known vLLM-origin fusion family first. |
|
||||
| vLLM-origin RMSNorm (+ residual add) + quant | `RMSNormQuantFusionPass`<br>`fused_add_rms_norm_static_fp8_quant`<br>`per_token_quant`<br>`per_group_quant` | `vllm/compilation/passes/fusion/rms_quant_fusion.py`<br>`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`<br>`SiluMulFp8*`<br>`Nvfp4`<br>`rocm_aiter` | `vllm/compilation/passes/fusion/act_quant_fusion.py`<br>`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`<br>`RocmAiterTritonAddRMSNormPadFusionPass`<br>`add_rmsnorm_pad` | `vllm/compilation/passes/fusion/rocm_aiter_fusion.py`<br>`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`<br>`AttnQuantFusionPass`<br>`merge_attn_states`<br>`output_scale`<br>`output_group_scale`<br>`output_block_scale` | `vllm/compilation/passes/fusion/attn_quant_fusion.py`<br>`vllm/v1/attention/ops/merge_attn_states.py`<br>`vllm/csrc/attention/merge_attn_states.cu`<br>`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`<br>`QKNormRoPEFusionPass`<br>`qk norm + rope` | `vllm/compilation/passes/fusion/qk_norm_rope_fusion.py`<br>`vllm/_custom_ops.py`<br>`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`<br>`triton_reshape_and_cache_flash`<br>`kv cache write` | `vllm/v1/attention/ops/triton_reshape_and_cache_flash.py`<br>`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`<br>`RopeKVCacheFusionPass`<br>`triton_rope_and_cache` | `vllm/compilation/passes/fusion/rope_kvcache_fusion.py`<br>`vllm/_aiter_ops.py`<br>`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. |
|
||||
| vLLM-origin fused MLA RoPE + concat/cache write | `concat_and_cache_mla_rope_fused`<br>`mla rope cache` | `vllm/_custom_ops.py`<br>`csrc/cache_kernels_fused.cu` | CUDA kernel fuses MLA-oriented RoPE preparation, concat, and cache write into a direct paged-store path | Treat MLA concat + cache-write ladders as a vLLM-origin precedent before calling them novel. |
|
||||
| vLLM-origin fused grouped top-k / biased grouped top-k router | `grouped_topk`<br>`biased_grouped_topk`<br>`grouped_topk_fused_kernel` | `vllm/_custom_ops.py`<br>`vllm/_aiter_ops.py`<br>`vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py`<br>`csrc/moe/grouped_topk_kernels.cu` | CUDA / ROCm router kernels fuse grouped score processing, top-k selection, and routed renorm / bias handling | Treat MoE router ladders as matching an upstream grouped-topk family first. |
|
||||
| vLLM-origin fused top-k softmax / sigmoid router | `topk_softmax`<br>`topk_sigmoid`<br>`topkGating`<br>`fused_topk` | `vllm/_custom_ops.py`<br>`vllm/_aiter_ops.py`<br>`vllm/model_executor/layers/fused_moe/router/fused_topk_router.py`<br>`vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py`<br>`csrc/moe/topk_softmax_kernels.cu` | CUDA and ROCm / AITER router kernels fuse score activation (`softmax` / `sigmoid`), top-k selection, optional bias correction, and routed renorm into one op instead of routing through grouped-topk or eager softmax-plus-topk ladders | Treat standalone score activation -> top-k -> bias / renorm chains as a known upstream fused router family first. |
|
||||
| vLLM-origin DSV3 router GEMM | `dsv3_router_gemm`<br>`allow_dsv3_router_gemm`<br>`router logits` | `vllm/_custom_ops.py`<br>`vllm/model_executor/layers/fused_moe/router/gate_linear.py`<br>`csrc/moe/dsv3_router_gemm_entry.cu`<br>`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`<br>`router gemm` | `vllm/_custom_ops.py`<br>`vllm/model_executor/layers/fused_moe/router/gate_linear.py`<br>`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`<br>`fused_qkv_a_proj`<br>`q_a_proj` | `vllm/model_executor/models/deepseek_v2.py`<br>`vllm/_custom_ops.py`<br>`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`<br>`MergedColumnParallelLinear`<br>`weights_proj` | `vllm/model_executor/models/deepseek_v2.py`<br>`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`<br>`minimax_allreduce_rmsnorm`<br>`MiniMax-M2.5`<br>`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`<br>`cutlass_scaled_mm_azp`<br>`scaled mm` | `vllm/_custom_ops.py`<br>`vllm/model_executor/kernels/linear/scaled_mm/cutlass.py`<br>`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`<br>`rocm_aiter_fused_moe`<br>`FusedMoE` | `vllm/model_executor/layers/fused_moe/layer.py`<br>`vllm/model_executor/layers/fused_moe/cpu_fused_moe.py`<br>`vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py`<br>`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`<br>`fused_moe_lora_fp8`<br>`w13_shrink`<br>`w2_expand` | `vllm/lora/ops/triton_ops/fused_moe_lora_op.py`<br>`vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py`<br>`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. |
|
||||
| vLLM-origin ViT fused bilinear position-embedding interpolation | `triton_pos_embed_interpolate`<br>`bilinear_pos_embed`<br>`pos_embed_interpolate_native` | `vllm/model_executor/models/qwen3_vl.py` | Triton kernel fuses bilinear interpolation and spatial-merge reorder for Qwen3-VL ViT position embeddings, replacing many tiny eager kernels | Treat VLM position-embedding ladders as an existing vLLM-origin Triton fusion family. |
|
||||
|
||||
## 15. vLLM-origin kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| vLLM-origin AsyncTP GEMM + collective overlap | `fuse_gemm_comms`<br>`fused_matmul_reduce_scatter`<br>`fused_all_gather_matmul` | `vllm/compilation/passes/fusion/collective_fusion.py`<br>`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`<br>`ReduceScatter`<br>`AllGather`<br>`SequenceParallelismPass` | `vllm/compilation/passes/fusion/sequence_parallelism.py`<br>`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`<br>`shared_experts_stream`<br>shared expert near router | `vllm/model_executor/layers/fused_moe/runner/shared_experts.py`<br>`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`<br>`all_to_all_single`<br>`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`<br>`wk`<br>`k_norm`<br>`aux_stream` | `PR #35968`<br>`vllm/model_executor/models/deepseek_v2.py`<br>`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`<br>`group_size=64`<br>`output_group_scale`<br>`per-group FP8` | `PR #37110`<br>`vllm/compilation/passes/fusion/attn_quant_fusion.py`<br>`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`<br>`MiniMax-M2`<br>`gate kernel` | `PR #38445`<br>`vllm/model_executor/layers/fused_moe/router/gate_linear.py`<br>`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`<br>`QK Norm + RoPE + Cache + Quant` | `PR #38621`<br>`csrc/fused_qk_norm_rope_cache_quant.cu`<br>`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 `#37646` ROCm AITER fused allreduce + RMSNorm | `rocm_aiter_fused_allreduce_rmsnorm`<br>`custom_fused_ar_rms`<br>`RocmAiterAllReduceFusionPass` | `PR #37646`<br>`vllm/_aiter_ops.py`<br>`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`<br>`flashinfer`<br>`NVFP4`<br>`rmsnorm + fp4 quant` | `PR #36413`<br>`vllm/compilation/passes/fusion/rms_quant_fusion.py`<br>`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 `#39301` GLM5 router GEMM with PDL overlap | `TRTLLM_ENABLE_PDL`<br>`router_gemm`<br>`GLM5`<br>`FI AR RMS fusion` | `PR #39301`<br>`vllm/model_executor/layers/fused_moe/router/gate_linear.py`<br>`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
|
||||
|
||||
| Toggle / env | Location | Effect on trace interpretation |
|
||||
| --- | --- | --- |
|
||||
| `enable_flashinfer_allreduce_fusion` | `python/sglang/srt/server_args.py` | Enables the FlashInfer TP allreduce fusion family. |
|
||||
| `enable_aiter_allreduce_fusion` | `python/sglang/srt/server_args.py` | Enables ROCm AITER TP allreduce fusion. |
|
||||
| `enable_deterministic_inference` | `python/sglang/srt/server_args.py` | Can intentionally disable or change some fast fusion paths, especially AITER allreduce fusion and some sampling / router choices, so split kernels may be expected. |
|
||||
| `enable_single_batch_overlap` | `python/sglang/srt/server_args.py` | Enables the SBO family. |
|
||||
| `enable_fused_moe_sum_all_reduce` | `python/sglang/srt/server_args.py` | Enables fused MoE sum-reduce in the down path. |
|
||||
| `SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO` | `python/sglang/srt/environ.py` | Alters how DeepSeek-style shared-expert overlap behaves on Blackwell. |
|
||||
| `SGLANG_NSA_FUSE_TOPK` | `python/sglang/srt/environ.py` | Gates NSA fused top-k transform / page-table build. |
|
||||
| `SGLANG_DISAGG_STAGING_BUFFER` | `python/sglang/srt/environ.py` | Enables the heterogeneous-TP staging-buffer family and its overlap windows. |
|
||||
| `SGLANG_STAGING_USE_TORCH` | `python/sglang/srt/disaggregation/common/staging_buffer.py` | Forces torch fallback for staging gather / scatter, so Triton staging kernels may disappear by design. |
|
||||
| `SGLANG_VIT_ENABLE_CUDA_GRAPH` | `python/sglang/srt/environ.py` | Can intentionally disable vision `aux_stream` overlap. |
|
||||
| `SGLANG_ENABLE_FUSED_QKNORM_ROPE` | `python/sglang/multimodal_gen/runtime/layers/layernorm.py` | Gates the diffusion fused qknorm+rope path. |
|
||||
| `enable_pdl` / `launch_with_pdl` | `flashinfer/norm/__init__.py`<br>`flashinfer/activation.py`<br>`flashinfer/rope.py`<br>`flashinfer/fused_moe/core.py`<br>`flashinfer/comm/allreduce.py` | Enables FlashInfer PDL across many kernels; launch grouping and same-stream overlap can change substantially when it is on. |
|
||||
| `trigger_completion_at_end` | `flashinfer/comm/allreduce.py` | `False` enables downstream PDL-aware overlap after FlashInfer allreduce fusion; `True` delays completion to kernel end and removes that overlap window. |
|
||||
| `use_cuda_graph` | `flashinfer/fused_moe/cute_dsl/fused_moe.py` | Enables the preallocated-buffer path and the safe aux-stream async-memset overlap in FlashInfer CuTeDSL MoE. |
|
||||
| `split_device_green_ctx*` | `flashinfer/green_ctx.py` | Changes trace shape by partitioning SMs into separate green contexts instead of overlapping full-device streams on the default context. |
|
||||
| `rmsnorm_backend` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Chooses whether AutoDeploy lowers RMSNorm to FlashInfer, so split norm ladders may reflect backend selection rather than a missing fuse. |
|
||||
| `insert_cached_attention.backend` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Selects the cached-attention backend; `flashinfer` enables the paged-KV cached-attention family. |
|
||||
| `insert_cached_mla_attention.backend` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Selects the cached MLA backend; `flashinfer_mla` enables the MLA prefill / decode family. |
|
||||
| `TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER` | `tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py` | Forces or guards the FlashInfer-backed TRTLLM-gen MoE family, so expert-kernel shape can change substantially when it is set. |
|
||||
| `multi_stream_moe` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Enables the TensorRT-LLM shared-expert vs routed-expert overlap family. |
|
||||
| `multi_stream_mla_attn` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Enables the TensorRT-LLM MLA Q-vs-KV branch overlap family. |
|
||||
| `multi_stream_gemm` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Enables generalized FP8 GEMM fork overlap in TensorRT-LLM AutoDeploy. |
|
||||
| `mlir_elementwise_fusion` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Can absorb merge adds into larger fused kernels, so missing explicit merge nodes in multi-stream traces may be intentional. |
|
||||
| `enable_torch_compile` | `python/sglang/srt/server_args.py`<br>`python/sglang/multimodal_gen/runtime/server_args.py` | Compiler-generated fusion / reordering can hide handwritten kernel names; absence of a custom kernel does not always mean absence of fusion. |
|
||||
| `enable_fused_grouped_gemm_combine` | `PR #21877` | In-flight path that intentionally disables SBO because combine is folded into down-GEMM. |
|
||||
| `PassConfig.fuse_allreduce_rms` | `vllm/config/compilation.py` | Enables vLLM's AllReduce -> RMSNorm (+ residual / quant) compile-time fusion family. |
|
||||
| `PassConfig.fuse_norm_quant` | `vllm/config/compilation.py` | Enables vLLM's RMSNorm(+residual add) -> FP8 / FP4 quant compile-time fusion family. |
|
||||
| `PassConfig.fuse_act_quant` | `vllm/config/compilation.py` | Enables vLLM's `SiLU+Mul -> quant` fusion family, plus ROCm AITER variants where applicable. |
|
||||
| `PassConfig.fuse_attn_quant` | `vllm/config/compilation.py` | Enables attention-epilogue quant fusion; requires the right backend / graph visibility, so split kernels may still be expected. |
|
||||
| `PassConfig.enable_qk_norm_rope_fusion` | `vllm/config/compilation.py` | Enables the compile-time QK RMSNorm + RoPE family on CUDA-like backends. |
|
||||
| `PassConfig.fuse_rope_kvcache` | `vllm/config/compilation.py` | Enables ROCm / AITER RoPE + KV-cache update fusion and is range-limited by token count. |
|
||||
| `PassConfig.enable_sp` | `vllm/config/compilation.py` | Rewrites all-reduce into sequence-parallel staging; this is often a prerequisite for the overlap family, not just a pure fuse toggle. |
|
||||
| `PassConfig.fuse_gemm_comms` | `vllm/config/compilation.py` | Enables AsyncTP GEMM + collective overlap and auto-enables `enable_sp` when valid. |
|
||||
| `TRTLLM_ENABLE_PDL` | `vllm/csrc/dsv3_fused_a_gemm.cu`<br>`vllm/csrc/moe/dsv3_router_gemm_utils.h` | Enables programmatic dependent launch for the DSV3 specialized CUDA kernels, which can change launch grouping and trace shape for router / QKV-A paths. |
|
||||
|
||||
## 18. 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.
|
||||
|
||||
```bash
|
||||
# Optional sibling checkouts used for comparative scanning:
|
||||
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|_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
|
||||
rg -n "single_batch_overlap|alt_stream|shared_expert|_comm_stream|scatter_stream|triton_mrope_fused|ring_attn|all_to_all_single|reorder_for_compute_comm_overlap|use_dual_stream" python/sglang
|
||||
git log --all --format='%h %s' | rg -i 'fused|fusion|overlap|cutedsl|triton|cuda|rope|topk|quant|combine|allreduce|all_to_all'
|
||||
rg -n "silu_and_mul|gelu_tanh_and_mul|gelu_and_mul|silu_and_mul_scaled_nvfp4_experts_quantize|rmsnorm_quant|fused_add_rmsnorm|fused_add_rmsnorm_quant|fused_rmsnorm_silu" "$FLASHINFER_REPO/flashinfer"
|
||||
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|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|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:
|
||||
# "fused OR overlap repo:sgl-project/sglang"
|
||||
# "triton OR cutedsl OR cuda fused repo:sgl-project/sglang"
|
||||
# "fused OR overlap repo:flashinfer-ai/flashinfer"
|
||||
# "pdl OR aux_stream OR green_ctx repo:flashinfer-ai/flashinfer"
|
||||
# "fused OR overlap repo:NVIDIA/TensorRT-LLM"
|
||||
# "flashinfer OR mla OR moe OR rmsnorm repo:NVIDIA/TensorRT-LLM"
|
||||
# "multi-stream OR aux_stream OR cudagraph repo:NVIDIA/TensorRT-LLM"
|
||||
# "fused OR overlap repo:vllm-project/vllm"
|
||||
# "triton OR cuda fused repo:vllm-project/vllm"
|
||||
```
|
||||
@@ -0,0 +1,119 @@
|
||||
# Overlap Heuristics
|
||||
|
||||
This analyzer is intentionally conservative.
|
||||
|
||||
## What Comes From Which Trace
|
||||
|
||||
### Mapping trace
|
||||
|
||||
Used for:
|
||||
|
||||
- `kernel -> cpu_op -> python scope`
|
||||
- launch-site call chains
|
||||
|
||||
This trace should be easier to read, even if it is not the exact final serving schedule.
|
||||
|
||||
### Formal trace
|
||||
|
||||
Used for:
|
||||
|
||||
- hidden ratio
|
||||
- exclusive ratio
|
||||
- overlap headroom
|
||||
- ASCII timelines
|
||||
|
||||
This trace should reflect the real serving shape.
|
||||
|
||||
## What It Treats As Hidden
|
||||
|
||||
A kernel is treated as hidden for a segment if:
|
||||
|
||||
- it is active during that segment
|
||||
- at least one kernel on a different stream is also active
|
||||
|
||||
If the overlapping kernel is compute-like, the analyzer separately records that it is hidden under compute.
|
||||
|
||||
## Category Heuristics
|
||||
|
||||
The analyzer classifies kernels by name:
|
||||
|
||||
- `compute`: GEMM, attention, cutlass, cublas, Triton matmul-like kernels
|
||||
- `communication`: NCCL, all-reduce, reduce-scatter, all-gather, DeepEP dispatch/combine
|
||||
- `elementwise`: sigmoid, top-k, gate, rmsnorm, layernorm, rope, casts
|
||||
- `memory`: memcpy, memset, fill, copy
|
||||
- `other`: everything else
|
||||
|
||||
These categories are for prioritization only.
|
||||
|
||||
## How To Read The Action Table
|
||||
|
||||
The overlap-opportunity table is intentionally not a full kernel dump.
|
||||
|
||||
It only keeps rows that already have an action-oriented label:
|
||||
|
||||
- `headroom`
|
||||
- `low-roi-hidden`
|
||||
|
||||
It also prunes very small `headroom` rows after prioritization.
|
||||
|
||||
- if a `headroom` row would end up as `P5` because it is below the default `1%` share bar, it is omitted from the table
|
||||
- `low-roi-hidden` rows can still remain even when they are small, because they are useful as "do not chase this first" signals
|
||||
|
||||
### `headroom`
|
||||
|
||||
Interpretation:
|
||||
|
||||
- the kernel still spends meaningful time exposed in the formal trace
|
||||
- the mapped Python scope is a good place to inspect scheduling or fusion opportunities
|
||||
- the dependency signal should still be checked before treating it as a serious overlap candidate
|
||||
|
||||
### `low-roi-hidden`
|
||||
|
||||
Interpretation:
|
||||
|
||||
- the kernel is already mostly hidden by another stream
|
||||
- optimizing it in isolation is less likely to move end-to-end latency
|
||||
- focus on fusion, launch reduction, or the surrounding schedule instead
|
||||
|
||||
## Dependency Signal
|
||||
|
||||
The table includes a dependency-oriented adjacency signal from the formal trace.
|
||||
|
||||
It is built from the nearest previous and next kernels on the same stream plus the mapping-trace source attribution.
|
||||
|
||||
Communication kernels are treated more conservatively than before:
|
||||
|
||||
- if a tight adjacent kernel looks like a likely producer or consumer, the table will raise the dependency risk even when the Python scope names differ
|
||||
- this avoids over-claiming that an all-reduce-like kernel is a clean overlap candidate just because its neighbors map to different functions
|
||||
|
||||
Typical labels:
|
||||
|
||||
- `serial risk low`: adjacent kernels do not look like a tight same-code serial chain
|
||||
- `prev-side serial risk`: the previous adjacent kernel looks tightly tied to the same code path
|
||||
- `next-side serial risk`: the next adjacent kernel looks tightly tied to the same code path
|
||||
- `both-side serial risk`: both sides look like a tight serial chain
|
||||
- `adjacency unclear`: the timing is tight but source attribution is too weak to trust a stronger claim
|
||||
|
||||
Treat this as a strong heuristic, not proof of dataflow.
|
||||
|
||||
The readable table compresses those into shorter labels:
|
||||
|
||||
- `low`
|
||||
- `high`
|
||||
- `unclear`
|
||||
|
||||
The recommendation labels are also intentionally short:
|
||||
|
||||
- `try overlap`
|
||||
- `try fusion`
|
||||
- `check deps`
|
||||
- `skip overlap`
|
||||
- `manual check`
|
||||
- `observe later`
|
||||
|
||||
## Important Limits
|
||||
|
||||
- A trace shows what overlapped, not what could legally overlap.
|
||||
- Two kernels on different streams do not prove they are dependency-free.
|
||||
- A mapped Python scope is a launch-site clue, not the only relevant code location.
|
||||
- A hidden kernel can still matter if it changes occupancy, launch count, or surrounding schedule.
|
||||
@@ -0,0 +1,180 @@
|
||||
# Overlap Catalog
|
||||
|
||||
This catalog is the overlap-only companion to
|
||||
`references/fuse-overlap-catalog.md`.
|
||||
|
||||
This revision is intentionally kernel-scoped. Keep rows here only when the
|
||||
overlap is visible in a profiler as GPU kernels, collective kernels, or
|
||||
streamed kernel families. Host-only scheduler, event-loop, executor, offload,
|
||||
and load-path overlaps are intentionally excluded.
|
||||
|
||||
Use it like this:
|
||||
|
||||
1. Start from the `overlap-opportunity table`.
|
||||
2. Match visible kernel windows, collective windows, or stream-level overlap
|
||||
against the rows below.
|
||||
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.
|
||||
5. Only call an overlap opportunity "new" when no row in this file or
|
||||
`fuse-overlap-catalog.md` fits.
|
||||
|
||||
The `vLLM-origin` sections below are comparative references. They are not
|
||||
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 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Single-batch overlap (SBO) | MoE combine, down-gemm, shared-expert work in nearby two-stream windows | `python/sglang/srt/batch_overlap/single_batch_overlap.py` | combine vs down-gemm overlap, combine vs shared-expert overlap, one-stream dispatch+shared overlap, explicit SM partitioning and events | If exposed MoE combine sits near neighboring compute, classify it against SBO before calling it new overlap. |
|
||||
| Q and K normalization on different streams | Q-side norm and K-side norm on different streams | `python/sglang/srt/models/utils.py::apply_qk_norm`<br>`python/sglang/srt/models/qwen3.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py` | Q stays on current stream, K can run on `alt_stream` in capture mode | Treat split Q / K norm as an existing overlap family when `alt_stream` is already wired. |
|
||||
| DeepSeek shared-expert / routed-expert overlap | shared-expert GEMMs near DeepEP dispatch / combine | `python/sglang/srt/models/deepseek_v2.py`<br>`python/sglang/srt/batch_overlap/single_batch_overlap.py` | shared experts on `alt_stream`, overlap with dispatch / combine and down-gemm, Blackwell-specific env gating | This is an established routed-vs-shared branch overlap pattern, not a novel idea. |
|
||||
| Llama4 shared branch vs routed branch overlap | shared expert branch plus routed MoE branch as adjacent windows | `python/sglang/srt/models/llama4.py` | shared expert on current stream, router + topk + routed experts on `alt_stream` | Use Llama4 as the first precedent for branch-level overlap in similar sparse models. |
|
||||
| ExaoneMoE shared experts vs router experts overlap | shared expert output and router-expert output form a two-branch window | `python/sglang/srt/models/exaone_moe.py::forward_normal_dual_stream` | shared experts on current stream, router + routed experts on `alt_stream`, explicit join before combine | This is an existing dual-stream MoE overlap family. |
|
||||
| Grok residual-MoE branch overlap | dense MLP and block-sparse MoE branches in parallel | `python/sglang/srt/models/grok.py::moe_with_rmoe` | dense MLP on current stream, MoE on `alt_stream`, fused dual residual RMSNorm around boundaries | Treat exposed Grok branch overlap as an existing pattern. |
|
||||
| NSA dual-stream overlap | Q-proj, K-proj, RoPE, cache-store, quantization in tight two-stream windows | `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Q / K projection split, RoPE split, cache-store vs quantization overlap | NSA already contains several dual-stream overlap precedents. |
|
||||
| MoriEP async dispatch / combine comm stream | `MoriEP`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. |
|
||||
| Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. |
|
||||
|
||||
## 2. Staging / communication kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Decode scatter on dedicated `scatter_stream` | `scatter_stream`<br>`_scatter_stream` | `python/sglang/srt/disaggregation/common/staging_handler.py` | staging scatter kernels are submitted to a dedicated stream so the decode thread does not block on the main forward stream | Treat decode-side staging scatter windows as an existing overlap pattern. |
|
||||
| Staging-buffer fused gather / scatter kernels | `_fused_gather_to_staging_kernel`<br>`_fused_scatter_from_staging_kernel` | `python/sglang/srt/disaggregation/common/staging_buffer.py` | Triton kernels gather KV slices into contiguous staging memory and scatter them back to KV cache | If heterogeneous-TP staging shows many small copy kernels, compare against this existing fused-plus-overlap family first. |
|
||||
|
||||
## 3. VLM / diffusion kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| Vision QK norm with aux stream | vision-side QK norm or norm-like kernels before attention | `python/sglang/srt/layers/attention/vision.py` | vision QK normalization can call shared `apply_qk_norm(...)`, with K-side work on `aux_stream` | If vision QK prep is split, first check this existing aux-stream path. |
|
||||
| ViT CUDA graph disables vision aux stream | expected vision overlap is absent under ViT graph | `python/sglang/srt/models/internvl.py`<br>`python/sglang/srt/layers/attention/vision.py`<br>`python/sglang/srt/environ.py::SGLANG_VIT_ENABLE_CUDA_GRAPH` | vision `aux_stream` is intentionally disabled when ViT CUDA graph is on | Missing vision overlap may be intentional, not a regression. |
|
||||
| Ulysses sequence-parallel attention | exposed `all_to_all` around attention blocks | `python/sglang/multimodal_gen/runtime/layers/attention/layer.py`<br>`python/sglang/multimodal_gen/runtime/distributed/communication_op.py` | head / sequence redistribution before and after attention | Treat sequence-parallel all-to-all as an existing distributed attention family. |
|
||||
| USP attention with all-to-all and ring attention | `all_to_all`, ring-attention comm, head / sequence reshards | `python/sglang/multimodal_gen/runtime/layers/attention/layer.py` | `_usp_input_all_to_all(...)`, `_usp_output_all_to_all(...)`, `ring_attn(...)` | This is the primary existing overlap / comm family for many diffusion models. |
|
||||
| Turbo-layer async all-to-all pipelining | pipelined A2A windows with explicit waits on a comm stream | `python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py` | looped `all_to_all_single(..., async_op=True)` plus staged postprocess on a comm stream | Treat exposed turbo A2A windows as an existing pipelined overlap pattern. |
|
||||
| TorchInductor compute / communication reorder | compiled traces with compute and comm partially interleaved | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py`<br>`python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py` | `torch._inductor.config.reorder_for_compute_comm_overlap = True` | Existing compile-time reordering may already explain partial overlap in diffusion traces. |
|
||||
| Dual-stream diffusion models | two nearby compute branches inside one DiT / UNet block | `python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py` | `use_dual_stream = True` | Treat dual-branch diffusion execution as an existing overlap family. |
|
||||
|
||||
## 4. PR-backed / in-flight kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PR `#21877` fused down-GEMM + combine superseding SBO | `enable_fused_grouped_gemm_combine`<br>`combine`<br>`down_gemm` | `PR #21877`<br>`python/sglang/srt/server_args.py`<br>`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. |
|
||||
|
||||
## 5. FlashInfer kernel-overlap families
|
||||
|
||||
These rows are comparative references from `flashinfer`. Use them when a trace
|
||||
looks like an upstream FlashInfer overlap family even if the current `sglang`
|
||||
checkout only calls part of that implementation.
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| FlashInfer PDL launch-overlap family | `enable_pdl`<br>`launch_with_pdl`<br>`cudaGridDependencySynchronize`<br>`cudaTriggerProgrammaticLaunchCompletion`<br>`trigger_completion_at_end=False`<br>`allreduce_fusion` | `flashinfer/norm/__init__.py`<br>`flashinfer/activation.py`<br>`flashinfer/rope.py`<br>`flashinfer/comm/allreduce.py`<br>`flashinfer/comm/trtllm_ar.py` | FlashInfer uses Programmatic Dependent Launch broadly, and the allreduce path can further advance completion so the next PDL-aware kernel overlaps on the same stream | Treat tight same-stream dependent windows and allreduce-followed-by-kernel windows as one existing FlashInfer launch-overlap family first. |
|
||||
| FlashInfer CuTeDSL MoE aux-stream async-memset overlap | `aux_stream`<br>`main_event`<br>`memset_event`<br>`use_async_memset` | `flashinfer/fused_moe/cute_dsl/fused_moe.py` | Preallocated MoE output is zeroed on an auxiliary CUDA stream while GEMM1 runs on the main stream, then both streams join before finalize | Treat GEMM1 vs output-zero windows as an existing FlashInfer multi-stream overlap family. |
|
||||
| FlashInfer green-context SM partition overlap | `split_device_green_ctx`<br>`split_device_green_ctx_by_sm_count`<br>`green_ctx` | `flashinfer/green_ctx.py` | CUDA green contexts partition SMs and create dedicated streams for concurrent kernel families on separate SM slices | Treat SM-partitioned concurrency as an existing FlashInfer overlap mechanism, not a novel scheduler idea. |
|
||||
|
||||
## 6. FlashInfer PR-backed / in-flight kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PR `#2720` PDL runtime-API migration | `cudaGridDependencySynchronize`<br>`cudaTriggerProgrammaticLaunchCompletion`<br>`inline PTX` | `PR #2720`<br>`include/flashinfer/comm/trtllm_allreduce_fusion.cuh`<br>`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
|
||||
|
||||
These rows are comparative references from `TensorRT-LLM`. Current mainline
|
||||
TensorRT-LLM overlap rows are mostly explicit auxiliary-stream rewrites in
|
||||
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`<br>`record_event_passthrough`<br>`_aux`<br>`wait_event` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`<br>`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`<br>`begin_aux_stream_passthrough`<br>`end_aux_stream_passthrough`<br>`wait_aux_stream_passthrough`<br>`mlir_elementwise_fusion`<br>`piecewise cudagraph`<br>`caller_stream.synchronize()` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`<br>`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`<br>`trtllm_finegrained_fp8_linear`<br>`record_event_passthrough`<br>`_aux` | `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_gemm.py`<br>`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. vLLM-origin kernel-overlap families
|
||||
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| vLLM-origin AsyncTP GEMM + collective overlap | `fuse_gemm_comms`<br>`fused_matmul_reduce_scatter`<br>`fused_all_gather_matmul` | `vllm/compilation/passes/fusion/collective_fusion.py`<br>`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`<br>`ReduceScatter`<br>`AllGather`<br>`SequenceParallelismPass` | `vllm/compilation/passes/fusion/sequence_parallelism.py`<br>`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`<br>`shared_experts_stream`<br>shared expert near router | `vllm/model_executor/layers/fused_moe/runner/shared_experts.py`<br>`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`<br>`all_to_all_single`<br>`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. |
|
||||
|
||||
## 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`<br>`wk`<br>`k_norm`<br>`aux_stream` | `PR #35968`<br>`vllm/model_executor/models/deepseek_v2.py`<br>`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`<br>`router_gemm`<br>`GLM5`<br>`FI AR RMS fusion` | `PR #39301`<br>`vllm/model_executor/layers/fused_moe/router/gate_linear.py`<br>`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. |
|
||||
|
||||
## 10. Important toggles and caveats
|
||||
|
||||
| Toggle / env | Location | Effect on trace interpretation |
|
||||
| --- | --- | --- |
|
||||
| `enable_single_batch_overlap` | `python/sglang/srt/server_args.py` | Enables the SBO family. |
|
||||
| `SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO` | `python/sglang/srt/environ.py` | Alters how DeepSeek-style shared-expert overlap behaves on Blackwell. |
|
||||
| `SGLANG_DISAGG_STAGING_BUFFER` | `python/sglang/srt/environ.py` | Enables the heterogeneous-TP staging-buffer family and its overlap windows. |
|
||||
| `SGLANG_STAGING_USE_TORCH` | `python/sglang/srt/disaggregation/common/staging_buffer.py` | Forces torch fallback for staging gather / scatter, so Triton staging kernels may disappear by design. |
|
||||
| `SGLANG_VIT_ENABLE_CUDA_GRAPH` | `python/sglang/srt/environ.py` | Can intentionally disable vision `aux_stream` overlap. |
|
||||
| `enable_pdl` / `launch_with_pdl` | `flashinfer/norm/__init__.py`<br>`flashinfer/activation.py`<br>`flashinfer/rope.py`<br>`flashinfer/fused_moe/core.py`<br>`flashinfer/comm/allreduce.py` | Enables FlashInfer PDL across many kernels; launch grouping and same-stream overlap can change substantially when it is on. |
|
||||
| `trigger_completion_at_end` | `flashinfer/comm/allreduce.py` | `False` enables downstream PDL-aware overlap after FlashInfer allreduce fusion; `True` delays completion to kernel end and removes that overlap window. |
|
||||
| `use_cuda_graph` | `flashinfer/fused_moe/cute_dsl/fused_moe.py` | Enables the preallocated-buffer path and the safe aux-stream async-memset overlap in FlashInfer CuTeDSL MoE. |
|
||||
| `split_device_green_ctx*` | `flashinfer/green_ctx.py` | Changes trace shape by partitioning SMs into separate green contexts instead of overlapping full-device streams on the default context. |
|
||||
| `multi_stream_moe` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Enables the TensorRT-LLM shared-expert vs routed-expert overlap family. |
|
||||
| `multi_stream_mla_attn` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Enables the TensorRT-LLM MLA Q-vs-KV branch overlap family. |
|
||||
| `multi_stream_gemm` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Enables generalized FP8 GEMM fork overlap in TensorRT-LLM AutoDeploy. |
|
||||
| `mlir_elementwise_fusion` | `tensorrt_llm/_torch/auto_deploy/config/default.yaml` | Can absorb merge adds into larger fused kernels, so missing explicit merge nodes in TensorRT-LLM multi-stream traces may be intentional. |
|
||||
| `enable_torch_compile` | `python/sglang/srt/server_args.py`<br>`python/sglang/multimodal_gen/runtime/server_args.py` | Compiler-generated reordering can hide or rename overlap windows. |
|
||||
| `enable_fused_grouped_gemm_combine` | `PR #21877` | In-flight path that intentionally disables SBO because combine is folded into down-GEMM. |
|
||||
| `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. |
|
||||
|
||||
## 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.
|
||||
|
||||
```bash
|
||||
# Optional sibling checkouts used for comparative scanning:
|
||||
FLASHINFER_REPO=${FLASHINFER_REPO:-../flashinfer}
|
||||
TRTLLM_REPO=${TRTLLM_REPO:-../TensorRT-LLM}
|
||||
VLLM_REPO=${VLLM_REPO:-../vllm}
|
||||
|
||||
rg -n "single_batch_overlap|alt_stream|shared_expert|scatter_stream|_fused_gather_to_staging_kernel|_fused_scatter_from_staging_kernel|async_op=True" python/sglang
|
||||
rg -n "apply_qk_norm|vision.py|ring_attn|all_to_all_single|reorder_for_compute_comm_overlap|use_dual_stream" python/sglang/multimodal_gen python/sglang/srt
|
||||
git log --all --format='%h %s' | rg -i 'fused|fusion|overlap|combine|all_to_all|ring attn|stream|triton|cutedsl|cuda'
|
||||
rg -n "enable_pdl|launch_with_pdl|trigger_completion_at_end|aux_stream|use_async_memset|split_device_green_ctx|split_device_green_ctx_by_sm_count" "$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 "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|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"
|
||||
# "triton OR cutedsl OR cuda overlap repo:sgl-project/sglang"
|
||||
# "fused OR overlap repo:flashinfer-ai/flashinfer"
|
||||
# "pdl OR aux_stream OR green_ctx repo:flashinfer-ai/flashinfer"
|
||||
# "fused OR overlap repo:NVIDIA/TensorRT-LLM"
|
||||
# "multi-stream OR aux_stream OR cudagraph repo:NVIDIA/TensorRT-LLM"
|
||||
# "mlir OR piecewise OR flashinfer repo:NVIDIA/TensorRT-LLM"
|
||||
# "fused OR overlap repo:vllm-project/vllm"
|
||||
# "triton OR cuda overlap repo:vllm-project/vllm"
|
||||
# "multi-stream OR aux_stream overlap repo:vllm-project/vllm"
|
||||
```
|
||||
@@ -0,0 +1,42 @@
|
||||
# Source Map
|
||||
|
||||
Use these upstream files when the workflow or behavior needs to be justified from SGLang source.
|
||||
|
||||
## Profiler entrypoints
|
||||
|
||||
- `python/sglang/profiler.py`
|
||||
- live profiler CLI
|
||||
- writes `server_args.json`
|
||||
- forwards `num_steps`, `profile_by_stage`, `merge_profiles`, and `profile_prefix`
|
||||
|
||||
- `python/sglang/test/send_one.py`
|
||||
- minimal request path that can trigger profiling from a single command
|
||||
|
||||
- `python/sglang/bench_serving.py`
|
||||
- profile-capable serving benchmark path
|
||||
- forwards `profile_activities`, `profile_by_stage`, `profile_stages`, and `profile_prefix`
|
||||
|
||||
## Scheduler-side trace writing
|
||||
|
||||
- `python/sglang/srt/managers/scheduler_profiler_mixin.py`
|
||||
- actual trace start/stop behavior
|
||||
- filename pattern for `TP/DP/PP/EP` and optional stage suffixes
|
||||
- `CUDA_PROFILER` and torch profiler handling
|
||||
|
||||
- `python/sglang/srt/utils/profile_merger.py`
|
||||
- merged distributed trace behavior
|
||||
- why merged traces should be treated differently from rank-local traces
|
||||
|
||||
- `python/sglang/srt/utils/profile_utils.py`
|
||||
- newer profile v2 manager path used for stage-based traces
|
||||
|
||||
## Documentation and tests
|
||||
|
||||
- `docs/developer_guide/benchmark_and_profiling.md`
|
||||
- canonical profiling docs
|
||||
|
||||
- `test/registered/profiling/test_start_profile.py`
|
||||
- validates `/start_profile` behavior, including `CUDA_PROFILER`
|
||||
|
||||
- `test/registered/profiling/test_profile_v2.py`
|
||||
- validates stage-scoped trace outputs under `SGLANG_PROFILE_V2`
|
||||
@@ -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:]))
|
||||
@@ -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:]))
|
||||
+132
@@ -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())
|
||||
@@ -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())
|
||||
@@ -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])
|
||||
@@ -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())
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user