[SKILL] add torch profiler analysis workflow (#22353)
This commit is contained in:
@@ -0,0 +1,223 @@
|
||||
---
|
||||
name: sglang-torch-profiler-analysis
|
||||
description: "Unified SGLang torch-profiler skill for trace generation, kernel/category breakdown, two-stage overlap analysis, and small Perfetto trace repair. Use when Codex should inspect an existing `trace.json(.gz)` or profile directory, trigger `sglang.profiler` against a live server, break down prefill/decode GPU time by kernel family, correlate a graph-off mapping trace with a graph-on formal trace to find overlap headroom tied back to Python code, or rewrite a trace so Perfetto renders overlapped events more reliably."
|
||||
---
|
||||
|
||||
# SGLang Torch Profiler Analysis
|
||||
|
||||
## Overview
|
||||
|
||||
Use this skill for all SGLang torch-profiler work. It replaces the old split between:
|
||||
|
||||
- kernel/category breakdown
|
||||
- overlap-specific diagnosis
|
||||
- small trace post-processing
|
||||
|
||||
Prefer the unified entrypoint:
|
||||
|
||||
- [scripts/analyze_sglang_torch_profile.py](scripts/analyze_sglang_torch_profile.py)
|
||||
|
||||
This entrypoint exposes four subcommands:
|
||||
|
||||
- `triage`: the default compact workflow that prints three main tables
|
||||
|
||||
- `breakdown`: one-trace kernel/category share analysis
|
||||
- `overlap`: required two-trace overlap analysis with source mapping
|
||||
- `perfetto-fix`: rewrite a trace when Perfetto drops some overlapped lanes
|
||||
|
||||
For normal use, prefer `triage`. It already collapses the result into three main tables:
|
||||
|
||||
- kernel table
|
||||
- overlap-opportunity table
|
||||
- fuse-opportunity table
|
||||
|
||||
Internal analyzers live here:
|
||||
|
||||
- [scripts/analyze_sglang_llm_torch_profile.py](scripts/analyze_sglang_llm_torch_profile.py)
|
||||
- [scripts/analyze_sglang_profiler_overlap.py](scripts/analyze_sglang_profiler_overlap.py)
|
||||
- [scripts/profile_common.py](scripts/profile_common.py)
|
||||
|
||||
## When To Use It
|
||||
|
||||
- inspect an SGLang torch profiler trace or profile directory
|
||||
- profile a live SGLang server and immediately analyze the output
|
||||
- quantify which kernel families dominate prefill or decode
|
||||
- compare communication, attention, MoE, quantization, norm, or memory share
|
||||
- map kernels back to Python code paths
|
||||
- judge whether a kernel still has overlap headroom in production shape
|
||||
- get a text table and a small ASCII timeline without opening Perfetto first
|
||||
- repair a trace so Perfetto can render overlapping events more faithfully
|
||||
|
||||
Do not use Nsight Systems as the default path for this workflow. This merged skill is torch-profiler-first.
|
||||
|
||||
## Main Commands
|
||||
|
||||
### 1. Compact triage from existing trace directories
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py triage \
|
||||
--mapping-input /path/to/graph_off_profile_dir \
|
||||
--formal-input /path/to/graph_on_profile_dir
|
||||
```
|
||||
|
||||
### 2. Compact triage from running servers
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py triage \
|
||||
--mapping-url http://127.0.0.1:31025 \
|
||||
--formal-url http://127.0.0.1:31026 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage
|
||||
```
|
||||
|
||||
### 3. Breakdown from an existing trace or profile dir
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--input /path/to/profile_dir
|
||||
```
|
||||
|
||||
### 4. Breakdown from a running server
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--url http://127.0.0.1:30000 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage \
|
||||
--table-only
|
||||
```
|
||||
|
||||
### 5. Two-stage overlap analysis
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py overlap \
|
||||
--mapping-input /path/to/graph_off_profile_dir \
|
||||
--formal-input /path/to/graph_on_profile_dir \
|
||||
--table-only
|
||||
```
|
||||
|
||||
Or profile both servers directly:
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py overlap \
|
||||
--mapping-url http://127.0.0.1:31025 \
|
||||
--formal-url http://127.0.0.1:31026 \
|
||||
--num-steps 5
|
||||
```
|
||||
|
||||
### 6. Perfetto-friendly trace rewrite
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py perfetto-fix \
|
||||
--input /path/to/trace.json.gz
|
||||
```
|
||||
|
||||
This small repair step is inspired by `torch_utils/src/convert_to_perfetto_compatible/convert_to_perfetto_compatible.py`.
|
||||
|
||||
## `profile_by_stage`
|
||||
|
||||
`profile_by_stage` is not only for PD disaggregation.
|
||||
|
||||
- On ordinary non-PD serving, it is still useful because prefill and decode usually have very different bottlenecks.
|
||||
- On the current profile-v2 path inside SGLang, stage-based profiling is effectively the normal path.
|
||||
- PD-disaggregated serving adds one extra rule: prefill workers and decode workers must be profiled separately. That is stricter than ordinary `profile_by_stage`.
|
||||
|
||||
## Which Mode To Choose
|
||||
|
||||
### `triage`
|
||||
|
||||
Use when you want the lowest-friction output:
|
||||
|
||||
- one kernel table
|
||||
- one overlap-opportunity table
|
||||
- one fuse-opportunity table
|
||||
- optional stage-aware rows when the trace directory includes both `EXTEND` and `DECODE`
|
||||
|
||||
This is the recommended default for final user-facing reports.
|
||||
|
||||
### `breakdown`
|
||||
|
||||
Use when you need:
|
||||
|
||||
- category share such as attention, communication, MoE, norm, quantize, memory
|
||||
- top kernels by cumulative GPU time
|
||||
- stage-aware prefill vs decode summaries
|
||||
- kernel tables keep full kernel names and full Python locations, already joined with CPU ops
|
||||
- conservative source-backed fusion opportunities
|
||||
|
||||
This mode works with one trace. A graph-off pre-pass plus `--kernel-map` is optional but recommended for the final polished report.
|
||||
|
||||
### `overlap`
|
||||
|
||||
Use when you need:
|
||||
|
||||
- a strong answer about which code paths still have overlap headroom
|
||||
- a table that says which kernels are already hidden and low ROI, with full kernel names and Python scopes
|
||||
- dependency-risk hints near adjacent kernels
|
||||
- an ASCII timeline around the most actionable windows
|
||||
|
||||
This mode requires two traces for a final answer:
|
||||
|
||||
1. mapping trace with `--disable-cuda-graph --disable-piecewise-cuda-graph`
|
||||
2. formal trace with the real serving optimizations enabled
|
||||
|
||||
Do not call the mapping pass a "fast profile". It exists to recover `kernel -> cpu_op -> python scope`.
|
||||
|
||||
### `perfetto-fix`
|
||||
|
||||
Use only when Perfetto fails to render obviously overlapping events cleanly. It is a post-processing utility, not the main analysis flow.
|
||||
|
||||
## Workflow
|
||||
|
||||
### One-trace breakdown workflow
|
||||
|
||||
1. If the user only wants kernel/category share, one trace is enough.
|
||||
2. Prefer rank-local `TP-0` traces over merged traces.
|
||||
3. For a live server, this skill can call `sglang.profiler` and automatically send a small probe request.
|
||||
4. Prefer `--profile-by-stage` even on standard serving unless the user explicitly wants an all-stage mixed trace.
|
||||
|
||||
### Two-trace overlap workflow
|
||||
|
||||
1. Produce a mapping trace first with graph disabled.
|
||||
2. Produce a formal trace second with graph enabled and the real serving flags kept on.
|
||||
3. Run `triage` for the compact three-table report, or `overlap` if you also want source context and ASCII timelines.
|
||||
4. Read the results in this order:
|
||||
- kernel table
|
||||
- overlap-opportunity table
|
||||
- fuse-opportunity table
|
||||
5. Before calling something a "new" optimization idea, compare the top rows against both [references/fuse-overlap-catalog.md](references/fuse-overlap-catalog.md) and [references/overlap-catalog.md](references/overlap-catalog.md). Always check the `PR-backed / in-flight` sections too. Prefer reporting:
|
||||
- an existing fused or overlap path that should already apply here
|
||||
- an existing path that appears disabled, unsupported, or regressed in this trace
|
||||
- an upstream PR-backed pattern that already exists but is not merged into the checked-out tree
|
||||
- a truly new opportunity only when no catalog entry fits
|
||||
6. Use the deeper `overlap` report only when you need source context or ASCII timelines beyond the compact three-table artifact.
|
||||
|
||||
## References
|
||||
|
||||
Load these only when needed:
|
||||
|
||||
- [references/source-map.md](references/source-map.md)
|
||||
- upstream SGLang profiler entrypoints and trace-writing source paths
|
||||
- [references/validated-workflows.md](references/validated-workflows.md)
|
||||
- validated two-pass examples for real SGLang models
|
||||
- [references/trace-workflow.md](references/trace-workflow.md)
|
||||
- practical guidance for mapping vs formal traces
|
||||
- [references/heuristics.md](references/heuristics.md)
|
||||
- overlap labels, dependency-risk interpretation, and limits
|
||||
- [references/fuse-overlap-catalog.md](references/fuse-overlap-catalog.md)
|
||||
- mixed source-backed catalog of existing fuse and overlap patterns, including PR-backed / in-flight rows
|
||||
- [references/overlap-catalog.md](references/overlap-catalog.md)
|
||||
- overlap-only lookup table across LLM, VLM, diffusion, disaggregation, HiSparse, and speculative scheduling
|
||||
|
||||
## Output Contract
|
||||
|
||||
### For `breakdown`
|
||||
|
||||
Return:
|
||||
|
||||
- trace path
|
||||
- model/server args when available
|
||||
- top categories
|
||||
- top kernels
|
||||
- one short conclusion about what dominates the run
|
||||
- any source-backed fusion opportunities worth checking
|
||||
@@ -0,0 +1,230 @@
|
||||
# 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 `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.
|
||||
|
||||
## 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>residual add right before norm | `python/sglang/srt/layers/layernorm.py`<br>`python/sglang/srt/layers/quantization/modelslim/modelslim.py` | Shared CUDA / ROCm / CPU / NPU fused add-RMSNorm implementations, including Gemma and NPU-bias variants | Treat split residual add + RMSNorm as an existing cross-backend fusion first, not a new idea. |
|
||||
| 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. |
|
||||
| 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. |
|
||||
| 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 merged 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 are intentionally not restricted to merged code. If the trace or
|
||||
user request is about upstream work, use these rows to avoid calling an
|
||||
already-known PR family "new".
|
||||
|
||||
| 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 `#21952` Gemma4 fused RMSNorm + residual + scalar | `gemma_rmsnorm_residual_scalar`<br>`_gemma_rmsnorm_residual_kernel`<br>`Gemma4` | `PR #21952`<br>`python/sglang/srt/layers/gemma4_fused_ops.py`<br>`python/sglang/srt/models/gemma4_causal.py` | Triton kernel fuses decoder post-FF RMSNorm, residual add, and per-layer scalar multiply into one pass | If Gemma4-style post-FF norm + residual + scalar steps appear split, treat them as an in-flight upstream Triton fuse family. |
|
||||
|
||||
## 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. 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>`output_scale`<br>`output_block_scale` | `vllm/compilation/passes/fusion/attn_quant_fusion.py`<br>`vllm/v1/attention/backends/`<br>`docs/design/fusions.md` | Compile-time fusion pushes FP8 / NVFP4 quantization into the attention epilogue on supported Triton / FlashInfer / ROCm / AITER backends | Treat attention-output quant kernels as a known upstream epilogue fusion family before calling them novel. |
|
||||
| vLLM-origin 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 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. |
|
||||
|
||||
## 9. 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/utils/torch_utils.py`<br>`vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py` | MoE shared experts can run on a dedicated aux stream and overlap with router-side work | Treat shared-expert vs router overlap as an existing upstream sparse-model family. |
|
||||
| vLLM-origin 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. |
|
||||
|
||||
## 10. 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` | Open PR overlaps the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. |
|
||||
| PR `#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 `#38684` DSV3.2 fused `wk + weights_proj` | `wk_weights_proj`<br>`MergedColumnParallelLinear`<br>`weights_proj` | `PR #38684`<br>`vllm/model_executor/models/deepseek_v2.py`<br>`vllm/model_executor/models/deepseek_mtp.py` | Merged PR fuses the DSV3.2 indexer `wk` and `weights_proj` projections into one GEMM; FP8 weight-loading caveats are being handled in follow-up `PR #38870` | Treat paired indexer projections as a concrete upstream fused linear family before calling the opportunity novel. |
|
||||
|
||||
## 11. 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_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. |
|
||||
|
||||
## 12. Suggested refresh commands
|
||||
|
||||
```bash
|
||||
rg -n "fused_add_rmsnorm|gemma_fused_add_rmsnorm|silu_and_mul|gelu_and_mul|fused_qk_rope_reshape_and_cache|fused_set_kv_buffer|fused_metadata_copy|normal_decode_set_metadata" python/sglang
|
||||
rg -n "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 "fused_add_rms_norm|fused_qk_norm_rope|grouped_topk|topk_softmax|topk_sigmoid|dsv3_router_gemm|dsv3_fused_a_gemm|concat_and_cache_mla_rope_fused|gpt_oss_router_gemm|cutlass_scaled_mm|cpu_fused_moe|fused_moe_lora|triton_pos_embed_interpolate" /Users/bbuf/工作目录/Common/vllm/vllm /Users/bbuf/工作目录/Common/vllm/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" /Users/bbuf/工作目录/Common/vllm/vllm /Users/bbuf/工作目录/Common/vllm/docs/design/fusions.md /Users/bbuf/工作目录/Common/vllm/csrc
|
||||
git -C /Users/bbuf/工作目录/Common/vllm 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: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,111 @@
|
||||
# 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.
|
||||
|
||||
## 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. 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/utils/torch_utils.py`<br>`vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py` | MoE shared experts can run on a dedicated aux stream and overlap with router-side work | Treat shared-expert vs router overlap as an existing upstream sparse-model family. |
|
||||
| vLLM-origin 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. |
|
||||
|
||||
## 6. 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` | Open PR overlaps the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. |
|
||||
|
||||
## 7. 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_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. |
|
||||
|
||||
## 8. Suggested refresh commands
|
||||
|
||||
```bash
|
||||
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 "fuse_gemm_comms|enable_sp|fused_matmul_reduce_scatter|fused_all_gather_matmul|shared_experts_stream|dcp_alltoall|async_op=True|aux_stream|maybe_execute_in_parallel" /Users/bbuf/工作目录/Common/vllm/vllm /Users/bbuf/工作目录/Common/vllm/docs/design/fusions.md
|
||||
git -C /Users/bbuf/工作目录/Common/vllm 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: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,119 @@
|
||||
# Trace Workflow
|
||||
|
||||
This skill is based on SGLang's existing profiling workflow.
|
||||
|
||||
Use:
|
||||
|
||||
- two traces for `triage`
|
||||
- one trace for `breakdown`
|
||||
- two traces for `overlap`
|
||||
- optional post-processing for `perfetto-fix`
|
||||
|
||||
`profile_by_stage` is still useful on normal non-PD serving because it separates prefill and decode. PD disaggregation adds an extra requirement beyond that: prefill workers and decode workers must be profiled separately.
|
||||
|
||||
## Existing SGLang Sources
|
||||
|
||||
- `sglang/.claude/skills/generate-profile/SKILL.md`
|
||||
- `sglang/docs/developer_guide/benchmark_and_profiling.md`
|
||||
- `sglang/docs/diffusion/performance/profiling.md`
|
||||
- `sglang/python/sglang/profiler.py`
|
||||
- `sglang/python/sglang/srt/utils/profile_utils.py`
|
||||
- `sglang/python/sglang/srt/utils/profile_merger.py`
|
||||
|
||||
## Required Two-Stage Flow
|
||||
|
||||
### Stage 1: Mapping trace
|
||||
|
||||
Collect a graph-off trace first.
|
||||
|
||||
Recommended properties:
|
||||
|
||||
- disable `cuda graph`
|
||||
- disable `piecewise cuda graph` if it would otherwise hide launch attribution
|
||||
- keep the same model, parallel shape, backend choices, and request pattern as much as possible
|
||||
|
||||
Purpose:
|
||||
|
||||
- preserve clean kernel launch attribution
|
||||
- recover `kernel -> cpu_op -> python scope`
|
||||
|
||||
### Stage 2: Formal trace
|
||||
|
||||
Collect a second trace with the real serving optimizations enabled.
|
||||
|
||||
Recommended properties:
|
||||
|
||||
- enable `cuda graph` if the real deployment uses it
|
||||
- enable `piecewise cuda graph` when the model normally captures it
|
||||
- keep the production MoE, attention, communication, and quantization backends
|
||||
|
||||
Purpose:
|
||||
|
||||
- measure real overlap under the real schedule
|
||||
- decide whether a code path still has overlap headroom
|
||||
|
||||
The final compact report should be built from:
|
||||
|
||||
- source attribution from stage 1
|
||||
- overlap conclusions from stage 2
|
||||
|
||||
The merged skill's `triage` command turns that into three tables:
|
||||
|
||||
- kernel table
|
||||
- overlap-opportunity table
|
||||
- fuse-opportunity table
|
||||
|
||||
## Ways To Produce A Trace
|
||||
|
||||
### Live server
|
||||
|
||||
```bash
|
||||
python3 -m sglang.profiler --url http://127.0.0.1:30000 --num-steps 5
|
||||
```
|
||||
|
||||
### One-shot request plus profile
|
||||
|
||||
```bash
|
||||
python3 -m sglang.test.send_one --profile
|
||||
```
|
||||
|
||||
### Bench serving
|
||||
|
||||
```bash
|
||||
export SGLANG_TORCH_PROFILER_DIR=/tmp/sglang-profile
|
||||
python3 -m sglang.bench_serving --backend sglang --num-prompts 10 --profile
|
||||
```
|
||||
|
||||
If you only call `python3 -m sglang.profiler`, remember that something still has to drive requests through the server while profiling is active. The merged skill's live URL flows handle this automatically by sending a small probe workload.
|
||||
|
||||
## Expected Output
|
||||
|
||||
Typical trace outputs are:
|
||||
|
||||
- `<profile_id>-TP-0.trace.json.gz`
|
||||
- `<profile_id>-TP-0-DP-0-PP-0-EP-0.trace.json.gz`
|
||||
- `merged-<profile_id>.trace.json.gz`
|
||||
- `server_args.json`
|
||||
|
||||
For overlap analysis, prefer a single-rank trace over a merged multi-rank trace.
|
||||
|
||||
## Optional Perfetto Repair
|
||||
|
||||
When Perfetto fails to render clearly overlapping events on the same logical lane, the unified script exposes:
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py perfetto-fix --input /path/to/trace.json.gz
|
||||
```
|
||||
|
||||
This is intentionally a narrow repair step. Do not make it part of the default profiling workflow unless rendering is actually broken.
|
||||
|
||||
## Why The Mapping Trace Must Exist
|
||||
|
||||
A single graph-on trace may still tell you that overlap is poor, but it is often not enough to say which Python code path owns the kernel.
|
||||
|
||||
That is why this skill requires:
|
||||
|
||||
- one trace for readable source attribution
|
||||
- one trace for real overlap behavior
|
||||
|
||||
Do not call the graph-off trace a "fast profile" in the final write-up. Its role is source mapping, not shortcutting the real analysis.
|
||||
@@ -0,0 +1,263 @@
|
||||
# Validated Workflows
|
||||
|
||||
These were the concrete workflows used to validate the skill design against real remote GPU environments.
|
||||
|
||||
## 1. Small single-GPU Qwen
|
||||
|
||||
Validated as a two-pass workflow on H100.
|
||||
|
||||
### Pass 1: no-CUDA-graph mapping pre-pass
|
||||
|
||||
Validated example:
|
||||
|
||||
```bash
|
||||
export SGLANG_TORCH_PROFILER_DIR=/tmp/sglang_torch_profile_qwen25_15b_map
|
||||
CUDA_VISIBLE_DEVICES=5 FLASHINFER_DISABLE_VERSION_CHECK=1 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen2.5-1.5B-Instruct \
|
||||
--host 127.0.0.1 \
|
||||
--port 32240 \
|
||||
--disable-cuda-graph \
|
||||
--disable-piecewise-cuda-graph
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--url http://127.0.0.1:32240 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage \
|
||||
--profile-prefix qwen25_15b_map \
|
||||
--export-kernel-map /tmp/qwen25_15b_kernel_map.json
|
||||
```
|
||||
|
||||
### Pass 2: final optimized profile
|
||||
|
||||
```bash
|
||||
export SGLANG_TORCH_PROFILER_DIR=/tmp/sglang_torch_profile_qwen25_15b_final
|
||||
CUDA_VISIBLE_DEVICES=5 FLASHINFER_DISABLE_VERSION_CHECK=1 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen2.5-1.5B-Instruct \
|
||||
--host 127.0.0.1 \
|
||||
--port 32241
|
||||
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 python3 -m sglang.profiler \
|
||||
--url http://127.0.0.1:32241 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage \
|
||||
--profile-prefix qwen25_15b_final
|
||||
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--input /tmp/sglang_torch_profile_qwen25_15b_final \
|
||||
--kernel-map /tmp/qwen25_15b_kernel_map.json
|
||||
```
|
||||
|
||||
## 2. Multi-GPU Qwen
|
||||
|
||||
Validated as a two-pass workflow on H100 with `Qwen/Qwen3-32B`.
|
||||
|
||||
### Pass 1: no-CUDA-graph mapping pre-pass
|
||||
|
||||
Validated target options:
|
||||
|
||||
- `Qwen/Qwen3-32B` on H100 with `--tp 2`
|
||||
- `Qwen/Qwen3-Next-80B-A3B-Instruct` on H200 with `--tp 4`
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
export SGLANG_TORCH_PROFILER_DIR=/tmp/sglang_torch_profile_qwen32b_map
|
||||
CUDA_VISIBLE_DEVICES=1,2 FLASHINFER_DISABLE_VERSION_CHECK=1 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen3-32B \
|
||||
--tp 2 \
|
||||
--host 127.0.0.1 \
|
||||
--port 32040 \
|
||||
--disable-cuda-graph \
|
||||
--disable-piecewise-cuda-graph
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--url http://127.0.0.1:32040 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage \
|
||||
--profile-prefix qwen32b_map \
|
||||
--export-kernel-map /tmp/qwen32b_kernel_map.json
|
||||
```
|
||||
|
||||
### Pass 2: final optimized profile
|
||||
|
||||
```bash
|
||||
export SGLANG_TORCH_PROFILER_DIR=/tmp/sglang_torch_profile_qwen32b_final
|
||||
CUDA_VISIBLE_DEVICES=1,2 FLASHINFER_DISABLE_VERSION_CHECK=1 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen3-32B \
|
||||
--tp 2 \
|
||||
--host 127.0.0.1 \
|
||||
--port 32041
|
||||
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 python3 -m sglang.profiler \
|
||||
--url http://127.0.0.1:32041 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage \
|
||||
--profile-prefix qwen32b_final
|
||||
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--input /tmp/sglang_torch_profile_qwen32b_final \
|
||||
--kernel-map /tmp/qwen32b_kernel_map.json
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- Prefer `TP-0` traces first for kernel share analysis and for the exported kernel map.
|
||||
- If the directory only contains merged traces, state that explicitly in the final conclusions.
|
||||
- For stage-aware comparisons, analyze `EXTEND` and `DECODE` separately before summarizing the overall model behavior.
|
||||
- On current SGLang builds, add `--disable-piecewise-cuda-graph` together with `--disable-cuda-graph` for the mapping pass, otherwise extend/prefill may still run under piecewise CUDA graph.
|
||||
- `Qwen/Qwen3-4B-Instruct-2507` was present in H200 cache during validation but did not work on the validated stack because of a `Qwen3Config.rope_parameters` compatibility issue.
|
||||
|
||||
## 3. Deliberately broken TP fusion rediscovery
|
||||
|
||||
Validated on B200 with `Qwen/Qwen2.5-0.5B-Instruct`, `TP=2`.
|
||||
|
||||
The validation intentionally commented out the fused TP all-reduce + RMSNorm path inside:
|
||||
|
||||
- `python/sglang/srt/layers/layernorm.py`
|
||||
|
||||
and forced the code to fall back to:
|
||||
|
||||
- plain `tensor_model_parallel_all_reduce`
|
||||
- then ordinary `norm_module.forward(...)`
|
||||
|
||||
### Mapping pass
|
||||
|
||||
Graph-off server:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=6,7 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--tp 2 \
|
||||
--host 127.0.0.1 \
|
||||
--port 32260 \
|
||||
--disable-cuda-graph \
|
||||
--disable-piecewise-cuda-graph
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py breakdown \
|
||||
--url http://127.0.0.1:32260 \
|
||||
--num-steps 5 \
|
||||
--profile-by-stage \
|
||||
--profile-prefix qwen25_tp2_map \
|
||||
--export-kernel-map /tmp/qwen25_tp2_map_kernel_map.json
|
||||
```
|
||||
|
||||
Observed result:
|
||||
|
||||
- the fuse table rediscovered `TP all-reduce + residual/RMSNorm`
|
||||
- it pointed back to `python/sglang/srt/layers/layernorm.py:89 _forward_with_allreduce_fusion`
|
||||
|
||||
### Formal pass
|
||||
|
||||
Graph-on server with the intentionally broken code still in place:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=6,7 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--tp 2 \
|
||||
--host 127.0.0.1 \
|
||||
--port 32261
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```bash
|
||||
python3 scripts/analyze_sglang_torch_profile.py triage \
|
||||
--mapping-input /tmp/sglang-torch-profile-_7gd033i/1775035260.3725493 \
|
||||
--formal-input /tmp/sglang-torch-profile-64oszwu2/1775035372.2416728
|
||||
```
|
||||
|
||||
Observed result:
|
||||
|
||||
- the kernel table showed `void sglang::cross_device_reduce_1stage<__nv_bfloat16, 2>` at roughly `31%` decode share
|
||||
- the overlap table surfaced that same communication kernel as a top headroom row
|
||||
- the fuse table again flagged `TP all-reduce + residual/RMSNorm`
|
||||
|
||||
This validation demonstrates that the skill can rediscover a real missing fusion path and also surface the exposed overlap opportunity that appears when the fusion is removed.
|
||||
|
||||
## 4. Dense Qwen3 QK-norm + RoPE rediscovery
|
||||
|
||||
Validated on B200 with `Qwen/Qwen3-32B`, `TP=2`.
|
||||
|
||||
### Mapping pass
|
||||
|
||||
Graph-off server:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1,2 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen3-32B \
|
||||
--tp 2 \
|
||||
--host 127.0.0.1 \
|
||||
--port 32320 \
|
||||
--disable-cuda-graph \
|
||||
--disable-piecewise-cuda-graph
|
||||
```
|
||||
|
||||
### Formal pass
|
||||
|
||||
Graph-on server:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1,2 python3 -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen3-32B \
|
||||
--tp 2 \
|
||||
--host 127.0.0.1 \
|
||||
--port 32321
|
||||
```
|
||||
|
||||
Observed result:
|
||||
|
||||
- on the current `sm100 + trtllm_mha` stack, `apply_qk_norm` and `RoPE` may collapse into a shared generic `void` kernel row
|
||||
- the useful source evidence still survives in the mapped Python locations:
|
||||
- `python/sglang/jit_kernel/rope.py:179 apply_rope_with_cos_sin_cache_inplace`
|
||||
- `python/sglang/srt/models/utils.py:204 apply_qk_norm`
|
||||
- the fuse detector therefore needs to treat a shared kernel row containing both `QK norm` and `RoPE` evidence as valid
|
||||
- after broadening the rule, triage produced:
|
||||
- `decode | Q/K RMSNorm + RoPE before attention | Conditional | 2.34 ms | 4.5%`
|
||||
|
||||
This validation demonstrates that dense-Qwen3 fuse detection should be source-evidence driven, not tied only to `norm` or `rope` kernel categories.
|
||||
|
||||
## 5. Single-GPU negative control
|
||||
|
||||
Validated on B200 with `Qwen/Qwen2.5-0.5B-Instruct`, single GPU, `TP=1`.
|
||||
|
||||
Observed result:
|
||||
|
||||
- the three main tables were still produced normally
|
||||
- no medium-confidence source-backed fusion opportunity was emitted
|
||||
- the skill did not incorrectly report:
|
||||
- `TP all-reduce + residual/RMSNorm`
|
||||
- `Q/K RMSNorm + RoPE before attention`
|
||||
|
||||
This validation is the negative control for avoiding TP-specific false positives when no TP communication exists.
|
||||
|
||||
## 6. MiniMax overlap-heavy trace validation
|
||||
|
||||
Validated on B200 using previously captured `MiniMaxAI/MiniMax-M2.5` mapping and formal traces.
|
||||
|
||||
Current note:
|
||||
|
||||
- a fresh launch on the current B200 `main` repo failed before profiling with:
|
||||
- `AttributeError: 'MiniMaxM2Config' object has no attribute 'rope_theta'`
|
||||
- that load-time issue is separate from the profiler skill itself
|
||||
|
||||
Using the existing B200 traces still validated the triage behavior:
|
||||
|
||||
- the kernel table was dominated by:
|
||||
- `void sglang::cross_device_reduce_1stage<__half, 4>`
|
||||
- `fused_moe_kernel`
|
||||
- the overlap table stayed compact and preserved only actionable overlap rows plus `low-roi-hidden` deprioritization rows
|
||||
- the fuse table flagged `TP all-reduce + residual/RMSNorm`
|
||||
|
||||
This validation demonstrates that the compact three-table artifact still stays readable on a communication-heavy MoE model.
|
||||
+1582
File diff suppressed because it is too large
Load Diff
+1768
File diff suppressed because it is too large
Load Diff
+559
@@ -0,0 +1,559 @@
|
||||
"""Unified entrypoint for SGLang torch-profiler analysis workflows."""
|
||||
|
||||
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 analyze_sglang_llm_torch_profile as breakdown_cli
|
||||
import analyze_sglang_profiler_overlap as overlap_cli
|
||||
from profile_common import (
|
||||
discover_trace_targets,
|
||||
load_server_args,
|
||||
load_trace_json,
|
||||
parse_stage,
|
||||
run_profiler,
|
||||
write_perfetto_compatible_trace,
|
||||
)
|
||||
|
||||
|
||||
def build_top_level_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="analyze_sglang_torch_profile.py",
|
||||
description=(
|
||||
"Unified torch-profiler entrypoint for SGLang. "
|
||||
"Use `breakdown` for kernel/category share analysis, "
|
||||
"`overlap` for two-trace overlap analysis, `triage` for the compact "
|
||||
"three-table workflow, or `perfetto-fix` to rewrite a trace into a "
|
||||
"more Perfetto-friendly form."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"command",
|
||||
nargs="?",
|
||||
choices=("breakdown", "overlap", "triage", "perfetto-fix"),
|
||||
help="Subcommand to run.",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def parse_perfetto_fix_args(argv: Sequence[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="analyze_sglang_torch_profile.py perfetto-fix",
|
||||
description="Rewrite a trace so overlapping kernel lanes render more reliably in Perfetto.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input", required=True, help="Input trace.json or trace.json.gz path."
|
||||
)
|
||||
parser.add_argument("--output", default=None, help="Optional output path.")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def parse_triage_args(argv: Sequence[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="analyze_sglang_torch_profile.py triage",
|
||||
description=(
|
||||
"Run the compact SGLang torch-profiler triage workflow. "
|
||||
"This prints three stage-aware tables: kernel mapping, overlap opportunities, "
|
||||
"and fuse opportunities."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mapping-input",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Graph-off mapping trace file or directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mapping-url",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Running graph-off SGLang server URL for the mapping trace.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--formal-input",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Formal graph-on trace file or directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--formal-url",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Running graph-on SGLang server URL for the formal trace.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mapping-output-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Trace output dir when using --mapping-url.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--formal-output-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Trace output dir when using --formal-url.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mapping-profile-prefix",
|
||||
type=str,
|
||||
default="mapping-trace",
|
||||
help="Profile prefix for the mapping trace.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--formal-profile-prefix",
|
||||
type=str,
|
||||
default="formal-trace",
|
||||
help="Profile prefix for the formal trace.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-steps",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Profiler steps when generating traces from URLs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-by-stage", action=argparse.BooleanOptionalAction, default=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--merge-profiles", action=argparse.BooleanOptionalAction, default=False
|
||||
)
|
||||
parser.add_argument("--probe-requests", type=int, default=1)
|
||||
parser.add_argument(
|
||||
"--probe-prompt",
|
||||
type=str,
|
||||
default=(
|
||||
"Repeat the word profiler many times with spaces so the server performs several decode steps. "
|
||||
"Do not add explanations."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--probe-max-new-tokens", type=int, default=None)
|
||||
parser.add_argument("--probe-delay", type=float, default=0.5)
|
||||
parser.add_argument(
|
||||
"--start-step",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Pass through to sglang.profiler when generating traces from URLs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pid-substring",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Restrict overlap analysis to PIDs containing this substring.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kernel-table-limit",
|
||||
type=int,
|
||||
default=0,
|
||||
help="How many kernel rows to print per stage. Use 0 for all kernels.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overlap-table-limit",
|
||||
type=int,
|
||||
default=0,
|
||||
help="How many overlap rows to print per stage. Use 0 for all kernels.",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
if bool(args.mapping_input) == bool(args.mapping_url):
|
||||
parser.error("Provide exactly one of --mapping-input or --mapping-url.")
|
||||
if bool(args.formal_input) == bool(args.formal_url):
|
||||
parser.error("Provide exactly one of --formal-input or --formal-url.")
|
||||
return args
|
||||
|
||||
|
||||
def resolve_profile_targets(
|
||||
*,
|
||||
label: str,
|
||||
input_path: Optional[str],
|
||||
url: Optional[str],
|
||||
output_dir: Optional[str],
|
||||
profile_prefix: Optional[str],
|
||||
args: argparse.Namespace,
|
||||
) -> Tuple[List[Path], Optional[dict]]:
|
||||
if bool(input_path) == bool(url):
|
||||
raise ValueError(f"{label} trace requires exactly one of input path or URL.")
|
||||
|
||||
if url:
|
||||
target_dir = run_profiler(
|
||||
url=url,
|
||||
output_dir=output_dir,
|
||||
num_steps=args.num_steps,
|
||||
profile_by_stage=args.profile_by_stage,
|
||||
merge_profiles=args.merge_profiles,
|
||||
profile_prefix=profile_prefix,
|
||||
probe_requests=max(0, args.probe_requests),
|
||||
probe_prompt=args.probe_prompt,
|
||||
probe_max_new_tokens=args.probe_max_new_tokens,
|
||||
probe_delay=args.probe_delay,
|
||||
start_step=args.start_step,
|
||||
)
|
||||
traces, server_args = discover_trace_targets(target_dir, all_traces=False)
|
||||
return traces, server_args
|
||||
|
||||
resolved = Path(input_path).resolve()
|
||||
traces, server_args = discover_trace_targets(resolved, all_traces=False)
|
||||
if server_args is None:
|
||||
server_args = load_server_args(resolved)
|
||||
return traces, server_args
|
||||
|
||||
|
||||
def build_mapping_kernel_map(trace_paths: Sequence[Path]) -> dict:
|
||||
# The graph-off mapping trace is only used to learn stable
|
||||
# kernel -> Python/CPU-op attribution. The final percentages still come from
|
||||
# the formal trace.
|
||||
stage_site_stats = defaultdict(
|
||||
lambda: defaultdict(lambda: defaultdict(breakdown_cli.MappingSiteAggregate))
|
||||
)
|
||||
stage_kernel_categories: Dict[str, Dict[str, str]] = defaultdict(dict)
|
||||
global_site_stats = defaultdict(
|
||||
lambda: defaultdict(breakdown_cli.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, _, _ = (
|
||||
breakdown_cli.extract_trace_data(trace)
|
||||
)
|
||||
cpu_ops_by_external_id = breakdown_cli.build_cpu_op_index(cpu_ops)
|
||||
launches_by_correlation = breakdown_cli.build_launch_index(launch_events)
|
||||
local_site_stats = breakdown_cli.aggregate_kernel_sites(
|
||||
kernels,
|
||||
cpu_ops_by_external_id,
|
||||
python_frames,
|
||||
launches_by_correlation=launches_by_correlation,
|
||||
)
|
||||
stage = parse_stage(trace_path)
|
||||
kernel_categories = {
|
||||
kernel.canonical_name: kernel.category for kernel in kernels
|
||||
}
|
||||
breakdown_cli.merge_site_stats(stage_site_stats[stage], local_site_stats)
|
||||
breakdown_cli.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: breakdown_cli.build_stage_payload(
|
||||
dict(site_stats), stage_kernel_categories.get(stage, {})
|
||||
)
|
||||
for stage, site_stats in stage_site_stats.items()
|
||||
}
|
||||
global_payload = breakdown_cli.build_stage_payload(
|
||||
dict(global_site_stats), global_kernel_categories
|
||||
)
|
||||
return {"stages": stage_payloads, "global": global_payload}
|
||||
|
||||
|
||||
def stage_index(stage: str) -> int:
|
||||
return {"extend": 0, "prefill": 0, "decode": 1, "all": 2}.get(stage, 99)
|
||||
|
||||
|
||||
def stage_display(stage: str) -> str:
|
||||
return breakdown_cli.stage_label(stage)
|
||||
|
||||
|
||||
def pick_trace_for_stage(stage_to_trace: Dict[str, Path], stage: str) -> Optional[Path]:
|
||||
if stage in stage_to_trace:
|
||||
return stage_to_trace[stage]
|
||||
if "all" in stage_to_trace:
|
||||
return stage_to_trace["all"]
|
||||
if len(stage_to_trace) == 1:
|
||||
return next(iter(stage_to_trace.values()))
|
||||
return None
|
||||
|
||||
|
||||
def build_stage_trace_map(trace_paths: Sequence[Path]) -> Dict[str, Path]:
|
||||
stage_map: Dict[str, Path] = {}
|
||||
for trace_path in sorted(
|
||||
trace_paths, key=lambda item: (stage_index(parse_stage(item)), item.name)
|
||||
):
|
||||
stage_map[parse_stage(trace_path)] = trace_path
|
||||
return stage_map
|
||||
|
||||
|
||||
def render_kernel_table(rows: Sequence[dict]) -> List[str]:
|
||||
lines = [
|
||||
"| Stage | Kernel | Category | GPU time | Share | Launches | Python location (site share) | CPU op |",
|
||||
"| --- | --- | --- | ---: | ---: | ---: | --- | --- |",
|
||||
]
|
||||
for row in rows:
|
||||
lines.append(
|
||||
"| {stage} | {kernel} | {category} | {gpu_time} | {share:.1f}% | {launches} | {location} | {cpu_op} |".format(
|
||||
stage=breakdown_cli.escape_md_cell(stage_display(row["stage"])),
|
||||
kernel=breakdown_cli.escape_md_cell(row["kernel"]),
|
||||
category=breakdown_cli.escape_md_cell(row["category"]),
|
||||
gpu_time=breakdown_cli.format_ms(row["total_us"]),
|
||||
share=row["share_pct"],
|
||||
launches=row["launches"],
|
||||
location=breakdown_cli.escape_md_cell(row["location"]),
|
||||
cpu_op=breakdown_cli.escape_md_cell(row["cpu_op"]),
|
||||
)
|
||||
)
|
||||
return lines
|
||||
|
||||
|
||||
def render_overlap_table(rows: Sequence[dict]) -> List[str]:
|
||||
lines = [
|
||||
"| Stage | Priority | Verdict | Kernel | Python scope | Formal signal | Dep risk | Recommendation |",
|
||||
"| --- | --- | --- | --- | --- | --- | --- | --- |",
|
||||
]
|
||||
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(
|
||||
[
|
||||
breakdown_cli.escape_md_cell(stage_display(row["stage"])),
|
||||
row["priority"],
|
||||
row["verdict"],
|
||||
breakdown_cli.escape_md_cell(row["kernel"]),
|
||||
breakdown_cli.escape_md_cell(row["python_scope"]),
|
||||
breakdown_cli.escape_md_cell(formal_signal),
|
||||
overlap_cli.dependency_risk_label(row["dependency_signal"]),
|
||||
row["recommendation"],
|
||||
]
|
||||
)
|
||||
+ " |"
|
||||
)
|
||||
return lines
|
||||
|
||||
|
||||
def render_fuse_table(rows: Sequence[dict]) -> List[str]:
|
||||
lines = [
|
||||
"| Stage | Pattern | Confidence | Related GPU time | Share | Evidence kernels | Current kernel Python location | Candidate fused Python path | Rationale |",
|
||||
"| --- | --- | --- | ---: | ---: | --- | --- | --- | --- |",
|
||||
]
|
||||
if not rows:
|
||||
lines.append(
|
||||
"| - | No medium-confidence source-backed fusion opportunity matched this trace. | - | - | - | - | - | - | - |"
|
||||
)
|
||||
return lines
|
||||
for row in rows:
|
||||
lines.append(
|
||||
"| {stage} | {pattern} | {confidence} | {gpu_time} | {share:.1f}% | {evidence} | {current_locations} | {candidate_path} | {rationale} |".format(
|
||||
stage=breakdown_cli.escape_md_cell(stage_display(row["stage"])),
|
||||
pattern=breakdown_cli.escape_md_cell(row["pattern"]),
|
||||
confidence=breakdown_cli.escape_md_cell(row["confidence"]),
|
||||
gpu_time=breakdown_cli.format_ms(row["related_us"]),
|
||||
share=row["share_pct"],
|
||||
evidence=breakdown_cli.escape_md_cell(row["evidence"]),
|
||||
current_locations=breakdown_cli.escape_md_cell(
|
||||
row["current_locations"]
|
||||
),
|
||||
candidate_path=breakdown_cli.escape_md_cell(row["candidate_path"]),
|
||||
rationale=breakdown_cli.escape_md_cell(row["rationale"]),
|
||||
)
|
||||
)
|
||||
return lines
|
||||
|
||||
|
||||
def run_triage(args: argparse.Namespace) -> int:
|
||||
mapping_traces, mapping_server_args = resolve_profile_targets(
|
||||
label="mapping",
|
||||
input_path=args.mapping_input,
|
||||
url=args.mapping_url,
|
||||
output_dir=args.mapping_output_dir,
|
||||
profile_prefix=args.mapping_profile_prefix,
|
||||
args=args,
|
||||
)
|
||||
formal_traces, formal_server_args = resolve_profile_targets(
|
||||
label="formal",
|
||||
input_path=args.formal_input,
|
||||
url=args.formal_url,
|
||||
output_dir=args.formal_output_dir,
|
||||
profile_prefix=args.formal_profile_prefix,
|
||||
args=args,
|
||||
)
|
||||
|
||||
mapping_kernel_map = build_mapping_kernel_map(mapping_traces)
|
||||
|
||||
kernel_rows_rendered: List[dict] = []
|
||||
fuse_rows_rendered: List[dict] = []
|
||||
|
||||
for formal_trace in formal_traces:
|
||||
trace = load_trace_json(formal_trace)
|
||||
kernels, _, _, _, _, _ = breakdown_cli.extract_trace_data(trace)
|
||||
if not kernels:
|
||||
continue
|
||||
stage = parse_stage(formal_trace)
|
||||
total_us = sum(kernel.dur for kernel in kernels)
|
||||
kernel_stats = breakdown_cli.aggregate(
|
||||
kernels, key_fn=lambda item: item.canonical_name
|
||||
)
|
||||
kernel_categories = {
|
||||
kernel.canonical_name: kernel.category for kernel in kernels
|
||||
}
|
||||
full_kernel_rows = breakdown_cli.build_kernel_rows(
|
||||
stage=stage,
|
||||
kernel_stats=kernel_stats,
|
||||
kernel_categories=kernel_categories,
|
||||
local_stage_payload=mapping_kernel_map.get("stages", {}).get(
|
||||
stage, {"kernels": {}}
|
||||
),
|
||||
external_kernel_map=mapping_kernel_map,
|
||||
)
|
||||
visible_kernel_rows = breakdown_cli.limit_kernel_rows(
|
||||
full_kernel_rows, args.kernel_table_limit
|
||||
)
|
||||
for row in visible_kernel_rows:
|
||||
kernel_rows_rendered.append(
|
||||
{
|
||||
"stage": stage,
|
||||
"kernel": row.name,
|
||||
"category": row.category,
|
||||
"total_us": row.total_us,
|
||||
"share_pct": breakdown_cli.pct(row.total_us, total_us),
|
||||
"launches": row.aggregate.count,
|
||||
"location": row.location,
|
||||
"cpu_op": row.cpu_op,
|
||||
}
|
||||
)
|
||||
for item in breakdown_cli.detect_fusion_opportunities(
|
||||
stage=stage,
|
||||
kernel_rows=full_kernel_rows,
|
||||
total_us=total_us,
|
||||
server_args=formal_server_args or mapping_server_args,
|
||||
):
|
||||
fuse_rows_rendered.append(
|
||||
{
|
||||
"stage": stage,
|
||||
"pattern": item.pattern,
|
||||
"confidence": item.confidence,
|
||||
"related_us": item.related_us,
|
||||
"share_pct": breakdown_cli.pct(item.related_us, total_us),
|
||||
"evidence": item.evidence,
|
||||
"current_locations": item.current_locations,
|
||||
"candidate_path": item.candidate_path,
|
||||
"rationale": item.rationale,
|
||||
}
|
||||
)
|
||||
|
||||
mapping_stage_map = build_stage_trace_map(mapping_traces)
|
||||
formal_stage_map = build_stage_trace_map(formal_traces)
|
||||
overlap_rows_rendered: List[dict] = []
|
||||
for stage in sorted(formal_stage_map, key=stage_index):
|
||||
formal_trace = formal_stage_map[stage]
|
||||
mapping_trace = pick_trace_for_stage(mapping_stage_map, stage)
|
||||
if mapping_trace is None:
|
||||
continue
|
||||
mapping_trace_json = load_trace_json(mapping_trace)
|
||||
mapping_events, mapping_pid = overlap_cli.extract_kernel_events(
|
||||
mapping_trace_json, args.pid_substring
|
||||
)
|
||||
if not mapping_events:
|
||||
continue
|
||||
formal_trace_json = load_trace_json(formal_trace)
|
||||
formal_events, formal_pid = overlap_cli.extract_kernel_events(
|
||||
formal_trace_json, args.pid_substring
|
||||
)
|
||||
if not formal_events:
|
||||
continue
|
||||
mapping_bundle = overlap_cli.TraceBundle(
|
||||
label=f"mapping-{stage}",
|
||||
trace_path=mapping_trace,
|
||||
server_args=mapping_server_args,
|
||||
raw_events=mapping_trace_json.get(
|
||||
"traceEvents",
|
||||
mapping_trace_json if isinstance(mapping_trace_json, list) else [],
|
||||
),
|
||||
events=mapping_events,
|
||||
pid=mapping_pid,
|
||||
)
|
||||
formal_bundle = overlap_cli.TraceBundle(
|
||||
label=f"formal-{stage}",
|
||||
trace_path=formal_trace,
|
||||
server_args=formal_server_args,
|
||||
raw_events=formal_trace_json.get(
|
||||
"traceEvents",
|
||||
formal_trace_json if isinstance(formal_trace_json, list) else [],
|
||||
),
|
||||
events=formal_events,
|
||||
pid=formal_pid,
|
||||
)
|
||||
formal_bundle.overlap_stats = overlap_cli.analyze_overlap(formal_bundle.events)
|
||||
aggregates = overlap_cli.aggregate_events(formal_bundle.events)
|
||||
source_map = overlap_cli.build_kernel_source_map(mapping_bundle)
|
||||
stage_rows = overlap_cli.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:
|
||||
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"Mapping traces: {', '.join(str(path) for path in mapping_traces)}")
|
||||
lines.append(f"Formal traces: {', '.join(str(path) for path in formal_traces)}")
|
||||
if formal_server_args or mapping_server_args:
|
||||
server_args = formal_server_args or mapping_server_args
|
||||
model = server_args.get("model_path") or server_args.get("model")
|
||||
if model:
|
||||
lines.append(f"Model: {model}")
|
||||
lines.append("")
|
||||
lines.append("Kernel Table")
|
||||
lines.extend(render_kernel_table(kernel_rows_rendered))
|
||||
lines.append("")
|
||||
lines.append("Overlap Opportunity Table")
|
||||
lines.extend(render_overlap_table(overlap_rows_rendered))
|
||||
lines.append("")
|
||||
lines.append("Fuse Opportunity Table")
|
||||
lines.extend(render_fuse_table(fuse_rows_rendered))
|
||||
print("\n".join(lines).rstrip())
|
||||
return 0
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
argv = list(argv or sys.argv[1:])
|
||||
top_parser = build_top_level_parser()
|
||||
|
||||
if not argv or argv[0] in {"-h", "--help"}:
|
||||
top_parser.print_help()
|
||||
return 0
|
||||
|
||||
command = argv[0]
|
||||
remainder = argv[1:]
|
||||
|
||||
if command == "breakdown":
|
||||
return breakdown_cli.main(remainder)
|
||||
if command == "overlap":
|
||||
return overlap_cli.main(remainder)
|
||||
if command == "triage":
|
||||
return run_triage(parse_triage_args(remainder))
|
||||
if command == "perfetto-fix":
|
||||
args = parse_perfetto_fix_args(remainder)
|
||||
output_path = write_perfetto_compatible_trace(
|
||||
input_path=Path(args.input),
|
||||
output_path=Path(args.output).resolve() if args.output else None,
|
||||
)
|
||||
print(f"Perfetto-friendly trace written to: {output_path}")
|
||||
return 0
|
||||
|
||||
top_parser.error(f"Unknown command: {command}")
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(sys.argv[1:]))
|
||||
@@ -0,0 +1,378 @@
|
||||
"""Shared helpers for SGLang torch-profiler skill scripts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import json
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from collections import Counter, defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
|
||||
from urllib import request
|
||||
|
||||
STAGE_ORDER = {"extend": 0, "prefill": 0, "decode": 1, "all": 2}
|
||||
TRACE_METADATA_NAMES = {
|
||||
"process_name",
|
||||
"thread_name",
|
||||
"process_sort_index",
|
||||
"thread_sort_index",
|
||||
}
|
||||
NON_KERNEL_TRACE_CATEGORIES = ("python_function", "cpu_op", "trace")
|
||||
PYTHON_SCOPE_NAME_PREFIXES = ("python/", "nn.module:")
|
||||
|
||||
|
||||
def normalize_text(value: object) -> str:
|
||||
return re.sub(r"\s+", " ", str(value)).strip()
|
||||
|
||||
|
||||
def normalize_repo_relative_path(path: object) -> str:
|
||||
text = normalize_text(path).replace("\\", "/")
|
||||
for marker in ("python/sglang/", "sgl_kernel/"):
|
||||
idx = text.find(marker)
|
||||
if idx != -1:
|
||||
return text[idx:].lstrip("/")
|
||||
idx = text.find("sglang/")
|
||||
if idx != -1:
|
||||
return ("python/" + text[idx:]).lstrip("/")
|
||||
return text.lstrip("/")
|
||||
|
||||
|
||||
def contains_any_keyword(text: str, keywords: Iterable[str]) -> bool:
|
||||
return any(keyword in text for keyword in keywords)
|
||||
|
||||
|
||||
def coerce_optional_int(value: object) -> Optional[int]:
|
||||
if value in (None, "", "None"):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
return int(value) if value.is_integer() else None
|
||||
try:
|
||||
return int(str(value))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def extract_trace_events(trace: object) -> Sequence[dict]:
|
||||
if isinstance(trace, dict):
|
||||
events = trace.get("traceEvents", [])
|
||||
return events if isinstance(events, list) else []
|
||||
if isinstance(trace, list):
|
||||
return trace
|
||||
return []
|
||||
|
||||
|
||||
def is_trace_metadata_name(name: object) -> bool:
|
||||
return str(name) in TRACE_METADATA_NAMES
|
||||
|
||||
|
||||
def is_complete_duration_event(event: dict) -> bool:
|
||||
if event.get("ph") != "X":
|
||||
return False
|
||||
dur = event.get("dur")
|
||||
ts = event.get("ts")
|
||||
if dur is None or ts is None:
|
||||
return False
|
||||
try:
|
||||
return float(dur) > 0
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def is_annotation_event(name: object, category: object) -> bool:
|
||||
lowered_name = normalize_text(name).lower()
|
||||
lowered_category = normalize_text(category).lower()
|
||||
return "annotation" in lowered_category or lowered_name.startswith("## call ")
|
||||
|
||||
|
||||
def is_non_kernel_trace_category(category: object) -> bool:
|
||||
lowered_category = normalize_text(category).lower()
|
||||
return any(token in lowered_category for token in NON_KERNEL_TRACE_CATEGORIES)
|
||||
|
||||
|
||||
def looks_like_python_scope_name(name: object) -> bool:
|
||||
lowered_name = normalize_text(name).lower()
|
||||
return ".py(" in lowered_name or lowered_name.startswith(PYTHON_SCOPE_NAME_PREFIXES)
|
||||
|
||||
|
||||
def has_stream_marker(args: Optional[dict]) -> bool:
|
||||
trace_args = args or {}
|
||||
return "stream" in trace_args or "cuda_stream" in trace_args
|
||||
|
||||
|
||||
def load_trace_json(path: Path) -> dict:
|
||||
if path.suffix == ".gz":
|
||||
with gzip.open(path, "rt", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
|
||||
|
||||
def load_server_args(path: Path) -> Optional[dict]:
|
||||
resolved = path.resolve()
|
||||
candidate_dirs: List[Path] = []
|
||||
if resolved.is_file():
|
||||
candidate_dirs.extend([resolved.parent, resolved.parent.parent])
|
||||
else:
|
||||
candidate_dirs.extend([resolved, resolved.parent])
|
||||
|
||||
seen: set[Path] = set()
|
||||
for candidate_dir in candidate_dirs:
|
||||
if candidate_dir in seen:
|
||||
continue
|
||||
seen.add(candidate_dir)
|
||||
candidate = candidate_dir / "server_args.json"
|
||||
if candidate.exists():
|
||||
with open(candidate, "r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
return None
|
||||
|
||||
|
||||
def parse_stage(path: Path) -> str:
|
||||
name = path.name.lower()
|
||||
if "-extend" in name or "-prefill" in name:
|
||||
return "extend"
|
||||
if "-decode" in name:
|
||||
return "decode"
|
||||
return "all"
|
||||
|
||||
|
||||
def parse_tp_rank(path: Path) -> Optional[int]:
|
||||
match = re.search(r"TP-(\d+)", path.name)
|
||||
return int(match.group(1)) if match else None
|
||||
|
||||
|
||||
def newest_trace_dir(path: Path) -> Path:
|
||||
if path.is_file():
|
||||
return path.parent
|
||||
direct = list(path.glob("*.trace.json")) + list(path.glob("*.trace.json.gz"))
|
||||
if direct:
|
||||
return path
|
||||
child_candidates = [item for item in path.rglob("*") if item.is_dir()]
|
||||
trace_dirs = [
|
||||
candidate
|
||||
for candidate in child_candidates
|
||||
if list(candidate.glob("*.trace.json"))
|
||||
or list(candidate.glob("*.trace.json.gz"))
|
||||
]
|
||||
if not trace_dirs:
|
||||
raise FileNotFoundError(f"No trace files found under {path}")
|
||||
trace_dirs.sort(key=lambda item: item.stat().st_mtime)
|
||||
return trace_dirs[-1]
|
||||
|
||||
|
||||
def discover_trace_targets(
|
||||
path: Path, all_traces: bool
|
||||
) -> Tuple[List[Path], Optional[dict]]:
|
||||
if path.is_file():
|
||||
return [path], load_server_args(path)
|
||||
|
||||
trace_dir = newest_trace_dir(path)
|
||||
traces = sorted(
|
||||
list(trace_dir.glob("*.trace.json")) + list(trace_dir.glob("*.trace.json.gz")),
|
||||
key=lambda item: item.stat().st_mtime,
|
||||
)
|
||||
if not traces:
|
||||
raise FileNotFoundError(f"No trace files found under {trace_dir}")
|
||||
|
||||
non_merged = [trace for trace in traces if not trace.name.startswith("merged-")]
|
||||
selected = non_merged or traces
|
||||
if not all_traces:
|
||||
ranks = sorted(
|
||||
{
|
||||
rank
|
||||
for rank in (parse_tp_rank(trace) for trace in selected)
|
||||
if rank is not None
|
||||
}
|
||||
)
|
||||
if ranks:
|
||||
rank = 0 if 0 in ranks else ranks[0]
|
||||
selected = [trace for trace in selected if parse_tp_rank(trace) == rank]
|
||||
grouped: Dict[str, List[Path]] = defaultdict(list)
|
||||
for trace in selected:
|
||||
grouped[parse_stage(trace)].append(trace)
|
||||
selected = [
|
||||
sorted(group, key=lambda item: item.stat().st_mtime)[-1]
|
||||
for group in grouped.values()
|
||||
]
|
||||
|
||||
selected.sort(key=lambda item: (STAGE_ORDER.get(parse_stage(item), 99), item.name))
|
||||
return selected, load_server_args(trace_dir)
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict, timeout: float = 60.0) -> Optional[dict]:
|
||||
req = request.Request(
|
||||
url=url,
|
||||
data=json.dumps(payload).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with request.urlopen(req, timeout=timeout) as response:
|
||||
raw = response.read()
|
||||
return json.loads(raw.decode("utf-8")) if raw else None
|
||||
|
||||
|
||||
def send_probe_request(
|
||||
url: str, prompt: str, max_new_tokens: int, sampling_seed: int
|
||||
) -> None:
|
||||
payload = {
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"sampling_seed": sampling_seed,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
},
|
||||
"stream": False,
|
||||
}
|
||||
post_json(url.rstrip("/") + "/generate", payload, timeout=300.0)
|
||||
|
||||
|
||||
def run_profiler(
|
||||
url: str,
|
||||
output_dir: Optional[str],
|
||||
num_steps: int,
|
||||
profile_by_stage: bool,
|
||||
merge_profiles: bool,
|
||||
profile_prefix: Optional[str],
|
||||
probe_requests: int,
|
||||
probe_prompt: str,
|
||||
probe_max_new_tokens: Optional[int],
|
||||
probe_delay: float,
|
||||
start_step: Optional[int] = None,
|
||||
) -> Path:
|
||||
if output_dir is None:
|
||||
output_dir = tempfile.mkdtemp(prefix="sglang-torch-profile-")
|
||||
output_path = Path(output_dir).resolve()
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"sglang.profiler",
|
||||
"--url",
|
||||
url,
|
||||
"--output-dir",
|
||||
str(output_path),
|
||||
"--num-steps",
|
||||
str(num_steps),
|
||||
"--cpu",
|
||||
"--gpu",
|
||||
"--merge-profiles" if merge_profiles else "--no-merge-profiles",
|
||||
"--profile-by-stage" if profile_by_stage else "--no-profile-by-stage",
|
||||
]
|
||||
if profile_prefix:
|
||||
cmd.extend(["--profile-prefix", profile_prefix])
|
||||
if start_step is not None:
|
||||
cmd.extend(["--start-step", str(start_step)])
|
||||
|
||||
profiler_proc = subprocess.Popen(cmd)
|
||||
try:
|
||||
if probe_requests > 0:
|
||||
time.sleep(max(0.0, probe_delay))
|
||||
effective_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8)
|
||||
for request_idx in range(probe_requests):
|
||||
send_probe_request(
|
||||
url=url,
|
||||
prompt=probe_prompt,
|
||||
max_new_tokens=effective_max_new_tokens,
|
||||
sampling_seed=request_idx,
|
||||
)
|
||||
if profiler_proc.poll() is not None:
|
||||
break
|
||||
return_code = profiler_proc.wait()
|
||||
finally:
|
||||
if profiler_proc.poll() is None:
|
||||
profiler_proc.kill()
|
||||
|
||||
if return_code != 0:
|
||||
raise subprocess.CalledProcessError(return_code, cmd)
|
||||
|
||||
deadline = time.time() + 15.0
|
||||
while time.time() < deadline:
|
||||
child_dirs = [path for path in output_path.iterdir() if path.is_dir()]
|
||||
if child_dirs:
|
||||
child_dirs.sort(key=lambda path: path.stat().st_mtime)
|
||||
newest_child = child_dirs[-1]
|
||||
if any(newest_child.glob("*.trace.json*")):
|
||||
return newest_child
|
||||
time.sleep(0.5)
|
||||
|
||||
child_dirs = [path for path in output_path.iterdir() if path.is_dir()]
|
||||
if child_dirs:
|
||||
child_dirs.sort(key=lambda path: path.stat().st_mtime)
|
||||
return child_dirs[-1]
|
||||
return output_path
|
||||
|
||||
|
||||
def select_heaviest_pid(
|
||||
events: Sequence[dict],
|
||||
event_filter: Callable[[dict], bool],
|
||||
pid_substring: Optional[str] = None,
|
||||
preferred_substrings: Iterable[str] = (),
|
||||
) -> Optional[str]:
|
||||
durations: Counter = Counter()
|
||||
for event in events:
|
||||
if not event_filter(event):
|
||||
continue
|
||||
pid = str(event.get("pid"))
|
||||
if pid_substring and pid_substring not in pid:
|
||||
continue
|
||||
durations[pid] += float(event["dur"])
|
||||
if not durations:
|
||||
return None
|
||||
|
||||
for substring in preferred_substrings:
|
||||
preferred = [pid for pid in durations if substring in pid]
|
||||
if preferred:
|
||||
return max(preferred, key=lambda pid: durations[pid])
|
||||
return max(durations, key=lambda pid: durations[pid])
|
||||
|
||||
|
||||
def write_perfetto_compatible_trace(
|
||||
input_path: Path, output_path: Optional[Path] = None
|
||||
) -> Path:
|
||||
resolved_input = input_path.resolve()
|
||||
if output_path is None:
|
||||
output_name = f"perfetto-compatible-{resolved_input.name}"
|
||||
output_path = resolved_input.with_name(output_name)
|
||||
|
||||
trace = load_trace_json(resolved_input)
|
||||
output = {key: value for key, value in trace.items() if key != "traceEvents"}
|
||||
output["traceEvents"] = _perfetto_fix_events(trace.get("traceEvents", []))
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if str(output_path).endswith(".gz"):
|
||||
with gzip.open(output_path, "wt", encoding="utf-8") as handle:
|
||||
json.dump(output, handle)
|
||||
else:
|
||||
with open(output_path, "w", encoding="utf-8") as handle:
|
||||
json.dump(output, handle)
|
||||
return output_path
|
||||
|
||||
|
||||
def _perfetto_fix_events(events: Sequence[dict]) -> List[dict]:
|
||||
fixed_events = [dict(event) for event in events]
|
||||
last_end_time_of_pid_tid: Dict[Tuple[str, str], float] = defaultdict(lambda: -1.0)
|
||||
|
||||
for event in fixed_events:
|
||||
if event.get("ph") != "X" or not _is_perfetto_overlap_interest_event(event):
|
||||
continue
|
||||
pid = str(event.get("pid"))
|
||||
tid = str(event.get("tid"))
|
||||
ts = float(event.get("ts", 0.0))
|
||||
dur = float(event.get("dur", 0.0))
|
||||
while ts < last_end_time_of_pid_tid[(pid, tid)]:
|
||||
tid = f"{tid}_hack"
|
||||
event["tid"] = tid
|
||||
last_end_time_of_pid_tid[(pid, tid)] = ts + dur
|
||||
return fixed_events
|
||||
|
||||
|
||||
def _is_perfetto_overlap_interest_event(event: dict) -> bool:
|
||||
args = event.get("args") or {}
|
||||
return "registers per thread" in args
|
||||
Reference in New Issue
Block a user