diff --git a/.claude/skills/sglang-torch-profiler-analysis/SKILL.md b/.claude/skills/sglang-torch-profiler-analysis/SKILL.md new file mode 100644 index 000000000..1567df1a7 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/SKILL.md @@ -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 diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/fuse-overlap-catalog.md b/.claude/skills/sglang-torch-profiler-analysis/references/fuse-overlap-catalog.md new file mode 100644 index 000000000..319fc2047 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/references/fuse-overlap-catalog.md @@ -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*`
`npu_add_rms_norm`
`add_rmsnorm_bias`
`gemma_fused_add_rmsnorm`
residual add right before norm | `python/sglang/srt/layers/layernorm.py`
`python/sglang/srt/layers/quantization/modelslim/modelslim.py` | Shared CUDA / ROCm / CPU / NPU fused add-RMSNorm implementations, including Gemma and NPU-bias variants | Treat split residual add + RMSNorm as an existing cross-backend fusion first, not a new idea. | +| FlashInfer unified `allreduce_fusion` | `cross_device_reduce_1stage*`
`all_reduce`
`FusedAddRMSNormKernel`
`rmsnorm*` | `python/sglang/srt/layers/flashinfer_comm_fusion.py`
`python/sglang/srt/layers/layernorm.py::forward_with_allreduce_fusion`
`python/sglang/srt/layers/communicator.py::apply_flashinfer_allreduce_fusion` | FlashInfer workspace creation plus `allreduce_fusion(..., pattern=AllReduceFusionPattern.kARResidualRMSNorm, ...)` | First suspect missing / disabled / unsupported FlashInfer allreduce fusion, not a brand new TP fusion idea. | +| AITER allreduce fusion | ROCm all-reduce plus RMSNorm still split | `python/sglang/srt/layers/layernorm.py::forward_with_allreduce_fusion`
`python/sglang/srt/distributed/communication_op.py::tensor_model_parallel_fused_allreduce_rmsnorm`
`python/sglang/srt/layers/communicator.py::apply_aiter_all_reduce_fusion` | ROCm-side fused TP all-reduce + RMSNorm with fallback to plain all-reduce plus norm | On AMD, rule out existing AITER fusion before proposing a new communication fusion. | +| Fused activation-and-mul (`SwiGLU` / `GeGLU`) | `silu_and_mul`
`gelu_and_mul`
`npu_swiglu` | `python/sglang/srt/layers/activation.py` | Single op covers activation plus elementwise multiply across CUDA / CPU / NPU / XPU backends | Treat separate activation + mul on packed MLP outputs as missing existing fusion. | +| Fused dual residual RMSNorm | residual add plus two RMSNorm-like kernels around Grok blocks | `python/sglang/srt/layers/elementwise.py::fused_dual_residual_rmsnorm`
`python/sglang/srt/models/grok.py` | One Triton kernel computes intermediate residual update and next RMSNorm output together | On Grok-like residual layouts, treat split residual + norm as missing existing fusion. | +| In-place QK RMSNorm | split `q_norm` / `k_norm` kernels | `python/sglang/srt/models/utils.py::apply_qk_norm`
`python/sglang/jit_kernel/norm.py::fused_inplace_qknorm` | In-place JIT QK norm plus optional `alt_stream` overlap for K | Check shape, dtype, deterministic mode, and in-place legality before proposing a new QK fuse. | +| MiniMax TP fused QK RMSNorm | `MiniMaxM2RMSNormTP`
`rms_sumsq_serial`
`rms_apply_serial`
`forward_qk` | `python/sglang/srt/models/minimax_m2.py` | Triton kernels compute Q / K sumsq together, TP all-reduces shared stats, then apply both RMSNorms together | On MiniMax traces, separate Q norm and K norm are usually a missed model-specific Triton fusion. | +| Fused QK RMSNorm + RoPE | `qknorm*` + `rope*` + `rotary*` as separate steps | `python/sglang/jit_kernel/fused_qknorm_rope.py`
`python/sglang/srt/models/qwen3_moe.py` | One JIT kernel applies QK RMSNorm and RoPE in-place on packed QKV | For compatible LLMs, classify split QK norm + RoPE as a missing existing fusion. | +| Fused QK RoPE reshape + KV cache write | `fused_qk_rope_reshape_and_cache*`
RoPE followed by reshape / cache DtoD | `python/sglang/srt/layers/attention/utils.py::fused_qk_rope_reshape_and_cache` | One Triton kernel applies RoPE to Q / K, reshapes cache layout, and writes K / V directly to paged cache | Treat separate RoPE + reshape + cache-write ladders as an existing attention-prep fusion family. | +| Fused RoPE + KV cache store | `fused_set_kv_buffer`
RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`
`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`
`cache_seqlens_int32`
`cu_seqlens_k`
`page_table`
`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`
`fused_metadata_copy_multi`
`fused_nsa_cache_seqlens`
`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`
`fused_qkv_a_proj_with_mqa`
`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`
`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`
`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`
`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`
`mrope`
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`
`python/sglang/srt/models/qwen3.py`
`python/sglang/srt/models/qwen3_moe.py`
`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`
`fp8 kv cache write`
`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*`
`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`
`python/sglang/srt/mem_cache/memory_pool.py` | MLA / NSA KV pool path can quantize K and write directly into KV storage without a separate concat-and-quant chain | Treat standalone quant + KV-buffer write on MLA paths as missing existing fusion first. | +| Fused MoE router / top-k / softcapping | `FusedMoeRouter`
`fused_moe_router*`
router GEMM + `topk` + `tanh` | `python/sglang/srt/layers/moe/router.py` | Single fused router kernel covers router matmul, softcapping, and top-k selection | Treat exposed router matmul + softcap + top-k chains as an existing MoE fusion family. | +| Fused MoE grouped-topk / gate kernels | `fused_topk_deepseek`
`moe_fused_gate`
`aiter_fused_topk`
`kimi_k2_moe_fused_gate` | `python/sglang/srt/layers/moe/topk.py` | CUDA / ROCm / FlashInfer kernels fuse bias, grouped-topk, renorm, and routed scaling into one gate op | Check backend / model eligibility before proposing a novel router-gate fusion. | +| Fused MoE dispatch / permute / combine | token permutation
dispatch / combine
grouped top-k
many small MoE support kernels | `python/sglang/srt/layers/moe/fused_moe_triton/layer.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | `FusedMoE` plus DeepEP / FlashInfer / FuseEP / standard dispatch backends and `permute_fusion=True` | First ask whether the model is missing an existing `FusedMoE`-style path or backend-specific dispatcher path. | +| Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. | +| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/jit_kernel/nvfp4.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | +| DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`
`fused_rms_mxfp4_quant`
`fused_flatten_fp8_group_quant`
`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
`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`
`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`
`act_quant`
`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`
`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`
`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`
`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`
`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*`
`qkvz_proj`
`ba_proj`
`qkvabz_proj`
`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`
`python/sglang/srt/models/qwen3_next.py`
`python/sglang/srt/models/qwen3_5.py`
`python/sglang/srt/models/kimi_linear.py`
`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`
`softplus`
`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`
`layer_norm_gated_fwd` | `python/sglang/srt/layers/attention/fla/fused_norm_gate.py`
`python/sglang/srt/models/qwen3_next.py`
`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`
`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`
`scaled_dot_kkt`
`solve_tril`
`recompute_w_u` | `python/sglang/srt/layers/attention/fla/chunk_fwd.py`
`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`
`fused_recurrent_gated_delta_rule_update`
`fused_kda_gate` | `python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py`
`python/sglang/srt/layers/attention/fla/fused_recurrent.py`
`python/sglang/srt/models/kimi_linear.py`
`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`
`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`
`_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`
`python/sglang/srt/models/qwen3.py`
`python/sglang/srt/models/qwen3_next.py`
`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`
`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`
`_comm_stream`
`dispatch`
`combine`
`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`
`_scatter_stream`
`staging` | `python/sglang/srt/disaggregation/common/staging_handler.py`
`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`
`qwen3_moe.py`
`glm4_moe.py`
`bailing_moe.py`
`llada2.py`
`grok.py`
`olmo2.py`
`step3p5.py`
`longcat_flash.py`
`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`
`python/sglang/srt/layers/attention/vision.py`
`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`
`multimodal_rotary_embedding_cpu`
`npu_mrope`
`MRotaryEmbedding` | `python/sglang/srt/layers/rotary_embedding/mrope.py`
`python/sglang/srt/layers/rotary_embedding/triton_kernels.py`
`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`
`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`
`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`
`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`
`python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` | `fused_add_rmsnorm(...)` and `triton_one_pass_rms_norm(...)` | For short hidden-size diffusion blocks, this is already an established fusion family. | +| Fused diffusion QK norm + RoPE | split QK norm and RoPE in diffusion attention blocks | `python/sglang/jit_kernel/diffusion/qknorm_rope.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py::apply_qk_norm_rope` | `fused_inplace_qknorm_rope(...)`, with fallback to QK norm plus `apply_flashinfer_rope_qk_inplace(...)` | Distinguish between missing fused qknorm + rope and the existing FlashInfer RoPE fallback. | +| Z-Image fused `norm(x) * tanh(scale) + shift` | `fused_norm_tanh_mul_add`
`tanh(gate) * rmsnorm(x)` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | CuTeDSL kernel plus runtime helper for Z-Image residual-form modulation | Treat split Z-Image residual-form modulation as a missing existing diffusion fusion, not a novel idea. | +| Z-Image fused residual modulation + next norm-scale | `fused_norm_tanh_mul_add_norm_scale`
`residual + tanh(gate) * rmsnorm(x)`
`ffn_norm1(x) * scale_mlp` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/models/dits/zimage.py` | One CuTeDSL kernel fuses the first residual-form modulation and the next normalization / scale stage | If you see this chain split in Z-Image traces, report it as a missing existing merged fusion family. | +| Nunchaku fused GELU MLP | `_fused_gelu_mlp`
`fused_gelu_mlp` | `python/sglang/multimodal_gen/runtime/models/dits/flux.py` | Nunchaku path fuses `fc1 GEMM + GELU + shift + re-quant + fc2.lora_down` before the second GEMM | Treat split GELU-MLP on Nunchaku checkpoints as an existing fused family, not a new discovery. | + +## 5. Diffusion kernel-overlap and async-communication families + +| 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`
`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`
`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`
`combine`
`fused grouped gemm combine` | `PR #21877`
`python/sglang/srt/layers/moe/ep_moe/flashinfer_cutedsl_moe.py`
`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`
`fp4 kv cache` | `PR #21889`
`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`
`WRITE_PT`
`dequant_fp4_paged_decode` | `PR #21889`
`python/sglang/srt/layers/attention/nsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. | +| PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`
`trtllm_fp8_block_scale_moe` | `PR #21491`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. | +| PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`
`per_token_quant_fp8` | `PR #22005`
`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`
`python/sglang/jit_kernel/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. | +| PR `#21952` Gemma4 fused RMSNorm + residual + scalar | `gemma_rmsnorm_residual_scalar`
`_gemma_rmsnorm_residual_kernel`
`Gemma4` | `PR #21952`
`python/sglang/srt/layers/gemma4_fused_ops.py`
`python/sglang/srt/models/gemma4_causal.py` | Triton kernel fuses decoder post-FF RMSNorm, residual add, and per-layer scalar multiply into one pass | If Gemma4-style post-FF norm + residual + scalar steps appear split, treat them as an in-flight upstream Triton fuse family. | + +## 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`
`combine`
`down_gemm` | `PR #21877`
`python/sglang/srt/server_args.py`
`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | Fused combine eliminates the standalone combine window, so SBO is intentionally disabled when this path is on | If the trace discussion is about combine overlap, first classify it as this upstream fused-overlap family. | + +## 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*`
residual add right before RMSNorm | `vllm/model_executor/layers/layernorm.py`
`vllm/_custom_ops.py`
`csrc/layernorm_kernels.cu`
`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`
`AllReduceFusionPass`
`allreduce + rmsnorm` | `vllm/compilation/passes/fusion/allreduce_rms_fusion.py`
`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`
`fused_add_rms_norm_static_fp8_quant`
`per_token_quant`
`per_group_quant` | `vllm/compilation/passes/fusion/rms_quant_fusion.py`
`vllm/compilation/passes/fusion/rocm_aiter_fusion.py` | Compile-time and ROCm AITER paths fuse RMSNorm or fused-add-RMSNorm with FP8 / FP4 quant output | Treat split norm/add + quant as an upstream fused family, not an unexplored direction. | +| vLLM-origin SiLU+Mul + quant | `ActivationQuantFusionPass`
`SiluMulFp8*`
`Nvfp4`
`rocm_aiter` | `vllm/compilation/passes/fusion/act_quant_fusion.py`
`vllm/compilation/passes/fusion/rocm_aiter_fusion.py` | Activation epilogues fuse `SiLU+Mul` with FP8 / NVFP4 / AITER group quant instead of materializing the BF16 activation first | Treat standalone activation then quant kernels as matching a vLLM-origin precedent. | +| vLLM-origin add + RMSNorm + pad | `fuse_act_padding`
`RocmAiterTritonAddRMSNormPadFusionPass`
`add_rmsnorm_pad` | `vllm/compilation/passes/fusion/rocm_aiter_fusion.py`
`docs/design/fusions.md` | ROCm / AITER path fuses residual add + RMSNorm directly into the padded layout expected by the next kernel | Treat norm-plus-padding ladders as an existing backend-specific fuse family first. | +| vLLM-origin attention + output quant | `fuse_attn_quant`
`AttnQuantFusionPass`
`output_scale`
`output_block_scale` | `vllm/compilation/passes/fusion/attn_quant_fusion.py`
`vllm/v1/attention/backends/`
`docs/design/fusions.md` | Compile-time fusion pushes FP8 / NVFP4 quantization into the attention epilogue on supported Triton / FlashInfer / ROCm / AITER backends | Treat attention-output quant kernels as a known upstream epilogue fusion family before calling them novel. | +| vLLM-origin fused QK RMSNorm + RoPE | `fused_qk_norm_rope`
`QKNormRoPEFusionPass`
`qk norm + rope` | `vllm/compilation/passes/fusion/qk_norm_rope_fusion.py`
`vllm/_custom_ops.py`
`csrc/fused_qknorm_rope_kernel.cu` | Compile-time and direct custom-op paths fuse per-head Q / K RMSNorm with RoPE | Treat split QK norm + RoPE as a clear vLLM-origin precedent. | +| vLLM-origin fused reshape + KV cache write | `reshape_and_cache`
`triton_reshape_and_cache_flash`
`kv cache write` | `vllm/v1/attention/ops/triton_reshape_and_cache_flash.py`
`vllm/v1/attention/backends/triton_attn.py` | Triton cache-update kernels reshape K / V into paged-cache layout and can include FP8 KV-cache scale/write logic | Treat reshape / transpose / cache-write ladders as an existing cache-store fusion family. | +| vLLM-origin fused RoPE + KV cache update | `fuse_rope_kvcache`
`RopeKVCacheFusionPass`
`triton_rope_and_cache` | `vllm/compilation/passes/fusion/rope_kvcache_fusion.py`
`vllm/_aiter_ops.py`
`docs/design/fusions.md` | ROCm / AITER compile-time fusion combines RoPE with paged KV cache update instead of launching them separately | Treat split RoPE + cache-store as a known upstream family, especially on ROCm-like paths. | +| vLLM-origin fused MLA RoPE + concat/cache write | `concat_and_cache_mla_rope_fused`
`mla rope cache` | `vllm/_custom_ops.py`
`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`
`biased_grouped_topk`
`grouped_topk_fused_kernel` | `vllm/_custom_ops.py`
`vllm/_aiter_ops.py`
`vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py`
`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`
`topk_sigmoid`
`topkGating`
`fused_topk` | `vllm/_custom_ops.py`
`vllm/_aiter_ops.py`
`vllm/model_executor/layers/fused_moe/router/fused_topk_router.py`
`vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py`
`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`
`allow_dsv3_router_gemm`
`router logits` | `vllm/_custom_ops.py`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`csrc/moe/dsv3_router_gemm_entry.cu`
`csrc/moe/dsv3_router_gemm_float_out.cu` | Hopper-class CUDA kernel specializes the DeepSeek router linear for small decode batches and can emit FP32 logits directly without a generic GEMM chain | Treat DeepSeek-style router linear paths as an existing upstream specialized fuse, distinct from grouped-topk itself. | +| vLLM-origin GPT-OSS router GEMM | `gpt_oss_router_gemm`
`router gemm` | `vllm/_custom_ops.py`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`csrc/moe/gpt_oss_router_gemm.cu` | Model-specific CUDA kernel replaces the router linear plus bias path with one specialized GEMM op | Treat GPT-OSS-style router linear chains as an existing upstream specialized fuse. | +| vLLM-origin DeepSeek min-latency fused QKV-A projection | `dsv3_fused_a_gemm`
`fused_qkv_a_proj`
`q_a_proj` | `vllm/model_executor/models/deepseek_v2.py`
`vllm/_custom_ops.py`
`csrc/dsv3_fused_a_gemm.cu` | Hopper-class CUDA kernel replaces the tiny-batch DeepSeek QKV-A projection path with one specialized min-latency GEMM instead of a generic linear launch | Treat small-batch DeepSeek QKV-A projection ladders as a known upstream fused kernel family first. | +| vLLM-origin CUTLASS scaled MM with scale / bias epilogue | `cutlass_scaled_mm`
`cutlass_scaled_mm_azp`
`scaled mm` | `vllm/_custom_ops.py`
`vllm/model_executor/kernels/linear/scaled_mm/cutlass.py`
`csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu` | CUTLASS kernels fuse activation scales, weight scales, matmul, and optional bias / AZP epilogues | Treat separate scale-mul + GEMM + bias ladders as a vLLM-origin fused linear family first. | +| vLLM-origin fused MoE expert execution | `cpu_fused_moe`
`rocm_aiter_fused_moe`
`FusedMoE` | `vllm/model_executor/layers/fused_moe/layer.py`
`vllm/model_executor/layers/fused_moe/cpu_fused_moe.py`
`vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py`
`vllm/_aiter_ops.py` | MoE backends on CUDA / ROCm / CPU already collapse packed expert execution into fused expert kernels rather than per-expert eager GEMMs | Treat exposed expert-side tiny GEMM ladders as matching an upstream fused-MoE family. | +| vLLM-origin fused MoE LoRA | `fused_moe_lora`
`fused_moe_lora_fp8`
`w13_shrink`
`w2_expand` | `vllm/lora/ops/triton_ops/fused_moe_lora_op.py`
`vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py`
`vllm/lora/layers/fused_moe.py` | Triton kernels fuse LoRA shrink / expand work into MoE expert execution, including FP8 variants | Treat MoE-LoRA adapter work as an upstream fused family before proposing a brand new kernel. | +| vLLM-origin ViT fused bilinear position-embedding interpolation | `triton_pos_embed_interpolate`
`bilinear_pos_embed`
`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`
`fused_matmul_reduce_scatter`
`fused_all_gather_matmul` | `vllm/compilation/passes/fusion/collective_fusion.py`
`docs/design/fusions.md` | AsyncTP overlaps GEMM with reduce-scatter / all-gather via symmetric-memory collectives | Treat GEMM+comm windows as a clear vLLM-origin overlap precedent first. | +| vLLM-origin Sequence Parallelism staging | `enable_sp`
`ReduceScatter`
`AllGather`
`SequenceParallelismPass` | `vllm/compilation/passes/fusion/sequence_parallelism.py`
`docs/design/fusions.md` | Sequence-parallel rewrites all-reduce into RS -> local norm -> AG so later passes can overlap comm and compute | Treat RS / AG staging around norm blocks as an upstream overlap-enabling family. | +| vLLM-origin shared-expert aux-stream overlap | `aux_stream`
`shared_experts_stream`
shared expert near router | `vllm/utils/torch_utils.py`
`vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py` | MoE shared experts can run on a dedicated aux stream and overlap with router-side work | Treat shared-expert vs router overlap as an existing upstream sparse-model family. | +| vLLM-origin DCP async all-to-all overlap | `dcp_alltoall`
`all_to_all_single`
`async_op=True` | `vllm/v1/attention/ops/dcp_alltoall.py` | Output / LSE exchange uses async all-to-all handles instead of serializing collective completion on the main path | Treat DCP all-to-all windows as an upstream async-collective family. | + +## 10. vLLM-origin PR-backed / in-flight fused-kernel and kernel-overlap families + +| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | +| --- | --- | --- | --- | --- | +| PR `#35968` DSV3.2 multi-stream indexer overlap | `weights_proj`
`wk`
`k_norm`
`aux_stream` | `PR #35968`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/utils/torch_utils.py` | Open PR overlaps the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. | +| PR `#37110` Triton attention + per-group FP8 dynamic quant | `group_size=128`
`group_size=64`
`output_group_scale`
`per-group FP8` | `PR #37110`
`vllm/compilation/passes/fusion/attn_quant_fusion.py`
`vllm/v1/attention/ops/triton_unified_attention.py` | In-flight Triton attention epilogue computes per-group FP8 scales and quantizes output directly instead of launching a separate group-quant kernel | Treat attention + per-group FP8 quant as a concrete upstream vLLM family, not a novel idea. | +| PR `#38445` MiniMax-M2 FP32 gate kernel | `fp32_router_gemm`
`MiniMax-M2`
`gate kernel` | `PR #38445`
`vllm/model_executor/layers/fused_moe/router/gate_linear.py`
`vllm/model_executor/models/minimax_m2.py` | Draft CUDA kernel fuses BF16->FP32 conversion and low-batch router GEMM for MiniMax-M2, replacing up to three kernels on the gate path | Treat MiniMax-M2 gate ladders as an in-flight upstream fused router family first. | +| PR `#38621` fused QK norm + RoPE + cache + quant | `fused_qk_norm_rope_cache_quant`
`QK Norm + RoPE + Cache + Quant` | `PR #38621`
`csrc/fused_qk_norm_rope_cache_quant.cu`
`vllm/compilation/passes/fusion/qk_norm_rope_cache_quant_fusion.py` | Draft CUDA kernel and compile-time pass try to fuse QK RMSNorm, RoPE, KV cache write, and optional FP8 quant for small-batch decode | Treat this as an in-flight upstream fusion family before calling a similar idea novel. | +| PR `#38684` DSV3.2 fused `wk + weights_proj` | `wk_weights_proj`
`MergedColumnParallelLinear`
`weights_proj` | `PR #38684`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/model_executor/models/deepseek_mtp.py` | Merged PR fuses the DSV3.2 indexer `wk` and `weights_proj` projections into one GEMM; FP8 weight-loading caveats are being handled in follow-up `PR #38870` | Treat paired indexer projections as a concrete upstream fused linear family before calling the opportunity novel. | + +## 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`
`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`
`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" +``` diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/heuristics.md b/.claude/skills/sglang-torch-profiler-analysis/references/heuristics.md new file mode 100644 index 000000000..e0f55ab4a --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/references/heuristics.md @@ -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. diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/overlap-catalog.md b/.claude/skills/sglang-torch-profiler-analysis/references/overlap-catalog.md new file mode 100644 index 000000000..76d6802f7 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/references/overlap-catalog.md @@ -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`
`python/sglang/srt/models/qwen3.py`
`python/sglang/srt/models/qwen3_next.py`
`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`
`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`
`_comm_stream`
`dispatch`
`combine`
`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`
`qwen3_moe.py`
`glm4_moe.py`
`bailing_moe.py`
`llada2.py`
`grok.py`
`olmo2.py`
`step3p5.py`
`longcat_flash.py`
`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`
`_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`
`_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`
`python/sglang/srt/layers/attention/vision.py`
`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`
`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`
`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`
`combine`
`down_gemm` | `PR #21877`
`python/sglang/srt/server_args.py`
`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | Fused combine eliminates the standalone combine window, so SBO is intentionally disabled when this path is on | If the trace discussion is about combine overlap, first classify it as this upstream fused-overlap family. | + +## 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`
`fused_matmul_reduce_scatter`
`fused_all_gather_matmul` | `vllm/compilation/passes/fusion/collective_fusion.py`
`docs/design/fusions.md` | AsyncTP overlaps GEMM with reduce-scatter / all-gather via symmetric-memory collectives | Treat GEMM+comm windows as a clear vLLM-origin overlap precedent first. | +| vLLM-origin Sequence Parallelism staging | `enable_sp`
`ReduceScatter`
`AllGather`
`SequenceParallelismPass` | `vllm/compilation/passes/fusion/sequence_parallelism.py`
`docs/design/fusions.md` | Sequence-parallel rewrites all-reduce into RS -> local norm -> AG so later passes can overlap comm and compute | Treat RS / AG staging around norm blocks as an upstream overlap-enabling family. | +| vLLM-origin shared-expert aux-stream overlap | `aux_stream`
`shared_experts_stream`
shared expert near router | `vllm/utils/torch_utils.py`
`vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py` | MoE shared experts can run on a dedicated aux stream and overlap with router-side work | Treat shared-expert vs router overlap as an existing upstream sparse-model family. | +| vLLM-origin DCP async all-to-all overlap | `dcp_alltoall`
`all_to_all_single`
`async_op=True` | `vllm/v1/attention/ops/dcp_alltoall.py` | Output / LSE exchange uses async all-to-all handles instead of serializing collective completion on the main path | Treat DCP all-to-all windows as an upstream async-collective family. | + +## 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`
`wk`
`k_norm`
`aux_stream` | `PR #35968`
`vllm/model_executor/models/deepseek_v2.py`
`vllm/utils/torch_utils.py` | Open PR overlaps the small `weights_proj` GEMM with `wk + k_norm` on a secondary CUDA stream for decode batches instead of serializing both on the default stream | Treat this as a concrete upstream decode-time kernel-overlap family when traces show underutilized projection overlap opportunities. | + +## 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`
`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" +``` diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/source-map.md b/.claude/skills/sglang-torch-profiler-analysis/references/source-map.md new file mode 100644 index 000000000..f460a45fa --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/references/source-map.md @@ -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` diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/trace-workflow.md b/.claude/skills/sglang-torch-profiler-analysis/references/trace-workflow.md new file mode 100644 index 000000000..26740c884 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/references/trace-workflow.md @@ -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: + +- `-TP-0.trace.json.gz` +- `-TP-0-DP-0-PP-0-EP-0.trace.json.gz` +- `merged-.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. diff --git a/.claude/skills/sglang-torch-profiler-analysis/references/validated-workflows.md b/.claude/skills/sglang-torch-profiler-analysis/references/validated-workflows.md new file mode 100644 index 000000000..c6219b03c --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/references/validated-workflows.md @@ -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. diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_llm_torch_profile.py b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_llm_torch_profile.py new file mode 100644 index 000000000..40c951f4c --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_llm_torch_profile.py @@ -0,0 +1,1582 @@ +"""Analyze SGLang LLM torch profiler traces into kernel/category shares.""" + +from __future__ import annotations + +import argparse +import json +import re +import time +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import DefaultDict, Dict, Iterable, List, Optional, Sequence, Tuple + +from profile_common import ( + coerce_optional_int, + contains_any_keyword, + discover_trace_targets, + extract_trace_events, + has_stream_marker, + is_annotation_event, + is_complete_duration_event, + is_non_kernel_trace_category, + is_trace_metadata_name, + load_trace_json, + looks_like_python_scope_name, + normalize_repo_relative_path, + normalize_text, + parse_stage, + run_profiler, + select_heaviest_pid, +) + +CATEGORY_PATTERNS: List[Tuple[str, Tuple[str, ...]]] = [ + ( + "hybrid_linear", + ( + "gdn", + "gated_delta", + "mamba", + "selective_scan", + "ssd", + "causal_conv", + "ssm", + ), + ), + ( + "attention", + ( + "flash_attn", + "flashattention", + "flash_attention", + "fmha", + "attention", + "mla", + "paged_attention", + "decode_attention", + ), + ), + ( + "moe", + ( + "fused_moe", + "grouped_mm", + "groupgemm", + "group_gemm", + "moe", + "expert", + "groupproblemshape", + ), + ), + ( + "gemm", + ( + "gemm", + "gemv", + "matmul", + "cublas", + "cutlass", + "wgmma", + "mma", + "bmm", + "nvjet", + ), + ), + ( + "norm", + ( + "rmsnorm", + "layernorm", + "_norm_", + " norm", + "normkernel", + ), + ), + ("rope", ("rotary", "rope", "mrope")), + ("softmax", ("softmax",)), + ("activation", ("silu", "gelu", "relu", "act_and_mul", "sigmoid")), + ("quantize", ("quant", "fp8", "mxfp", "nvfp4", "dequant", "cvt")), + ( + "reduce_topk", + ("topk", "reduce", "argmax", "argtopk", "sampling", "multinomial"), + ), + ( + "sampling_io", + ( + "prepare_inputs", + "write_req_to", + "catarraybatched", + "prepare_next", + "copy_next", + ), + ), + ( + "elementwise", + ( + "elementwise", + "vectorized_elementwise_kernel", + "unrolled_elementwise_kernel", + "gpu_kernel_impl", + "binary_internal", + "unaryfunctor", + "add_kernel", + "sub_kernel", + "mul_kernel", + "div_", + "floor_kernel", + "log_kernel", + "neg_kernel", + ), + ), +] + +COMMUNICATION_STRONG_KEYWORDS = ( + "nccl", + "allreduce", + "all_reduce", + "reduce_scatter", + "allgather", + "all_gather", + "alltoall", + "all_to_all", + "cross_device_reduce", + "deepep", + "mooncake", +) + +COMMUNICATION_WEAK_KEYWORDS = ( + "broadcast", + "dispatch", + "combine", +) + +MEMORY_STRONG_KEYWORDS = ( + "memcpy", + "memset", + "dma", + "prefetch", +) + +MEMORY_WEAK_KEYWORDS = ( + "copy", + "fill", +) + +COMPUTE_HINT_KEYWORDS = ( + "gemm", + "gemv", + "matmul", + "cublas", + "cutlass", + "wgmma", + "mma", + "bmm", + "nvjet", + "fmha", + "attention", + "flash_attn", + "flashattention", + "flash_attention", + "grouped_mm", + "groupgemm", + "moe", + "expert", +) + +NOISE_FRAME_PREFIXES = ( + "threading.py(", + "multiprocessing/", + "contextlib.py(", + "torch/utils/_contextlib.py(", + "runpy.py(", + "asyncio/", + "selectors.py(", + "queue.py(", + "socket.py(", + "tqdm/_monitor.py(", + "(", + "python/sglang/srt/distributed/communication_op.py:21 " + "tensor_model_parallel_fused_allreduce_rmsnorm" +) +QWEN3_QK_ROPE_FUSION_PATH = ( + "python/sglang/srt/models/qwen3.py:141 forward_prepare_native" + "
python/sglang/srt/models/utils.py:230 apply_qk_norm" + "
python/sglang/jit_kernel/fused_qknorm_rope.py:34 fused_qk_norm_rope_out" + "
python/sglang/srt/models/qwen3_moe.py:592 apply_qk_norm_rope" +) + + +@dataclass +class KernelEvent: + name: str + canonical_name: str + category: str + pid: str + tid: str + ts: float + dur: float + external_id: Optional[int] + correlation: Optional[int] = None + + +@dataclass +class CpuOpEvent: + name: str + pid: str + tid: str + ts: float + dur: float + external_id: int + + +@dataclass +class LaunchEvent: + name: str + pid: str + tid: str + ts: float + dur: float + correlation: int + + +@dataclass +class PythonFrame: + name: str + normalized_name: str + pid: str + tid: str + ts: float + dur: float + python_id: Optional[int] + parent_id: Optional[int] + + @property + def end_ts(self) -> float: + return self.ts + self.dur + + +@dataclass +class Aggregate: + total_us: float = 0.0 + count: int = 0 + max_us: float = 0.0 + + @property + def avg_us(self) -> float: + return self.total_us / self.count if self.count else 0.0 + + +@dataclass +class MappingSiteAggregate: + total_us: float = 0.0 + count: int = 0 + cpu_ops: Counter = field(default_factory=Counter) + stacks: Counter = field(default_factory=Counter) + + +@dataclass +class KernelRow: + name: str + category: str + aggregate: Aggregate + location: str + cpu_op: str + entry: Optional[dict] + + @property + def total_us(self) -> float: + return self.aggregate.total_us + + +@dataclass +class FusionOpportunity: + pattern: str + confidence: str + related_us: float + evidence: str + current_locations: str + candidate_path: str + rationale: str + + +def short_name(name: str, max_len: int = 96) -> str: + text = normalize_text(name) + if len(text) <= max_len: + return text + return text[: max_len - 3] + "..." + + +def canonicalize_name(name: str) -> str: + text = normalize_text(name) + text = re.sub(r"0x[0-9a-fA-F]+", "0xADDR", text) + if text.startswith("void ") and text.endswith(")"): + depth = 0 + split_idx: Optional[int] = None + for idx in range(len(text) - 1, -1, -1): + char = text[idx] + if char == ")": + depth += 1 + elif char == "(": + depth -= 1 + if depth == 0: + split_idx = idx + break + if split_idx is not None: + text = text[:split_idx] + return text + + +def classify_kernel(name: str) -> str: + # Keep the matching order explicit: strong communication/memory signals win + # first, then we fall back to weaker category hints. + lowered = name.lower() + if contains_any_keyword(lowered, COMMUNICATION_STRONG_KEYWORDS): + return "communication" + if contains_any_keyword(lowered, MEMORY_STRONG_KEYWORDS): + return "memory" + looks_compute_like = contains_any_keyword(lowered, COMPUTE_HINT_KEYWORDS) + if contains_any_keyword(lowered, MEMORY_WEAK_KEYWORDS) and not looks_compute_like: + return "memory" + for category, keywords in CATEGORY_PATTERNS: + if contains_any_keyword(lowered, keywords): + return category + if ( + contains_any_keyword(lowered, COMMUNICATION_WEAK_KEYWORDS) + and not looks_compute_like + ): + return "communication" + return "other" + + +def normalize_source_location(name: str) -> str: + text = normalize_text(name) + match = re.match(r"(?P.+?)\((?P\d+)\): (?P.+)$", text) + if not match: + return text + path = normalize_repo_relative_path(match.group("path")) + return f"{path}:{match.group('line')} {match.group('func')}" + + +def source_location_priority(location: str) -> int: + text = str(location).strip() + if not text or text == "unresolved": + return -100 + if text.startswith("python/sglang/"): + return 300 + if text.startswith("sglang/"): + return 290 + if text.startswith("sgl_kernel/"): + return 260 + if text.startswith("python/"): + return 180 + if text.startswith("torch/") or "/torch/" in text: + return 20 + if ".py:" in text: + return 120 + return 0 + + +def is_preferred_source_location(location: str) -> bool: + text = str(location).strip() + return ( + text.startswith("python/sglang/") + or text.startswith("sglang/") + or text.startswith("sgl_kernel/") + ) + + +def extract_preferred_stack_location(stack: Optional[str]) -> Optional[str]: + if not stack: + return None + parts = [str(part).strip() for part in str(stack).split("->")] + ranked: List[Tuple[int, int, str]] = [] + for index, part in enumerate(parts): + normalized = normalize_source_location(part) + priority = source_location_priority(normalized) + if priority <= 0: + continue + ranked.append((priority, index, normalized)) + if not ranked: + return None + ranked.sort(key=lambda item: (item[0], item[1]), reverse=True) + return ranked[0][2] + + +def site_display_location(site: dict) -> str: + location = str(site.get("location") or "unresolved").strip() + if is_preferred_source_location(location): + return location + stack_location = extract_preferred_stack_location(site.get("stack")) + if stack_location: + return stack_location + return location + + +def choose_best_location(locations: Dict[str, MappingSiteAggregate]) -> str: + if not locations: + return "unresolved" + ranked = sorted( + locations.items(), + key=lambda pair: ( + source_location_priority(pair[0]), + pair[1].total_us, + pair[1].count, + ), + reverse=True, + ) + return ranked[0][0] + + +def frame_priority(frame_name: str) -> int: + text = str(frame_name).strip() + if text.startswith(NOISE_FRAME_PREFIXES): + return -20 + if text.startswith("/data/") or text.startswith("/Users/"): + if "/sglang/" in text: + return 120 + return 100 + if ".py(" in text and "/sglang/" in text: + return 110 + if ".py(" in text and ("site-packages" in text or text.startswith("torch/")): + return 45 + if ".py(" in text: + return 35 + if text.startswith(" str: + if stage == "extend": + return "extend/prefill" + return stage + + +def stage_aliases(stage: str) -> List[str]: + if stage == "extend": + return ["extend", "prefill", "all"] + if stage == "prefill": + return ["prefill", "extend", "all"] + if stage == "decode": + return ["decode", "all"] + return [stage, "all"] + + +def escape_md_cell(text: str) -> str: + return str(text).replace("|", "\\|").replace("\n", "
") + + +def pct(part: float, whole: float) -> float: + return 100.0 * part / whole if whole else 0.0 + + +def format_ms(value_us: float) -> str: + return f"{value_us / 1000.0:.2f} ms" + + +def is_cuda_launch_event(name: str, cat: str) -> bool: + lowered_name = normalize_text(name).lower() + lowered_cat = normalize_text(cat).lower() + if lowered_cat not in {"cuda_runtime", "cuda_driver"}: + return False + return "launch" in lowered_name + + +def is_gpu_kernel_event(event: dict) -> bool: + # Be conservative here: first drop trace metadata / Python scopes / + # annotations, then only accept entries with clear GPU-kernel markers. + if not is_complete_duration_event(event): + return False + name = normalize_text(event.get("name", "")) + if is_trace_metadata_name(name): + return False + cat = normalize_text(event.get("cat", "")).lower() + args = event.get("args") or {} + if is_non_kernel_trace_category(cat): + return False + if is_annotation_event(name, cat): + return False + if "kernel" in cat or cat.startswith("gpu_"): + return True + if looks_like_python_scope_name(name): + return False + return has_stream_marker(args) + + +def extract_trace_data( + trace: dict, +) -> Tuple[ + List[KernelEvent], + List[CpuOpEvent], + Dict[Tuple[str, str], List[PythonFrame]], + List[LaunchEvent], + Optional[str], + float, +]: + # Build the basic trace views in one pass so later stages can stay simple: + # GPU kernels for ranking, CPU ops for External-id mapping, Python frames for + # source attribution, and CUDA launch calls for correlation-based fallback. + raw_events = extract_trace_events(trace) + correlation_external = build_correlation_external_lookup(raw_events) + chosen_pid = select_heaviest_pid( + raw_events, + is_gpu_kernel_event, + preferred_substrings=("TP00", "TP-0"), + ) + + kernels: List[KernelEvent] = [] + cpu_ops: List[CpuOpEvent] = [] + launches: List[LaunchEvent] = [] + python_frames: DefaultDict[Tuple[str, str], List[PythonFrame]] = defaultdict(list) + min_ts = None + max_end = None + + for event in raw_events: + if event.get("ph") != "X": + 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)) + cat = str(event.get("cat", "")) + args = event.get("args") or {} + name = str(event.get("name", "")) + + if cat == "python_function": + python_frames[(pid, tid)].append( + PythonFrame( + name=name, + normalized_name=normalize_source_location(name), + pid=pid, + tid=tid, + ts=ts, + dur=dur, + python_id=coerce_optional_int(args.get("Python id")), + parent_id=coerce_optional_int(args.get("Python parent id")), + ) + ) + + correlation = coerce_optional_int(args.get("correlation")) + external_id = coerce_optional_int(args.get("External id")) + if external_id is None and correlation is not None: + external_id = correlation_external.get(correlation) + if cat == "cpu_op" and external_id is not None: + cpu_ops.append( + CpuOpEvent( + name=name, + pid=pid, + tid=tid, + ts=ts, + dur=dur, + external_id=external_id, + ) + ) + if is_cuda_launch_event(name, cat) and correlation is not None: + launches.append( + LaunchEvent( + name=name, + pid=pid, + tid=tid, + ts=ts, + dur=dur, + correlation=correlation, + ) + ) + + if chosen_pid is None or not is_gpu_kernel_event(event) or pid != chosen_pid: + continue + + min_ts = ts if min_ts is None else min(min_ts, ts) + max_end = ts + dur if max_end is None else max(max_end, ts + dur) + kernels.append( + KernelEvent( + name=name, + canonical_name=canonicalize_name(name), + category=classify_kernel(name), + pid=pid, + tid=tid, + ts=ts, + dur=dur, + external_id=external_id, + correlation=correlation, + ) + ) + + for frames in python_frames.values(): + frames.sort(key=lambda item: (item.ts, item.end_ts)) + + window_us = 0.0 if min_ts is None or max_end is None else max_end - min_ts + return kernels, cpu_ops, dict(python_frames), launches, chosen_pid, window_us + + +def build_correlation_external_lookup(raw_events: Sequence[dict]) -> Dict[int, int]: + lookup: Dict[int, int] = {} + for event in raw_events: + args = event.get("args", {}) or {} + correlation = coerce_optional_int(args.get("correlation")) + external_id = coerce_optional_int(args.get("External id")) + if correlation is not None and external_id is not None: + lookup[correlation] = external_id + return lookup + + +def build_cpu_op_index(cpu_ops: Sequence[CpuOpEvent]) -> Dict[int, List[CpuOpEvent]]: + output: DefaultDict[int, List[CpuOpEvent]] = defaultdict(list) + for cpu_op in cpu_ops: + output[cpu_op.external_id].append(cpu_op) + for items in output.values(): + items.sort(key=lambda item: item.ts) + return dict(output) + + +def match_cpu_op( + kernel: KernelEvent, cpu_ops_by_external_id: Dict[int, List[CpuOpEvent]] +) -> Optional[CpuOpEvent]: + if kernel.external_id is None: + return None + return match_timed_event( + cpu_ops_by_external_id.get(kernel.external_id, []), kernel.ts + ) + + +def build_launch_index( + launch_events: Sequence[LaunchEvent], +) -> Dict[int, List[LaunchEvent]]: + output: DefaultDict[int, List[LaunchEvent]] = defaultdict(list) + for launch in launch_events: + output[launch.correlation].append(launch) + for items in output.values(): + items.sort(key=lambda item: item.ts) + return dict(output) + + +def match_launch_event( + kernel: KernelEvent, launches_by_correlation: Dict[int, List[LaunchEvent]] +) -> Optional[LaunchEvent]: + if kernel.correlation is None: + return None + return match_timed_event( + launches_by_correlation.get(kernel.correlation, []), kernel.ts + ) + + +def match_timed_event(events: Sequence, probe_ts: float): + if not events: + return None + earlier = [item for item in events if item.ts <= probe_ts + 1e-3] + if earlier: + return min(earlier, key=lambda item: abs((item.ts + item.dur) - probe_ts)) + return min(events, key=lambda item: abs(item.ts - probe_ts)) + + +def find_active_python_frames( + cpu_op: CpuOpEvent, + python_frames: Dict[Tuple[str, str], List[PythonFrame]], +) -> List[PythonFrame]: + frames = python_frames.get((cpu_op.pid, cpu_op.tid), []) + if not frames: + return [] + probe_ts = cpu_op.ts + min(cpu_op.dur * 0.5, 1.0) + active = [item for item in frames if item.ts <= probe_ts <= item.end_ts] + active.sort(key=lambda item: (item.ts, item.end_ts)) + return active + + +def find_active_python_frames_at_ts( + *, + pid: str, + tid: str, + ts: float, + python_frames: Dict[Tuple[str, str], List[PythonFrame]], +) -> List[PythonFrame]: + frames = python_frames.get((pid, tid), []) + if not frames: + return [] + active = [item for item in frames if item.ts <= ts <= item.end_ts] + active.sort(key=lambda item: (item.ts, item.end_ts)) + return active + + +def render_kernel_site( + active_frames: Sequence[PythonFrame], cpu_op_name: str +) -> Tuple[str, str, str]: + chosen_frame = choose_mapping_frame(active_frames) + if chosen_frame is None: + return "unresolved", "", cpu_op_name + return chosen_frame.normalized_name, build_stack_display(active_frames), cpu_op_name + + +def resolve_kernel_site_context( + kernel: KernelEvent, + cpu_ops_by_external_id: Dict[int, List[CpuOpEvent]], + python_frames: Dict[Tuple[str, str], List[PythonFrame]], + launches_by_correlation: Dict[int, List[LaunchEvent]], +) -> Tuple[str, str, str]: + # Prefer the normal External-id path first. If the kernel dropped that link, + # fall back to the correlated CUDA launch and reuse the Python frames that + # were active when the launch happened. + cpu_op = match_cpu_op(kernel, cpu_ops_by_external_id) + if cpu_op is not None: + active_frames = find_active_python_frames(cpu_op, python_frames) + if active_frames: + return render_kernel_site(active_frames, cpu_op.name) + + launch_event = match_launch_event(kernel, launches_by_correlation) + if launch_event is not None: + active_frames = find_active_python_frames_at_ts( + pid=launch_event.pid, + tid=launch_event.tid, + ts=launch_event.ts, + python_frames=python_frames, + ) + if active_frames: + cpu_op_name = cpu_op.name if cpu_op is not None else launch_event.name + return render_kernel_site(active_frames, cpu_op_name) + return "unresolved", "", launch_event.name + + cpu_op_name = cpu_op.name if cpu_op is not None else "" + return "unresolved", "", cpu_op_name + + +def choose_mapping_frame(active_frames: Sequence[PythonFrame]) -> Optional[PythonFrame]: + if not active_frames: + return None + ranked = sorted( + active_frames, + key=lambda item: (frame_priority(item.name), item.ts, -item.dur), + ) + return ranked[-1] + + +def build_stack_display(active_frames: Sequence[PythonFrame]) -> str: + if not active_frames: + return "" + filtered = [ + item.normalized_name for item in active_frames if frame_priority(item.name) > 0 + ] + if not filtered: + filtered = [active_frames[-1].normalized_name] + return " -> ".join(filtered[-4:]) + + +def aggregate(events: Iterable[KernelEvent], key_fn) -> Dict[str, Aggregate]: + output: Dict[str, Aggregate] = defaultdict(Aggregate) + for event in events: + key = key_fn(event) + item = output[key] + item.total_us += event.dur + item.count += 1 + item.max_us = max(item.max_us, event.dur) + return output + + +def aggregate_kernel_sites( + kernels: Sequence[KernelEvent], + cpu_ops_by_external_id: Dict[int, List[CpuOpEvent]], + python_frames: Dict[Tuple[str, str], List[PythonFrame]], + launches_by_correlation: Optional[Dict[int, List[LaunchEvent]]] = None, +) -> Dict[str, Dict[str, MappingSiteAggregate]]: + # Each kernel is mapped independently so the fallback behavior stays easy to + # reason about and easy to regression-test. + output: DefaultDict[str, DefaultDict[str, MappingSiteAggregate]] = defaultdict( + lambda: defaultdict(MappingSiteAggregate) + ) + launch_index = launches_by_correlation or {} + for kernel in kernels: + location, stack, cpu_op_name = resolve_kernel_site_context( + kernel, + cpu_ops_by_external_id, + python_frames, + launch_index, + ) + + item = output[kernel.canonical_name][location] + item.total_us += kernel.dur + item.count += 1 + if cpu_op_name: + item.cpu_ops[cpu_op_name] += 1 + if stack: + item.stacks[stack] += 1 + return {kernel_name: dict(locations) for kernel_name, locations in output.items()} + + +def merge_site_stats( + destination: DefaultDict[str, DefaultDict[str, MappingSiteAggregate]], + source: Dict[str, Dict[str, MappingSiteAggregate]], +) -> None: + for kernel_name, locations in source.items(): + for location, aggregate_item in locations.items(): + target = destination[kernel_name][location] + target.total_us += aggregate_item.total_us + target.count += aggregate_item.count + target.cpu_ops.update(aggregate_item.cpu_ops) + target.stacks.update(aggregate_item.stacks) + + +def build_stage_payload( + site_stats: Dict[str, Dict[str, MappingSiteAggregate]], + kernel_categories: Dict[str, str], +) -> Dict[str, dict]: + kernels_payload: Dict[str, dict] = {} + for kernel_name, locations in sorted(site_stats.items()): + total_us = sum(item.total_us for item in locations.values()) + sites = [] + for location, aggregate_item in sorted( + locations.items(), + key=lambda pair: pair[1].total_us, + reverse=True, + ): + sites.append( + { + "location": location, + "display_location": extract_preferred_stack_location( + aggregate_item.stacks.most_common(1)[0][0] + if aggregate_item.stacks + else None + ) + or location, + "launches": aggregate_item.count, + "total_us": round(aggregate_item.total_us, 3), + "share_pct_within_kernel": round( + pct(aggregate_item.total_us, total_us), 3 + ), + "top_cpu_op": ( + aggregate_item.cpu_ops.most_common(1)[0][0] + if aggregate_item.cpu_ops + else None + ), + "stack": ( + aggregate_item.stacks.most_common(1)[0][0] + if aggregate_item.stacks + else None + ), + } + ) + sites.sort( + key=lambda site: ( + source_location_priority(site_display_location(site)), + float(site.get("total_us", 0.0)), + int(site.get("launches", 0)), + ), + reverse=True, + ) + kernels_payload[kernel_name] = { + "category": kernel_categories.get(kernel_name, "other"), + "sites": sites, + "best_location": ( + site_display_location(sites[0]) + if sites + else choose_best_location(locations) + ), + } + return {"kernels": kernels_payload} + + +def load_kernel_map(path: Path) -> dict: + with open(path, "r", encoding="utf-8") as handle: + return json.load(handle) + + +def relaxed_kernel_entry_lookup( + kernels: Dict[str, dict], kernel_name: str +) -> Optional[dict]: + if kernel_name in kernels: + return kernels[kernel_name] + lowered = kernel_name.lower() + best_key = None + best_score = -1 + for candidate_key in kernels: + candidate_lowered = candidate_key.lower() + if candidate_lowered.startswith(lowered) or lowered.startswith( + candidate_lowered + ): + score = min(len(candidate_lowered), len(lowered)) + elif candidate_lowered in lowered or lowered in candidate_lowered: + score = min(len(candidate_lowered), len(lowered)) // 2 + else: + continue + if score > best_score: + best_key = candidate_key + best_score = score + return kernels.get(best_key) if best_key else None + + +def lookup_kernel_map_entry( + kernel_map: dict, stage: str, kernel_name: str +) -> Optional[dict]: + stage_map = kernel_map.get("stages", {}) + for candidate_stage in stage_aliases(stage): + entry = relaxed_kernel_entry_lookup( + stage_map.get(candidate_stage, {}).get("kernels", {}), + kernel_name, + ) + if entry: + return entry + return relaxed_kernel_entry_lookup( + kernel_map.get("global", {}).get("kernels", {}), kernel_name + ) + + +def best_site_summary(kernel_entry: Optional[dict]) -> Tuple[str, str]: + if not kernel_entry: + return "unresolved", "-" + sites = kernel_entry.get("sites") or [] + if not sites: + return kernel_entry.get("best_location", "unresolved"), "-" + preferred_sites = [ + site + for site in sites + if is_preferred_source_location(site_display_location(site)) + ] + candidate_sites = preferred_sites or sites + rendered_locations = [] + rendered_cpu_ops = [] + for site in candidate_sites[:2]: + location = site_display_location(site) + share = site.get("share_pct_within_kernel") + if len(candidate_sites) > 1 and share is not None: + rendered_locations.append(f"{location} (site share {share:.0f}%)") + else: + rendered_locations.append(location) + cpu_op = site.get("top_cpu_op") + if cpu_op: + rendered_cpu_ops.append(cpu_op) + return "
".join(rendered_locations), ( + "
".join(rendered_cpu_ops) if rendered_cpu_ops else "-" + ) + + +def resolve_kernel_entry( + stage: str, + kernel_name: str, + local_stage_payload: dict, + external_kernel_map: Optional[dict], +) -> Optional[dict]: + if external_kernel_map: + kernel_entry = lookup_kernel_map_entry(external_kernel_map, stage, kernel_name) + if kernel_entry: + return kernel_entry + return local_stage_payload.get("kernels", {}).get(kernel_name) + + +def build_kernel_rows( + stage: str, + kernel_stats: Dict[str, Aggregate], + kernel_categories: Dict[str, str], + local_stage_payload: dict, + external_kernel_map: Optional[dict], +) -> List[KernelRow]: + rows: List[KernelRow] = [] + for kernel_name, aggregate_item in sorted( + kernel_stats.items(), + key=lambda pair: pair[1].total_us, + reverse=True, + ): + kernel_entry = resolve_kernel_entry( + stage, kernel_name, local_stage_payload, external_kernel_map + ) + location, cpu_op = best_site_summary(kernel_entry) + rows.append( + KernelRow( + name=kernel_name, + category=kernel_categories.get(kernel_name, "other"), + aggregate=aggregate_item, + location=location, + cpu_op=cpu_op, + entry=kernel_entry, + ) + ) + return rows + + +def limit_kernel_rows(rows: Sequence[KernelRow], table_limit: int) -> List[KernelRow]: + if table_limit <= 0: + return list(rows) + return list(rows[:table_limit]) + + +def entry_sites(kernel_entry: Optional[dict]) -> List[dict]: + if not kernel_entry: + return [] + sites = kernel_entry.get("sites") or [] + return [site for site in sites if site.get("location")] + + +def ordered_unique(values: Iterable[str], limit: int = 4) -> List[str]: + output: List[str] = [] + seen = set() + for value in values: + item = str(value).strip() + if not item or item in seen: + continue + seen.add(item) + output.append(item) + if len(output) >= limit: + break + return output + + +def kernel_row_locations(row: KernelRow, limit: int = 4) -> List[str]: + values = [site_display_location(site) for site in entry_sites(row.entry)] + if not values and row.location and row.location != "unresolved": + values = [fragment.strip() for fragment in row.location.split("
")] + return ordered_unique(values, limit=limit) + + +def format_location_for_fusion_display(location: str) -> str: + text = normalize_text(location) + match = re.match(r"(?P.+?):(?P\d+)\s+(?P.+)$", text) + if not match: + return text + return f"{match.group('func')} @ {match.group('path')}:{match.group('line')}" + + +def row_matches(row: KernelRow, *needles: str) -> bool: + lowered = " ".join([row.name, row.location, row.cpu_op]).lower() + return any(needle in lowered for needle in needles) + + +def summarize_text(values: Iterable[str], limit: int = 4) -> str: + items = ordered_unique(values, limit=limit) + return "
".join(items) if items else "-" + + +def summarize_locations(values: Iterable[str], limit: int = 4) -> str: + items = ordered_unique( + (format_location_for_fusion_display(value) for value in values), + limit=limit, + ) + return "
".join(items) if items else "-" + + +def summarize_evidence( + rows: Sequence[KernelRow], total_us: float, limit: int = 3 +) -> str: + items = [] + for row in rows[:limit]: + items.append(f"{row.name} ({pct(row.total_us, total_us):.1f}%)") + return "
".join(items) if items else "-" + + +def model_path_from_server_args(server_args: Optional[dict]) -> str: + if not isinstance(server_args, dict): + return "" + return str(server_args.get("model_path") or server_args.get("model") or "") + + +def detect_fusion_opportunities( + stage: str, + kernel_rows: Sequence[KernelRow], + total_us: float, + server_args: Optional[dict], +) -> List[FusionOpportunity]: + opportunities: List[FusionOpportunity] = [] + if total_us <= 0: + return opportunities + + model_path = model_path_from_server_args(server_args).lower() + tp_size = 1 + if isinstance(server_args, dict): + tp_size = int(server_args.get("tp_size") or 1) + + comm_rows = [ + row + for row in kernel_rows + if row.category == "communication" + and ( + row_matches( + row, + "cross_device_reduce", + "allreduce", + "all_reduce", + "custom_all_reduce_ops.py", + ) + ) + ] + comm_us = sum(row.total_us for row in comm_rows) + if comm_rows and ( + tp_size > 1 + or any(row_matches(row, "custom_all_reduce_ops.py") for row in comm_rows) + ): + opportunities.append( + FusionOpportunity( + pattern="TP all-reduce + residual/RMSNorm", + confidence="Likely" if pct(comm_us, total_us) >= 4.0 else "Conditional", + related_us=comm_us, + evidence=summarize_evidence(comm_rows, total_us), + current_locations=summarize_locations( + location + for row in comm_rows + for location in kernel_row_locations(row) + ), + candidate_path=ALLREDUCE_FUSION_PATH, + rationale=( + f"TP communication already consumes {pct(comm_us, total_us):.1f}% of cumulative GPU kernel " + "time, and SGLang already exposes a fused allreduce+RMSNorm path for the residual/norm " + "boundary." + ), + ) + ) + + is_qwen3_dense = ( + "qwen3" in model_path + and "moe" not in model_path + and "qwen3-next" not in model_path + ) + qwen3_qk_rows = [ + row + for row in kernel_rows + if row_matches( + row, + "apply_qk_norm", + "fused_inplace_qknorm", + "qknorm", + ) + ] + qwen3_rope_rows = [ + row + for row in kernel_rows + if row_matches( + row, + "apply_rope", + "rope.py", + "rope_inplace", + "rotary", + ) + ] + qwen3_rows: List[KernelRow] = [] + seen_qwen3_keys = set() + for row in qwen3_qk_rows + qwen3_rope_rows: + row_key = (row.name, row.location, row.cpu_op) + if row_key in seen_qwen3_keys: + continue + seen_qwen3_keys.add(row_key) + qwen3_rows.append(row) + qwen3_related_us = sum(row.total_us for row in qwen3_rows) + if ( + is_qwen3_dense + and qwen3_qk_rows + and qwen3_rope_rows + and pct(qwen3_related_us, total_us) >= 1.0 + ): + opportunities.append( + FusionOpportunity( + pattern="Q/K RMSNorm + RoPE before attention", + confidence="Conditional", + related_us=qwen3_related_us, + evidence=summarize_evidence(qwen3_rows, total_us), + current_locations=summarize_locations( + location + for row in qwen3_rows + for location in kernel_row_locations(row) + ), + candidate_path=QWEN3_QK_ROPE_FUSION_PATH, + rationale=( + "Dense Qwen3 still prepares Q/K with `apply_qk_norm` and then `rotary_emb` in separate source " + "steps, while SGLang already ships a fused QK-norm+RoPE kernel path in its JIT/MoE stack." + ), + ) + ) + + opportunities.sort(key=lambda item: item.related_us, reverse=True) + return opportunities + + +def generate_takeaways( + stage: str, + total_us: float, + window_us: float, + category_stats: Dict[str, Aggregate], + resolved_us: float, + server_args: Optional[dict], + fusion_opportunities: Sequence[FusionOpportunity], +) -> List[str]: + items = sorted( + category_stats.items(), key=lambda pair: pair[1].total_us, reverse=True + ) + if not items: + return ["No GPU kernel events were found in the selected trace."] + + takeaways: List[str] = [] + top_name, top_agg = items[0] + takeaways.append( + f"{stage_label(stage)} is dominated by `{top_name}` at {pct(top_agg.total_us, total_us):.1f}% of cumulative GPU kernel time." + ) + if len(items) > 1: + second_name, second_agg = items[1] + combined = pct(top_agg.total_us + second_agg.total_us, total_us) + takeaways.append( + f"The top two categories are `{top_name}` + `{second_name}` at {combined:.1f}% combined." + ) + + comm_share = pct( + category_stats.get("communication", Aggregate()).total_us, total_us + ) + if comm_share >= 10.0: + tp = server_args.get("tp_size") if isinstance(server_args, dict) else None + if tp and tp > 1: + takeaways.append( + f"`communication` already accounts for {comm_share:.1f}% of cumulative GPU time in this TP={tp} run." + ) + else: + takeaways.append( + f"`communication` shows up at {comm_share:.1f}% even without an obvious large-TP context." + ) + + if pct(resolved_us, total_us) >= 70.0: + takeaways.append( + f"Kernel-to-Python mapping covers {pct(resolved_us, total_us):.1f}% of cumulative GPU time, so the table is representative enough for code triage." + ) + + if window_us > 0: + parallelism = total_us / window_us + if parallelism >= 1.15: + takeaways.append( + f"Summed kernel time is {parallelism:.2f}x the GPU time window, so these percentages are cumulative launch share rather than wall time share." + ) + if fusion_opportunities: + top_pattern = fusion_opportunities[0] + takeaways.append( + f"The strongest source-backed fuse candidate is `{top_pattern.pattern}` with {pct(top_pattern.related_us, total_us):.1f}% related GPU time in this stage." + ) + return takeaways + + +def print_mapping_table( + kernel_rows: Sequence[KernelRow], + total_us: float, + table_limit: int, +) -> float: + resolved_us = 0.0 + rendered_rows = limit_kernel_rows(kernel_rows, table_limit) + label = "all kernels" if table_limit <= 0 else f"first {len(rendered_rows)} kernels" + print(f"\nKernel-to-Python mapping (Markdown, {label}):") + print( + "| Kernel | Category | GPU time | Share | Launches | Python location (site share) | CPU op |" + ) + print("| --- | --- | ---: | ---: | ---: | --- | --- |") + for row in rendered_rows: + if row.location != "unresolved": + resolved_us += row.total_us + print( + "| {kernel} | {category} | {gpu_time} | {share:.1f}% | {launches} | {location} | {cpu_op} |".format( + kernel=escape_md_cell(row.name), + category=escape_md_cell(row.category), + gpu_time=format_ms(row.total_us), + share=pct(row.total_us, total_us), + launches=row.aggregate.count, + location=escape_md_cell(row.location), + cpu_op=escape_md_cell(row.cpu_op), + ) + ) + return resolved_us + + +def print_fusion_opportunity_table( + opportunities: Sequence[FusionOpportunity], + total_us: float, +) -> None: + print("\nKernel fuse opportunities (Markdown):") + print( + "| Pattern | Confidence | Related GPU time | Share | Evidence kernels | Current kernel Python location | Candidate fused Python path | Rationale |" + ) + print("| --- | --- | ---: | ---: | --- | --- | --- | --- |") + if not opportunities: + print( + "| No medium-confidence source-backed fusion opportunity matched this trace. | - | - | - | - | - | - | - |" + ) + return + for item in opportunities: + print( + "| {pattern} | {confidence} | {gpu_time} | {share:.1f}% | {evidence} | {current_locations} | {candidate_path} | {rationale} |".format( + pattern=escape_md_cell(item.pattern), + confidence=escape_md_cell(item.confidence), + gpu_time=format_ms(item.related_us), + share=pct(item.related_us, total_us), + evidence=escape_md_cell(item.evidence), + current_locations=escape_md_cell(item.current_locations), + candidate_path=escape_md_cell(item.candidate_path), + rationale=escape_md_cell(item.rationale), + ) + ) + + +def print_report( + trace_path: Path, + server_args: Optional[dict], + kernels: List[KernelEvent], + chosen_pid: Optional[str], + window_us: float, + local_stage_payload: dict, + external_kernel_map: Optional[dict], + top_k: int, + kernel_table_limit: int, + table_only: bool, +) -> None: + stage = parse_stage(trace_path) + total_us = sum(kernel.dur for kernel in kernels) + print(f"Trace: {trace_path}") + print(f"Stage: {stage_label(stage)}") + if chosen_pid: + print(f"Selected PID: {chosen_pid}") + + if server_args: + model_path = server_args.get("model_path") or server_args.get("model") + tp_size = server_args.get("tp_size") + dp_size = server_args.get("dp_size") + print(f"Model: {model_path}") + if tp_size or dp_size: + print(f"Parallelism: tp={tp_size or 1} dp={dp_size or 1}") + + if not kernels: + print("No GPU kernel events found.\n") + return + + print( + f"GPU kernels: {len(kernels)} | cumulative kernel time: {format_ms(total_us)} | " + f"GPU window: {format_ms(window_us)} | avg parallelism: {total_us / window_us:.2f}x" + if window_us + else f"GPU kernels: {len(kernels)} | cumulative kernel time: {format_ms(total_us)}" + ) + + category_stats = aggregate(kernels, key_fn=lambda item: item.category) + kernel_stats = aggregate(kernels, key_fn=lambda item: item.canonical_name) + kernel_categories = {kernel.canonical_name: kernel.category for kernel in kernels} + kernel_rows = build_kernel_rows( + stage=stage, + kernel_stats=kernel_stats, + kernel_categories=kernel_categories, + local_stage_payload=local_stage_payload, + external_kernel_map=external_kernel_map, + ) + fusion_opportunities = detect_fusion_opportunities( + stage=stage, + kernel_rows=kernel_rows, + total_us=total_us, + server_args=server_args, + ) + + if not table_only: + print("\nTop categories by cumulative GPU kernel time:") + for idx, (name, aggregate_item) in enumerate( + sorted( + category_stats.items(), key=lambda pair: pair[1].total_us, reverse=True + )[:8], + start=1, + ): + print( + f" {idx}. {name:<16} {format_ms(aggregate_item.total_us):>10} " + f"{pct(aggregate_item.total_us, total_us):>5.1f}% launches={aggregate_item.count}" + ) + + print("\nTop kernels by cumulative GPU kernel time:") + for idx, (name, aggregate_item) in enumerate( + sorted( + kernel_stats.items(), key=lambda pair: pair[1].total_us, reverse=True + )[:top_k], + start=1, + ): + print( + f" {idx}. {short_name(name, 76):<76} {format_ms(aggregate_item.total_us):>10} " + f"{pct(aggregate_item.total_us, total_us):>5.1f}% launches={aggregate_item.count} avg={format_ms(aggregate_item.avg_us)}" + ) + + resolved_us = print_mapping_table( + kernel_rows=kernel_rows, + total_us=total_us, + table_limit=kernel_table_limit, + ) + print_fusion_opportunity_table(fusion_opportunities, total_us) + + if not table_only: + print("\nTakeaways:") + for takeaway in generate_takeaways( + stage, + total_us, + window_us, + category_stats, + resolved_us, + server_args, + fusion_opportunities, + ): + print(f" - {takeaway}") + print() + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Analyze SGLang LLM torch profiler traces into kernel/category shares." + ) + parser.add_argument("--input", type=str, help="Trace file or profile directory.") + parser.add_argument("--url", type=str, help="Running SGLang server URL.") + parser.add_argument( + "--output-dir", type=str, default=None, help="Output root for live profiling." + ) + parser.add_argument( + "--num-steps", type=int, default=5, help="Profiler steps for live mode." + ) + parser.add_argument( + "--profile-by-stage", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--merge-profiles", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument("--profile-prefix", type=str, default=None) + parser.add_argument("--probe-requests", type=int, default=1) + parser.add_argument( + "--probe-prompt", + type=str, + default="Explain what tensor parallelism is in one short paragraph.", + ) + parser.add_argument("--probe-max-new-tokens", type=int, default=96) + parser.add_argument("--probe-delay", type=float, default=1.0) + parser.add_argument( + "--top-k", + type=int, + default=12, + help="How many top kernels to summarize above the tables.", + ) + parser.add_argument( + "--kernel-table-limit", + type=int, + default=0, + help="How many kernels to include in the Markdown kernel table. Use 0 for all kernels.", + ) + parser.add_argument( + "--kernel-map", + type=str, + default=None, + help="Existing kernel_map.json from a no-CUDA-graph torch pre-pass.", + ) + parser.add_argument( + "--export-kernel-map", + type=str, + default=None, + help="Write kernel-to-Python mapping JSON to this file.", + ) + parser.add_argument( + "--all-traces", + action="store_true", + help="Analyze every matching trace in the selected run directory.", + ) + parser.add_argument( + "--table-only", + action="store_true", + help="Print the trace header plus the kernel and fuse-opportunity tables only.", + ) + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + if not args.input and not args.url: + raise SystemExit("Provide either --input or --url.") + + input_path = Path(args.input).resolve() if args.input else None + if args.url: + profile_dir = run_profiler( + url=args.url, + output_dir=args.output_dir, + num_steps=args.num_steps, + profile_by_stage=args.profile_by_stage, + merge_profiles=args.merge_profiles, + profile_prefix=args.profile_prefix, + probe_requests=args.probe_requests, + probe_prompt=args.probe_prompt, + probe_max_new_tokens=args.probe_max_new_tokens, + probe_delay=args.probe_delay, + start_step=None, + ) + print(f"Generated profile directory: {profile_dir}\n") + input_path = profile_dir + + external_kernel_map = ( + load_kernel_map(Path(args.kernel_map).resolve()) if args.kernel_map else None + ) + traces, server_args = discover_trace_targets(input_path, all_traces=args.all_traces) + + stage_site_stats: DefaultDict[ + str, DefaultDict[str, DefaultDict[str, MappingSiteAggregate]] + ] = defaultdict(lambda: defaultdict(lambda: defaultdict(MappingSiteAggregate))) + stage_kernel_categories: DefaultDict[str, Dict[str, str]] = defaultdict(dict) + global_site_stats: DefaultDict[str, DefaultDict[str, MappingSiteAggregate]] = ( + defaultdict(lambda: defaultdict(MappingSiteAggregate)) + ) + global_kernel_categories: Dict[str, str] = {} + reports = [] + + for trace_path in traces: + trace = load_trace_json(trace_path) + kernels, cpu_ops, python_frames, launch_events, chosen_pid, window_us = ( + extract_trace_data(trace) + ) + cpu_ops_by_external_id = build_cpu_op_index(cpu_ops) + launches_by_correlation = build_launch_index(launch_events) + local_site_stats = 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 + } + + merge_site_stats(stage_site_stats[stage], local_site_stats) + merge_site_stats(global_site_stats, local_site_stats) + stage_kernel_categories[stage].update(kernel_categories) + global_kernel_categories.update(kernel_categories) + reports.append( + (trace_path, kernels, chosen_pid, window_us, stage, kernel_categories) + ) + + stage_payloads = { + stage: build_stage_payload( + dict(site_stats), stage_kernel_categories.get(stage, {}) + ) + for stage, site_stats in stage_site_stats.items() + } + global_payload = build_stage_payload( + dict(global_site_stats), global_kernel_categories + ) + + if args.export_kernel_map: + export_path = Path(args.export_kernel_map).resolve() + export_path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "version": 2, + "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "generated_from": str(input_path), + "notes": "Kernel-to-Python mapping extracted from torch profiler traces. Prefer generating this from a --disable-cuda-graph --disable-piecewise-cuda-graph pre-pass.", + "server_args": server_args, + "stages": stage_payloads, + "global": global_payload, + } + with open(export_path, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, ensure_ascii=False) + print(f"Exported kernel map: {export_path}\n") + + for trace_path, kernels, chosen_pid, window_us, stage, _ in reports: + print_report( + trace_path=trace_path, + server_args=server_args, + kernels=kernels, + chosen_pid=chosen_pid, + window_us=window_us, + local_stage_payload=stage_payloads.get(stage, {"kernels": {}}), + external_kernel_map=external_kernel_map, + top_k=args.top_k, + kernel_table_limit=args.kernel_table_limit, + table_only=args.table_only, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_profiler_overlap.py b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_profiler_overlap.py new file mode 100644 index 000000000..b2da3d445 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_profiler_overlap.py @@ -0,0 +1,1768 @@ +"""Analyze SGLang PyTorch profiler traces with a two-stage workflow. + +This script correlates two traces: + +1. A mapping trace, usually collected with CUDA graph disabled, to keep + `kernel -> cpu_op -> python scope` attribution readable. +2. A formal trace, collected with the real serving optimizations enabled, to + measure actual overlap and critical-path exposure. + +It prints: +- mapping/formal trace summaries +- an action table that maps formal overlap findings back to Python code +- a few ASCII timelines from the formal trace +""" + +from __future__ import annotations + +import argparse +import math +import re +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +from profile_common import ( + coerce_optional_int, + contains_any_keyword, + extract_trace_events, + has_stream_marker, + is_annotation_event, + is_complete_duration_event, + is_non_kernel_trace_category, + is_trace_metadata_name, + load_server_args, + load_trace_json, + looks_like_python_scope_name, + normalize_repo_relative_path, + normalize_text, +) +from profile_common import run_profiler as shared_run_profiler +from profile_common import ( + select_heaviest_pid, +) + +COMMUNICATION_STRONG_KEYWORDS = ( + "allreduce", + "all_reduce", + "reduce_scatter", + "allgather", + "all_gather", + "nccl", + "cross_device_reduce", + "deepep", + "a2a", + "alltoall", + "allreduce_fusion", + "mooncake", +) + +COMMUNICATION_WEAK_KEYWORDS = ( + "broadcast", + "dispatch", + "combine", +) + +MEMORY_STRONG_KEYWORDS = ( + "memcpy", + "memset", + "dma", + "prefetch", +) + +MEMORY_WEAK_KEYWORDS = ( + "fill", + "copy", +) + +ELEMENTWISE_KEYWORDS = ( + "sigmoid", + "silu", + "gelu", + "relu", + "softmax", + "layernorm", + "rmsnorm", + "norm", + "rotary", + "rope", + "topk", + "gate", + "bias", + "_cast", + "index", + "gather", + "scatter", + "masked", + "elementwise", + "activation", +) + +COMPUTE_KEYWORDS = ( + "cublas", + "cudnn", + "cutlass", + "triton", + "gemm", + "gemv", + "matmul", + "grouped_mm", + "flash", + "attention", + "fmha", + "marlin", + "fused_moe", + "moe_kernel", + "groupgemm", + "mma", + "wgmma", + "conv", + "bmm", + "mm_kernel", +) + +CATEGORY_CHARS = { + "compute": "#", + "communication": "=", + "elementwise": "~", + "memory": "+", + "other": "*", +} + +CATEGORY_PRIORITY = { + "compute": 4, + "communication": 3, + "memory": 2, + "elementwise": 1, + "other": 0, +} + +PYTHON_SCOPE_IGNORE_PREFIXES = ( + "threading.py(", + "selectors.py(", + "contextlib.py(", + "queue.py(", + "logging/", + "logging/__init__.py(", + "socket.py(", + "asyncio/", + "concurrent/futures/", + "tqdm/", + "uvicorn/", + "fastapi/", + "starlette/", + "http/", + "torch/_ops.py(", + "torch/nn/modules/module.py(", + "torch/utils/_contextlib.py(", + "torch/autograd/", + "torch/_tensor.py(", + "torch/distributed/", + "torch/_dynamo/", + "torch/_inductor/", +) +KERNEL_NAME_HINTS = ( + COMMUNICATION_STRONG_KEYWORDS + + COMMUNICATION_WEAK_KEYWORDS + + MEMORY_STRONG_KEYWORDS + + MEMORY_WEAK_KEYWORDS + + COMPUTE_KEYWORDS +) + + +@dataclass +class KernelEvent: + idx: int + name: str + canonical_name: str + category: str + pid: str + tid: str + stream: str + ts: float + dur: float + end: float + external_id: Optional[int] = None + correlation: Optional[int] = None + hidden_us: float = 0.0 + exclusive_us: float = 0.0 + hidden_by_compute_us: float = 0.0 + overlap_with: Counter = field(default_factory=Counter) + + +@dataclass +class AggregateStats: + name: str + category: str + count: int = 0 + total_us: float = 0.0 + hidden_us: float = 0.0 + exclusive_us: float = 0.0 + hidden_by_compute_us: float = 0.0 + overlap_with: Counter = field(default_factory=Counter) + representative_idx: Optional[int] = None + representative_score: float = -1.0 + + @property + def hidden_ratio(self) -> float: + return self.hidden_us / self.total_us if self.total_us else 0.0 + + @property + def exclusive_ratio(self) -> float: + return self.exclusive_us / self.total_us if self.total_us else 0.0 + + +@dataclass +class PythonScope: + name: str + normalized_name: str + pid: str + tid: str + ts: float + dur: float + end: float + + +@dataclass +class CPUOpContext: + external_id: int + cpu_op_name: str + pid: str + tid: str + ts: float + dur: float + end: float + scope_chain: Tuple[str, ...] + + +@dataclass +class KernelSourceStats: + name: str + total_count: int = 0 + mapped_count: int = 0 + scope_counter: Counter = field(default_factory=Counter) + chain_counter: Counter = field(default_factory=Counter) + launch_op_counter: Counter = field(default_factory=Counter) + + @property + def mapping_ratio(self) -> float: + return self.mapped_count / self.total_count if self.total_count else 0.0 + + @property + def best_scope(self) -> Optional[str]: + return self.scope_counter.most_common(1)[0][0] if self.scope_counter else None + + @property + def best_chain(self) -> Optional[str]: + return self.chain_counter.most_common(1)[0][0] if self.chain_counter else None + + @property + def best_launch_op(self) -> Optional[str]: + return ( + self.launch_op_counter.most_common(1)[0][0] + if self.launch_op_counter + else None + ) + + +@dataclass +class TraceBundle: + label: str + trace_path: Path + server_args: Optional[dict] + raw_events: Sequence[dict] + events: List[KernelEvent] + pid: Optional[str] + overlap_stats: Optional[Dict[str, float]] = None + + +@dataclass +class ActionRow: + priority: str + verdict: str + kernel: str + category: str + total_us: float + share_pct: float + exclusive_ratio: float + hidden_ratio: float + python_scope: str + launch_op: str + mapping_ratio: float + dependency_signal: str + prev_neighbor: str + next_neighbor: str + recommendation: str + suggestion: str + representative_idx: Optional[int] + + +def short_name(name: str, max_len: int = 80) -> str: + name = normalize_text(name) + if len(name) <= max_len: + return name + return name[: max_len - 3] + "..." + + +def canonicalize_name(name: str) -> str: + name = normalize_text(name) + name = re.sub(r"0x[0-9a-fA-F]+", "0xADDR", name) + if name.startswith("void ") and name.endswith(")"): + depth = 0 + split_idx: Optional[int] = None + for idx in range(len(name) - 1, -1, -1): + char = name[idx] + if char == ")": + depth += 1 + elif char == "(": + depth -= 1 + if depth == 0: + split_idx = idx + break + if split_idx is not None: + name = name[:split_idx] + return name + + +def canonicalize_python_scope_name(name: str) -> str: + name = normalize_text(name) + name = re.sub(r"0x[0-9a-fA-F]+", "0xADDR", name) + match = re.match(r"(?P.+?)\((?P\d+)\): (?P.+)$", name) + if match: + path = normalize_repo_relative_path(match.group("path")) + name = f"{path}({match.group('line')}): {match.group('func')}" + return name + + +def canonicalize_cpu_op_name(name: str) -> str: + return short_name(normalize_text(name), max_len=100) + + +def classify_kernel(name: str) -> str: + # This script only needs broad overlap buckets, so keep the precedence small + # and deterministic: memory/communication first, then compute/elementwise. + lowered = name.lower() + looks_compute_like = contains_any_keyword(lowered, COMPUTE_KEYWORDS) + if contains_any_keyword(lowered, MEMORY_STRONG_KEYWORDS): + return "memory" + if contains_any_keyword(lowered, COMMUNICATION_STRONG_KEYWORDS): + return "communication" + if contains_any_keyword(lowered, COMPUTE_KEYWORDS): + return "compute" + if contains_any_keyword(lowered, ELEMENTWISE_KEYWORDS): + return "elementwise" + if contains_any_keyword(lowered, MEMORY_WEAK_KEYWORDS) and not looks_compute_like: + return "memory" + if ( + contains_any_keyword(lowered, COMMUNICATION_WEAK_KEYWORDS) + and not looks_compute_like + ): + return "communication" + if lowered.startswith("void "): + return "other" + return "other" + + +def is_kernel_event(event: dict) -> bool: + # The overlap script prefers a slightly broader kernel detector than the + # breakdown script, but it still rejects annotations and Python frames up + # front so the later overlap math only sees real GPU work. + if not is_complete_duration_event(event): + return False + name = normalize_text(event.get("name", "")) + if is_trace_metadata_name(name): + return False + cat = normalize_text(event.get("cat", "")).lower() + args = event.get("args", {}) or {} + if is_non_kernel_trace_category(cat): + return False + if is_annotation_event(name, cat): + return False + if "kernel" in cat or cat.startswith("gpu_"): + return True + lowered = name.lower() + if looks_like_python_scope_name(name): + return False + if has_stream_marker(args) and ( + lowered.startswith("void ") + or lowered.startswith("ampere_") + or lowered.startswith("sm80_") + or lowered.startswith("sm90_") + or contains_any_keyword(lowered, KERNEL_NAME_HINTS) + ): + return True + return False + + +def is_meaningful_python_scope(name: str) -> bool: + normalized = canonicalize_python_scope_name(name) + if not normalized: + return False + if normalized.startswith(" bool: + normalized = canonicalize_python_scope_name(name) + if ( + not normalized + or normalized.startswith(" Dict[Tuple[str, str], str]: + mapping: Dict[Tuple[str, str], str] = {} + for event in events: + if event.get("ph") != "M" or event.get("name") != "thread_name": + continue + pid = str(event.get("pid")) + tid = str(event.get("tid")) + thread_name = str((event.get("args") or {}).get("name", "")) + if thread_name: + mapping[(pid, tid)] = thread_name + return mapping + + +def build_correlation_external_lookup(raw_events: Sequence[dict]) -> Dict[int, int]: + lookup: Dict[int, int] = {} + for event in raw_events: + args = event.get("args", {}) or {} + correlation = coerce_optional_int(args.get("correlation")) + external_id = coerce_optional_int(args.get("External id")) + if correlation is not None and external_id is not None: + lookup[correlation] = external_id + return lookup + + +def extract_kernel_events( + trace: dict, pid_substring: Optional[str] +) -> Tuple[List[KernelEvent], Optional[str]]: + # We first build a clean kernel list from the chosen TP rank, then later + # overlap analysis can stay focused on stream timing instead of trace noise. + raw_events = extract_trace_events(trace) + thread_names = extract_thread_names(raw_events) + correlation_external = build_correlation_external_lookup(raw_events) + chosen_pid = select_heaviest_pid( + raw_events, + is_kernel_event, + pid_substring=pid_substring, + preferred_substrings=(() if pid_substring else ("TP00",)), + ) + kernel_events: List[KernelEvent] = [] + if chosen_pid is None: + return kernel_events, None + + idx = 0 + for event in raw_events: + if not is_kernel_event(event): + continue + pid = str(event.get("pid")) + if pid != chosen_pid: + continue + tid = str(event.get("tid")) + args = event.get("args", {}) or {} + stream = ( + args.get("stream") + or args.get("cuda_stream") + or thread_names.get((pid, tid)) + or f"tid={tid}" + ) + correlation = coerce_optional_int(args.get("correlation")) + external_id = coerce_optional_int(args.get("External id")) + if external_id is None and correlation is not None: + external_id = correlation_external.get(correlation) + name = str(event["name"]) + dur = float(event["dur"]) + ts = float(event["ts"]) + kernel_events.append( + KernelEvent( + idx=idx, + name=name, + canonical_name=canonicalize_name(name), + category=classify_kernel(name), + pid=pid, + tid=tid, + stream=str(stream), + ts=ts, + dur=dur, + end=ts + dur, + external_id=external_id, + correlation=correlation, + ) + ) + idx += 1 + return kernel_events, chosen_pid + + +def dominant_overlap_name( + event: KernelEvent, active_events: Iterable[KernelEvent] +) -> Optional[str]: + candidates = [ + other + for other in active_events + if other.idx != event.idx and other.stream != event.stream + ] + if not candidates: + return None + candidates.sort( + key=lambda other: (CATEGORY_PRIORITY.get(other.category, 0), other.dur), + reverse=True, + ) + return candidates[0].canonical_name + + +def analyze_overlap(events: Sequence[KernelEvent]) -> Dict[str, float]: + # Sweep line over kernel start/end points. For each active time slice we + # decide whether a kernel was exposed on the critical path or hidden by work + # on other streams. + points: List[Tuple[float, int, int]] = [] + event_map = {event.idx: event for event in events} + for event in events: + points.append((event.ts, 1, event.idx)) + points.append((event.end, 0, event.idx)) + points.sort(key=lambda item: (item[0], item[1])) + + total_busy = 0.0 + total_overlap = 0.0 + max_concurrent = 0 + active: Dict[int, KernelEvent] = {} + prev_time: Optional[float] = None + + for time_point, is_start, event_idx in points: + if prev_time is not None and time_point > prev_time and active: + segment = time_point - prev_time + active_events = list(active.values()) + distinct_streams = {event.stream for event in active_events} + total_busy += segment + max_concurrent = max(max_concurrent, len(distinct_streams)) + if len(distinct_streams) >= 2: + total_overlap += segment + for event in active_events: + overlapping_events = [ + other + for other in active_events + if other.idx != event.idx and other.stream != event.stream + ] + if overlapping_events: + event.hidden_us += segment + if any(other.category == "compute" for other in overlapping_events): + event.hidden_by_compute_us += segment + overlap_name = dominant_overlap_name(event, active_events) + if overlap_name: + event.overlap_with[overlap_name] += segment + else: + event.exclusive_us += segment + + if is_start == 0: + active.pop(event_idx, None) + else: + active[event_idx] = event_map[event_idx] + prev_time = time_point + + return { + "total_busy_us": total_busy, + "total_overlap_us": total_overlap, + "max_concurrent_streams": float(max_concurrent), + } + + +def aggregate_events( + events: Sequence[KernelEvent], +) -> Dict[Tuple[str, str], AggregateStats]: + aggregates: Dict[Tuple[str, str], AggregateStats] = {} + for event in events: + key = (event.canonical_name, event.category) + if key not in aggregates: + aggregates[key] = AggregateStats( + name=event.canonical_name, category=event.category + ) + stats = aggregates[key] + stats.count += 1 + stats.total_us += event.dur + stats.hidden_us += event.hidden_us + stats.exclusive_us += event.exclusive_us + stats.hidden_by_compute_us += event.hidden_by_compute_us + stats.overlap_with.update(event.overlap_with) + score = event.hidden_us + event.exclusive_us + if score > stats.representative_score: + stats.representative_score = score + stats.representative_idx = event.idx + return aggregates + + +def top_hidden_low_roi( + aggregates: Dict[Tuple[str, str], AggregateStats], +) -> List[AggregateStats]: + candidates = [ + stats + for stats in aggregates.values() + if stats.category in {"elementwise", "memory"} + and stats.total_us >= 5.0 + and stats.hidden_ratio >= 0.65 + ] + candidates.sort( + key=lambda stats: ( + stats.hidden_us + * (1.0 + stats.hidden_by_compute_us / max(stats.hidden_us, 1.0)), + stats.hidden_ratio, + ), + reverse=True, + ) + return candidates[:5] + + +def top_overlap_opportunities( + aggregates: Dict[Tuple[str, str], AggregateStats], +) -> List[AggregateStats]: + category_weight = { + "communication": 1.3, + "memory": 1.15, + "elementwise": 1.0, + "compute": 0.35, + "other": 0.8, + } + candidates = [ + stats + for stats in aggregates.values() + if stats.total_us >= 5.0 and stats.exclusive_ratio >= 0.45 + ] + primary = [stats for stats in candidates if stats.category != "compute"] + fallback = [stats for stats in candidates if stats.category == "compute"] + primary.sort( + key=lambda stats: stats.exclusive_us * category_weight.get(stats.category, 1.0), + reverse=True, + ) + fallback.sort( + key=lambda stats: stats.exclusive_us * category_weight.get(stats.category, 1.0), + reverse=True, + ) + return (primary + fallback)[:5] + + +def choose_window_events( + events: Sequence[KernelEvent], + representative_idx: int, + window_us: Optional[float], +) -> Tuple[float, float, List[KernelEvent]]: + center = next(event for event in events if event.idx == representative_idx) + span = window_us if window_us is not None else max(40.0, center.dur * 6.0) + start = max(0.0, center.ts - span * 0.35) + end = center.end + span * 0.65 + window_events = [ + event for event in events if event.end >= start and event.ts <= end + ] + return start, end, window_events + + +def render_ascii_timeline( + events: Sequence[KernelEvent], + representative_idx: int, + window_us: Optional[float], + width: int, +) -> str: + start, end, window_events = choose_window_events( + events, representative_idx, window_us + ) + if not window_events: + return "No events found in the selected window." + + streams = sorted( + {event.stream for event in window_events}, key=lambda item: (len(item), item) + ) + symbol_map: Dict[int, str] = {} + legend_events = sorted(window_events, key=lambda event: event.dur, reverse=True)[:8] + symbol_alphabet = list( + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + ) + for index, event in enumerate(legend_events): + symbol_map[event.idx] = symbol_alphabet[index] + + label_width = max(len(stream) for stream in streams) + lines = [] + marker_positions = [0, width // 4, width // 2, (3 * width) // 4, width - 1] + header = [" "] * width + for position in marker_positions: + header[position] = "|" + lines.append("time(us) " + "".join(header)) + + time_line = [" "] * width + markers = [ + start, + start + (end - start) * 0.25, + start + (end - start) * 0.5, + start + (end - start) * 0.75, + end, + ] + for position, value in zip(marker_positions, markers): + text = f"{value:.1f}" + begin = min(max(position - len(text) // 2, 0), max(0, width - len(text))) + for offset, char in enumerate(text): + time_line[begin + offset] = char + lines.append(" " + "".join(time_line)) + + for stream in streams: + row = ["."] * width + row_events = [event for event in window_events if event.stream == stream] + for event in row_events: + char = symbol_map.get(event.idx, CATEGORY_CHARS[event.category]) + left = int((event.ts - start) / max(end - start, 1.0) * (width - 1)) + right = int( + math.ceil((event.end - start) / max(end - start, 1.0) * (width - 1)) + ) + right = max(left + 1, min(right, width - 1)) + for pos in range(max(0, left), min(width, right + 1)): + row[pos] = char + lines.append(f"{stream:<{label_width}} " + "".join(row)) + + if legend_events: + lines.append("legend:") + for event in legend_events: + symbol = symbol_map[event.idx] + lines.append( + f" {symbol} [{event.category[:4]}] {short_name(event.canonical_name, 72)} ({event.dur:.1f} us)" + ) + return "\n".join(lines) + + +def choose_best_scope(scope_chain: Sequence[str]) -> Optional[str]: + ranked: List[Tuple[float, str]] = [] + for index, scope in enumerate(scope_chain): + score = float(index) + if scope.startswith("python/sglang/"): + score += 50.0 + elif scope.startswith("sglang/"): + score += 48.0 + elif scope.startswith("sgl_kernel/"): + score += 30.0 + elif ".py(" in scope: + score += 10.0 + if "utils.py" in scope and "__call__" in scope: + score -= 15.0 + if "scheduler_profiler_mixin.py" in scope: + score -= 20.0 + ranked.append((score, scope)) + return max(ranked, key=lambda item: item[0])[1] if ranked else None + + +def scope_chain_key(scope_chain: Sequence[str]) -> Optional[str]: + if not scope_chain: + return None + trimmed = list(scope_chain[-4:]) + return " -> ".join(trimmed) + + +def extract_cpu_launch_contexts( + raw_events: Sequence[dict], +) -> Dict[int, List[CPUOpContext]]: + # Rebuild `External id -> CPU op -> active Python scopes` so formal-trace + # kernels can be mapped back to readable Python locations from the mapping + # trace even when launches are interleaved on the same thread. + scopes_by_thread: Dict[Tuple[str, str], List[PythonScope]] = defaultdict(list) + cpu_ops_by_thread: Dict[Tuple[str, str], List[CPUOpContext]] = defaultdict(list) + + for event in raw_events: + if not is_complete_duration_event(event): + continue + cat = str(event.get("cat", "")) + pid = str(event.get("pid")) + tid = str(event.get("tid")) + ts = float(event.get("ts", 0.0)) + dur = float(event.get("dur", 0.0)) + args = event.get("args", {}) or {} + if cat == "python_function": + name = canonicalize_python_scope_name(event.get("name", "")) + scopes_by_thread[(pid, tid)].append( + PythonScope( + name=str(event.get("name", "")), + normalized_name=name, + pid=pid, + tid=tid, + ts=ts, + dur=dur, + end=ts + dur, + ) + ) + elif cat == "cpu_op": + external_id = coerce_optional_int(args.get("External id")) + if external_id is None: + continue + cpu_ops_by_thread[(pid, tid)].append( + CPUOpContext( + external_id=external_id, + cpu_op_name=str(event.get("name", "")), + pid=pid, + tid=tid, + ts=ts, + dur=dur, + end=ts + dur, + scope_chain=(), + ) + ) + + contexts_by_external_id: Dict[int, List[CPUOpContext]] = defaultdict(list) + for thread_key in set(scopes_by_thread) | set(cpu_ops_by_thread): + scopes = scopes_by_thread.get(thread_key, []) + cpu_ops = cpu_ops_by_thread.get(thread_key, []) + timeline = [] + for scope in scopes: + timeline.append((scope.ts, 0, scope)) + timeline.append((scope.end, 2, scope)) + for cpu_op in cpu_ops: + timeline.append((cpu_op.ts, 1, cpu_op)) + timeline.sort(key=lambda item: (item[0], item[1])) + + active_scopes: List[PythonScope] = [] + for _, kind, payload in timeline: + if kind == 0: + active_scopes.append(payload) + elif kind == 1: + normalized_chain = [scope.normalized_name for scope in active_scopes] + meaningful = [ + scope + for scope in normalized_chain + if is_meaningful_python_scope(scope) + ] + fallback = [ + scope + for scope in normalized_chain + if is_fallback_python_scope(scope) + ] + chosen_chain = tuple((meaningful or fallback)[-6:]) + contexts_by_external_id[payload.external_id].append( + CPUOpContext( + external_id=payload.external_id, + cpu_op_name=payload.cpu_op_name, + pid=payload.pid, + tid=payload.tid, + ts=payload.ts, + dur=payload.dur, + end=payload.end, + scope_chain=chosen_chain, + ) + ) + else: + if payload in active_scopes: + active_scopes.remove(payload) + return contexts_by_external_id + + +def choose_cpu_context( + contexts: Sequence[CPUOpContext], kernel_ts: float +) -> Optional[CPUOpContext]: + if not contexts: + return None + return min(contexts, key=lambda context: (abs(context.ts - kernel_ts), context.dur)) + + +def extract_meaningful_python_scopes(raw_events: Sequence[dict]) -> List[PythonScope]: + scopes: List[PythonScope] = [] + for event in raw_events: + if not is_complete_duration_event(event): + continue + if str(event.get("cat", "")) != "python_function": + continue + ts = float(event.get("ts", 0.0)) + dur = float(event.get("dur", 0.0)) + normalized_name = canonicalize_python_scope_name(event.get("name", "")) + if not is_meaningful_python_scope(normalized_name): + continue + scopes.append( + PythonScope( + name=str(event.get("name", "")), + normalized_name=normalized_name, + pid=str(event.get("pid")), + tid=str(event.get("tid")), + ts=ts, + dur=dur, + end=ts + dur, + ) + ) + return scopes + + +def choose_temporal_scope_chain( + scopes: Sequence[PythonScope], kernel_ts: float +) -> Tuple[str, ...]: + matches = [scope for scope in scopes if scope.ts <= kernel_ts <= scope.end] + if not matches: + return () + matches.sort(key=lambda scope: (scope.ts, -scope.dur, scope.normalized_name)) + chain = [] + seen = set() + for scope in matches: + if scope.normalized_name in seen: + continue + seen.add(scope.normalized_name) + chain.append(scope.normalized_name) + return tuple(chain[-6:]) + + +def build_kernel_source_map( + mapping_bundle: TraceBundle, +) -> Dict[str, KernelSourceStats]: + contexts_by_external_id = extract_cpu_launch_contexts(mapping_bundle.raw_events) + temporal_scopes = extract_meaningful_python_scopes(mapping_bundle.raw_events) + source_map: Dict[str, KernelSourceStats] = {} + for event in mapping_bundle.events: + stats = source_map.setdefault( + event.canonical_name, KernelSourceStats(name=event.canonical_name) + ) + stats.total_count += 1 + cpu_context = None + if event.external_id is not None: + cpu_context = choose_cpu_context( + contexts_by_external_id.get(event.external_id, []), event.ts + ) + + launch_op = None + scope_chain: Tuple[str, ...] = () + if cpu_context is not None: + launch_op = canonicalize_cpu_op_name(cpu_context.cpu_op_name) + scope_chain = cpu_context.scope_chain + else: + scope_chain = choose_temporal_scope_chain(temporal_scopes, event.ts) + if scope_chain: + launch_op = "time-window fallback" + + if not scope_chain: + continue + + stats.mapped_count += 1 + best_scope = choose_best_scope(scope_chain) + if best_scope: + stats.scope_counter[best_scope] += 1 + chain = scope_chain_key(scope_chain) + if chain: + stats.chain_counter[chain] += 1 + if launch_op: + stats.launch_op_counter[launch_op] += 1 + return source_map + + +def format_overlap_counter(counter: Counter, limit: int = 2) -> str: + if not counter: + return "n/a" + parts = [] + for name, duration in counter.most_common(limit): + parts.append(f"{short_name(name, 48)} ({duration:.1f} us)") + return ", ".join(parts) + + +def build_headroom_suggestion(stats: AggregateStats) -> str: + if stats.category == "communication": + return "Exposed comm path. Check whether this code path can overlap with nearby compute." + if stats.category in {"elementwise", "memory"}: + return "Still exposed. Try to fuse it or move it under a nearby compute-heavy window." + return "Meaningful exposed time remains. Inspect stream placement and surrounding dependencies." + + +def build_hidden_suggestion(stats: AggregateStats) -> str: + overlap = format_overlap_counter(stats.overlap_with, limit=1) + if overlap != "n/a": + return f"Mostly hidden under {overlap}. Standalone tuning is probably low ROI." + return ( + "Mostly hidden already. Optimize it only if you also change fusion or schedule." + ) + + +def build_other_suggestion(stats: AggregateStats) -> str: + if stats.exclusive_ratio >= 0.6: + return "Some exposed time remains, but it did not rank among the strongest headroom rows." + if stats.hidden_ratio >= 0.6: + return "Often hidden already. Usually secondary unless it also drives launch count." + return "Mixed exposure and overlap. Inspect after the stronger rows above." + + +def parse_scope_signature(scope: str) -> Tuple[str, str]: + if not scope or scope in {"unmapped", "n/a"}: + return "", "" + match = re.match(r"(.+?)\(\d+\):\s*(.+)$", scope) + if match: + return match.group(1), match.group(2) + return scope, "" + + +def same_scope_family(left: str, right: str) -> bool: + left_path, left_func = parse_scope_signature(left) + right_path, right_func = parse_scope_signature(right) + if not left_path or not right_path: + return False + if left_path == right_path: + return True + return bool(left_func and right_func and left_func == right_func) + + +def is_neighbor_dependency_like( + current: KernelEvent, neighbor: Optional[KernelEvent] +) -> bool: + if neighbor is None: + return False + if current.category == "communication": + return neighbor.category in {"compute", "elementwise", "memory", "other"} + if current.category in {"elementwise", "memory"}: + return neighbor.category in { + "compute", + "communication", + "elementwise", + "memory", + } + return False + + +def build_stream_neighbor_index( + events: Sequence[KernelEvent], +) -> Dict[int, Tuple[Optional[KernelEvent], Optional[KernelEvent]]]: + by_stream: Dict[str, List[KernelEvent]] = defaultdict(list) + for event in events: + by_stream[event.stream].append(event) + + index: Dict[int, Tuple[Optional[KernelEvent], Optional[KernelEvent]]] = {} + for stream_events in by_stream.values(): + stream_events.sort(key=lambda event: (event.ts, event.end, event.idx)) + for pos, event in enumerate(stream_events): + prev_event = stream_events[pos - 1] if pos > 0 else None + next_event = ( + stream_events[pos + 1] if pos + 1 < len(stream_events) else None + ) + index[event.idx] = (prev_event, next_event) + return index + + +def describe_neighbor( + neighbor: Optional[KernelEvent], + gap_us: Optional[float], + source_map: Dict[str, KernelSourceStats], +) -> str: + if neighbor is None: + return "none" + source = source_map.get(neighbor.canonical_name) + scope = source.best_scope if source and source.best_scope else "unmapped" + if gap_us is not None: + gap_us = max(gap_us, 0.0) + gap_text = f"{gap_us:.1f} us" + else: + gap_text = "n/a" + return ( + f"{short_name(neighbor.canonical_name, 28)} " + f"@ {short_name(scope, 28)} " + f"(gap {gap_text})" + ) + + +def classify_dependency_signal( + current: KernelEvent, + source: Optional[KernelSourceStats], + prev_event: Optional[KernelEvent], + next_event: Optional[KernelEvent], + source_map: Dict[str, KernelSourceStats], +) -> Tuple[str, str, str]: + current_scope = source.best_scope if source and source.best_scope else "unmapped" + current_launch = ( + source.best_launch_op if source and source.best_launch_op else "n/a" + ) + + prev_gap = current.ts - prev_event.end if prev_event is not None else None + next_gap = next_event.ts - current.end if next_event is not None else None + prev_source = ( + source_map.get(prev_event.canonical_name) if prev_event is not None else None + ) + next_source = ( + source_map.get(next_event.canonical_name) if next_event is not None else None + ) + prev_scope = ( + prev_source.best_scope if prev_source and prev_source.best_scope else "unmapped" + ) + next_scope = ( + next_source.best_scope if next_source and next_source.best_scope else "unmapped" + ) + prev_launch = ( + prev_source.best_launch_op + if prev_source and prev_source.best_launch_op + else "n/a" + ) + next_launch = ( + next_source.best_launch_op + if next_source and next_source.best_launch_op + else "n/a" + ) + + if prev_gap is not None: + prev_gap = max(prev_gap, 0.0) + if next_gap is not None: + next_gap = max(next_gap, 0.0) + + tight_gap_threshold = max(2.0, min(20.0, current.dur * 0.15)) + prev_tight = prev_gap is not None and prev_gap <= tight_gap_threshold + next_tight = next_gap is not None and next_gap <= tight_gap_threshold + + prev_risk = prev_tight and ( + same_scope_family(current_scope, prev_scope) + or (current_launch != "n/a" and current_launch == prev_launch) + or is_neighbor_dependency_like(current, prev_event) + ) + next_risk = next_tight and ( + same_scope_family(current_scope, next_scope) + or (current_launch != "n/a" and current_launch == next_launch) + or is_neighbor_dependency_like(current, next_event) + ) + + prev_unclear = ( + prev_tight + and not prev_risk + and (current_scope == "unmapped" or prev_scope == "unmapped") + ) + next_unclear = ( + next_tight + and not next_risk + and (current_scope == "unmapped" or next_scope == "unmapped") + ) + + if prev_risk and next_risk: + signal = "both-side serial risk" + elif prev_risk: + signal = "prev-side serial risk" + elif next_risk: + signal = "next-side serial risk" + elif prev_unclear or next_unclear: + signal = "adjacency unclear" + else: + signal = "serial risk low" + + prev_desc = describe_neighbor(prev_event, prev_gap, source_map) + next_desc = describe_neighbor(next_event, next_gap, source_map) + return signal, prev_desc, next_desc + + +def dependency_risk_label(signal: str) -> str: + mapping = { + "serial risk low": "low", + "prev-side serial risk": "high", + "next-side serial risk": "high", + "both-side serial risk": "high", + "adjacency unclear": "unclear", + } + return mapping.get(signal, signal) + + +def build_priority_and_recommendation( + verdict: str, + category: str, + dependency_signal: str, + stats: AggregateStats, + share_pct: float, +) -> Tuple[str, str]: + dep_label = dependency_risk_label(dependency_signal) + if share_pct < 1.0: + return "P5", "skip overlap" + + if verdict == "headroom": + if dep_label == "low": + if category == "communication": + return "P1", "try overlap" + return "P1", "try fusion" + return "P2", "check deps" + + if verdict == "low-roi-hidden": + return "P4", "skip overlap" + + if stats.exclusive_ratio >= 0.85 and dep_label == "low": + return "P3", "observe later" + if stats.hidden_ratio >= 0.7: + return "P5", "skip overlap" + if dep_label == "high": + return "P4", "check deps" + if dep_label == "unclear": + return "P4", "manual check" + return "P4", "observe later" + + +def make_action_row( + stats: AggregateStats, + verdict: str, + suggestion: str, + source_map: Dict[str, KernelSourceStats], + formal_events: Sequence[KernelEvent], + neighbor_index: Dict[int, Tuple[Optional[KernelEvent], Optional[KernelEvent]]], + total_busy_us: float, +) -> ActionRow: + source = source_map.get(stats.name) + representative_idx = stats.representative_idx + dependency_signal = "adjacency unclear" + prev_neighbor = "none" + next_neighbor = "none" + share_pct = (stats.total_us / total_busy_us * 100.0) if total_busy_us > 0 else 0.0 + if representative_idx is not None: + current_event = next( + (event for event in formal_events if event.idx == representative_idx), None + ) + if current_event is not None: + prev_event, next_event = neighbor_index.get( + representative_idx, (None, None) + ) + dependency_signal, prev_neighbor, next_neighbor = ( + classify_dependency_signal( + current=current_event, + source=source, + prev_event=prev_event, + next_event=next_event, + source_map=source_map, + ) + ) + priority, recommendation = build_priority_and_recommendation( + verdict=verdict, + category=stats.category, + dependency_signal=dependency_signal, + stats=stats, + share_pct=share_pct, + ) + + return ActionRow( + priority=priority, + verdict=verdict, + kernel=stats.name, + category=stats.category, + total_us=stats.total_us, + share_pct=share_pct, + exclusive_ratio=stats.exclusive_ratio, + hidden_ratio=stats.hidden_ratio, + python_scope=source.best_scope if source and source.best_scope else "unmapped", + launch_op=source.best_launch_op if source and source.best_launch_op else "n/a", + mapping_ratio=source.mapping_ratio if source else 0.0, + dependency_signal=dependency_signal, + prev_neighbor=prev_neighbor, + next_neighbor=next_neighbor, + recommendation=recommendation, + suggestion=suggestion, + representative_idx=representative_idx, + ) + + +def build_action_rows( + aggregates: Dict[Tuple[str, str], AggregateStats], + source_map: Dict[str, KernelSourceStats], + formal_events: Sequence[KernelEvent], + total_busy_us: float, + table_limit: int, +) -> List[ActionRow]: + rows: List[ActionRow] = [] + seen: set[str] = set() + neighbor_index = build_stream_neighbor_index(formal_events) + + for stats in top_overlap_opportunities(aggregates): + row = make_action_row( + stats=stats, + verdict="headroom", + suggestion=build_headroom_suggestion(stats), + source_map=source_map, + formal_events=formal_events, + neighbor_index=neighbor_index, + total_busy_us=total_busy_us, + ) + if row.priority == "P5": + continue + rows.append(row) + seen.add(stats.name) + + for stats in top_hidden_low_roi(aggregates): + if stats.name in seen: + continue + rows.append( + make_action_row( + stats=stats, + verdict="low-roi-hidden", + suggestion=build_hidden_suggestion(stats), + source_map=source_map, + formal_events=formal_events, + neighbor_index=neighbor_index, + total_busy_us=total_busy_us, + ) + ) + seen.add(stats.name) + + if table_limit > 0: + return rows[:table_limit] + return rows + + +def render_action_table(rows: Sequence[ActionRow]) -> List[str]: + lines = [ + "| Priority | Verdict | Kernel | Python scope | Formal signal | Dep risk | Recommendation |", + "| --- | --- | --- | --- | --- | --- | --- |", + ] + if not rows: + lines.append( + "| - | No actionable overlap rows stood out from the formal trace. | - | - | - | - | - |" + ) + return lines + for row in rows: + formal_signal = ( + f"share {row.share_pct:.1f}%, " + f"excl {row.exclusive_ratio * 100:.1f}% / " + f"hid {row.hidden_ratio * 100:.1f}%" + ) + lines.append( + "| " + + " | ".join( + [ + row.priority, + row.verdict, + row.kernel, + row.python_scope, + f"{row.total_us:.1f} us, {formal_signal}", + dependency_risk_label(row.dependency_signal), + row.recommendation, + ] + ) + + " |" + ) + return lines + + +def trace_summary_line(bundle: TraceBundle) -> str: + events = bundle.events + streams = sorted({event.stream for event in events}) + if bundle.overlap_stats is None: + return f"{bundle.label}: {len(events)} kernel events, {len(streams)} streams" + overlap_ratio = ( + bundle.overlap_stats["total_overlap_us"] / bundle.overlap_stats["total_busy_us"] + if bundle.overlap_stats["total_busy_us"] + else 0.0 + ) + return ( + f"{bundle.label}: {len(events)} kernel events, {len(streams)} streams, " + f"busy={bundle.overlap_stats['total_busy_us']:.1f} us, " + f"2+ stream overlap={bundle.overlap_stats['total_overlap_us']:.1f} us " + f"({overlap_ratio * 100:.1f}%), " + f"peak_concurrency={int(bundle.overlap_stats['max_concurrent_streams'])}" + ) + + +def launch_summary(server_args: Optional[dict]) -> Optional[str]: + if not server_args: + return None + model_path = server_args.get("model_path") or server_args.get("model") + shape_bits = [] + if model_path: + shape_bits.append(f"model={model_path}") + for key in ("tp_size", "dp_size", "pp_size", "ep_size", "enable_dp_attention"): + if key in server_args: + shape_bits.append(f"{key}={server_args[key]}") + return ", ".join(shape_bits) if shape_bits else None + + +def build_report( + mapping_bundle: TraceBundle, + formal_bundle: TraceBundle, + source_map: Dict[str, KernelSourceStats], + aggregates: Dict[Tuple[str, str], AggregateStats], + rows: Sequence[ActionRow], + window_us: Optional[float], + timeline_count: int, + width: int, + table_only: bool, +) -> str: + lines: List[str] = [] + lines.append(f"Mapping Trace: {mapping_bundle.trace_path}") + mapping_launch = launch_summary(mapping_bundle.server_args) + if mapping_launch: + lines.append(f"Mapping Launch: {mapping_launch}") + if mapping_bundle.pid: + lines.append(f"Mapping PID slice: {mapping_bundle.pid}") + lines.append(trace_summary_line(mapping_bundle)) + + lines.append("") + lines.append(f"Formal Trace: {formal_bundle.trace_path}") + formal_launch = launch_summary(formal_bundle.server_args) + if formal_launch: + lines.append(f"Formal Launch: {formal_launch}") + if formal_bundle.pid: + lines.append(f"Formal PID slice: {formal_bundle.pid}") + lines.append(trace_summary_line(formal_bundle)) + + mapped_kernels = sum(1 for stats in source_map.values() if stats.mapped_count > 0) + table_mapped = sum(1 for row in rows if row.python_scope != "unmapped") + lines.append("") + lines.append( + "Source Map Coverage: " + f"{mapped_kernels}/{len(source_map)} mapping-trace kernels found a Python scope, " + f"{table_mapped}/{len(rows)} table rows were mapped back to code." + ) + + lines.append("") + lines.append("Action Table") + lines.extend(render_action_table(rows)) + + if not table_only: + detail_lookup = {row.kernel: row for row in rows} + focus_rows = list(rows) + lines.append("") + lines.append("Source Context") + if not focus_rows: + lines.append(" No source-mapped rows to expand.") + else: + for index, row in enumerate(focus_rows, start=1): + stats = source_map.get(row.kernel) + lines.append( + f" {index}. {short_name(row.kernel, 88)} [{row.priority}, {row.verdict}, {row.category}] " + f"mapping={row.mapping_ratio * 100:.1f}%" + ) + lines.append(f" time share: {row.share_pct:.1f}%") + lines.append(f" python scope: {row.python_scope}") + lines.append(f" launch op: {row.launch_op}") + lines.append( + f" dependency signal: {dependency_risk_label(row.dependency_signal)}" + ) + lines.append(f" prev neighbor: {row.prev_neighbor}") + lines.append(f" next neighbor: {row.next_neighbor}") + lines.append(f" recommendation: {row.recommendation}") + if stats and stats.best_chain: + lines.append( + f" call chain: {short_name(stats.best_chain, 132)}" + ) + lines.append(f" conclusion: {row.suggestion}") + + timeline_targets: List[int] = [] + for row in rows: + if ( + row.representative_idx is not None + and row.representative_idx not in timeline_targets + ): + timeline_targets.append(row.representative_idx) + timeline_targets = timeline_targets[:timeline_count] + + if timeline_targets: + lines.append("") + lines.append("ASCII Timelines") + for index, representative_idx in enumerate(timeline_targets, start=1): + event = next( + event + for event in formal_bundle.events + if event.idx == representative_idx + ) + mapped_scope = ( + detail_lookup.get(event.canonical_name).python_scope + if event.canonical_name in detail_lookup + else "unmapped" + ) + lines.append( + f" Window {index}: {short_name(event.canonical_name, 90)} " + f"[{event.category}] ts={event.ts:.1f} us dur={event.dur:.1f} us" + ) + lines.append(f" mapped scope: {short_name(mapped_scope, 120)}") + lines.append( + render_ascii_timeline( + formal_bundle.events, representative_idx, window_us, width + ) + ) + lines.append("") + + lines.append("Notes") + lines.append( + " - The mapping trace should be graph-off so kernel-to-code attribution stays readable." + ) + lines.append( + " - The formal trace should keep the real serving optimizations enabled; overlap conclusions come from this trace." + ) + lines.append( + " - A mapped Python scope is a launch-site clue, not proof that the code is dependency-free to reorder." + ) + return "\n".join(lines).rstrip() + + +def discover_trace_file(path: Path) -> Tuple[Path, Optional[dict]]: + if path.is_file(): + return path, load_server_args(path) + + traces = sorted( + path.glob("*.trace.json.gz"), key=lambda candidate: candidate.stat().st_mtime + ) + traces.extend( + sorted( + [ + candidate + for candidate in path.glob("*.trace.json") + if candidate.name not in {trace.name[:-3] for trace in traces} + ], + key=lambda candidate: candidate.stat().st_mtime, + ) + ) + if not traces: + child_dirs = sorted( + [candidate for candidate in path.iterdir() if candidate.is_dir()], + key=lambda candidate: candidate.stat().st_mtime, + ) + for child_dir in reversed(child_dirs): + child_traces = list(child_dir.glob("*.trace.json.gz")) + list( + child_dir.glob("*.trace.json") + ) + if child_traces: + return discover_trace_file(child_dir) + if not traces: + raise FileNotFoundError(f"No trace files found under {path}") + + non_merged = [trace for trace in traces if not trace.name.startswith("merged-")] + tp0 = [ + trace for trace in non_merged if "-TP-0" in trace.name or "TP-0" in trace.name + ] + chosen = tp0[-1] if tp0 else (non_merged[-1] if non_merged else traces[-1]) + return chosen, load_server_args(path) + + +def resolve_trace_source( + label: str, + input_path: Optional[str], + url: Optional[str], + output_dir: Optional[str], + profile_prefix: Optional[str], + args: argparse.Namespace, +) -> TraceBundle: + if bool(input_path) == bool(url): + raise ValueError(f"{label} trace requires exactly one of input path or URL.") + + if url: + target_dir = shared_run_profiler( + url=url, + output_dir=output_dir, + num_steps=args.num_steps, + profile_by_stage=args.profile_by_stage, + merge_profiles=args.merge_profiles, + profile_prefix=profile_prefix, + probe_requests=max(0, args.probe_requests), + probe_prompt=args.probe_prompt, + probe_max_new_tokens=args.probe_max_new_tokens, + probe_delay=args.probe_delay, + start_step=args.start_step, + ) + trace_path, server_args = discover_trace_file(target_dir) + else: + trace_path, server_args = discover_trace_file(Path(input_path).resolve()) + + trace = load_trace_json(trace_path) + raw_events = trace.get("traceEvents", trace if isinstance(trace, list) else []) + events, pid = extract_kernel_events(trace, args.pid_substring) + if not events: + raise RuntimeError(f"No GPU kernel events found in {trace_path}") + return TraceBundle( + label=label, + trace_path=trace_path, + server_args=server_args, + raw_events=raw_events, + events=events, + pid=pid, + ) + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Analyze SGLang profiler overlap with mapping-trace correlation." + ) + 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("--input", type=str, default=None, help=argparse.SUPPRESS) + parser.add_argument("--url", type=str, default=None, help=argparse.SUPPRESS) + + 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( + "--start-step", + type=int, + default=None, + help="Pass through to sglang.profiler when generating traces from URLs.", + ) + parser.add_argument( + "--num-steps", + type=int, + default=5, + help="Number of steps to profile when generating traces from URLs.", + ) + parser.add_argument( + "--profile-by-stage", + action="store_true", + help="Pass through to sglang.profiler.", + ) + parser.add_argument( + "--merge-profiles", action="store_true", help="Pass through to sglang.profiler." + ) + parser.add_argument( + "--probe-requests", + type=int, + default=1, + help="When generating traces from a URL, send this many probe requests so the profiler captures real work.", + ) + parser.add_argument( + "--probe-max-new-tokens", + type=int, + default=None, + help="Override max_new_tokens for the synthetic probe workload. Defaults to max(64, num_steps * 8).", + ) + parser.add_argument( + "--probe-prompt", + type=str, + default=( + "Repeat the word overlap many times with spaces so the server performs several decode steps. " + "Do not add analysis or explanations." + ), + help="Prompt used for the synthetic probe request when profiling a live URL.", + ) + parser.add_argument( + "--probe-delay", + type=float, + default=0.5, + help="Seconds to wait after starting the profiler before sending the synthetic probe request.", + ) + parser.add_argument( + "--pid-substring", + type=str, + default=None, + help="Restrict analysis to PIDs containing this substring.", + ) + parser.add_argument( + "--table-limit", + type=int, + default=0, + help="Maximum number of rows in the action table. Use 0 to print all kernels.", + ) + parser.add_argument( + "--timeline-count", + type=int, + default=3, + help="Number of ASCII windows to render from the formal trace.", + ) + parser.add_argument( + "--timeline-width", type=int, default=96, help="ASCII timeline width." + ) + parser.add_argument( + "--window-us", + type=float, + default=None, + help="Fixed timeline window size in microseconds.", + ) + parser.add_argument( + "--table-only", + action="store_true", + help="Print trace summaries plus the overlap action table only.", + ) + args = parser.parse_args(argv) + + if args.formal_input is None and args.input is not None: + args.formal_input = args.input + if args.formal_url is None and args.url is not None: + args.formal_url = args.url + + 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 main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + + mapping_bundle = resolve_trace_source( + 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_bundle = resolve_trace_source( + 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, + ) + + formal_bundle.overlap_stats = analyze_overlap(formal_bundle.events) + aggregates = aggregate_events(formal_bundle.events) + source_map = build_kernel_source_map(mapping_bundle) + rows = build_action_rows( + aggregates, + source_map, + formal_bundle.events, + formal_bundle.overlap_stats["total_busy_us"], + table_limit=max(0, args.table_limit), + ) + report = build_report( + mapping_bundle=mapping_bundle, + formal_bundle=formal_bundle, + source_map=source_map, + aggregates=aggregates, + rows=rows, + window_us=args.window_us, + timeline_count=max(0, args.timeline_count), + width=max(40, args.timeline_width), + table_only=args.table_only, + ) + print(report) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py new file mode 100644 index 000000000..b47693fc8 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/scripts/analyze_sglang_torch_profile.py @@ -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:])) diff --git a/.claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py b/.claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py new file mode 100644 index 000000000..7d2ecf7f8 --- /dev/null +++ b/.claude/skills/sglang-torch-profiler-analysis/scripts/profile_common.py @@ -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