From db938484f24781b4077a1e2e5f4e9fc8e4656a6c Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:26:01 -0700 Subject: [PATCH] docs(skills): add GPU performance analysis and inference optimization skills Add two skills covering the profile-to-optimization workflow for Megatron-Core inference performance work: - nsight-system-analysis: analyze nsys traces to find the per-iteration gap between two implementations or break down a single profile. Covers iteration anchoring, GPU busy/idle split, per-category kernel breakdown, module slicing between GEMM anchors, exposed communication, and single-forward-pass kernel composition, with scripts for each step. - optimize-inference-siddharth: what to change, in what order, and what never to break when an architecture is slower than vLLM. Covers CUDA-graph scope and bucket coverage, the inference_optimized MoE stack, Triton production hygiene, per-step host overhead, decision gates that quantify a lever's ceiling before any code is written, competitor-trace differential analysis, and the A/B protocol needed to make sub-1% wins falsifiable. The skill is self-maintaining: it documents how to extend, correct, and prune itself as campaigns learn more. Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> --- skills/nsight-system-analysis/SKILL.md | 245 +++++++++ .../assets/report_template.md | 194 ++++++++ .../references/pitfalls.md | 153 ++++++ .../references/sql_recipes.md | 114 +++++ .../references/taxonomy_template.yml | 74 +++ skills/nsight-system-analysis/scripts/_lib.py | 218 ++++++++ .../scripts/busy_idle.py | 224 +++++++++ .../scripts/categorize.py | 230 +++++++++ .../scripts/exposed_comm.py | 238 +++++++++ .../scripts/forward_pass.py | 328 ++++++++++++ .../scripts/iter_anchor.py | 381 ++++++++++++++ .../scripts/module_diff.py | 126 +++++ .../scripts/module_slice.py | 338 +++++++++++++ .../nsight-system-analysis/scripts/run_all.py | 328 ++++++++++++ skills/optimize-inference-siddharth/SKILL.md | 466 ++++++++++++++++++ .../assets/review-checklist.md | 160 ++++++ .../references/commit-log.md | 121 +++++ .../references/cuda-graphs.md | 302 ++++++++++++ .../references/decision-gates.md | 237 +++++++++ .../references/host-path.md | 225 +++++++++ .../references/mamba-and-triton.md | 321 ++++++++++++ .../references/measuring.md | 359 ++++++++++++++ .../references/moe-inference.md | 382 ++++++++++++++ .../references/updating-this-skill.md | 193 ++++++++ .../references/vllm-differential.md | 140 ++++++ 25 files changed, 6097 insertions(+) create mode 100644 skills/nsight-system-analysis/SKILL.md create mode 100644 skills/nsight-system-analysis/assets/report_template.md create mode 100644 skills/nsight-system-analysis/references/pitfalls.md create mode 100644 skills/nsight-system-analysis/references/sql_recipes.md create mode 100644 skills/nsight-system-analysis/references/taxonomy_template.yml create mode 100644 skills/nsight-system-analysis/scripts/_lib.py create mode 100644 skills/nsight-system-analysis/scripts/busy_idle.py create mode 100644 skills/nsight-system-analysis/scripts/categorize.py create mode 100644 skills/nsight-system-analysis/scripts/exposed_comm.py create mode 100644 skills/nsight-system-analysis/scripts/forward_pass.py create mode 100644 skills/nsight-system-analysis/scripts/iter_anchor.py create mode 100644 skills/nsight-system-analysis/scripts/module_diff.py create mode 100644 skills/nsight-system-analysis/scripts/module_slice.py create mode 100644 skills/nsight-system-analysis/scripts/run_all.py create mode 100644 skills/optimize-inference-siddharth/SKILL.md create mode 100644 skills/optimize-inference-siddharth/assets/review-checklist.md create mode 100644 skills/optimize-inference-siddharth/references/commit-log.md create mode 100644 skills/optimize-inference-siddharth/references/cuda-graphs.md create mode 100644 skills/optimize-inference-siddharth/references/decision-gates.md create mode 100644 skills/optimize-inference-siddharth/references/host-path.md create mode 100644 skills/optimize-inference-siddharth/references/mamba-and-triton.md create mode 100644 skills/optimize-inference-siddharth/references/measuring.md create mode 100644 skills/optimize-inference-siddharth/references/moe-inference.md create mode 100644 skills/optimize-inference-siddharth/references/updating-this-skill.md create mode 100644 skills/optimize-inference-siddharth/references/vllm-differential.md diff --git a/skills/nsight-system-analysis/SKILL.md b/skills/nsight-system-analysis/SKILL.md new file mode 100644 index 00000000000..cf197eedd68 --- /dev/null +++ b/skills/nsight-system-analysis/SKILL.md @@ -0,0 +1,245 @@ +--- +name: nsight-system-analysis +description: Analyze NVIDIA Nsight Systems (nsys) GPU profiles to find the per-iteration performance gap between two implementations of the same workload, or to break down a single profile. Use whenever the user provides one or two `.nsys-rep` / `.sqlite` traces and asks "why is X slower than Y?", "compare these two nsys profiles", "investigate this nsys trace", "find the perf gap", "analyze this GPU profile", "look at one forward pass / one decode step", "how many kernels per step", or anything similar. Also use when the user wants per-iteration time, GPU busy vs idle split, per-category kernel breakdown (GEMM/Conv/MHA/Norm), module-slicing between GEMM anchors, exposed communication time, single-forward-pass kernel-count/composition comparison, or source-level root-cause analysis of perf differences. Covers single-GPU and multi-GPU, training and inference. Do not use for non-GPU profiling, non-nsys profilers, or anything that isn't an nsys trace. +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# Nsight Systems Performance Analysis + +This skill produces rigorous per-iteration performance analysis of nsys traces — either comparing two profiles to find the gap, or breaking down a single profile. + +## Three workflows + +**Workflow A — comparative gap**: two `.sqlite` profiles of the same workload. Output: a structured report (`assets/report_template.md`) attributing the wall-clock per-iter Δ to specific kernels and (if source paths are provided) specific code changes. + +**Workflow B — single-profile breakdown**: one `.sqlite` profile. Same template, but with absolute numbers instead of Δs and an "optimization priority" section instead of root causes. + +**Workflow C — single forward pass (decode step)**: the fast, focused workflow for **inference** traces where only the steady-state decode loop matters and everything else (load, warmup, prefill, client-wait idle) is noise. One command isolates *one* forward pass and reports its wall time, kernel count, and per-category composition — for one profile, or side-by-side for two. **Prefer this over Workflow A/B when the user asks "look at one forward pass / one decode step", "how many kernels per step", "why is one step slower", or is comparing two inference engines (e.g. mcore vs vLLM).** See the dedicated section below. + +The user's question tells you which workflow: "compare A and B" or "why is X slower" (training / whole-iter) → A. "Analyze this profile", "what's expensive in this run" → B. "One forward pass", "one decode step", "kernels per step", inference engine comparison → **C**. + +## Inputs you should ensure you have + +Before starting, confirm with the user: +1. **Profile path(s)** (one or two `.sqlite` or `.nsys-rep` files). If only `.nsys-rep` is provided, ask for or generate the `.sqlite` via `nsys export --type sqlite`. +2. **Approximate iteration count** (optional — the iter-anchor script will auto-detect, but a user-provided count makes warmup/cooldown trimming unambiguous). +3. **For Step 6 (root cause)**: source code paths or access for both implementations. Without source access, the analysis stops cleanly after Step 5 and reports phenomena only. + +## The protocol + +Follow Steps 1 → 5 (or 1 → 6 if source is available) **in order**. Every later step builds on the prior step's numbers; skipping or reordering breaks the attribution arithmetic. After each step, write that section of the report immediately — don't batch. + +### Fast path: `scripts/run_all.py` + +For routine analysis, run the integrated pipeline: + +``` +python scripts/run_all.py \ + --profile-a flat.sqlite [--profile-b nonflat.sqlite] \ + --yaml taxonomy.yml \ + --out out/ +``` + +It executes Steps 1, 2, 3, 5 on each profile, decides Step 4 mode (op-group vs +module-slicing) from the auto-computed `fused_share_of_residual_pct`, then +runs the chosen Step 4 script. All intermediates (windows, busy/idle, categorize, +module-slice, exposed-comm) are written to `/` as separate JSON files. +A consolidated `summary.json` plus arithmetic invariants are written and printed. + +**Always start with the fast path.** Then read the per-step intermediate JSON +to fill the report template. The per-step sections below describe each script +in case you need to re-run an individual step with different flags. + +## Workflow C: single forward pass (decode step) — `scripts/forward_pass.py` + +Use this when the user cares about **one forward pass** of an inference engine and wants everything else treated as noise. It answers, in one command: + +- What is the forward-pass (per-decode-step) wall time? +- How many kernels launch in one forward pass? +- Where does the GPU-time go, per category (attention / MoE GEMM / MoE routing / activation / norm / dense GEMM / comm / elementwise / …)? +- **When comparing two engines**: is any individual kernel slower on one side (`µs/kernel`), or does one side simply launch **more** kernels for the same work (`#`)? This is the crux for mcore-vs-vLLM-style gaps, where the slower engine is usually not running slower kernels — it's running ~3× as many small, unfused ones. + +**Run (one profile):** + +``` +python scripts/forward_pass.py mcore.sqlite +``` + +**Run (compare two — label order preserved):** + +``` +python scripts/forward_pass.py mcore.sqlite vllm.sqlite --label-a mcore --label-b vllm +``` + +**How it isolates one pass (fully automatic):** +1. Picks the **busiest GPU/rank** (one `deviceId`) so ranks don't interleave. +2. Finds the **steady-state decode region** = densest contiguous 1 s-binned kernel window (drops load/warmup/prefill and client-wait idle). +3. Finds a **per-step anchor**: the kernel that fires ~once per forward pass (embedding / LM-head / sampling / metadata). Per-step kernels have the *fewest* launches among recurring kernels (per-layer kernels fire `n_layers×` more) and the lowest inter-arrival jitter. Its median inter-arrival period **is** the forward-pass wall time — cross-check it against TPOT. +4. Extracts **one representative step** (a typical anchor-to-anchor interval near the middle of the decode region) and reports launches, GPU-busy (interval union), idle, and the per-category breakdown. + +**Overrides** (when auto-detection is wrong or the user already found the boundaries in the GUI): +- `--anchor ` — force the once-per-step marker (e.g. `--anchor _fused_metadata_kernel` for mcore, `--anchor triton_red_fused__to_copy_embedding_rms_norm` for vLLM). +- `--step-window t0,t1` — force the window in **seconds** (Nsight GUI timeline time = `raw_ns / 1e9`), e.g. `--step-window 156.9498,156.9543`. +- `--device N` — force a specific `deviceId`. +- `--yaml taxonomy.yml` — replace the built-in inference taxonomy with the skill's regex taxonomy. +- `--json` — machine-readable output. + +**Interpreting the comparison table:** +- `Δus > 0` on a row ⇒ that category costs more in profile A. Sort by Δ to get the ranked optimization targets. +- Same `µs/kernel` but higher `#` ⇒ the lever is **fusion / fewer launches**, not a faster kernel. (e.g. a separate SwiGLU activation kernel per layer, FC1 and FC2 as two grouped-GEMM launches, or a "storm" of tiny routing/permute kernels vs a single fused routing kernel.) +- Higher `µs/kernel` on the same shape ⇒ a genuine kernel-selection / tile / dtype finding — drill into it with Step 6. + +**Caveats:** +- `GPU-time (Σ durations)` per category **sums** kernel durations and so over-counts wall time when streams overlap — that's intentional for composition. Use the reported **GPU-busy (interval union)** for the true single-step wall figure, and the anchor **period** for the forward-pass time. +- The built-in taxonomy is tuned for MoE decode; verify the category assignments once (as in Step 3's YAML verification) if a large "other/misc" bucket appears, and extend via `--yaml`. + +For deeper attribution (source-level root cause, exposed comm, module-slicing) fall through to Steps 1–6 below, restricting the windows to the decode region this workflow identified. + +### Step 1: Find a per-iter anchor and measure per-iter time + +This is **always first**. Without a correct per-iter time, every later number is wrong. + +**Run**: `python scripts/iter_anchor.py [--n-iters N] [--anchor PATTERN]` — emits the chosen anchor, detected iter count, the windows, and per-iter timing stats as JSON to stdout. **Drops iter 1 and iter N by default** (warmup + cooldown); pass `--keep-warmup-cooldown` to keep them. + +The script auto-detects the anchor in priority order: +1. NCCL collectives (AllGather, ReduceScatter, AllReduce) — best for distributed training. +2. CUDA stream/device sync — good when present (often absent in graph-captured loops). +3. Optimizer step kernels (`AdamFunctor`, `AdamCapturable`) — training only. +4. HtoD memcpy — fallback; async/non-blocking, less precise. +5. Densest recurring kernel — final fallback for inference / fusion-graph workloads. + +The script picks the candidate with constant integer N per iter and lowest within-iter jitter. **If detection is ambiguous (no candidate yields constant N), it exits with an error listing the candidates** — ask the user which to use, or ask for the iteration count. + +**Cross-check**: if two anchors agree on per-iter time within ~1 ms, the anchor is reliable. The script reports this when possible. + +**Drop warmup/cooldown**: drop iter 1 and iter N when reporting median. Report median + min + max over the remaining iters. + +**Report**: per-iter median + range for each profile, absolute Δ and Δ%. This is the headline. + +### Step 2: GPU busy vs GPU idle = CPU-bound time + +GPU busy = wall time during which ≥1 stream is running a kernel or memcpy (interval union). GPU idle = `iter_time − busy`, and **is** the CPU-bound portion (host dispatch, Python overhead, sync wait). + +**Run**: `python scripts/busy_idle.py --yaml taxonomy.yml --windows windows.json` — emits per-iter busy, idle, and per-stream union JSON. **Always pass `--yaml`** so the script can compute the non-NCCL stream variant (otherwise it warns and the non-NCCL fields are identical to the all-kernel fields). + +Why this matters: it tells you whether the gap is GPU-side (kernel composition, fusion, comm) or CPU-side (launch latency, host code). Pursue the dominant axis in later steps; if both contribute, break each down. + +**Include memcpys in the union.** Kernels alone undercount GPU busy when there's significant HtoD/DtoH activity. + +**Per-stream union**: the longest single-stream union is the critical-path stream. Compute the non-NCCL variant too (NCCL on a co-located stream contaminates the metric — see references/pitfalls.md). + +**Counterintuitive case to expect**: the *faster* profile may have the *longer* single-stream union. That happens when the faster profile concentrates compute on one main stream while the slower one fans the same work across many parallel streams. In that case the slower profile's longest-single-stream is shorter — but its total (union across all streams) is larger. This is a finding to note, not a contradiction; the wall-clock gap is measured by total iter time (Step 1), not by the longest single stream. + +### Step 3: Heavy-compute categories — GEMM / Conv / MHA + +Only these three are reported here. They are model-level ops that are essentially never fused into anything else (same matmul shape on both sides should produce the same nvjet/cutlass call). If they differ, that's a real algorithmic / tile-heuristic / kernel-selection finding. + +**Norm goes in Step 4**, not here. Norm is commonly fused into neighbors (rope, AdaLN, gating); reporting a "norm gap" in Step 3 mistakes fusion-placement for an algorithmic finding. + +**Run**: `python scripts/categorize.py --yaml taxonomy.yml --windows windows.json` — emits per-category time and a list of uncategorized kernels above a threshold. + +**Verify the YAML before reporting numbers** (this is non-negotiable — see "YAML verification" below). The YAML at `references/taxonomy_template.yml` is a **starting point, not a golden reference**. + +### Step 4: Per-op breakdown — op-group first, module-slicing as fallback + +Default: **op-group** breakdown of the residual (everything outside GEMM / Conv / MHA / NCCL). Norm is included here. Categories are emitted by `categorize.py` using the full taxonomy from the YAML. + +Op-group is more interpretable for vanilla workloads: rows are recognizable operators (rmsnorm, rope, elementwise, fp8_cast). Each row shows "how much time does this operator take on each side", which is directly actionable. + +**Decision rule — when to switch to module-slicing**: op-group breaks down when the two implementations have different fusion patterns. Concrete check, automated by the script: compute the share of non-anchor compute time spent in **custom-fused** kernels (names matching `_*_fused_*`, names containing 2+ op-root tokens like `_qkv_split_norm_rope` or `_fused_ln_adaln`, or anything matching a YAML category named `*fused*`). If that share is **>10%** of non-anchor time in *either* profile, op-group attribution becomes unreliable (a single misallocated `add` or `clone` can swing rows by tens of ms — see references/pitfalls.md). Switch to module-slicing. + +**Run for op-group (default)**: `python scripts/categorize.py --yaml taxonomy.yml --windows windows.json --residual-only` — same script, residual mode. + +**Run for module-slicing (fallback)**: `python scripts/module_slice.py --yaml taxonomy.yml --windows windows.json --anchor-categories gemm,mha` — emits per-window times grouped by anchor-pair signature. + +**For comparative module-slicing**, after running module_slice on both profiles, also run `python scripts/module_diff.py mod_a.json mod_b.json` to get a paired diff. The script handles the case where the same window signature has different per-iter counts on each side — it surfaces both the total Δ and a per-call Δ (normalized by min count) so you can see per-call cost differences regardless of count mismatch. + +Report which mode was used and **why**. The decision rule's input (custom-fused share) goes in the report so the choice is auditable. + +#### Module-slicing details (when used) + +- Anchors are all GEMM + MHA kernels across **all compute streams** (NCCL excluded by name), ordered globally by start time. +- Anchor counts on both sides should be approximately equal (same model = same matmul count). If they differ significantly, that's a structural finding to report. +- Anchors may overlap (different streams): when `anchor[k+1].start < anchor[k].end`, the window has zero or negative width — skip and set the next window's left boundary to `max(anchor[k].end, anchor[k+1].end)`. +- Windows are grouped by `(left_anchor_category, right_anchor_category)` signature so the same logical region across N blocks aggregates into one row. +- Each window's time is the **interval-union** of non-NCCL work in that window, not the duration sum (sum overcounts parallel-stream overlap). +- Both sides are non-zero on every row by construction — that's the property op-group can't always give. + +### Step 5: Communication — volume and exposed time + +Skip if the profile has no NCCL kernels (single-GPU, inference, etc.). The script exits cleanly with `{"comm": "none"}` in that case. + +**Run**: `python scripts/exposed_comm.py --yaml taxonomy.yml --windows windows.json` — emits per-collective counts and times, total comm time, exposed time (union of NCCL minus union of non-NCCL), and hidden percentage. + +Two metrics: +- **Comm volume** (count × avg per collective type) — diagnoses bucketing / sharding differences. +- **Exposed comm time** — the wall-clock contribution. Most comm is hidden by overlap; only the exposed portion adds to iter time. + +**Reconcile arithmetic**: `exposed_comm + non_NCCL_union ≈ GPU_busy` from Step 2 within ~0.5 ms. If not, you double-counted or mis-clipped — go back. + +### Step 6: Root cause (only if source paths provided for all profiles being analyzed) + +This is where the report becomes a finding rather than a measurement. Skip if source is unavailable; output a clear note that Step 6 needs source access. + +**Coverage rule**: every Δ row in the final attribution table needs **all three** of: +1. **Source-level evidence on both sides** (file:line each, citing the actual code path that runs at runtime). +2. **Mechanism**: what specifically in the code causes the kernel selection / fusion / dispatch / scheduling to differ. Not "the heuristic picks differently" — *why* it picks differently. +3. **Actionable change**: the code or config change that would close the Δ, named at file:line. + +If you can't produce all three for a given Δ, mark it "unverified — needs follow-up" rather than ship an unverified guess as a root cause. + +**Drill-down templates** (use the one that matches each Δ): + +- **GEMM/MHA tile or kernel-selection Δ**: list per-shape kernel-name + count + time. Diff. The same matmul shape with different kernel-name suffix means the same heuristic took different inputs — find which arg differs (epilogue, dtype, accumulate, alignment, fp8 recipe, layout flags). Read the actual call sites on both sides. +- **Op-group or module-slice window Δ**: dump the full ordered kernel list inside the top-N differing windows / rows on both sides. Identify the operator(s) each kernel implements. Trace back to source: is the difference a custom fused kernel on one side, a `torch.compile` boundary on the other, a framework-level decomposition, a precision/cast inserted by autocast? +- **Exposed comm Δ**: which collective is unhidden? What compute *should* have overlapped? Look for `cudaStreamWaitEvent`, sync-point insertion, bucket-size config (`bucket_size`, `align_param_gather`, `overlap_grad_reduce`). +- **GPU idle Δ**: count `cudaLaunchKernel` events on each side (CUPTI_ACTIVITY_KIND_RUNTIME). Diff. Identify what each extra launch is *for* (optimizer? grad clip? metric logging?). Look for CUDA-graph scope (capture region might be narrower on one side). + +**Phenomenon vs cause traps** — see references/pitfalls.md: +- "The heuristic picks differently" is not a cause. Find the input that differs. +- "More launches" is not a cause. Find what the extra launches do. +- "Different bucketing" is not a cause. Find the config or env var that sets it. + +**Collapse phenomena to fewer causes**: several Δ rows often trace to a single root cause (e.g. three different windows all slower because the same fused custom kernel is missing on one side). Final report should have 3–6 root causes, not 20 phenomena. + +**Verification status table**: end the report with a table listing every claim and whether it's *verified from source*, *inferred from trace*, or *unverified — empirical follow-up needed*. This is required — without it, readers can't tell which findings are solid. + +## Hard rules (apply throughout) + +1. **Union, never sum, for wall time.** Multiple streams overlap; summing kernel durations overstates wall time 2–4×. Use interval union (`scripts/_lib.py:union_intervals`). +2. **YAML is a reference, not golden.** Always verify per-category matches by hand-inspecting the kernel names before reporting numbers. See "YAML verification" below. +3. **Apples-to-apples comparison.** In op-group mode, every row must be non-zero on both sides if the workload actually performs that op. In module-slicing mode, anchor counts must match (structural alignment). +4. **Phenomenon ≠ root cause.** Step 6 requires source-level evidence with file:line, mechanism, and actionable change. +5. **Every Δ must reconcile.** Final attribution table sum must match measured iter Δ to ≤0.5 ms. If it doesn't, find what you missed — don't round and ship. + +## YAML verification (Step 3 prerequisite) + +Before reporting any Step 3 numbers: + +1. Run `categorize.py` and inspect the matched-kernel list per category. Confirm every matched kernel name actually belongs to its category. Common false positives: cuBLAS epilogue helpers caught by a loose `gemm` regex; unrelated kernels with "norm" in the name; Triton catch-all swallowing a fused norm or fused matmul. +2. Inspect the uncategorized list. For each name with non-trivial time (>1% of iter), decide: extend a regex to include it, or add a new category. Pay attention to fused-op kernels (`_qkv_split_norm_rope_kernel`, `_fused_ln_adaln_*` style) — they often need their own custom category. +3. Iterate: re-run, re-check, fix until matched-kernels-per-category-look-correct AND uncategorized-list-is-empty-of-significant-time. Save the iterated YAML as part of the deliverable. + +For comparative mode, the YAML must produce verified matches in both profiles — kernels that exist in one profile may not exist in the other, but the regexes must be sound for whichever kernels do appear. + +## Output: the report + +Fill `assets/report_template.md`. Section structure is fixed (Steps 1–6 + Inputs + verification-status table). Numbers go in the tables; prose goes between them. + +Two arithmetic invariants the final report must satisfy: +- **Step 2 sanity**: `iter_time = GPU_busy + GPU_idle` on each side, within ~0.2 ms. +- **Step 6 sanity** (comparative mode): `Σ(Δ rows in attribution table) = measured iter Δ`, within ~0.5 ms. + +If either invariant fails, the analysis is incomplete or wrong — fix before finalizing. + +**Common arithmetic trap when attributing the gap**: GPU busy is an *interval union*, but Step 4 window-work is sometimes a *sum* across parallel streams. When non-anchor kernels run concurrently on multiple compute streams, the sum overcounts wall-time by the overlap. Always use the **union** column for attribution (`module_slice.py` emits both `*_union_*` and `*_sum_*`; use union). If you find a leftover ~ms-scale residual between `Σ(Δ rows)` and measured iter Δ that you can't account for, this is usually the culprit — recompute with unions throughout. + +**A second small (~1–3 ms) residual is expected** between (Step 4 anchor + Step 4 window-union + Step 5 exposed-comm) and (Step 2 GPU-busy) when NCCL kernels are co-located on the same streams as compute. Step 4 excludes NCCL by name from windows, but if a tiny non-NCCL kernel happens to be straddled by NCCL clipping it can slip out of the accounting. Flag this in the report; do not paper over it. + +## Two-page references + +- `references/taxonomy_template.yml` — minimal generic regex taxonomy. Copy and extend per workload. +- `references/sql_recipes.md` — paste-ready SQL queries for ad-hoc inspection. +- `references/pitfalls.md` — common traps and rejected approaches (sum-overcounting, single-stream anchor, op-group with heavy fusion, single-rank assumption under TP/PP, YAML-as-golden, "the heuristic flipped"). diff --git a/skills/nsight-system-analysis/assets/report_template.md b/skills/nsight-system-analysis/assets/report_template.md new file mode 100644 index 00000000000..1e9350e3452 --- /dev/null +++ b/skills/nsight-system-analysis/assets/report_template.md @@ -0,0 +1,194 @@ + + +# — Performance + +## Inputs + + +- **Profile A** (fast / baseline): `` +- **Profile B** (slow / variant): `` +- Same config: +- YAML taxonomy: `` + +## Step 1: Per-Iteration Time + +### Anchor + + + +| Anchor | A count | B count | Per-iter N (A / B) | Notes | +|---|---|---|---|---| +| | … | … | … | … | +| Cross-check anchor | … | … | … | agree to ms | + +### Per-iteration time + +| Profile | Median (ms) | Min (ms) | Max (ms) | +|---|---|---|---| +| A | … | … | … | +| B | … | … | … | +| **Δ** | **… ms** | — | — | +| **Δ %** | **… %** | — | — | + +**Headline**: + +## Step 2: GPU Busy vs GPU Idle + +GPU busy = interval-union of kernels + memcpys across all streams within each iter +window. GPU idle = `iter_time − busy` and is the CPU-bound portion. + +| Metric | A | B | +|---|---|---| +| iter time (ms) | … | … | +| GPU busy (union) | … | … | +| GPU idle | … | … | +| GPU idle % | … | … | +| Longest single-stream union (any kernel) | … | … | +| Longest single-stream union (non-NCCL) | … | … | + +### Δ (B − A) + +| Metric | Δ | Share of iter Δ | +|---|---|---| +| iter time | … | 100% | +| GPU busy | … | … | +| GPU idle | … | … | + +**Interpretation**: + +## Step 3: Heavy-Compute Categories (GEMM / Conv / MHA) + +Scope: only model-level compute that is essentially never fused into anything +else. Norm is in Step 4 (commonly fused). + +| Category | A (ms/iter) | B (ms/iter) | Δ ms | Δ % | +|---|---|---|---|---| +| gemm | … | … | … | … | +| conv | … | … | … | … | +| mha (flash + cudnn_sdpa) | … | … | … | … | +| **Sum** | … | … | … | … | + +**Interpretation**: + +### YAML verification + +YAML at `` was iterated. Final state: +- Matched-kernel-list per category inspected; all matches confirmed. +- Uncategorized list contains 1% of iter>. + +## Step 4: + + + +**Decision metric**: custom-fused kernel share of non-anchor compute time = +(threshold: 10%). Mode chosen: **** because . + +### Op-group mode (when fusion is light) — DELETE this section if using module-slicing + +| Operator | A (ms/iter) | B (ms/iter) | Δ ms | Notes | +|---|---|---|---|---| +| | … | … | … | … | +| **Sum (residual)** | … | … | … | | + +### Module-slicing mode (when fusion is heavy) — DELETE this section if using op-group + +Global anchor sequence: A had **** anchors/iter, B had **** (within +%). Anchor overlaps (windows skipped): . + +| Window (signature) | Count/iter | A union (ms) | B union (ms) | Δ ms | Notes | +|---|---|---|---|---|---| +| | … | … | … | … | | +| ... | | | | | | + +## Step 5: Communication (Volume + Exposed) + + + +### Volume + +| Op | A count | A ms | A avg | B count | B ms | B avg | Δ ms | +|---|---|---|---|---|---|---|---| +| AllGather | … | … | … | … | … | … | … | +| ReduceScatter | … | … | … | … | … | … | … | +| AllReduce | … | … | … | … | … | … | … | +| **Total** | … | … | — | … | … | — | … | + +### Exposed comm + +| Profile | Comm kernel (ms) | Exposed (ms) | Hidden % | +|---|---|---|---| +| A | … | … | … | +| B | … | … | … | +| **Δ** | … | **… ms** | — | + +**Reconciliation**: `exposed + non_nccl_union ≈ GPU_busy` from Step 2. A: vs + (within ms). B: similar. + +**Interpretation**: + +## Step 6: Root Causes + + + +### Architectural picture (verified) + +- **A**: +- **B**: same. + +### Root Cause 1 — + +| | A | B | +|---|---|---| +| Mechanism | … | … | +| Source evidence | | | +| Owned phenomena | | + +**Net contribution**: <Δ ms> + +### Root Cause N — + + + +### Δ ownership table + +| Phenomenon (Steps 2–5) | Δ ms | Owning root cause(s) | +|---|---|---| +| Step 2 GPU idle | … | RCx | +| Step 3 anchor-time (GEMM+MHA) | … | RCy | +| Step 4 residual (op-group / window-work) | … | RCz | +| Step 5 exposed comm | … | RCw | +| **Sum** | **…** | **vs measured … (within … ms ✓)** | + +### Recommendations (ordered by recoverable impact) + +1. +2. ... + +## Verification status + +| Claim | Status | +|---|---| +| Step 1 anchor + per-iter time | Verified from sqlite | +| Step 2 GPU busy/idle | Verified from sqlite | +| Step 3 GEMM/MHA equality | Verified from sqlite + YAML | +| Step 4 breakdown | Verified from sqlite | +| Step 5 comm volume + exposed | Verified from sqlite | +| RC1 mechanism | | +| RC2 mechanism | … | +| RC | … | + +### Open follow-ups (not yet verified) + +1. +2. ... diff --git a/skills/nsight-system-analysis/references/pitfalls.md b/skills/nsight-system-analysis/references/pitfalls.md new file mode 100644 index 00000000000..4fe7e010371 --- /dev/null +++ b/skills/nsight-system-analysis/references/pitfalls.md @@ -0,0 +1,153 @@ + + +# Pitfalls and rejected approaches + +Things that look reasonable but produce wrong numbers. Documented here so future-you +doesn't re-discover them. + +## 1. Dividing wall-window by N_iter + +``` +per_iter ≈ (last_kernel.end − first_kernel.start) / num_iters ← WRONG +``` + +The wall window includes warmup, profiler startup, post-iter teardown, and any +idle time at the head/tail. This can underestimate per-iter time and dilute Δ. +Real case: this method reported a 5.7% gap when the actual gap was 19%. + +**Use**: an iter-anchor (NCCL/sync/optimizer kernel timestamps) and measure +`anchor[i+1].start − anchor[i].start`. See `scripts/iter_anchor.py`. + +## 2. Summing kernel durations to estimate wall time + +``` +total_busy ≈ SUM(end - start) over all kernels ← WRONG +``` + +Modern training/inference uses multiple streams (compute + comm, fwd/bwd +overlap, multi-stream fan-out). Summing overstates wall time 2–4× because it +double-counts intervals when ≥2 streams are busy simultaneously. + +**Use**: interval union (`_lib.union_total_ns`). Sum is fine for *per-category +work volume* (Step 3) but never for wall-time attribution (Steps 2, 5, 6). + +## 3. Same iter count ≠ same op count per iter + +Two implementations of the same model can issue different numbers of kernels +per iter (different fusion, different parallelism, different bucketing). The +user telling you "both ran 10 iters" doesn't mean each iter has the same +internal structure. Real case: 18 vs 23 AllGathers per iter from different +distributed-optimizer bucket sizes. + +**Use**: verify anchor count constancy per iter, on each profile +independently. If it differs between profiles, that's a finding to investigate +(Step 6 RC) rather than a methodology bug. + +## 4. Single-main-stream anchoring under stream fan-out + +If implementation A places all compute on one big stream and implementation B +fans out across many small streams, anchoring on "the main compute stream" +will yield mismatched anchor counts (e.g. 1138 vs 540) — module-slicing then +can't pair up the windows. + +**Use**: global anchor sequence across **all compute streams** (NCCL excluded +by name, not by stream — comm streams sometimes carry non-comm work too). +Anchor counts then match because they reflect the model's matmul count, not +the dispatcher's stream choice. See `scripts/module_slice.py`. + +## 5. Op-group attribution under heavy fusion + +Op-group is the natural Step 4 default: "rmsnorm took X ms on flat, Y ms on +non-flat". But if implementation A has a single custom kernel +`_qkv_split_norm_rope_kernel` doing 3 ops while B runs separate rmsnorm + rope ++ clone, then: + +- A's `rmsnorm` row shows 0 ms (it's hidden in the fused kernel). +- B's `rmsnorm` row shows the full cost. +- Reporting "A saves X ms in rmsnorm" misrepresents the gap. + +Op-group breakdown also has a hard problem of attributing small ambiguous +kernels (a `vectorized_add` could belong to multiple op-groups; a single +misallocation can swing a row by 10–40 ms). + +**Use**: op-group as the default; switch to module-slicing when the share of +non-anchor compute time in *custom-fused kernels* (heuristic: names matching +`_*_fused_*` or 2+ op tokens in `triton_*_fused_*`) exceeds 10% in either +profile. `scripts/categorize.py` emits `fused_share_of_residual_pct` and +`module_slicing_recommended` to automate this. + +## 6. Phenomenon ≠ root cause + +"The cuBLAS heuristic picks a different kernel" is not a root cause — it's a +phenomenon restatement. Library heuristics are deterministic: same inputs → +same output. If two callers get different results, at least one input differs. +**Find which input.** + +Similarly: "more kernel launches", "different bucketing", "different SDPA +backend" — all phenomena. The cause is the specific config/env-var/code-path +that produces the different choice. + +**Use**: for each Δ, drill until the answer is "implementation A calls X with +Y; B calls X with Z; the heuristic at picks differently because of +Y vs Z". Step 6 of the SKILL prompts for source-level evidence, mechanism, and +actionable change. + +## 7. YAML as golden reference + +It's tempting to ship a "correct" YAML taxonomy. There isn't one. The +taxonomy must match the workload — Triton kernel names differ between +torch.compile versions; custom kernels have idiosyncratic names. A regex +that's right for workload A may mis-classify workload B. + +**Use**: treat the bundled `taxonomy_template.yml` as a starting point. For +every Step 3 / 4 run, inspect the matched-kernel list per category (false +positives) and the uncategorized list (false negatives), then iterate the +regexes. Save the final YAML alongside the report. + +## 8. Naming `_h_badd_` as "bias-add" + +The nvjet/cuBLAS kernel-family suffixes (`_h_bz_`, `_h_badd_`, `_NTT`, `_TNT`, +etc.) are internal cuBLAS heuristic names; their meaning is **not publicly +documented**. Treating them as if they encode op semantics (e.g. "badd = +bias-add epilogue") can be wrong — for example, a wgrad GEMM with `bias=None` +and `accumulate=True` may pick `_h_badd_*` for the FP32-RMW accumulator path, +not for a bias epilogue. + +**Use**: the suffix flip is a *symptom*. The cause is whichever cuBLAS input +differs between the two callers (epilogue arg, output dtype, accumulate flag, +alignment, recipe). Read the actual `cublasLtMatmul` call args, not the kernel +name. + +## 9. Treating "graph break" as proven from kernel pattern alone + +If implementation A has one big fused Triton kernel and B has many small +Triton + TE kernels, the natural conclusion is "torch.compile graph-breaks at +TE C++ calls in B." That's plausible — but the kernel pattern alone doesn't +prove it. To verify, read `transformer_engine/pytorch/jit.py` for the +`no_torch_dynamo` / `torch._dynamo.disable` decorators applied to TE modules; +or capture `TORCH_LOGS=graph_breaks` from a runtime invocation. + +**Use**: in Step 6, include source decorator evidence or runtime log evidence; +otherwise mark the mechanism as "inferred from kernel pattern" in the +verification-status table. + +## 10. Forgetting the verification-status table at the end of Step 6 + +Without it, readers can't tell which claims are source-verified, which are +trace-inferred, and which are unverified guesses. Every RC's mechanism needs +a status row. Required, not optional. + +## 11. Single-rank assumption under TP/PP/SP + +With tensor/pipeline/sequence parallelism, the profiled rank only sees the +collectives and compute that rank participates in. An anchor that exists on +rank 0 may not exist on rank 1. Per-iter time on rank 0 may differ from rank +1 (bubble, imbalance). The user telling you "rank 0" doesn't mean rank 0's +view is representative. + +**Use**: check the anchor type's frequency. If it's `~N` per iter with low +jitter, fine. If it's wildly variable or missing, the profile may be +non-representative or the anchor is wrong for this rank. diff --git a/skills/nsight-system-analysis/references/sql_recipes.md b/skills/nsight-system-analysis/references/sql_recipes.md new file mode 100644 index 00000000000..bb397ba8c23 --- /dev/null +++ b/skills/nsight-system-analysis/references/sql_recipes.md @@ -0,0 +1,114 @@ + + +# Paste-ready SQL for nsys-exported sqlite + +All recipes assume the nsys sqlite schema. Open with sqlite3 or query from Python. + +## Convert .nsys-rep → .sqlite + +```bash +nsys export --type sqlite -o profile.sqlite profile.nsys-rep +``` + +## Top kernels by total time + +```sql +SELECT + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') AS name, + COUNT(*) AS count, + SUM(k.end - k.start)/1e6 AS total_ms +FROM CUPTI_ACTIVITY_KIND_KERNEL k +GROUP BY name +ORDER BY total_ms DESC +LIMIT 30; +``` + +## Per-stream summary (with NCCL-on-stream detection) + +```sql +SELECT + k.streamId, + COUNT(*) AS n, + SUM(k.end - k.start)/1e6 AS total_ms, + (MAX(k.end) - MIN(k.start))/1e6 AS span_ms +FROM CUPTI_ACTIVITY_KIND_KERNEL k +GROUP BY k.streamId +ORDER BY total_ms DESC; +``` + +## NCCL collective timestamps (anchor candidates) + +```sql +SELECT + k.start, k.end, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') AS name +FROM CUPTI_ACTIVITY_KIND_KERNEL k +WHERE name LIKE '%AllGather%' + OR name LIKE '%ReduceScatter%' + OR name LIKE '%AllReduce%' +ORDER BY k.start; +``` + +## CUDA synchronization events + +```sql +SELECT start, end, syncType FROM CUPTI_ACTIVITY_KIND_SYNCHRONIZATION +ORDER BY start; +``` + +## Memcpy summary (by copyKind: 1=HtoD, 2=DtoH, 8=DtoD) + +```sql +SELECT copyKind, COUNT(*), SUM(end-start)/1e6 AS ms +FROM CUPTI_ACTIVITY_KIND_MEMCPY +GROUP BY copyKind; +``` + +## Host-side launch dispatch (idle-investigation) + +```sql +-- Count cudaLaunchKernel events per iter (host time) +SELECT COUNT(*), SUM(end-start)/1e6 AS host_api_ms +FROM CUPTI_ACTIVITY_KIND_RUNTIME r +JOIN StringIds s ON r.nameId = s.id +WHERE s.value = 'cudaLaunchKernel'; +``` + +## Kernel name shape filter (per-GEMM-shape diff) + +```sql +SELECT + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') AS name, + COUNT(*), + SUM(k.end - k.start)/1e6 AS total_ms +FROM CUPTI_ACTIVITY_KIND_KERNEL k +WHERE name LIKE 'nvjet_%' +GROUP BY name +ORDER BY total_ms DESC; +``` + +## Kernels in a specific time window (e.g. between two anchors) + +```sql +SELECT k.start, k.end, k.streamId, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') AS name +FROM CUPTI_ACTIVITY_KIND_KERNEL k +WHERE k.end > :lo AND k.start < :hi +ORDER BY k.start; +``` + +## Number of distinct streams (compute vs comm) + +```sql +SELECT COUNT(DISTINCT streamId) FROM CUPTI_ACTIVITY_KIND_KERNEL; +``` + +## Distinct demangled-name count (for sanity-checking the YAML) + +```sql +SELECT COUNT(DISTINCT s.value) +FROM CUPTI_ACTIVITY_KIND_KERNEL k JOIN StringIds s ON k.demangledName = s.id; +``` diff --git a/skills/nsight-system-analysis/references/taxonomy_template.yml b/skills/nsight-system-analysis/references/taxonomy_template.yml new file mode 100644 index 00000000000..c54f655ef04 --- /dev/null +++ b/skills/nsight-system-analysis/references/taxonomy_template.yml @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Minimal generic kernel-classification taxonomy for nsys traces. +# Order matters: first match wins. Put SPECIFIC categories before catch-alls. +# +# This is a STARTING POINT, not a golden reference. After running categorize.py +# you MUST inspect the matched-kernel list per category and the uncategorized +# list, then iterate the regexes until both look correct for your workload. +# See SKILL.md "YAML verification" section. +# +# Categories typically reported in the SKILL workflow: +# Step 3 (heavy compute, never fused into anything else): +# gemm, conv, flash_attn, cudnn_sdpa +# Step 4 (op-group breakdown of the residual): +# rmsnorm, layer_norm, group_norm, rope, elementwise, fp8_cast, fp8_swizzle, +# reduce, concat, triton, mta_adam, mta_l2norm, mta_scale, sincos, inf_check, +# and any workload-specific custom-fused category you add. +# Step 5 (comm): nccl + +Overall: + # ---- Heavy compute (Step 3) ---- + gemm: 'xmma_gemm|cutlass::Kernel|cublasLt::|cutlass3x(?:.*?)(? list[tuple[int, int]]: + """Merge overlapping intervals. Returns sorted non-overlapping list.""" + ivals = sorted(intervals) + if not ivals: + return [] + out = [list(ivals[0])] + for s, e in ivals[1:]: + if s <= out[-1][1]: + if e > out[-1][1]: + out[-1][1] = e + else: + out.append([s, e]) + return [(s, e) for s, e in out] + + +def union_total_ns(intervals: Iterable[tuple[int, int]]) -> int: + """Total ns covered by at least one interval.""" + return sum(e - s for s, e in union_intervals(intervals)) + + +def subtract_union( + a: Iterable[tuple[int, int]], + b_union: list[tuple[int, int]], +) -> list[tuple[int, int]]: + """Return union(a) minus b_union as a list of intervals. + + b_union must be already merged (non-overlapping, sorted). Result is sorted + non-overlapping. + """ + a_union = union_intervals(a) + out = [] + for s, e in a_union: + cursor = s + for bs, be in b_union: + if be <= cursor: + continue + if bs >= e: + break + if bs > cursor: + out.append((cursor, min(bs, e))) + cursor = max(cursor, be) + if cursor >= e: + break + if cursor < e: + out.append((cursor, e)) + return out + + +def clip( + intervals: Iterable[tuple[int, int]], + lo: int, + hi: int, +) -> list[tuple[int, int]]: + """Clip intervals to [lo, hi).""" + out = [] + for s, e in intervals: + s2, e2 = max(s, lo), min(e, hi) + if e2 > s2: + out.append((s2, e2)) + return out + + +# ---------- taxonomy / YAML ---------- + + +def load_taxonomy(path: str) -> dict[str, re.Pattern]: + """Load a YAML taxonomy and compile its regexes. + + The YAML may have a top-level `Overall:` map or be a flat map; we accept + either. Order is preserved (Python 3.7+ dict ordering). + """ + try: + import yaml # type: ignore + except ImportError: + sys.stderr.write( + "ERROR: PyYAML required. Install with `pip install pyyaml` or run " + "in an env that has it.\n" + ) + sys.exit(2) + with open(path) as f: + data = yaml.safe_load(f) + if ( + isinstance(data, dict) + and "Overall" in data + and isinstance(data["Overall"], dict) + ): + data = data["Overall"] + if not isinstance(data, dict): + sys.stderr.write(f"ERROR: taxonomy YAML at {path} is not a mapping.\n") + sys.exit(2) + return {name: re.compile(pattern) for name, pattern in data.items()} + + +def classify(name: str, taxonomy: dict[str, re.Pattern]) -> str | None: + """First-match-wins classification. Returns category name or None.""" + for cat, rx in taxonomy.items(): + if rx.search(name): + return cat + return None + + +# Heuristic detection of "custom-fused" kernels — used by Step 4 to decide +# between op-group and module-slicing. The criteria are deliberately broad; +# false positives bias toward "use module-slicing", which is the safe default. + +_CUSTOM_FUSED_HINTS = re.compile( + r"_fused_[A-Za-z]|fused_.*kernel|" + r"(?:qkv|rope|norm|adaln|gate|residual|silu|gelu|tanh|layernorm|rmsnorm)" + r"_(?:.*_)?(?:qkv|rope|norm|adaln|gate|residual|silu|gelu|tanh|layernorm|rmsnorm|split|join|cat|sum|fwd|bwd)", + re.IGNORECASE, +) + + +def looks_custom_fused(name: str) -> bool: + """Heuristic: does this kernel name suggest fusion of 2+ ops? + + Used by the op-group-vs-module-slicing decision in Step 4. Matches kernel + names like `_qkv_split_norm_rope_kernel`, `_fused_ln_adaln_fwd_kernel`, + `triton_red_fused__to_copy_mul_native_layer_norm_*`, etc. Pure GEMM and + pure MHA/SDPA names should not match. + """ + if not name: + return False + # Triton fused-* kernels with 2+ op tokens + if "triton_" in name and "fused_" in name: + # Pull out the bit after `fused_` and look for 2+ op tokens + tail = name.split("fused_", 1)[1] + op_tokens = re.findall( + r"(?:add|mul|sub|div|tanh|gelu|silu|relu|sigmoid|exp|log|" + r"native_layer_norm|rms_norm|layer_norm|sum|view|cat|copy|clone|" + r"rope|qkv|norm|adaln|gate)", + tail, + ) + if len(op_tokens) >= 2: + return True + return bool(_CUSTOM_FUSED_HINTS.search(name)) + + +# ---------- output helpers ---------- + + +def write_json(obj, fp=None): + """Write JSON to stdout (default) or a file-like object.""" + fp = fp or sys.stdout + json.dump(obj, fp, indent=2, sort_keys=False) + fp.write("\n") + + +def err(msg: str) -> None: + sys.stderr.write(msg.rstrip() + "\n") + + +def die(msg: str, code: int = 1): + err(f"ERROR: {msg}") + sys.exit(code) diff --git a/skills/nsight-system-analysis/scripts/busy_idle.py b/skills/nsight-system-analysis/scripts/busy_idle.py new file mode 100644 index 00000000000..66db54e6310 --- /dev/null +++ b/skills/nsight-system-analysis/scripts/busy_idle.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compute per-iter GPU busy / idle and per-stream union, clipped to iter windows. + +GPU busy = interval-union of (kernels + memcpys) across all streams within +each iter window. GPU idle = iter_time − GPU_busy and is the CPU-bound portion. + +Usage: + python busy_idle.py --windows windows.json + python busy_idle.py --windows windows.json --exclude-nccl-regex 'ncclDev' + python busy_idle.py --windows windows.json --exclude-nccl-yaml taxonomy.yml + +Inputs: + windows.json: either the raw output of iter_anchor.py, or a JSON + `{"windows_ns": [[start, end], ...]}` or just `[[start, end], ...]`. + +Output (stdout JSON): + { + "per_iter": [ + {"i": 0, "iter_ms": ..., "busy_ms": ..., "idle_ms": ..., "idle_pct": ...}, + ... + ], + "median": {"iter_ms": ..., "busy_ms": ..., "idle_ms": ..., "idle_pct": ...}, + "per_stream_union_ms_median": {stream_id: ms, ...}, + "per_stream_union_ms_median_excluding_nccl": {stream_id: ms, ...}, + "longest_single_stream_union_ms_median": { + "all": ms, "non_nccl": ms + } + } +""" + +from __future__ import annotations + +import argparse +import json +import re +from statistics import median + +from _lib import ( + die, + err, + load_taxonomy, + open_sqlite, + union_total_ns, + write_json, +) + + +def load_windows(path: str) -> list[tuple[int, int]]: + with open(path) as f: + data = json.load(f) + if isinstance(data, dict) and "windows_ns" in data: + data = data["windows_ns"] + if not isinstance(data, list): + die("windows file must be a list or contain 'windows_ns'") + return [(int(w[0]), int(w[1])) for w in data] + + +def get_nccl_classifier(args) -> callable[[str], bool]: + if args.exclude_nccl_yaml: + taxo = load_taxonomy(args.exclude_nccl_yaml) + nccl_rx = taxo.get("nccl") + if nccl_rx is None: + die(f"YAML {args.exclude_nccl_yaml!r} has no `nccl` category") + return lambda name: bool(nccl_rx.search(name)) + if args.exclude_nccl_regex: + rx = re.compile(args.exclude_nccl_regex) + return lambda name: bool(rx.search(name)) + return lambda name: False # do not exclude anything + + +def fetch_intervals_for_window(con, lo: int, hi: int): + """Fetch kernels + memcpys whose [start, end) overlaps [lo, hi). + + Returns list of (start, end, stream_id, name) tuples. Memcpys get a + synthetic name 'memcpy'. + """ + out = [] + for s, e, sid, name in con.execute( + """ + SELECT k.start, k.end, k.streamId, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + FROM CUPTI_ACTIVITY_KIND_KERNEL k + WHERE k.end > ? AND k.start < ? + """, + (lo, hi), + ): + out.append((max(s, lo), min(e, hi), sid, name)) + for s, e, sid in con.execute( + """ + SELECT start, end, streamId + FROM CUPTI_ACTIVITY_KIND_MEMCPY + WHERE end > ? AND start < ? + """, + (lo, hi), + ): + out.append((max(s, lo), min(e, hi), sid, "memcpy")) + return out + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("sqlite") + p.add_argument( + "--windows", required=True, help="iter windows JSON from iter_anchor.py" + ) + p.add_argument( + "--yaml", + help=( + "Taxonomy YAML used to identify NCCL kernels for the non-NCCL split. " + "If not provided, the script falls back to --exclude-nccl-regex; " + "if neither is given, the non-NCCL split is the same as the all-kernel " + "split (warned in output)." + ), + ) + p.add_argument( + "--exclude-nccl-regex", + help="Regex (alternative to --yaml) to identify NCCL kernels to " + "exclude in the non-NCCL stats", + ) + args = p.parse_args() + # Adapt to the classifier-loader's argument names. + args.exclude_nccl_yaml = args.yaml + + windows = load_windows(args.windows) + is_nccl = get_nccl_classifier(args) + nccl_filter_provided = bool(args.exclude_nccl_yaml or args.exclude_nccl_regex) + + per_iter = [] + per_stream_ms = [] # list of dicts per iter + per_stream_ms_no_nccl = [] + + with open_sqlite(args.sqlite) as con: + for i, (lo, hi) in enumerate(windows): + iter_ns = hi - lo + ivals = fetch_intervals_for_window(con, lo, hi) + all_ivals = [(s, e) for s, e, _, _ in ivals] + busy_ns = union_total_ns(all_ivals) + idle_ns = max(0, iter_ns - busy_ns) + per_iter.append( + { + "i": i, + "iter_ms": round(iter_ns / 1e6, 3), + "busy_ms": round(busy_ns / 1e6, 3), + "idle_ms": round(idle_ns / 1e6, 3), + "idle_pct": round(100 * idle_ns / iter_ns, 3) if iter_ns else 0, + } + ) + # Per-stream union + streams: dict[int, list[tuple[int, int]]] = {} + streams_nn: dict[int, list[tuple[int, int]]] = {} + for s, e, sid, name in ivals: + streams.setdefault(sid, []).append((s, e)) + if not is_nccl(name): + streams_nn.setdefault(sid, []).append((s, e)) + per_stream_ms.append( + {sid: union_total_ns(iv) / 1e6 for sid, iv in streams.items()} + ) + per_stream_ms_no_nccl.append( + {sid: union_total_ns(iv) / 1e6 for sid, iv in streams_nn.items()} + ) + + # Medians + def med_key(values): + return round(median(values), 3) if values else 0.0 + + iter_med = { + "iter_ms": med_key([x["iter_ms"] for x in per_iter]), + "busy_ms": med_key([x["busy_ms"] for x in per_iter]), + "idle_ms": med_key([x["idle_ms"] for x in per_iter]), + "idle_pct": med_key([x["idle_pct"] for x in per_iter]), + } + + # Per-stream medians (union of stream IDs across iters) + all_sids = sorted({sid for d in per_stream_ms for sid in d}) + per_stream_med = { + sid: med_key([d.get(sid, 0.0) for d in per_stream_ms]) for sid in all_sids + } + all_sids_nn = sorted({sid for d in per_stream_ms_no_nccl for sid in d}) + per_stream_med_nn = { + sid: med_key([d.get(sid, 0.0) for d in per_stream_ms_no_nccl]) + for sid in all_sids_nn + } + + longest_all = max(per_stream_med.values()) if per_stream_med else 0.0 + longest_nn = max(per_stream_med_nn.values()) if per_stream_med_nn else 0.0 + + out = { + "per_iter": per_iter, + "median": iter_med, + "per_stream_union_ms_median": { + str(k): round(v, 3) + for k, v in sorted(per_stream_med.items(), key=lambda kv: -kv[1]) + }, + "per_stream_union_ms_median_non_nccl": { + str(k): round(v, 3) + for k, v in sorted(per_stream_med_nn.items(), key=lambda kv: -kv[1]) + }, + "longest_single_stream_union_ms_median": { + "all": round(longest_all, 3), + "non_nccl": round(longest_nn, 3), + }, + "nccl_filter": ( + "yaml" + if args.exclude_nccl_yaml + else ("regex" if args.exclude_nccl_regex else "NONE — non_nccl == all") + ), + } + if not nccl_filter_provided: + err( + "WARNING: no --yaml or --exclude-nccl-regex provided. The 'non_nccl' " + "fields are identical to the 'all' fields. Pass --yaml " + "to get a meaningful non-NCCL split." + ) + write_json(out) + + +if __name__ == "__main__": + main() diff --git a/skills/nsight-system-analysis/scripts/categorize.py b/skills/nsight-system-analysis/scripts/categorize.py new file mode 100644 index 00000000000..3fcc830018d --- /dev/null +++ b/skills/nsight-system-analysis/scripts/categorize.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Apply a YAML taxonomy of regexes to kernels, emit per-category time + uncategorized. + +Usage: + python categorize.py --yaml taxonomy.yml --windows windows.json + python categorize.py --yaml taxonomy.yml --windows windows.json \ + --residual-only # excludes gemm,conv,mha,nccl (= Step 4 op-group mode) + python categorize.py --yaml taxonomy.yml --windows windows.json \ + --report-categories gemm,conv,mha # only emit these (= Step 3) + +Inputs: + windows.json: from iter_anchor.py. + taxonomy.yml: ordered category → regex map. First match wins. + +Output (stdout JSON): + { + "per_category": { + "gemm": {"median_ms_per_iter": 184.98, "count_per_iter": 1245, + "unique_kernel_names": N}, + ... + }, + "uncategorized": [ + {"name": "...", "median_ms_per_iter": ..., "count_per_iter": ...}, + ... + ], + "fused_share_of_residual_pct": 12.3, # heuristic for Step 4 decision + "matched_kernels": { + "gemm": ["name1", "name2", ...], + ... + } + } + +The `fused_share_of_residual_pct` is the share of non-anchor non-NCCL time +spent in kernels that look_custom_fused (per the heuristic in _lib.py). If +this is >10%, prefer module_slice.py over op-group attribution in Step 4. +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict + +from _lib import ( + classify, + die, + load_taxonomy, + looks_custom_fused, + open_sqlite, + write_json, +) + + +def load_windows(path: str) -> list[tuple[int, int]]: + with open(path) as f: + data = json.load(f) + if isinstance(data, dict) and "windows_ns" in data: + data = data["windows_ns"] + return [(int(w[0]), int(w[1])) for w in data] + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("sqlite") + p.add_argument("--yaml", required=True, help="Taxonomy YAML") + p.add_argument("--windows", required=True, help="iter windows JSON") + p.add_argument( + "--residual-only", + action="store_true", + help="Exclude gemm/conv/mha/nccl categories from the report " + "(Step 4 op-group mode)", + ) + p.add_argument( + "--report-categories", + help="Comma-separated list of categories to report " + "(e.g. gemm,conv,mha for Step 3)", + ) + p.add_argument( + "--uncategorized-threshold-pct", + type=float, + default=1.0, + help="Min %% of total iter time to surface an uncategorized kernel " + "(default: 1.0)", + ) + args = p.parse_args() + + taxonomy = load_taxonomy(args.yaml) + windows = load_windows(args.windows) + if not windows: + die("No windows in windows.json") + + n_iters = len(windows) + total_iter_ns = sum(hi - lo for lo, hi in windows) + + # Per-category accumulators across all windows. + cat_total_ns: dict[str, int] = defaultdict(int) + cat_count: dict[str, int] = defaultdict(int) + cat_matched_names: dict[str, set] = defaultdict(set) + uncat_total_ns: dict[str, int] = defaultdict(int) + uncat_count: dict[str, int] = defaultdict(int) + fused_total_ns: int = 0 + + nccl_rx = taxonomy.get("nccl") + + with open_sqlite(args.sqlite) as con: + for lo, hi in windows: + rows = con.execute( + """ + SELECT k.start, k.end, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + FROM CUPTI_ACTIVITY_KIND_KERNEL k + WHERE k.end > ? AND k.start < ? + """, + (lo, hi), + ).fetchall() + for s, e, name in rows: + if not name: + name = "(unnamed)" + # Clip to window + cs, ce = max(s, lo), min(e, hi) + if ce <= cs: + continue + dur = ce - cs + cat = classify(name, taxonomy) + if cat is None: + uncat_total_ns[name] += dur + uncat_count[name] += 1 + else: + cat_total_ns[cat] += dur + cat_count[cat] += 1 + cat_matched_names[cat].add(name) + # Fused-share accumulator: only over non-anchor non-NCCL work + is_anchor = cat in ("gemm", "conv", "mha", "flash_attn", "cudnn_sdpa") + is_nccl = cat == "nccl" or (nccl_rx and nccl_rx.search(name)) + if not is_anchor and not is_nccl and looks_custom_fused(name): + fused_total_ns += dur + + # Compute per-category medians by re-scanning per-iter (simpler: divide by + # n_iters since we're already at median sense over windows). + # We'll convert totals → per-iter by dividing by n_iters (this is the + # "mean per iter", close to median for steady-state). + def per_iter_ms(total_ns: int) -> float: + return total_ns / n_iters / 1e6 + + skip_cats = set() + if args.residual_only: + skip_cats = {"gemm", "conv", "mha", "flash_attn", "cudnn_sdpa", "nccl"} + + report_cats = None + if args.report_categories: + report_cats = { + c.strip() for c in args.report_categories.split(",") if c.strip() + } + + per_category = {} + for cat in taxonomy: + if cat in skip_cats: + continue + per_category[cat] = { + "median_ms_per_iter": round(per_iter_ms(cat_total_ns.get(cat, 0)), 3), + "count_per_iter": round(cat_count.get(cat, 0) / n_iters, 3), + "unique_kernel_names": len(cat_matched_names.get(cat, set())), + } + + # Always build a synthetic "mha" row when flash_attn + cudnn_sdpa are present — + # the Step 3 view conventionally combines them. + if "flash_attn" in taxonomy and "cudnn_sdpa" in taxonomy: + mha_ns = cat_total_ns.get("flash_attn", 0) + cat_total_ns.get("cudnn_sdpa", 0) + mha_count = cat_count.get("flash_attn", 0) + cat_count.get("cudnn_sdpa", 0) + mha_names = cat_matched_names.get("flash_attn", set()) | cat_matched_names.get( + "cudnn_sdpa", set() + ) + if mha_ns and "mha" not in skip_cats: + per_category["mha"] = { + "median_ms_per_iter": round(per_iter_ms(mha_ns), 3), + "count_per_iter": round(mha_count / n_iters, 3), + "unique_kernel_names": len(mha_names), + "constituents": ["flash_attn", "cudnn_sdpa"], + } + + # Now apply --report-categories filter (so synthetic mha is included if requested). + if report_cats is not None: + per_category = {k: v for k, v in per_category.items() if k in report_cats} + + # Uncategorized: surface only those above threshold. + threshold_ns = total_iter_ns * args.uncategorized_threshold_pct / 100 + uncat = [] + for name, total_ns in uncat_total_ns.items(): + if total_ns >= threshold_ns: + uncat.append( + { + "name": name, + "median_ms_per_iter": round(per_iter_ms(total_ns), 3), + "count_per_iter": round(uncat_count[name] / n_iters, 3), + } + ) + uncat.sort(key=lambda x: -x["median_ms_per_iter"]) + + # Fused share of residual (non-anchor non-NCCL): the decision metric. + residual_ns = 0 + for cat, total_ns in cat_total_ns.items(): + if cat in {"gemm", "conv", "mha", "flash_attn", "cudnn_sdpa", "nccl"}: + continue + residual_ns += total_ns + residual_ns += sum(uncat_total_ns.values()) # uncategorized counts as residual + fused_share_pct = ( + round(100 * fused_total_ns / residual_ns, 2) if residual_ns > 0 else 0.0 + ) + + out = { + "per_category": per_category, + "uncategorized_above_threshold": uncat, + "uncategorized_threshold_pct": args.uncategorized_threshold_pct, + "fused_share_of_residual_pct": fused_share_pct, + "module_slicing_recommended": fused_share_pct > 10.0, + "matched_kernels": { + cat: sorted(names) for cat, names in cat_matched_names.items() + }, + } + write_json(out) + + +if __name__ == "__main__": + main() diff --git a/skills/nsight-system-analysis/scripts/exposed_comm.py b/skills/nsight-system-analysis/scripts/exposed_comm.py new file mode 100644 index 00000000000..2f4af161069 --- /dev/null +++ b/skills/nsight-system-analysis/scripts/exposed_comm.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compute communication volume and exposed time per iter. + +Exposed comm time = portion of NCCL kernel wall-time during which no non-NCCL +kernel is in flight. It's the comm contribution to wall-clock iter time; the +rest is hidden by overlap. + +Usage: + python exposed_comm.py --yaml taxonomy.yml --windows windows.json + +Output (stdout JSON), or `{"comm": "none"}` if no NCCL kernels found: + { + "per_collective_type": { + "AllGather": {"count_per_iter": 18, "median_ms_per_iter": 44.7, + "avg_ms_per_call": 2.48}, + "ReduceScatter": {"count_per_iter": 18, "median_ms_per_iter": 60.7, ...}, + "AllReduce": {...}, + "Other": {...} + }, + "totals": { + "comm_kernel_ms_median_per_iter": 105.3, + "comm_stream_union_ms_median_per_iter": 105.3, + "exposed_ms_median_per_iter": 6.4, + "hidden_pct_median": 94.0 + }, + "reconciliation": { + "exposed_plus_non_nccl_union_ms": ..., + "gpu_busy_union_ms": ... + } + } +""" + +from __future__ import annotations + +import argparse +import json +import re +from collections import defaultdict +from statistics import median + +from _lib import ( + die, + load_taxonomy, + open_sqlite, + subtract_union, + union_intervals, + union_total_ns, + write_json, +) + + +def load_windows(path: str) -> list[tuple[int, int]]: + with open(path) as f: + data = json.load(f) + if isinstance(data, dict) and "windows_ns" in data: + data = data["windows_ns"] + return [(int(w[0]), int(w[1])) for w in data] + + +COLLECTIVE_PATTERNS = [ + ("AllGather", re.compile(r"AllGather")), + ("ReduceScatter", re.compile(r"ReduceScatter")), + ("AllReduce", re.compile(r"AllReduce")), + ("Broadcast", re.compile(r"Broadcast")), + ("AllToAll", re.compile(r"AllToAll")), + ("Send", re.compile(r"\bSend\b")), + ("Recv", re.compile(r"\bRecv\b")), +] + + +def collective_type(name: str) -> str: + for typ, rx in COLLECTIVE_PATTERNS: + if rx.search(name): + return typ + return "Other" + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("sqlite") + p.add_argument("--yaml", required=True) + p.add_argument("--windows", required=True) + args = p.parse_args() + + taxonomy = load_taxonomy(args.yaml) + nccl_rx = taxonomy.get("nccl") + if not nccl_rx: + die("YAML has no `nccl` category — cannot identify NCCL kernels.") + + windows = load_windows(args.windows) + + per_type_count_per_iter: dict[str, list[int]] = defaultdict(list) + per_type_ms_per_iter: dict[str, list[float]] = defaultdict(list) + total_comm_ms_per_iter: list[float] = [] + comm_union_ms_per_iter: list[float] = [] + exposed_ms_per_iter: list[float] = [] + hidden_pct_per_iter: list[float] = [] + non_nccl_union_ms_per_iter: list[float] = [] + gpu_busy_union_ms_per_iter: list[float] = [] + iter_ms: list[float] = [] + + with open_sqlite(args.sqlite) as con: + for lo, hi in windows: + rows = con.execute( + """ + SELECT k.start, k.end, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + FROM CUPTI_ACTIVITY_KIND_KERNEL k + WHERE k.end > ? AND k.start < ? + """, + (lo, hi), + ).fetchall() + memcpy = con.execute( + """ + SELECT start, end FROM CUPTI_ACTIVITY_KIND_MEMCPY + WHERE end > ? AND start < ? + """, + (lo, hi), + ).fetchall() + + nccl_ivals = [] + non_nccl_ivals = [] + per_type_count_this_iter: dict[str, int] = defaultdict(int) + per_type_ms_this_iter: dict[str, float] = defaultdict(float) + + for s, e, name in rows: + cs, ce = max(s, lo), min(e, hi) + if ce <= cs: + continue + if name and nccl_rx.search(name): + nccl_ivals.append((cs, ce)) + typ = collective_type(name) + per_type_count_this_iter[typ] += 1 + per_type_ms_this_iter[typ] += (ce - cs) / 1e6 + else: + non_nccl_ivals.append((cs, ce)) + for s, e in memcpy: + cs, ce = max(s, lo), min(e, hi) + if ce <= cs: + continue + non_nccl_ivals.append((cs, ce)) + + iter_ns = hi - lo + iter_ms.append(iter_ns / 1e6) + if not nccl_ivals: + # No NCCL: short-circuit at end after gathering all iters. + total_comm_ms_per_iter.append(0.0) + comm_union_ms_per_iter.append(0.0) + exposed_ms_per_iter.append(0.0) + hidden_pct_per_iter.append(100.0) + non_nccl_union_ms_per_iter.append(union_total_ns(non_nccl_ivals) / 1e6) + gpu_busy_union_ms_per_iter.append(union_total_ns(non_nccl_ivals) / 1e6) + continue + + total_comm_ns = sum(e - s for s, e in nccl_ivals) + comm_union = union_intervals(nccl_ivals) + comm_union_ns = sum(e - s for s, e in comm_union) + non_nccl_union = union_intervals(non_nccl_ivals) + non_nccl_union_ns = sum(e - s for s, e in non_nccl_union) + exposed = subtract_union(nccl_ivals, non_nccl_union) + exposed_ns = sum(e - s for s, e in exposed) + gpu_busy_union_ns = union_total_ns(nccl_ivals + non_nccl_ivals) + + total_comm_ms_per_iter.append(total_comm_ns / 1e6) + comm_union_ms_per_iter.append(comm_union_ns / 1e6) + exposed_ms_per_iter.append(exposed_ns / 1e6) + hidden_pct_per_iter.append( + 100 * (1 - exposed_ns / comm_union_ns) if comm_union_ns else 100.0 + ) + non_nccl_union_ms_per_iter.append(non_nccl_union_ns / 1e6) + gpu_busy_union_ms_per_iter.append(gpu_busy_union_ns / 1e6) + for typ in set( + list(per_type_count_this_iter) + list(per_type_ms_this_iter) + ): + per_type_count_per_iter[typ].append( + per_type_count_this_iter.get(typ, 0) + ) + per_type_ms_per_iter[typ].append(per_type_ms_this_iter.get(typ, 0.0)) + + # Decide if NCCL was found at all. + has_nccl = sum(total_comm_ms_per_iter) > 0 + if not has_nccl: + write_json({"comm": "none"}) + return + + def med(xs): + return round(median(xs), 3) if xs else 0.0 + + per_collective_type = {} + for typ in sorted(set(list(per_type_count_per_iter) + list(per_type_ms_per_iter))): + counts = per_type_count_per_iter.get(typ, []) + mss = per_type_ms_per_iter.get(typ, []) + # Pad shorter list with zeros (for iters where this type didn't appear) + target_len = len(windows) + while len(counts) < target_len: + counts.append(0) + while len(mss) < target_len: + mss.append(0.0) + c_median = med(counts) + ms_median = med(mss) + avg = round(ms_median / c_median, 4) if c_median else 0.0 + per_collective_type[typ] = { + "count_per_iter": c_median, + "median_ms_per_iter": ms_median, + "avg_ms_per_call": avg, + } + + out = { + "per_collective_type": per_collective_type, + "totals": { + "comm_kernel_ms_median_per_iter": med(total_comm_ms_per_iter), + "comm_stream_union_ms_median_per_iter": med(comm_union_ms_per_iter), + "exposed_ms_median_per_iter": med(exposed_ms_per_iter), + "hidden_pct_median": med(hidden_pct_per_iter), + }, + "reconciliation": { + "non_nccl_union_ms_median_per_iter": med(non_nccl_union_ms_per_iter), + "gpu_busy_union_ms_median_per_iter": med(gpu_busy_union_ms_per_iter), + "exposed_plus_non_nccl_union_ms_median": round( + med(exposed_ms_per_iter) + med(non_nccl_union_ms_per_iter), 3 + ), + "note": ( + "exposed + non_nccl_union should approximately equal gpu_busy_union " + "from busy_idle.py — if not, double-counting or clip mismatch" + ), + }, + } + write_json(out) + + +if __name__ == "__main__": + main() diff --git a/skills/nsight-system-analysis/scripts/forward_pass.py b/skills/nsight-system-analysis/scripts/forward_pass.py new file mode 100644 index 00000000000..68aa44ae70b --- /dev/null +++ b/skills/nsight-system-analysis/scripts/forward_pass.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-forward-pass (decode-step) analysis for nsys traces. + +Everything in an inference trace except the steady-state decode loop is noise: +model load, warmup, prefill, and long idle stretches while the server waits for +the client. This script throws all of that away and reports insights about +*one* forward pass (one decode step), for one or two profiles. + +What it does, fully automatically: + 1. Picks the busiest GPU/rank (one deviceId) so ranks don't interleave. + 2. Finds the steady-state decode region (densest contiguous kernel window). + 3. Finds a per-step "anchor" kernel: one that fires exactly once per forward + pass (embedding / LM-head / sampling / metadata). Its inter-arrival period + IS the forward-pass wall time (and matches TPOT). + 4. Extracts exactly one representative step (a typical anchor-to-anchor + interval) and reports: + - forward-pass wall time, GPU-busy (interval union), idle gaps + - total kernel launches in the step + - per-category breakdown: launches, GPU-time (Σ durations), µs/kernel + +With two profiles it prints a side-by-side comparison and, per category, the +delta in launches and time, so you can answer the two questions that matter: + (a) is any individual kernel slower in one engine, or + (b) does one engine just launch MORE kernels for the same work? + +Usage: + # one profile + python scripts/forward_pass.py mcore.sqlite + + # compare two (label order is preserved in the output) + python scripts/forward_pass.py mcore.sqlite vllm.sqlite \ + --label-a mcore --label-b vllm + + # override auto-detection (seconds match the Nsight GUI timeline; raw_ns/1e9) + python scripts/forward_pass.py mcore.sqlite \ + --anchor _fused_metadata_kernel + python scripts/forward_pass.py vllm.sqlite --step-window 156.9498,156.9543 + + # machine-readable + python scripts/forward_pass.py mcore.sqlite vllm.sqlite --json + +Notes: + - "GPU-time (Σ durations)" per category is the sum of kernel durations; it + intentionally over-counts wall time when kernels overlap across streams + (that overlap is the point of comparing composition). Use the reported + GPU-busy (interval union) for the true wall figure. + - The default categories are a sensible inference taxonomy. Pass --yaml + to override with the skill's regex taxonomy instead. +""" + +from __future__ import annotations + +import argparse +import statistics +import sys + +from _lib import classify, clip, err, load_taxonomy, open_sqlite, union_total_ns, write_json + +# Ordered, first-match-wins default taxonomy (case-insensitive substring match). +# Tuned for MoE decode inference (mcore + vLLM). Override with --yaml. +_DEFAULT_CATEGORIES: list[tuple[str, tuple[str, ...]]] = [ + ("comm (EP dispatch/combine)", ("multimem", "nccl", "all_gather", "allgather", + "reduce_scatter", "reducescatter", "alltoall", "sendrecv")), + ("attention", ("flash", "fmha", "attention", "attn", "paged", "mha")), + ("MoE expert GEMM", ("_fused_moe_kernel", "bmm_")), + ("MoE routing/permute", ("routing", "topk", "gathertopk", "moe_sum", "count_local", + "metadata", "finalize", "scatter", "permute", "sort", + "cumsum", "histogram", "softmax")), + ("activation (SwiGLU)", ("silu", "act_and_mul", "swiglu", "glu", "gelu")), + ("norm (RMSNorm)", ("rmsnorm", "rms_norm", "layernorm", "layer_norm", "norm_fwd", "rsqrt")), + ("dense GEMM (qkv/o/router/lm-head)", ("nvjet", "cutlass", "sgemm", "gemm", + "sm100_tst", "sm90", "cublas", "ampere")), + ("sampling/logits", ("reduce_kernel", "argmax", "multinomial", "sample", "catarray", "cat_")), + ("embedding", ("embedding",)), + ("elementwise/copy/cast", ("elementwise", "vectorized", "copy", "index_", "triton_poi", + "triton_red", "fused", "cast", "_add")), +] + + +def _categorize(name: str, taxonomy) -> str: + if taxonomy is not None: + return classify(name, taxonomy) or "other/misc" + low = name.lower() + for cat, pats in _DEFAULT_CATEGORIES: + if any(p in low for p in pats): + return cat + return "other/misc" + + +def _kernels(con, dev=None): + """(start, end, deviceId, name) ordered by start, optionally one device.""" + where = "1=1" if dev is None else f"k.deviceId={int(dev)}" + cur = con.execute( + f"""SELECT k.start, k.end, k.deviceId, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + FROM CUPTI_ACTIVITY_KIND_KERNEL k WHERE {where} ORDER BY k.start""" + ) + return cur.fetchall() + + +def _busiest_device(con) -> int: + row = con.execute( + "SELECT deviceId, COUNT(*) c FROM CUPTI_ACTIVITY_KIND_KERNEL " + "GROUP BY deviceId ORDER BY c DESC LIMIT 1" + ).fetchone() + if row is None: + raise SystemExit("ERROR: no kernels in trace") + return row[0] + + +def _decode_window(starts: list[int], bin_ns: int = 1_000_000_000) -> tuple[int, int]: + """Densest contiguous region of kernel launches = steady-state decode.""" + mn, mx = starts[0], starts[-1] + nb = (mx - mn) // bin_ns + 1 + cnt = [0] * nb + for s in starts: + cnt[(s - mn) // bin_ns] += 1 + pk = max(range(nb), key=lambda i: cnt[i]) + thr = 0.3 * cnt[pk] + lo = hi = pk + while lo > 0 and cnt[lo - 1] >= thr: + lo -= 1 + while hi < nb - 1 and cnt[hi + 1] >= thr: + hi += 1 + return mn + lo * bin_ns, mn + (hi + 1) * bin_ns + + +# Kernels that fire ~once per forward pass (step boundaries), across engines. +_STEP_HINTS = ("metadata", "to_copy_embedding", "embedding", "lm_head", "logits", + "argmax", "multinomial", "sample") + + +def _gap_stats(ss): + ss = sorted(ss) + gaps = [b - a for a, b in zip(ss, ss[1:])] + m = statistics.median(gaps) + cv = (statistics.pstdev(gaps) / m) if m else 9e9 + return ss, m, cv + + +def _find_anchor(win_kernels, window_ns, forced=None): + """Return (anchor_name, sorted_starts_on_one_device, period_ns). + + A per-step anchor fires ~once per forward pass. Per-layer kernels fire + n_layers× more often, so among *low-jitter* recurring kernels the per-step + marker has the *longest* inter-arrival period. We therefore: + 1. prefer a kernel whose name matches a known step-boundary marker + (embedding / lm-head / sampling / metadata), else + 2. pick the longest-period, low-jitter recurring kernel (a sane [0.5 ms, + window/3] band, so it spans several steps but isn't per-layer). + + win_kernels spans all devices (the marker may live on a different rank than + the busiest compute device), but the *period* is always measured on a + SINGLE device — otherwise N ranks firing the same marker once per step, + staggered, would collapse the apparent period by ~N×. + """ + from collections import defaultdict + + dev_starts = defaultdict(lambda: defaultdict(list)) # name -> device -> [starts] + for s, _e, d, name in win_kernels: + dev_starts[name][d].append(s) + + # For each name, evaluate it on the single device where it fires most. + stats = {} # name -> (starts_on_best_dev, period_ns, cv) + for name, per_dev in dev_starts.items(): + best_dev = max(per_dev, key=lambda d: len(per_dev[d])) + ss = per_dev[best_dev] + if len(ss) < 15: + continue + s, m, cv = _gap_stats(ss) + stats[name] = (s, m, cv) + if not stats: + raise SystemExit("ERROR: no recurring kernel found in decode window; pass --anchor/--step-window") + + if forced is not None: + cands = {n: v for n, v in stats.items() if forced.lower() in n.lower()} + if not cands: + raise SystemExit(f"ERROR: --anchor '{forced}' matched no recurring kernel") + name = min(cands, key=lambda n: cands[n][2]) # lowest jitter + s, m, _ = cands[name] + return name, s, m + + lo, hi = 0.5e6, window_ns / 3 # plausible forward-pass period band (ns) + periodic = {n: v for n, v in stats.items() if v[2] < 0.5 and lo <= v[1] <= hi} + if not periodic: # relax the band if nothing qualifies + periodic = {n: v for n, v in stats.items() if v[1] > 0} + + hinted = {n: v for n, v in periodic.items() if any(h in n.lower() for h in _STEP_HINTS)} + pool = hinted or periodic + # longest low-jitter period == fires least often == once per step + name = max(pool, key=lambda n: pool[n][1]) + s, m, _ = pool[name] + return name, s, m + + +def analyze(path: str, taxonomy, forced_anchor=None, forced_dev=None, forced_window=None): + with open_sqlite(path) as con: + if forced_window is not None: + w0, w1 = forced_window + dev = forced_dev + if dev is None: + row = con.execute( + "SELECT deviceId, COUNT(*) c FROM CUPTI_ACTIVITY_KIND_KERNEL " + "WHERE start>=? AND start [count, dur_ns] + for s, e, name in step: + c = _categorize(name, taxonomy) + cats[c][0] += 1 + cats[c][1] += e - s + + return { + "profile": path, + "device": dev, + "decode_window_s": [w0 / 1e9, w1 / 1e9], + "steps_in_window": n_steps, + "anchor": anchor, + "forward_pass_ms": period / 1e6, + "step_wall_ms": wall / 1e6, + "gpu_busy_ms": busy / 1e6, + "gpu_idle_ms": (wall - busy) / 1e6, + "n_kernels": len(step), + "categories": { + c: {"launches": v[0], "gpu_time_us": v[1] / 1e3, + "us_per_kernel": (v[1] / 1e3 / v[0]) if v[0] else 0.0} + for c, v in cats.items() + }, + } + + +def _print_one(a: dict, label: str): + print(f"\n=== {label} ({a['profile'].split('/')[-1]}) ===") + print(f" device={a['device']} decode window {a['decode_window_s'][0]:.1f}-{a['decode_window_s'][1]:.1f}s" + f" (~{a['steps_in_window']} steps) anchor: {a['anchor'][:48]}") + print(f" forward pass = {a['forward_pass_ms']:.3f} ms step wall = {a['step_wall_ms']:.3f} ms" + f" GPU-busy = {a['gpu_busy_ms']:.3f} ms idle = {a['gpu_idle_ms']:.3f} ms") + print(f" kernels in one forward pass = {a['n_kernels']}") + rows = sorted(a["categories"].items(), key=lambda kv: -kv[1]["gpu_time_us"]) + print(f" {'category':34} {'#':>5} {'GPU us':>9} {'us/kern':>8}") + for c, v in rows: + print(f" {c:34} {v['launches']:>5} {v['gpu_time_us']:>9.1f} {v['us_per_kernel']:>8.1f}") + + +def _print_compare(a: dict, b: dict, la: str, lb: str): + print(f"\n=== COMPARISON: {la} vs {lb} (one forward pass) ===") + print(f" forward pass : {a['forward_pass_ms']:.3f} ms ({la}) vs {b['forward_pass_ms']:.3f} ms ({lb})" + f" -> {la} is {a['forward_pass_ms']/b['forward_pass_ms']:.2f}x") + print(f" kernels/step : {a['n_kernels']} ({la}) vs {b['n_kernels']} ({lb})" + f" -> {la} launches {a['n_kernels']/max(b['n_kernels'],1):.2f}x") + allc = list(dict.fromkeys(list(a["categories"]) + list(b["categories"]))) + def g(d, c): return d["categories"].get(c, {"launches": 0, "gpu_time_us": 0.0, "us_per_kernel": 0.0}) + allc.sort(key=lambda c: -(g(a, c)["gpu_time_us"] - g(b, c)["gpu_time_us"])) + print(f"\n {'category':34} | {lb+' #':>7} {lb+' us':>9} {'us/k':>6} | {la+' #':>7} {la+' us':>9} {'us/k':>6} | {'Δus':>8}") + print(" " + "-" * 104) + for c in allc: + av, bv = g(a, c), g(b, c) + print(f" {c:34} | {bv['launches']:>7} {bv['gpu_time_us']:>9.1f} {bv['us_per_kernel']:>6.1f} |" + f" {av['launches']:>7} {av['gpu_time_us']:>9.1f} {av['us_per_kernel']:>6.1f} |" + f" {av['gpu_time_us']-bv['gpu_time_us']:>8.1f}") + print("\n Read: 'Δus' > 0 = category costs more in %s. Compare 'us/k' to see if a kernel is\n" + " individually slower, vs '#' to see if it's just more launches for the same work." % la) + + +def main(argv=None): + ap = argparse.ArgumentParser(description="Single forward-pass (decode step) analysis for nsys sqlite traces.") + ap.add_argument("profiles", nargs="+", help="1 or 2 .sqlite profiles (A then optional B).") + ap.add_argument("--anchor", default=None, help="Substring of a once-per-step kernel to force the anchor.") + ap.add_argument("--device", type=int, default=None, help="Force GPU deviceId (default: busiest).") + ap.add_argument("--step-window", default=None, + help="Force window 't0,t1' in seconds (Nsight GUI time = raw_ns/1e9).") + ap.add_argument("--yaml", default=None, help="Override default taxonomy with a regex YAML.") + ap.add_argument("--label-a", default="A") + ap.add_argument("--label-b", default="B") + ap.add_argument("--json", action="store_true", help="Emit JSON instead of tables.") + args = ap.parse_args(argv) + + if len(args.profiles) > 2: + err("WARNING: more than 2 profiles given; only the first two are used.") + taxonomy = load_taxonomy(args.yaml) if args.yaml else None + fw = None + if args.step_window: + t0, t1 = (float(x) for x in args.step_window.split(",")) + fw = (int(t0 * 1e9), int(t1 * 1e9)) + + a = analyze(args.profiles[0], taxonomy, args.anchor, args.device, fw) + b = analyze(args.profiles[1], taxonomy, args.anchor, args.device, fw) if len(args.profiles) >= 2 else None + + if args.json: + write_json({"a": a, "b": b}) + return + _print_one(a, args.label_a) + if b is not None: + _print_one(b, args.label_b) + _print_compare(a, b, args.label_a, args.label_b) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/skills/nsight-system-analysis/scripts/iter_anchor.py b/skills/nsight-system-analysis/scripts/iter_anchor.py new file mode 100644 index 00000000000..b2a27dc9991 --- /dev/null +++ b/skills/nsight-system-analysis/scripts/iter_anchor.py @@ -0,0 +1,381 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Find a per-iteration anchor and emit per-iter timing. + +Auto-detects an anchor kernel pattern by trying candidates in priority order +and picking the one whose timestamps cluster into a constant integer N per +iter with the lowest jitter. Exits with an error (and a list of candidates) +if no clean anchor is found, leaving the caller to specify `--anchor`. + +Usage: + python iter_anchor.py + python iter_anchor.py --anchor 'ncclDev.*AllGather' + python iter_anchor.py --n-iters 10 + python iter_anchor.py --drop-warmup-cooldown # drop first+last iter + +Output (JSON, stdout): + { + "anchor": {"name": "AllGather", "pattern": "...", "count": 198}, + "iter_count_detected": 10, + "iter_count_used": 10, + "anchors_per_iter": 18, + "windows_ns": [[start, end], ...], + "per_iter_ms": [368.7, 369.3, ...], + "median_ms": 369.0, + "min_ms": 366.3, + "max_ms": 373.3, + "cross_check": {"second_anchor": "...", "agreement_ms": 0.4} # if available + } +""" + +from __future__ import annotations + +import argparse +import re +import sys +from statistics import median + +from _lib import ( + die, + err, + open_sqlite, + write_json, +) + +# Anchor candidates in priority order. Each entry: (name, sql LIKE pattern, regex). +CANDIDATES = [ + # Distributed training: NCCL collectives. + ("AllGather", "%AllGather%", r"AllGather"), + ("ReduceScatter", "%ReduceScatter%", r"ReduceScatter"), + ("AllReduce", "%AllReduce%", r"AllReduce"), + # CUDA synchronization events (often absent under CUDA Graphs). + # We search the kernel table for sync-like names; for the CUPTI sync + # table see _attempt_sync_anchor. + # Optimizer step (training). + ("Adam", "%Adam%", r"Adam(Functor|Capturable|CapturableFunctor)"), + # HtoD memcpy (input batches). + # Handled separately via the memcpy table. +] + + +def _fetch_timestamps_by_pattern(con, like_pattern: str, regex: str) -> list[int]: + """Get start_ns of kernels matching the like_pattern, refined by regex.""" + rx = re.compile(regex) + rows = con.execute( + """ + SELECT k.start, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + FROM CUPTI_ACTIVITY_KIND_KERNEL k + WHERE COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + LIKE ? + ORDER BY k.start + """, + (like_pattern,), + ).fetchall() + return [start for start, name in rows if rx.search(name or "")] + + +def _fetch_sync_timestamps(con) -> list[int]: + """CUPTI sync events (cudaStreamSynchronize, cudaDeviceSynchronize).""" + try: + rows = con.execute( + "SELECT start FROM CUPTI_ACTIVITY_KIND_SYNCHRONIZATION ORDER BY start" + ).fetchall() + return [r[0] for r in rows] + except Exception: + return [] + + +def _fetch_htod_memcpy_timestamps(con) -> list[int]: + try: + rows = con.execute( + """ + SELECT start FROM CUPTI_ACTIVITY_KIND_MEMCPY + WHERE copyKind=1 -- HtoD + ORDER BY start + """ + ).fetchall() + return [r[0] for r in rows] + except Exception: + return [] + + +def _try_densest_recurring_kernel(con) -> tuple[str, list[int]] | None: + """Final fallback: pick the most numerous kernel that has > N_min instances + and cluster well into iters.""" + rows = con.execute( + """ + SELECT COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + AS name, + COUNT(*) AS c + FROM CUPTI_ACTIVITY_KIND_KERNEL k + WHERE name <> '' + GROUP BY name + ORDER BY c DESC + LIMIT 5 + """ + ).fetchall() + for name, _ in rows: + ts = con.execute( + """ + SELECT start FROM CUPTI_ACTIVITY_KIND_KERNEL + WHERE COALESCE((SELECT value FROM StringIds WHERE id=demangledName),'') = ? + ORDER BY start + """, + (name,), + ).fetchall() + ts = [r[0] for r in ts] + if len(ts) >= 20: + return (name, ts) + return None + + +def _cluster_iters(timestamps: list[int], gap_multiplier: float = 5.0) -> list[int]: + """Return indices where each iter starts. The first index is always 0. + + Uses the heuristic: an iter boundary is a gap > gap_multiplier * median_gap. + """ + if len(timestamps) < 2: + return [0] + gaps = [timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)] + g_sorted = sorted(gaps) + med = g_sorted[len(g_sorted) // 2] + if med == 0: + return [0] + boundaries = [0] + [i + 1 for i, g in enumerate(gaps) if g > gap_multiplier * med] + return boundaries + + +def _evaluate_candidate(timestamps: list[int]) -> dict | None: + """Score an anchor candidate. Returns None if unusable.""" + if len(timestamps) < 4: + return None + boundaries = _cluster_iters(timestamps) + n_iters = len(boundaries) + if n_iters < 2: + return None + per_iter_counts = [ + boundaries[i + 1] - boundaries[i] for i in range(len(boundaries) - 1) + ] + # Add the last partial cluster + per_iter_counts.append(len(timestamps) - boundaries[-1]) + # Drop the first cluster (may be partial warmup) before checking constancy + constant_n = ( + len(set(per_iter_counts[1:-1])) <= 1 if len(per_iter_counts) > 3 else False + ) + # Per-iter durations: anchor[boundaries[i]] -> anchor[boundaries[i+1]] + iter_times_ns = [ + timestamps[boundaries[i + 1]] - timestamps[boundaries[i]] + for i in range(len(boundaries) - 1) + ] + if not iter_times_ns: + return None + iter_ms = [t / 1e6 for t in iter_times_ns] + med = median(iter_ms) + jitter = max(iter_ms) - min(iter_ms) if len(iter_ms) > 1 else 0.0 + return { + "boundaries": boundaries, + # Number of complete iter windows (= number of boundary pairs). + "n_iters_detected": max(0, len(boundaries) - 1), + "per_iter_counts": per_iter_counts, + "constant_n_per_iter": constant_n, + "anchors_per_iter": per_iter_counts[1] if constant_n else None, + "iter_times_ms": iter_ms, + "median_ms": med, + "jitter_ms": jitter, + } + + +def _build_windows( + timestamps: list[int], boundaries: list[int] +) -> list[tuple[int, int]]: + """Iter window i = [anchor[boundaries[i]], anchor[boundaries[i+1]]).""" + out = [] + for i in range(len(boundaries) - 1): + out.append((timestamps[boundaries[i]], timestamps[boundaries[i + 1]])) + return out + + +def auto_detect(con, user_n_iters: int | None) -> tuple[dict, str, list[int]] | None: + """Try anchor candidates and return the best one. + + Returns (result_dict, anchor_name, timestamps) or None if no candidate + yields a constant per-iter count. + """ + tried = [] + best = None + for name, like, regex in CANDIDATES: + ts = _fetch_timestamps_by_pattern(con, like, regex) + if not ts: + tried.append({"name": name, "count": 0}) + continue + ev = _evaluate_candidate(ts) + tried.append( + { + "name": name, + "count": len(ts), + "constant_n_per_iter": ev["constant_n_per_iter"] if ev else False, + "n_iters_detected": ev["n_iters_detected"] if ev else None, + "jitter_ms": round(ev["jitter_ms"], 3) if ev else None, + } + ) + if ev and ev["constant_n_per_iter"]: + if best is None or ev["jitter_ms"] < best[0]["jitter_ms"]: + best = (ev, name, ts) + + # Try CUPTI sync events + sync_ts = _fetch_sync_timestamps(con) + if sync_ts: + ev = _evaluate_candidate(sync_ts) + tried.append( + { + "name": "cudaSynchronization", + "count": len(sync_ts), + "constant_n_per_iter": ev["constant_n_per_iter"] if ev else False, + "jitter_ms": round(ev["jitter_ms"], 3) if ev else None, + } + ) + if ev and ev["constant_n_per_iter"]: + if best is None or ev["jitter_ms"] < best[0]["jitter_ms"]: + best = (ev, "cudaSynchronization", sync_ts) + + # Try HtoD memcpy + h2d = _fetch_htod_memcpy_timestamps(con) + if h2d: + ev = _evaluate_candidate(h2d) + tried.append( + { + "name": "HtoD memcpy", + "count": len(h2d), + "constant_n_per_iter": ev["constant_n_per_iter"] if ev else False, + "jitter_ms": round(ev["jitter_ms"], 3) if ev else None, + } + ) + if ev and ev["constant_n_per_iter"]: + if best is None or ev["jitter_ms"] < best[0]["jitter_ms"]: + best = (ev, "HtoD memcpy", h2d) + + # Final fallback: densest recurring kernel + if best is None: + densest = _try_densest_recurring_kernel(con) + if densest: + name, ts = densest + ev = _evaluate_candidate(ts) + tried.append( + { + "name": name[:60], + "count": len(ts), + "constant_n_per_iter": ev["constant_n_per_iter"] if ev else False, + "jitter_ms": round(ev["jitter_ms"], 3) if ev else None, + } + ) + if ev and ev["constant_n_per_iter"]: + best = (ev, name, ts) + + if best is None: + err("\nAuto-detection failed. Candidates tried:") + for t in tried: + err(f" {t}") + err( + "\nNo candidate produced a constant per-iter count. Specify --anchor " + "explicitly (e.g. --anchor 'ncclDev.*AllGather') or provide --n-iters.\n" + ) + return None + + return best + + +def cross_check(con, primary_name: str, primary_iter_ms: list[float]) -> dict | None: + """Try a second anchor and report agreement, if possible.""" + for name, like, regex in CANDIDATES: + if name == primary_name: + continue + ts = _fetch_timestamps_by_pattern(con, like, regex) + if not ts: + continue + ev = _evaluate_candidate(ts) + if not ev or not ev["constant_n_per_iter"]: + continue + # Compare medians + agreement = abs(ev["median_ms"] - median(primary_iter_ms)) + return { + "second_anchor": name, + "second_median_ms": round(ev["median_ms"], 3), + "agreement_ms": round(agreement, 3), + } + return None + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("sqlite") + p.add_argument("--anchor", help="Override regex for the anchor kernel name") + p.add_argument("--n-iters", type=int, help="Expected total iter count (for sanity)") + p.add_argument( + "--keep-warmup-cooldown", + action="store_true", + help="Keep iter 1 and iter N in the windows (default: drop both).", + ) + args = p.parse_args() + # Default behavior: drop warmup + cooldown. Flag inverts. + args.drop_warmup_cooldown = not args.keep_warmup_cooldown + + with open_sqlite(args.sqlite) as con: + if args.anchor: + ts = _fetch_timestamps_by_pattern(con, "%", args.anchor) + if not ts: + die(f"No kernels match --anchor {args.anchor!r}") + ev = _evaluate_candidate(ts) + if not ev: + die("Anchor produced too few events to cluster.") + picked_name = args.anchor + else: + best = auto_detect(con, args.n_iters) + if best is None: + sys.exit(2) + ev, picked_name, ts = best + + # Build windows + windows = _build_windows(ts, ev["boundaries"]) + iter_ms = ev["iter_times_ms"] + + if args.drop_warmup_cooldown and len(iter_ms) >= 3: + iter_ms = iter_ms[1:-1] + windows = windows[1:-1] + + # Sanity vs --n-iters + if args.n_iters and ev["n_iters_detected"] != args.n_iters: + err( + f"WARNING: detected {ev['n_iters_detected']} iters but user said " + f"{args.n_iters}. Using detected count." + ) + + xc = cross_check(con, picked_name, iter_ms) if not args.anchor else None + + out = { + "anchor": { + "name": picked_name, + "anchor_count_total": len(ts), + "anchors_per_iter": ev["anchors_per_iter"], + }, + "iter_count_detected": ev["n_iters_detected"], + "iter_count_used": len(iter_ms), + "windows_ns": [list(w) for w in windows], + "per_iter_ms": [round(x, 3) for x in iter_ms], + "median_ms": round(median(iter_ms), 3), + "min_ms": round(min(iter_ms), 3), + "max_ms": round(max(iter_ms), 3), + "warmup_cooldown_dropped": args.drop_warmup_cooldown, + } + if xc: + out["cross_check"] = xc + write_json(out) + + +if __name__ == "__main__": + main() diff --git a/skills/nsight-system-analysis/scripts/module_diff.py b/skills/nsight-system-analysis/scripts/module_diff.py new file mode 100644 index 00000000000..c386f6fca69 --- /dev/null +++ b/skills/nsight-system-analysis/scripts/module_diff.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Diff two module_slice.py outputs side-by-side: pair window signatures, show +A ms, B ms, and Δ. Tolerates count mismatches between profiles for the same +signature by emitting the per-iter counts as a separate column. + +Usage: + python module_diff.py mod_a.json mod_b.json [--top N] + +Output (stdout JSON): + { + "anchor_count_a": 1530, "anchor_count_b": 1524, + "anchor_count_match_pct": 99.6, + "iter_total_anchor_ms": {"a": ..., "b": ..., "delta": ...}, + "iter_total_window_union_ms": {"a": ..., "b": ..., "delta": ...}, + "signatures": [ + { + "signature": "...", + "count_per_iter_a": 38, "count_per_iter_b": 38, + "union_ms_a": ..., "union_ms_b": ..., "delta_union_ms": ..., + "delta_per_call_us": ... + }, + ... + ] + } + +The `delta_per_call_us` field normalizes Δ by the lower of the two per-iter +counts, exposing the per-call gap regardless of count mismatch. +""" + +from __future__ import annotations + +import argparse +import json +import sys + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("a_json") + p.add_argument("b_json") + p.add_argument( + "--top", type=int, default=20, help="Top N signatures by |delta_union_ms|" + ) + args = p.parse_args() + + with open(args.a_json) as f: + A = json.load(f) + with open(args.b_json) as f: + B = json.load(f) + + a_by_sig = {g["signature"]: g for g in A["grouped_windows"]} + b_by_sig = {g["signature"]: g for g in B["grouped_windows"]} + all_sigs = set(a_by_sig) | set(b_by_sig) + + rows = [] + for sig in all_sigs: + ga = a_by_sig.get(sig, {}) + gb = b_by_sig.get(sig, {}) + ca = ga.get("count_per_iter", 0) + cb = gb.get("count_per_iter", 0) + ua = ga.get("median_union_ms_per_iter", 0.0) + ub = gb.get("median_union_ms_per_iter", 0.0) + delta = ub - ua + # Per-call: normalize by min count if both >0; else fall back to + # whichever exists. + min_count = min(ca, cb) if ca > 0 and cb > 0 else max(ca, cb) + per_call_us = round(1000 * delta / min_count, 3) if min_count else None + rows.append( + { + "signature": sig, + "count_per_iter_a": ca, + "count_per_iter_b": cb, + "union_ms_a": round(ua, 3), + "union_ms_b": round(ub, 3), + "delta_union_ms": round(delta, 3), + "delta_per_call_us": per_call_us, + "only_in": ( + None + if sig in a_by_sig and sig in b_by_sig + else ("a" if sig in a_by_sig else "b") + ), + } + ) + rows.sort(key=lambda r: -abs(r["delta_union_ms"])) + + ac = A.get("anchor_count_per_iter_first", 0) + bc = B.get("anchor_count_per_iter_first", 0) + out = { + "anchor_count_a": ac, + "anchor_count_b": bc, + "anchor_count_match_pct": ( + round(100 * min(ac, bc) / max(ac, bc), 2) if max(ac, bc) else 0 + ), + "iter_total_anchor_ms": { + "a": A.get("iter_total_anchor_ms_median", 0), + "b": B.get("iter_total_anchor_ms_median", 0), + "delta": round( + B.get("iter_total_anchor_ms_median", 0) + - A.get("iter_total_anchor_ms_median", 0), + 3, + ), + }, + "iter_total_window_union_ms": { + "a": A.get("iter_total_window_union_ms_median", 0), + "b": B.get("iter_total_window_union_ms_median", 0), + "delta": round( + B.get("iter_total_window_union_ms_median", 0) + - A.get("iter_total_window_union_ms_median", 0), + 3, + ), + }, + "signatures": rows[: args.top], + "signatures_total": len(rows), + } + json.dump(out, sys.stdout, indent=2) + sys.stdout.write("\n") + + +if __name__ == "__main__": + main() diff --git a/skills/nsight-system-analysis/scripts/module_slice.py b/skills/nsight-system-analysis/scripts/module_slice.py new file mode 100644 index 00000000000..c4fee11ecf9 --- /dev/null +++ b/skills/nsight-system-analysis/scripts/module_slice.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Module-slice a profile: anchor on every GEMM+MHA kernel globally across all +compute streams, then sum (or union) the non-NCCL work in each inter-anchor window. + +Designed for Step 4 when op-group attribution is unreliable due to heavy fusion. +Both implementations should produce roughly equal anchor counts (same model → +same number of matmuls); the resulting windows are apples-to-apples. + +Usage: + python module_slice.py --yaml taxonomy.yml --windows windows.json + python module_slice.py --yaml taxonomy.yml --windows windows.json \ + --anchor-categories gemm,mha,flash_attn,cudnn_sdpa + +Output (stdout JSON): + { + "anchor_count_total": 1530, + "anchor_count_per_iter": [1530, 1530, ...], + "anchor_count_per_iter_constant": true, + "anchor_overlap_pct": 33.2, + "grouped_windows": [ + { + "signature": "gemm_TNT -> gemm_TNT", + "count_per_iter": 38, + "median_total_union_ms_per_iter": 30.4, + "median_total_sum_ms_per_iter": 30.6 + }, + ... + ], + "iter_total_anchor_ms_median": 222.5, + "iter_total_window_union_ms_median": 51.9, + "iter_total_window_sum_ms_median": 53.5 + } + +The `signature` is +` -> `, +where the "shape" is the part of the kernel name after the category-defining +substring. This keeps related anchors grouped while distinguishing different +matmul shapes. If shapes are too noisy, use `--signature-mode category` for the +coarser "gemm -> gemm" / "gemm -> mha" grouping. +""" + +from __future__ import annotations + +import argparse +import json +import re +from collections import defaultdict +from statistics import median + +from _lib import ( + classify, + clip, + load_taxonomy, + open_sqlite, + union_total_ns, + write_json, +) + + +def load_windows(path: str) -> list[tuple[int, int]]: + with open(path) as f: + data = json.load(f) + if isinstance(data, dict) and "windows_ns" in data: + data = data["windows_ns"] + return [(int(w[0]), int(w[1])) for w in data] + + +def anchor_signature(name: str, cat: str, mode: str) -> str: + """Build a signature string for an anchor kernel. + + mode='category': just the category (e.g. 'gemm', 'mha'). + mode='shape': category + the kernel-shape tail (e.g. 'gemm_128x256_TNT'). + """ + if mode == "category": + return cat + # Pull out a shape hint from the kernel name. + # nvjet: nvjet_sm103_qqtst_128x256_128x6_2x2f_2cta_h_bz_..._NTT + m = re.search(r"(\d+x\d+(?:_\d+x\d+)?)", name) + shape = m.group(1) if m else "" + # Pull out NTT/NNT/TNT/TNN suffix if present + m2 = re.search(r"_([NT]{2,4})(?:\b|$)", name) + suffix = m2.group(1) if m2 else "" + parts = [cat] + if shape: + parts.append(shape) + if suffix: + parts.append(suffix) + return "_".join(parts) + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("sqlite") + p.add_argument("--yaml", required=True) + p.add_argument("--windows", required=True) + p.add_argument( + "--anchor-categories", + default="gemm,flash_attn,cudnn_sdpa,mha,conv", + help="Comma-separated YAML categories whose kernels serve as anchors", + ) + p.add_argument( + "--signature-mode", + choices=["category", "shape"], + default="shape", + help="Granularity of anchor-pair signatures used for window grouping", + ) + args = p.parse_args() + + taxonomy = load_taxonomy(args.yaml) + windows = load_windows(args.windows) + anchor_cats = {c.strip() for c in args.anchor_categories.split(",") if c.strip()} + nccl_rx = taxonomy.get("nccl") + + anchor_count_per_iter: list[int] = [] + overlaps_per_iter: list[int] = [] + iter_anchor_ms: list[float] = [] + iter_window_union_ms: list[float] = [] + iter_window_sum_ms: list[float] = [] + + # signature -> list[ (window_union_ns, window_sum_ns) ]; one entry per + # (window instance) accumulated across iters. + sig_window_union: dict[str, list[int]] = defaultdict(list) + sig_window_sum: dict[str, list[int]] = defaultdict(list) + sig_count_per_iter: dict[str, list[int]] = defaultdict(list) + + with open_sqlite(args.sqlite) as con: + for lo, hi in windows: + # Pull all kernels in the iter window. + rows = con.execute( + """ + SELECT k.start, k.end, k.streamId, + COALESCE((SELECT value FROM StringIds WHERE id=k.demangledName), '') + FROM CUPTI_ACTIVITY_KIND_KERNEL k + WHERE k.end > ? AND k.start < ? + ORDER BY k.start + """, + (lo, hi), + ).fetchall() + + anchors = [] # list of (start, end, category, name) + non_anchor_intervals_by_stream: dict[int, list[tuple[int, int]]] = ( + defaultdict(list) + ) + + for s, e, sid, name in rows: + cs, ce = max(s, lo), min(e, hi) + if ce <= cs: + continue + cat = classify(name, taxonomy) or "_uncat_" + is_nccl = cat == "nccl" or (nccl_rx and nccl_rx.search(name)) + if is_nccl: + continue # NCCL excluded entirely from module-slicing + if cat in anchor_cats: + anchors.append((cs, ce, cat, name)) + else: + non_anchor_intervals_by_stream[sid].append((cs, ce)) + + anchors.sort(key=lambda a: a[0]) + anchor_count_per_iter.append(len(anchors)) + + # Anchor time + window times. + anchor_total_ns = sum(e - s for s, e, _, _ in anchors) + iter_anchor_ms.append(anchor_total_ns / 1e6) + + # Group non-anchor intervals across all streams into one flat list + # (we want the union across streams = wall-time GPU busy in the window). + all_non_anchor = [] + for sid, ivs in non_anchor_intervals_by_stream.items(): + all_non_anchor.extend(ivs) + + overlaps = 0 + window_union_ns_iter = 0 + window_sum_ns_iter = 0 + sig_count_this_iter: dict[str, int] = defaultdict(int) + + # Iterate inter-anchor windows. To avoid losing time inside + # overlap-collapsed window groups, we maintain a "cursor" which is + # the running max(end) of the most recent anchor group. The next + # window begins at the cursor, not at the immediately preceding + # anchor's end. This way: if anchors k..k+M all overlap, we + # collapse them into a single anchor group ending at max(end_k..k+M), + # and the window from there to anchor[k+M+1].start carries any + # non-anchor work without being dropped. + n = len(anchors) + cursor = anchors[0][1] if n else 0 + k = 0 + while k < n - 1: + a_left = anchors[k] + # Advance cursor past any overlapping run of anchors that + # collectively span the next anchor's start. + end_of_group = a_left[1] + while k + 1 < n and anchors[k + 1][0] < end_of_group: + end_of_group = max(end_of_group, anchors[k + 1][1]) + overlaps += 1 + k += 1 + if k + 1 >= n: + break + a_right = anchors[k + 1] + w_lo = max(end_of_group, cursor) + w_hi = a_right[0] + if w_hi > w_lo: + sig = ( + anchor_signature(a_left[3], a_left[2], args.signature_mode) + + " -> " + + anchor_signature(a_right[3], a_right[2], args.signature_mode) + ) + clipped = clip(all_non_anchor, w_lo, w_hi) + w_sum_ns = sum(e - s for s, e in clipped) + w_union_ns = union_total_ns(clipped) + sig_window_sum[sig].append(w_sum_ns) + sig_window_union[sig].append(w_union_ns) + sig_count_this_iter[sig] += 1 + window_sum_ns_iter += w_sum_ns + window_union_ns_iter += w_union_ns + cursor = max(end_of_group, a_right[1]) + k += 1 + + # Tail: include any non-anchor work between the last anchor and + # the iter window end. This is the "post-last-anchor" region. + if n: + last_end = anchors[-1][1] + w_lo = max(last_end, cursor) + w_hi = hi + if w_hi > w_lo: + sig = ( + anchor_signature( + anchors[-1][3], anchors[-1][2], args.signature_mode + ) + + " -> [iter_end]" + ) + clipped = clip(all_non_anchor, w_lo, w_hi) + w_sum_ns = sum(e - s for s, e in clipped) + w_union_ns = union_total_ns(clipped) + sig_window_sum[sig].append(w_sum_ns) + sig_window_union[sig].append(w_union_ns) + sig_count_this_iter[sig] += 1 + window_sum_ns_iter += w_sum_ns + window_union_ns_iter += w_union_ns + # Pre-first-anchor region: also include. + if n: + w_lo = lo + w_hi = anchors[0][0] + if w_hi > w_lo: + sig = "[iter_start] -> " + anchor_signature( + anchors[0][3], anchors[0][2], args.signature_mode + ) + clipped = clip(all_non_anchor, w_lo, w_hi) + w_sum_ns = sum(e - s for s, e in clipped) + w_union_ns = union_total_ns(clipped) + sig_window_sum[sig].append(w_sum_ns) + sig_window_union[sig].append(w_union_ns) + sig_count_this_iter[sig] += 1 + window_sum_ns_iter += w_sum_ns + window_union_ns_iter += w_union_ns + + for sig in sig_count_this_iter: + sig_count_per_iter[sig].append(sig_count_this_iter[sig]) + overlaps_per_iter.append(overlaps) + iter_window_union_ms.append(window_union_ns_iter / 1e6) + iter_window_sum_ms.append(window_sum_ns_iter / 1e6) + + n_iters = len(windows) + n_anchors = anchor_count_per_iter[0] if anchor_count_per_iter else 0 + constant = len(set(anchor_count_per_iter)) == 1 + overlap_pct = ( + 100 + * sum(overlaps_per_iter) + / max(1, sum(a - 1 for a in anchor_count_per_iter if a > 0)) + if any(anchor_count_per_iter) + else 0.0 + ) + + # Aggregate per-signature: sum across windows in iter, then median across iters. + grouped = [] + for sig in sig_window_union: + total_union_per_iter = [] + total_sum_per_iter = [] + # We accumulated per-window samples; convert back to per-iter by chunking + # using sig_count_per_iter. + counts = sig_count_per_iter[sig] + idx = 0 + for c in counts: + u_slice = sig_window_union[sig][idx : idx + c] + s_slice = sig_window_sum[sig][idx : idx + c] + idx += c + total_union_per_iter.append(sum(u_slice)) + total_sum_per_iter.append(sum(s_slice)) + # Pad with zeros if some iters have no instance of this sig. + while len(total_union_per_iter) < n_iters: + total_union_per_iter.append(0) + total_sum_per_iter.append(0) + grouped.append( + { + "signature": sig, + "count_per_iter": round(median(counts) if counts else 0, 2), + "median_union_ms_per_iter": round( + median(total_union_per_iter) / 1e6, 3 + ), + "median_sum_ms_per_iter": round(median(total_sum_per_iter) / 1e6, 3), + } + ) + grouped.sort(key=lambda g: -g["median_union_ms_per_iter"]) + + out = { + "anchor_count_per_iter": anchor_count_per_iter, + "anchor_count_constant_across_iters": constant, + "anchor_count_per_iter_first": n_anchors, + "anchor_overlap_pct": round(overlap_pct, 2), + "iter_total_anchor_ms_median": round( + median(iter_anchor_ms) if iter_anchor_ms else 0, 3 + ), + "iter_total_window_union_ms_median": round( + median(iter_window_union_ms) if iter_window_union_ms else 0, 3 + ), + "iter_total_window_sum_ms_median": round( + median(iter_window_sum_ms) if iter_window_sum_ms else 0, 3 + ), + "grouped_windows": grouped, + "notes": [ + "Use median_union_ms_per_iter as the primary number for " + "wall-time attribution.", + "median_sum_ms_per_iter is shown for comparison; sum overcounts " + "when non-anchor kernels run in parallel on multiple compute streams.", + "NCCL kernels are excluded entirely (by name match against the " + "YAML's `nccl` category).", + ], + } + write_json(out) + + +if __name__ == "__main__": + main() diff --git a/skills/nsight-system-analysis/scripts/run_all.py b/skills/nsight-system-analysis/scripts/run_all.py new file mode 100644 index 00000000000..230bf327343 --- /dev/null +++ b/skills/nsight-system-analysis/scripts/run_all.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the full Steps 1–5 pipeline on one or two profiles and write all JSON +intermediates to an output directory. + +Usage: + # Comparative mode + python run_all.py --profile-a flat.sqlite --profile-b nonflat.sqlite \ + --yaml taxonomy.yml --out /tmp/analysis/ + + # Single-profile mode + python run_all.py --profile-a profile.sqlite --yaml taxonomy.yml --out /tmp/analysis/ + +The script: + 1. Runs iter_anchor.py on each profile → windows_.json + 2. Runs busy_idle.py on each → busy_.json + 3. Runs categorize.py on each → cat_.json (full taxonomy + Step 3 view) + 4. Decides Step 4 mode (op-group vs module-slicing) from the Step 3 output + `module_slicing_recommended` flag (uses OR across both profiles). + 5. If module-slicing: runs module_slice.py on each → mod_.json + If op-group: re-runs categorize with --residual-only → opgroup_.json + 6. Runs exposed_comm.py on each → comm_.json (skips cleanly if no NCCL). + 7. Writes a top-level summary.json with all key numbers plus the Step 4 decision. + +This is a convenience wrapper. For more control, invoke individual scripts. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path + +SCRIPT_DIR = Path(__file__).parent + + +def run(cmd: list[str], capture: bool = True) -> dict | str: + """Run a script, return parsed JSON from stdout (or text).""" + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + sys.stderr.write(f"ERROR running: {' '.join(cmd)}\nstderr:\n{result.stderr}\n") + sys.exit(result.returncode) + if result.stderr.strip(): + # Forward warnings to the user, but don't fail. + sys.stderr.write(result.stderr) + if not capture: + return result.stdout + try: + return json.loads(result.stdout) + except json.JSONDecodeError: + sys.stderr.write( + f"Could not parse JSON from {cmd[0]}:\n{result.stdout[:1000]}\n" + ) + sys.exit(2) + + +def analyze_profile( + profile: Path, yaml: Path, out_dir: Path, tag: str, n_iters: int | None +) -> dict: + """Run Steps 1–5 on a single profile. Returns the assembled summary.""" + summary: dict = {"profile": str(profile), "tag": tag} + py = sys.executable + + # Step 1: anchor + windows + cmd = [py, str(SCRIPT_DIR / "iter_anchor.py"), str(profile)] + if n_iters: + cmd += ["--n-iters", str(n_iters)] + anchor = run(cmd) + windows_path = out_dir / f"windows_{tag}.json" + windows_path.write_text(json.dumps(anchor)) + summary["step1"] = { + "anchor": anchor["anchor"]["name"], + "iter_count_used": anchor["iter_count_used"], + "median_ms": anchor["median_ms"], + "min_ms": anchor["min_ms"], + "max_ms": anchor["max_ms"], + "cross_check": anchor.get("cross_check"), + } + + # Step 2: busy/idle + busy = run( + [ + py, + str(SCRIPT_DIR / "busy_idle.py"), + str(profile), + "--windows", + str(windows_path), + "--yaml", + str(yaml), + ] + ) + (out_dir / f"busy_{tag}.json").write_text(json.dumps(busy)) + summary["step2"] = { + "median": busy["median"], + "longest_single_stream_union_ms_median": busy[ + "longest_single_stream_union_ms_median" + ], + } + + # Step 3 view: gemm/conv/mha + cat3 = run( + [ + py, + str(SCRIPT_DIR / "categorize.py"), + str(profile), + "--yaml", + str(yaml), + "--windows", + str(windows_path), + "--report-categories", + "gemm,conv,mha", + ] + ) + (out_dir / f"cat_step3_{tag}.json").write_text(json.dumps(cat3)) + summary["step3"] = { + "per_category": cat3["per_category"], + "fused_share_of_residual_pct": cat3["fused_share_of_residual_pct"], + "module_slicing_recommended": cat3["module_slicing_recommended"], + } + + # Full categorize for uncategorized inspection + cat_full = run( + [ + py, + str(SCRIPT_DIR / "categorize.py"), + str(profile), + "--yaml", + str(yaml), + "--windows", + str(windows_path), + ] + ) + (out_dir / f"cat_full_{tag}.json").write_text(json.dumps(cat_full)) + summary["uncategorized_above_1pct"] = cat_full["uncategorized_above_threshold"][:10] + + # Step 5: comm + comm = run( + [ + py, + str(SCRIPT_DIR / "exposed_comm.py"), + str(profile), + "--yaml", + str(yaml), + "--windows", + str(windows_path), + ] + ) + (out_dir / f"comm_{tag}.json").write_text(json.dumps(comm)) + summary["step5"] = comm + + return summary + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--profile-a", required=True, help="Path to first .sqlite") + p.add_argument("--profile-b", help="Path to second .sqlite (comparative mode)") + p.add_argument("--yaml", required=True, help="Taxonomy YAML") + p.add_argument("--out", required=True, help="Output directory for intermediates") + p.add_argument("--n-iters", type=int, help="Expected iteration count (optional)") + args = p.parse_args() + + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + yaml_path = Path(args.yaml) + + summary: dict = {"mode": "comparative" if args.profile_b else "single"} + summary["a"] = analyze_profile( + Path(args.profile_a), yaml_path, out_dir, "a", args.n_iters + ) + if args.profile_b: + summary["b"] = analyze_profile( + Path(args.profile_b), yaml_path, out_dir, "b", args.n_iters + ) + + # Step 4 decision + fused_a = summary["a"]["step3"]["fused_share_of_residual_pct"] + fused_b = ( + summary.get("b", {}).get("step3", {}).get("fused_share_of_residual_pct", 0.0) + ) + use_module_slicing = fused_a > 10.0 or fused_b > 10.0 + summary["step4_mode"] = "module-slicing" if use_module_slicing else "op-group" + summary["step4_decision"] = { + "fused_share_a_pct": fused_a, + "fused_share_b_pct": fused_b if args.profile_b else None, + "threshold_pct": 10.0, + "use_module_slicing": use_module_slicing, + } + + py = sys.executable + for tag, prof in [("a", args.profile_a)] + ( + [("b", args.profile_b)] if args.profile_b else [] + ): + windows_path = out_dir / f"windows_{tag}.json" + if use_module_slicing: + mod = run( + [ + py, + str(SCRIPT_DIR / "module_slice.py"), + str(prof), + "--yaml", + str(yaml_path), + "--windows", + str(windows_path), + "--signature-mode", + "shape", + ] + ) + mod_path = out_dir / f"mod_{tag}.json" + mod_path.write_text(json.dumps(mod)) + summary[tag]["step4_module_slice"] = { + "anchor_count_per_iter": mod["anchor_count_per_iter_first"], + "anchor_count_constant": mod["anchor_count_constant_across_iters"], + "anchor_overlap_pct": mod["anchor_overlap_pct"], + "iter_total_anchor_ms": mod["iter_total_anchor_ms_median"], + "iter_total_window_union_ms": mod["iter_total_window_union_ms_median"], + "iter_total_window_sum_ms": mod["iter_total_window_sum_ms_median"], + "top_windows": mod["grouped_windows"][:10], + } + else: + opgroup = run( + [ + py, + str(SCRIPT_DIR / "categorize.py"), + str(prof), + "--yaml", + str(yaml_path), + "--windows", + str(windows_path), + "--residual-only", + ] + ) + (out_dir / f"opgroup_{tag}.json").write_text(json.dumps(opgroup)) + summary[tag]["step4_op_group"] = opgroup["per_category"] + + # Arithmetic invariants (top-level sanity) + invariants = {} + for tag in ["a"] + (["b"] if args.profile_b else []): + s = summary[tag] + med = s["step2"]["median"] + invariants[tag] = { + "iter_eq_busy_plus_idle": round( + med["iter_ms"] - med["busy_ms"] - med["idle_ms"], 3 + ), + } + if "step5" in s and "totals" in s["step5"]: + recon = s["step5"]["totals"] + exposed = recon["exposed_ms_median_per_iter"] + non_nccl = s["step5"]["reconciliation"]["non_nccl_union_ms_median_per_iter"] + invariants[tag]["exposed_plus_non_nccl_minus_busy_ms"] = round( + exposed + non_nccl - med["busy_ms"], 3 + ) + summary["invariants"] = invariants + + # Comparative Δ + if args.profile_b: + med_a = summary["a"]["step2"]["median"] + med_b = summary["b"]["step2"]["median"] + summary["delta"] = { + "iter_ms": round(med_b["iter_ms"] - med_a["iter_ms"], 3), + "busy_ms": round(med_b["busy_ms"] - med_a["busy_ms"], 3), + "idle_ms": round(med_b["idle_ms"] - med_a["idle_ms"], 3), + } + + # Comparative module-slice diff if both profiles + module-slicing mode. + if args.profile_b and use_module_slicing: + diff = run( + [ + py, + str(SCRIPT_DIR / "module_diff.py"), + str(out_dir / "mod_a.json"), + str(out_dir / "mod_b.json"), + ] + ) + (out_dir / "mod_diff.json").write_text(json.dumps(diff)) + summary["step4_module_diff_top"] = diff["signatures"] + + summary_path = out_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2)) + + # Compact stdout summary — agent reads summary.json for the rest. + print(f"\nFull summary JSON: {summary_path}") + print(f"Intermediates in: {out_dir}/") + print() + print(f"Mode: {summary['mode']}") + a = summary["a"] + print( + f" A ({Path(a['profile']).name}): " + f"anchor={a['step1']['anchor']}, iters_used={a['step1']['iter_count_used']} " + f"(after dropping warmup+cooldown), " + f"per_iter_median={a['step1']['median_ms']:.2f} ms" + ) + if args.profile_b: + b = summary["b"] + print( + f" B ({Path(b['profile']).name}): " + f"anchor={b['step1']['anchor']}, " + f"iters_used={b['step1']['iter_count_used']} " + f"(after dropping warmup+cooldown), " + f"per_iter_median={b['step1']['median_ms']:.2f} ms" + ) + d = summary["delta"] + print( + f" Δ iter={d['iter_ms']:+.2f} ms, " + f"busy={d['busy_ms']:+.2f}, idle={d['idle_ms']:+.2f}" + ) + print() + print( + f"Step 4 mode: {summary['step4_mode']} " + f"(fused_share A={fused_a:.1f}%" + + (f", B={fused_b:.1f}%" if args.profile_b else "") + + ", threshold=10%)" + ) + print() + print("Invariants (should be near zero):") + for tag, inv in summary["invariants"].items(): + print(f" {tag}: {inv}") + + +if __name__ == "__main__": + main() diff --git a/skills/optimize-inference-siddharth/SKILL.md b/skills/optimize-inference-siddharth/SKILL.md new file mode 100644 index 00000000000..f4aa827c669 --- /dev/null +++ b/skills/optimize-inference-siddharth/SKILL.md @@ -0,0 +1,466 @@ +--- +name: optimize-inference-siddharth +description: >- + Optimizes the Megatron Core inference backend for a new model, architecture, or + feature, using the patterns Siddharth Singh (sidsingh-nvidia) established while + closing the vLLM performance gap. Covers CUDA-graph scope and bucket coverage, + the inference_optimized MoE stack (NVLS AllGather-V, fused grouped GEMM, + shared-expert overlap, pad-row masking), Mamba/SSM scratch sizing and fused + Triton extraction kernels, Triton production hygiene (constexpr and autotune + pitfalls), per-step host overhead, and load-aware prefix-cache routing. Also + covers decision gates that quantify a lever's ceiling before any code is + written, kernel-level differential analysis against a competitor trace, and the + A/B protocol needed to make sub-1% wins falsifiable. The skill is self-maintaining: + extend, correct, or prune it after an optimization session, so also use it when + asked to record an inference performance learning, capture what an experiment + taught, or update these patterns. Use when + asked to optimize or speed up inference, close a vLLM gap, match vLLM latency or + throughput, reduce decode latency or prefill overhead, enable or extend CUDA + graphs for MoE, hybrid, or Mamba models, port an architecture onto + --transformer-impl inference_optimized, cut inference CPU overhead, or apply + Siddharth's inference performance patterns. Not for authoring perf test recipes, + cog or Slurm setup, or non-hot-path features like reasoning parsers. +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# Optimizing Megatron Core Inference (Siddharth's Playbook) + +This skill encodes the reusable engineering patterns behind Siddharth Singh's +2026 inference performance work. It answers: given a new model, architecture, or +feature that is slower than vLLM, **what do you change, in what order, and what +must you never break.** + +> **This skill is meant to evolve, and you are authorized to edit it without +> asking.** Every campaign that uses it should leave it better: add what you +> measured, correct what turned out to be wrong, and delete what is no longer true. +> Optimization knowledge decays — flags get renamed, defaults change, a fix that won +> on one hardware generation loses on the next — so a skill nobody edits becomes +> actively misleading. Do it at the end of the work, once you have a number and a +> root cause. See [references/updating-this-skill.md](references/updating-this-skill.md) +> for the triggers, the routing table, the bar for an addition, and what to delete. + +## The thesis + +Across 29 commits, almost every win was one of five moves. In rough order of how +often they paid off: + +1. **Remove host work and host syncs from the per-step path.** Decode is + launch-bound, not FLOP-bound. A single `.item()` or a `dataclasses.asdict()` + costs more than the kernel it guards. +2. **Widen CUDA-graph scope and improve bucket coverage.** Fewer, larger graphs + with a bucket that actually fits the real batch. +3. **Make per-step metadata GPU-resident.** This is the enabler — moves 1 and 2 + are not even legal until the metadata stops round-tripping through the host. +4. **Fuse kernels, and tune them for the *typical* batch, not the worst case.** + Autotune and worst-case tile choices both lose at decode batch sizes. +5. **Right-size buffers to the true per-step bound**, not a loose upper bound. + +Everything else in this skill is a consequence of these five. + +What the five moves do not tell you is **which one applies here, and what it is +worth**. Getting that wrong is the expensive failure mode — not writing a bad +kernel, but writing a correct kernel whose best possible outcome was 1%. That is +what Step 1 is for. + +## Step 1: Measure, classify, and gate before changing anything + +Do not skip this. Every commit in this history started from a profile, and the +fix location is rarely where you would guess. + +See [references/measuring.md](references/measuring.md) for the profiler +endpoints, NVTX ranges, built-in counters, honest idle accounting, and — read this +before adding host visibility — **the trace flags that deadlock nsys finalization +on MoE decode workloads**. For trace analysis use the `nsight-system-analysis` +skill; for the throughput harness use `run-inference-performance-tests`. + +Classify the dominant signal, then jump to the matching section: + +| Dominant signal in the profile | Where to work | +|---|---| +| GPU idle gaps between kernels; host ahead of device | Step 2, CUDA graphs | +| Host time after the forward (serialize, detokenize, ZMQ) | Step 5, host path | +| Exposed NCCL, or EP ranks waiting on each other | Step 3, NVLS AllGather-V and `ep_consensus_interval` | +| Periodic multi-ms stalls, or slow first steps | Step 4, Triton recompilation and autotune | +| Grouped GEMM inefficient at small batch | Step 3, grouped-GEMM backend and tile heuristic | +| OOM, or a prefix cache far smaller than expected | Step 4, scratch sizing math | +| Many sub-microsecond kernels in a row | Step 3, fusion — but gate it first | + +### Then gate the lever: share is not headroom + +A category's share of device time is not what you can win back, for two reasons that +have each burned real weeks. A category can be 33% of device time with **2% of +headroom**, because it is already moving the bytes it has to move. And **device time +is not wall time** — kernels overlap across streams, so the step's wall clock is the +critical path through the per-layer dependency chain; work on a side stream +contributes less than its share, while a tiny kernel on the serial chain costs its +duration *plus* the dispatch gap behind it, every layer. + +So before writing anything that takes more than about a day, compute the ceiling: +establish the floor from measured machine constants, measure the current cost under +**graph replay** at the decode token count, take the ratio, subtract what the fix +itself costs (added launches, grid syncs, atomics, extra passes), and express the +result as a fraction of the step. Then write down **proceed** or **gated out**. + +Full method, the per-launch fixed costs you need, and three case studies where a +gate killed a multi-week effort — including a hand-written grouped GEMM whose +entire ceiling was 1.45×, most of it reachable by tile tuning alone — are in +[references/decision-gates.md](references/decision-gates.md). + +Skip the gate for cheap reversible changes: flag flips, tile retunes, backend enum +swaps. Gating those costs more than trying them. + +### If the target is a competitor, diff against its trace + +When the goal is "match vLLM," its trace is a specification, not just a scoreboard. +Take matched windows of one forward block from each and answer one question: **is +any individual kernel slower, or does mcore just launch more of them?** The fixes +are opposite. On Qwen3-30B the answer was ~467 kernels against ~1784 for the same +block, which pointed at fusion and away from kernel rewrites. + +Method, window anchoring, and the traps in adopting a competitor's kernel are in +[references/vllm-differential.md](references/vllm-differential.md). + +### Open a ledger before the first change + +For anything spanning more than one session, keep an append-only ledger: a fixed +protocol table (hardware, model, batch, output length, parallelism, warmup/timed +counts), one row per experiment **including every rejection with its root cause**, +a running distance-to-target, and a next-levers list re-derived after each profile. +Never edit a recorded result; supersede it. See +[references/measuring.md](references/measuring.md). The negative results are the +higher-value half — they are what stops the next person re-deriving a dead end. + +## Step 2: CUDA-graph the largest region that is safe + +Full detail in [references/cuda-graphs.md](references/cuda-graphs.md). + +Set `cuda_graph_impl="local"` and `inference_cuda_graph_scope=block` to capture +the whole decoder block in one graph. For hybrid models, graph ownership lives on +`HybridModel`, not the stack, so the embedding and output layers land inside the +same capture — that widening was itself a measurable win. + +Then fix coverage. A step whose shape matches no captured bucket silently falls +back to eager, so a "CUDA graphs enabled" run can still be launch-bound. Check +`num_cuda_graphs` (`-1` auto-sizes), `cuda_graph_max_tokens` (512 by default, so +prefill and mixed steps up to 512 tokens get a graph), and +`cuda_graph_sizing_distribution`. + +Know what the wide capture costs, though — it is a trade, not a free win. Under +full-iteration inference capture, the flashinfer sampling backend cannot run, async +scheduling is guarded off for EP (and deadlocks if you open the guard), and +multi-stream comm/compute overlap is bounded by the graph's structure rather than by +hardware queues, so `CUDA_DEVICE_MAX_CONNECTIONS` does nothing. Details and error +signatures are under *What a wide capture costs you* in +[references/cuda-graphs.md](references/cuda-graphs.md). + +## Step 3: Enable the model-path stack + +**MoE** — see [references/moe-inference.md](references/moe-inference.md). +Use `--transformer-impl inference_optimized`, which swaps in +`InferenceTopKRouter` and `InferenceGroupedMLP` and picks a dispatcher from +`inference_moe_token_dispatcher_type`. Prefer `nvls` (the default): the NCCL +fallback requires equal token counts across EP ranks, which forces decode-only +graphs. Then check the grouped-GEMM backend, shared-expert overlap, and that +padding rows route to no expert. + +**Hybrid / Mamba** — see +[references/mamba-and-triton.md](references/mamba-and-triton.md). +Size the extraction scratch to the real per-step bound before the durable prefix +cache claims the rest, and use fused gather-plus-scatter kernels gated on a +runtime count instead of materializing intermediates. + +**Dense GPT** — Steps 2, 4, and 5 usually cover it. Still apply the padding and +dtype contracts to any fused kernel you add. + +## Step 4: Kernel and Triton hygiene + +Full detail in [references/mamba-and-triton.md](references/mamba-and-triton.md). +The two rules that cause the most damage when violated are the `tl.constexpr` +specialization rule and the no-autotune-in-production rule, both in the hard +rules below. + +## Step 5: Trim the host path + +Full detail in [references/host-path.md](references/host-path.md). The per-step +host critical section is the `bookkeeping` to `detokenization` to +`coordinator_communication` span in +[dynamic_engine.py](megatron/core/inference/engines/dynamic_engine.py). Anything +there runs once per step and blocks the next one. + +## Step 6: Re-measure, keep a kill switch, and feed the result back + +Re-run the same measurement from Step 1 with the identical config. Then confirm +correctness with `run-inference-functional-tests`, and regenerate golden values +if you changed bucketization (padding changes shift outputs — expected, not a +bug). + +Re-measure **in the same allocation, back to back, arms alternating** — identical +configs drifted by up to 1.6% between sessions on this workload, which is larger +than most individual wins, so a cross-session comparison will mislead you in both +directions. Accept on **distribution separation** (slowest ON beats fastest OFF), +not on mean delta. Protocol and worked examples: *The noise floor is bigger than +your win* in [references/measuring.md](references/measuring.md). + +Expect the end-to-end result to differ from the kernel-level result, often by +several times in either direction. Work on the serial dependency chain converts at +more than 1:1 (a ~1% microbench ceiling delivered +2.9%, because removing a launch +also removes a graph node and a dispatch gap ×48 layers); already-overlapped work +converts at a third to a half; work off the critical path converts at roughly zero +(a 1.25× kernel win delivered a wash). If the measured conversion is far from your +prediction, you mis-identified where the work sits. + +Every optimization here shipped with a way to turn it off +(`inference_disable_triton_nvls_kernels`, +`inference_moe_disable_fused_quant_kernels`, the backend enums). Add one. It is +how the next person A/Bs your change instead of reverting it. + +Then close the loop: append the result to the campaign ledger — **including +rejections, with their root cause** — and promote whatever generalizes into this +skill. The promotion test is whether it would change what someone does on a +*different* model; if it only describes this one, it stays in the ledger. Mechanics +in [references/updating-this-skill.md](references/updating-this-skill.md). + +--- + +## Hard rules + +These are non-negotiable. Each one exists because violating it broke something. + +### 1. No host sync on the per-step path + +No `.item()`, `.tolist()`, `.cpu()`, or data-dependent Python branch in code that +runs every step. To publish a per-step scalar, `fill_` it into a **preallocated, +fixed-address** GPU tensor and read it inside the kernel: + +```python +# megatron/core/inference/contexts/attention_context/mamba_metadata.py +# fill_ is async (no host sync) and keeps the tensor at the same address +# captured graphs reference. +self._intermediate_real_count_buffer.fill_(self.intermediate_count) +``` + +The fixed address is the whole point: a captured graph records pointers, so the +*value* may change between replays but the *address* may not. + +### 2. `tl.constexpr` only for values fixed for the process lifetime + +A `tl.constexpr` parameter is baked into the compiled kernel, so **every distinct +value triggers a fresh JIT compile.** Marking a per-step batch or token count +`constexpr` means recompiling on every step. That was the entirety of commit +`f29c747` — a four-line fix worth a large latency spike. + +Constexpr is correct for block and tile sizes, `tl.arange` bounds, and unroll +counts. It is wrong for anything that varies per step. Note that even plain +`int` arguments get specialized on divisibility-by-16 and `== 1`; if such a value +varies, use `@triton.jit(do_not_specialize=["batch"])`. + +### 3. Graph-safe kernel shape: worst-case grid, data-conditional body + +Size the launch grid at capture time to the maximum, then let padded programs +exit immediately: + +```python +# megatron/core/ssm/ops/intermediate_extraction.py +real_count = tl.load(real_count_ptr).to(tl.int32) +if pid_slot >= real_count: + return +``` + +This keeps the grid static for replay while making padded slots nearly free. + +### 4. Size per-step work to the matched bucket, not the global max + +Once a step matches a CUDA-graph bucket, all metadata updates and scratch writes +should be bounded by that bucket, never by global `max_tokens` or `max_requests`. +Commit `9b4074b` was largely this one change applied to Mamba prefill. + +### 5. Padding must not do real work + +Three concrete forms: + +- CUDA-graph pad rows get routing index `-1` (via `mask_routing_padding`) so they + activate no expert. +- Pad tokens point at a reserved `dummy_block_idx` so KV-append writes somewhere + valid and throwaway. +- Do not *zero* rows past `valid_tokens`. Downstream only reads the valid prefix, + so a zeroing pass is pure wasted bandwidth — just don't write them. + +### 6. Tune for the expected batch; never autotune in production + +Pick tile, warp, and stage counts from a host-side heuristic keyed on the typical +token count. Commit `20f09364` replaced 25 `@triton.autotune` configs with +vLLM's `_get_default_config` heuristic, which cut compile time and picked better +tiles for decode-sized batches. Route any remaining autotuning through +`autotune_configs` in +[determinism.py](megatron/core/ssm/ops/determinism.py). + +### 7. Preserve Transformer Engine `nn.Parameter` identity + +Grouped GEMM wants one stacked `[num_experts, out, in]` weight; TE stores +per-expert parameters. Redirect `param.data` to a view into the stacked buffer +rather than replacing the `Parameter` object, and build it **lazily on first +forward**, after checkpoint load: + +```python +# megatron/core/transformer/moe/experts.py +# Redirect param.data to view into contiguous buffer. +# The nn.Parameter object stays the same - TE's internal state is preserved. +fc1_param.data = _fc1_weight[i] +``` + +Replacing the object corrupts TE's FP8 and bookkeeping state; building eagerly at +`__init__` reads weights that do not exist yet. Commit `905c0e38` fixed both +after they broke the RL integration. + +### 8. Hold the dtype contract at every kernel boundary + +Token ids are `int64` throughout the pipeline. FlashInfer's sampling kernels +return `int32`, so their results are cast with `.long()` at the boundary. +Normalize dtype where the external kernel enters, not by letting it propagate — +the cast is free once, a mismatch inserts conversion kernels on every step. + +Related, and a sharper trap: **not everything belongs in a graph.** FlashInfer +sampling is deliberately left eager, because its kernel choice is data-dependent +*and* FlashInfer bakes the philox RNG state into a graph as a by-value constant at +capture — so a captured sampler replays identical random numbers every step. Check +for baked-in state before widening a graph over anything stateful or random. + +### 9. Guard correctness with a reference implementation + +The testing pattern throughout: compare the fast kernel against an obvious +PyTorch or plain-loop reference with `torch.testing.assert_close`; assert the +gating behavior explicitly by prefilling the output with a sentinel and checking +padded slots still hold it; and re-derive any sizing formula independently in the +test so a changed bound fails loudly. See `add-inference-unit-tests`. + +### 10. A sub-1% claim needs same-session, back-to-back, separated arms + +Session-to-session drift on identical configs reached 1.6% — larger than most +individual wins once the easy ones are gone. So a number compared against last +session's baseline is not evidence. + +Re-baseline the current best config **in the allocation you will test in**, run the +OFF and ON arms **back to back**, repeat the pair, and accept only if the arms do +not overlap: + +``` +min(ON iterations) > max(OFF iterations) +``` + +Report pairwise deltas rather than one average, and treat the first timed iteration +as a suspected cold outlier. A +1% mean with overlapping arms is not a result; a ++0.9% with separated arms is. + +### 11. If it is not bit-exact, say so and justify it differently + +Bit-exactness is achievable more often than assumed — retuning tiles while holding +`BLOCK_SIZE_K` fixed preserves the fp32 reduction order, and a masked slot adding an +exact 0.0 keeps a fused reduction bit-exact. **Check whether a bit-exact formulation +exists before accepting drift.** + +When none does — norm and reduction fusions frequently land one ulp off, since TE's +internal rsqrt and reduction order differ — the acceptance argument changes shape and +must be made explicitly: bound the deviation in ulps against the reference across +several token counts and seeds (`max_rel ≤ 7.9e-3` is about one bf16 ulp); diff +fixed temperature-0 coherence output against the gate-OFF arm; and for any prompt +that diverges, **inspect where**. A divergence at a genuinely low-confidence branch +with both continuations fluent and factually correct is acceptable; one that degrades +fluency or correctness is not. Record which prompts diverged — "not bit-exact" must +be lookup-able, not rediscovered. + +Never let an ulp-level deviation ride on `assert_close` alone: the tolerance that +passes it also passes a real bug. + +--- + +## Code entry points + +| Area | Paths | +|---|---| +| MoE dispatchers | [token_dispatcher_inference.py](megatron/core/transformer/moe/token_dispatcher_inference.py) — `NCCLAllGatherDispatcher`, `NVLSAllGatherVDispatcher` | +| MoE experts / router | [experts.py](megatron/core/transformer/moe/experts.py) `InferenceGroupedMLP`, [router.py](megatron/core/transformer/moe/router.py) `InferenceTopKRouter`, [moe_layer.py](megatron/core/transformer/moe/moe_layer.py) | +| Fused MoE kernels | [megatron/core/inference/moe/](megatron/core/inference/moe/) — `fused_moe.py`, `vllm_fused_moe.py`, `permute.py`, `activations.py`, `metadata.py` | +| Pad-row masking | [inference_routing_mask_kernel.py](megatron/core/transformer/moe/inference_routing_mask_kernel.py) | +| Backend selection | [backends.py](megatron/core/models/backends.py), [moe_module_specs.py](megatron/core/models/gpt/moe_module_specs.py) | +| CUDA-graph buckets | [batch_dimensions_utils.py](megatron/core/inference/batch_dimensions_utils.py) | +| Graph scope hooks | [enums.py](megatron/core/transformer/enums.py), [transformer_block.py](megatron/core/transformer/transformer_block.py), [hybrid_model.py](megatron/core/models/hybrid/hybrid_model.py) | +| Per-step context state | [dynamic_context.py](megatron/core/inference/contexts/dynamic_context.py), [gpu_view.py](megatron/core/inference/contexts/gpu_view.py) | +| Mamba / SSM | [mamba_mixer.py](megatron/core/ssm/mamba_mixer.py), [intermediate_extraction.py](megatron/core/ssm/ops/intermediate_extraction.py), [mamba_metadata.py](megatron/core/inference/contexts/attention_context/mamba_metadata.py), [mamba_slot_allocator.py](megatron/core/inference/contexts/mamba_slot_allocator.py) | +| Symmetric memory / NVLS | [symmetric_memory.py](megatron/core/inference/symmetric_memory.py), [torch_symm_triton/](megatron/core/inference/communication/torch_symm_triton/), [inference_layers.py](megatron/core/tensor_parallel/inference_layers.py) | +| Engine / host path | [dynamic_engine.py](megatron/core/inference/engines/dynamic_engine.py), [inference_request.py](megatron/core/inference/inference_request.py), [data_parallel_inference_coordinator/](megatron/core/inference/data_parallel_inference_coordinator/) | + +## Flags + +Verified against the current tree. Defaults matter — several are already the +tuned value, so the useful move is often *checking* rather than changing them. + +| Flag | Default | Effect | +|---|---|---| +| `transformer_impl` | `transformer_engine` | `inference_optimized` swaps in the whole inference MoE stack | +| `inference_moe_token_dispatcher_type` | `'nvls'` | `nccl` fallback forces equal EP token counts and decode-only graphs | +| `inference_grouped_gemm_backend` | `"vllm"` | `vllm`, `flashinfer`, or `torch`; MXFP8 needs `torch` | +| `inference_disable_triton_nvls_kernels` | `False` | Kill switch: fall back to NCCL collectives | +| `inference_moe_disable_fused_quant_kernels` | `False` | Kill switch: unfuse activation + quantize | +| `moe_router_dtype` | — | Must be `fp32` for `inference_optimized`, to avoid per-decode dtype conversion | +| `cuda_graph_impl` | `"none"` | Inference graphs need `local` | +| `inference_cuda_graph_scope` | derived | `none`, `layer`, or `block`; `local` derives `layer`, so set `block` explicitly | +| `cuda_graph_max_tokens` | `512` | Token ceiling for prefill and mixed graphs | +| `num_cuda_graphs` | `16` | `-1` auto-sizes from `log2(max_tokens)` | +| `cuda_graph_sizing_distribution` | `EXPONENTIAL` | `LINEAR` gives the dense small-batch ladder | +| `ep_consensus_interval` | `20` | Skip EP consensus all-reduces while busy | +| `prefix_caching_routing_alpha` | `0.5` | 0 is pure load balance, 1 is pure cache affinity | +| `SamplingParams.return_prompt_tokens` | `False` | Opt in to echoing prompt ids over the wire | +| `moe_enable_routing_replay` | `False` | Record per-token expert choices for imbalance analysis | + +### Flags that look like free wins and are not + +Measured and rejected on Qwen3-30B-A3B EP4 under `inference_optimized` + block +scope. Each is the sort of thing you would reasonably try first. + +| Flag | Result | +|---|---| +| `--moe-router-fusion`, `--moe-permute-fusion` | **Crash**: `AssertionError: hidden_size mismatch: 128 vs 8`. Wired to the training MoE path; TE's fused router emits a dense `num_experts` map, not the dense top-k the inference dispatcher needs. | +| `--inference-dynamic-batching-sampling-backend flashinfer` | **Crash** under graph capture: `Generator not registered with the capturing graph`. | +| `--inference-dynamic-batching-async-sched-mode serial` | Guarded off for EP; opening the guard deadlocks the first single-request decode step. Gains +0.85% at batch — the post-sampling chain is a genuine data dependency, not a scheduling artifact. | +| `CUDA_DEVICE_MAX_CONNECTIONS=8` | Flat. Overlap is bounded by graph structure, not queue count. | +| `--inference-moe-token-dispatcher-type nccl` | **0.66×**. Pads to the worst-case per-rank count, roughly doubling comm volume. Correctness fallback only. | +| `--inference-grouped-gemm-backend torch` | **0.82×** for BF16. Use only when MXFP8 forces it. | + +## References + +- [references/decision-gates.md](references/decision-gates.md) — quantify a lever's ceiling before building it; per-launch fixed costs; three gates that killed multi-week efforts +- [references/vllm-differential.md](references/vllm-differential.md) — kernel-level comparison against a competitor trace; slower kernels vs more kernels +- [references/cuda-graphs.md](references/cuda-graphs.md) — scope, bucketing, graph-safety invariants, idle EP ranks, what a wide capture forecloses +- [references/moe-inference.md](references/moe-inference.md) — dispatchers, grouped GEMM, overlap, pad masking +- [references/mamba-and-triton.md](references/mamba-and-triton.md) — Triton rules, fused extraction, scratch sizing +- [references/host-path.md](references/host-path.md) — serialization, IPC payloads, DP routing +- [references/measuring.md](references/measuring.md) — profiler endpoints, NVTX, built-in counters, nsys flags that deadlock, union-busy idle accounting, A/B protocol, kernel-to-e2e conversion, ledger +- [references/commit-log.md](references/commit-log.md) — every commit mapped to its pattern, plus superseded APIs +- [references/updating-this-skill.md](references/updating-this-skill.md) — how to extend, correct, and prune this skill as you learn +- [assets/review-checklist.md](assets/review-checklist.md) — pre-PR checklist + +## Out of scope + +- Authoring perf or functional test recipes — use `add-inference-performance-test` +- Cluster and container bootstrap — use `cog-setup-and-help` +- Non-hot-path product features (reasoning parsers, chat-template retention) + +## Keeping this skill current + +Edit this skill at the end of a piece of work, once you have a number and a root +cause — never mid-experiment on a hypothesis. Triggers: an accepted optimization, a +**rejection** (the highest-value entry type), a negative decision gate, a tooling +failure that cost over an hour, a measurement that contradicts something here, a +flag or default that differs from what is documented, or an invariant whose +violation broke something. + +Route by subject into the reference that owns it — only new invariants, flag +behavior, and routing lines belong in this file, because `SKILL.md` is the router +and it stops working once it needs its own router. An addition must be measured, +root-caused, scoped, and actionable. Delete what is false; never delete a measured +negative result. + +Triggers, routing table, house style, deletion policy, size budgets, and the +revision log: [references/updating-this-skill.md](references/updating-this-skill.md). diff --git a/skills/optimize-inference-siddharth/assets/review-checklist.md b/skills/optimize-inference-siddharth/assets/review-checklist.md new file mode 100644 index 00000000000..f57d33e339d --- /dev/null +++ b/skills/optimize-inference-siddharth/assets/review-checklist.md @@ -0,0 +1,160 @@ +# Inference Optimization Pre-PR Checklist + +Copy this into the working notes and walk it before opening the PR. Most items map +to a hard rule in [SKILL.md](../SKILL.md); the rest are review feedback these +commits actually received. + +## Decision gate (before writing the code, not before the PR) + +``` +- [ ] Floor established from MEASURED machine constants, not datasheet peaks +- [ ] Current cost measured under graph replay, at the real decode token count +- [ ] Gross ceiling stated as a ratio AND as a % of the step +- [ ] Net ceiling subtracts what the fix costs (added launches, grid syncs, + atomics, extra passes) +- [ ] Mechanism identified (bytes / FLOPs / occupancy / CTA utilization / + dependent load / pure launch overhead) — not just the magnitude +- [ ] Target located on the serial chain vs already-overlapped vs off-path, so the + expected kernel-to-e2e conversion is predicted before measuring +- [ ] Verdict recorded in the ledger, including gates that came out negative +``` + +## Measurement + +``` +- [ ] Baseline captured before any change, with the config recorded +- [ ] Baseline re-run in the SAME allocation as the test arm (cross-session drift + reached 1.6%, larger than most wins) +- [ ] OFF and ON arms run back to back, and the pair repeated at least once +- [ ] Arms do not overlap: slowest ON beats fastest OFF (report pairwise deltas, + not one average) +- [ ] First timed iteration checked for cold-start outlier before averaging +- [ ] After-measurement uses the identical batch size, sequence length, parallelism +- [ ] Warmup ran long enough that CUDA-graph capture is excluded from both numbers +- [ ] The win is attributed to a specific kernel or host span, not just end-to-end +- [ ] Measured conversion is consistent with the predicted one; if not, the target + was mis-located +- [ ] One mechanism per PR, so the measurement attributes cleanly +- [ ] Profiler flags known-good (no osrt, no process-tree sampling, no NVTX under + graph capture) or the capture will not finalize +``` + +## Host path + +``` +- [ ] No .item() / .tolist() / .cpu() added to per-step code +- [ ] No dataclasses.asdict() on any object that can hold a tensor +- [ ] No torch.save / pickle on the IPC path +- [ ] New per-request fields on the wire are scalars, or opt-in behind a flag +- [ ] Work not needed for the next step is off the step loop +- [ ] New host work on the step loop has an NVTX range around it +``` + +## CUDA graph safety + +``` +- [ ] Per-step scalars go into preallocated fixed-address GPU tensors via fill_/copy_ +- [ ] No new buffer allocation inside a captured region +- [ ] Grow-only buffers pre-sized to the worst case BEFORE capture +- [ ] No Python-level assignment in forward() relied on at replay time (use copy_) +- [ ] Nothing stateful or RNG-bearing newly captured without checking what gets + frozen by value +- [ ] Per-step work bounded by the matched bucket, not global max_tokens/max_requests +- [ ] Bucket coverage verified: confirm real steps match a graph rather than + silently falling back to eager +``` + +## Padding + +``` +- [ ] Pad rows route to expert -1 (no expert activated) +- [ ] Pad tokens point at reserved dummy storage, not real blocks +- [ ] Kernels skip rows where permutation_map == -1 +- [ ] No zeroing pass over rows past valid_tokens (don't write them at all) +- [ ] Idle EP ranks use the lightweight dummy path, never add_request +``` + +## Triton kernels + +``` +- [ ] tl.constexpr only on values fixed for the process lifetime +- [ ] Per-step ints either do_not_specialize or passed as a 0-d GPU tensor +- [ ] Grid sized to worst case, body gated by `if pid >= real_count: return` +- [ ] Any autotune goes through autotune_configs(), with a short config list +- [ ] Tile/warp/stage choice driven by a typical-batch hint, not buffer capacity +- [ ] torch.empty rather than torch.zeros where the kernel overwrites every row +``` + +## MoE specifics + +``` +- [ ] moe_router_dtype is fp32 (hard requirement for inference_optimized) +- [ ] NVLS eligibility checked via are_tensors_nvls_eligible, not re-derived +- [ ] NCCL fallback path still correct for non-NVLink / non-bf16 / unaligned shapes +- [ ] TE nn.Parameter identity preserved (param.data views, built lazily after + checkpoint load) +- [ ] Shared-expert overlap re-measured on THIS model; CTA caps not copied forward +- [ ] Unsupported config combinations rejected with a clear error, not silently wrong +``` + +## Buffer sizing + +``` +- [ ] Sized from the true per-step bound, not a loose upper bound +- [ ] The bound is a single shared value, not duplicated across modules +- [ ] Scratch reserved before durable pools +- [ ] Over-budget raises a config-time error naming the knob to reduce +- [ ] Assert on overrun rather than trusting the caller +``` + +## Correctness + +``` +- [ ] Reference implementation test (plain PyTorch or loop) with assert_close +- [ ] Gating asserted with a sentinel: padded slots still hold the sentinel after + the kernel runs +- [ ] Sizing formulas re-derived independently in the test, covering BOTH regimes + of any min()/max() +- [ ] Boundary cases covered (window before position 0, sub-block prefills, + non-default chunk sizes) +- [ ] Error paths tested (too-small budget raises, not OOM later) +- [ ] Golden values regenerated if bucketization or padding changed, and the + generated text sanity-checked for coherence +- [ ] Inference and training paths still agree where they should (e.g. router + picks the same experts) +- [ ] Checked whether a BIT-EXACT formulation exists (hold BLOCK_SIZE_K fixed; + masked slots adding exact 0.0) before accepting any deviation +- [ ] If not bit-exact: deviation bounded in ulps across several token counts and + seeds, temperature-0 coherence diffed against gate-OFF, every divergence + inspected for low-confidence branch + fluent + factually correct, and the + divergence recorded in the ledger +``` + +## Shipping + +``` +- [ ] Kill switch exists so the next person can A/B without reverting +- [ ] New flags documented with defaults in arguments.py and the config docstring +- [ ] Imports run through `uv run isort` +- [ ] Ledger entry appended (hypothesis, changed files/flags, throughput, delta, + correctness, run id, conclusion) — including for rejected attempts +- [ ] PR opened as a draft, commits signed with both -s and -S +``` + +## Skill maintenance (after the work, not before the PR) + +``` +- [ ] Anything that generalizes beyond this model promoted into the skill, routed + to the right reference rather than piled into SKILL.md +- [ ] Rejections and negative gate verdicts recorded with their root cause +- [ ] Anything this work proved wrong in the skill corrected in place; contradictions + scoped by hardware/version rather than overwritten +- [ ] Content that is now false deleted; measured negative results kept +- [ ] Links resolve, frontmatter parses, revision log appended +``` + +See [../references/updating-this-skill.md](../references/updating-this-skill.md). + +See `add-inference-unit-tests` for where new tests belong, `mcore-testing` for +recipe structure, and `run-inference-functional-tests` for verification on the +cluster. diff --git a/skills/optimize-inference-siddharth/references/commit-log.md b/skills/optimize-inference-siddharth/references/commit-log.md new file mode 100644 index 00000000000..b6fa11ebdee --- /dev/null +++ b/skills/optimize-inference-siddharth/references/commit-log.md @@ -0,0 +1,121 @@ +# Commit Provenance + +Every commit by `sidsingh@nvidia.com` on `main` since 2026-01-01, grouped by the +pattern it demonstrates. Use this to read the primary example of any technique: + +```bash +git show +git show --stat # scope first +git log -1 --format=%B # message and co-authors +``` + +Regenerate this list: + +```bash +git log --format='%h|%ad|%s' --date=short --since=2026-01-01 --author='sidsingh' main +``` + +## MoE inference stack + +| SHA | PR | Pattern | +|---|---|---| +| `7d1c01685` | #3496 | **Foundational.** Forks a parallel inference MoE hierarchy: `InferenceTopKRouter`, `InferenceGroupedMLP`, graph-safe AllGather dispatcher, GPU-resident offsets, fused NVLS collectives, centralized `are_tensors_nvls_eligible`. | +| `905c0e386` | #3851 | Lazy `SymmetricMemoryManager`; lazy non-destructive weight concatenation via `param.data` views to preserve TE `Parameter` identity. Fixes the RL integration. | +| `589cd9e12` | #3858 | `megatron/core/inference/moe/` package: `mcore_fused_moe`, Triton permute, padding-aware activations, fused activation+MXFP8-quantize+swizzle. Backend enum replaces a boolean. | +| `bfd45740c` | #4258 | **Architectural rewrite.** Variable-count AllGather-V / ReduceScatter-V so EP ranks carry different token counts; `fused_metadata_update` collapses 5 kernels into 1; buffers allocated once at init; dispatcher swap moves into `MoELayer.train()`. | +| `442a936a1` | #4570 | Shared-expert overlap on `SharedExpertMLP.stream`, with the AGV capped to 16 CTAs so it does not starve the side stream. | +| `c817dad28` | #4587 | Cuts the EP graph-size sync back to the minimum (max token count + is-anyone-non-decode) after it accumulated dead complexity. | +| `20f09364e` | #4603 | Replaces 25 `@triton.autotune` configs with vLLM's host-side `_get_default_config` heuristic keyed on a token-count *hint*; persistent-grid `moe_sum` that stops zeroing rows past `valid_tokens`; latent-MoE shared-expert overlap; `ep_consensus_interval`. | +| `3a253ac5c` | #4922 | `mask_routing_padding` writes `-1` into padded rows' topk slots so CUDA-graph padding activates no expert. Fixed-address `real_token_count` scalar. | + +Start with #3496 for the structure and #4258 for the current design. Read them in +that order — #4258 only makes sense as a response to #3496's equal-token-count +limitation. + +## CUDA graphs + +| SHA | PR | Pattern | +|---|---|---| +| `32efeffd2` | #3250 | Splits inference graph scope from the training `full_iteration` scope; adds block-scope graphs for the Mamba block. Establishes the `_should_call_local_cudagraph` predicate. | +| `fde3b90a8` | #3527 | `num_cuda_graphs=-1` auto-sizing with a dense small-batch ladder, `[1,2,4] + range(8,256,8) + range(256,max,16)` — fine granularity where decode batches actually live. | +| `60a25aa67` | #3525 | `add_dummy_requests_for_expert_parallel_step`: idle EP ranks `fill_` preallocated tensors instead of constructing request objects. Rewrote 1448 lines of golden values because bucketization changed. | +| `35f76df3f` | #4440 | Moves hybrid graph ownership from `HybridStack` up to `HybridModel` so embedding and output layers fall inside the capture. Widening scope as an optimization. | +| `740c16e6b` | #5797 | `cuda_graph_max_tokens` default 512, so prefill and mixed steps up to 512 tokens get a graph instead of falling back to eager. | + +## Mamba / SSM and Triton + +| SHA | PR | Pattern | +|---|---|---| +| `ab2b33d54` | #4397 | `do_not_specialize=["batch"]`; `fast_exp` via `exp2`; capability-specialized `BLOCK_SIZE_M`/`num_warps` for Blackwell; `torch.zeros` to `torch.empty` where the kernel overwrites every row. | +| `9b4074b51` | #4764 | Bounds per-step Mamba work by the padded bucket rather than the global max; retunes the varlen conv autotune list to two regimes; removes the AGV CTA cap from #4570 after measuring it net-negative. | +| `f29c747f2` | #5608 | **Four lines, high value.** Demotes per-step batch sizes from `tl.constexpr` to runtime args, ending per-step JIT recompilation. | +| `411a5d8b2` | #5863 | Scratch sized `min(ceil(max_tokens / block_size_tokens), 3 * max_requests)` — an order of magnitude at low concurrency. Also fixes a hardcoded `mamba_chunk_size = 128`. | +| `648bc011f` | #5866 | Fused gather-plus-conditional-scatter Triton kernels replacing dense gather + `copy_`; runtime `real_count` gate; transpose folded into the write address. | + +## Host path, scheduling, observability + +| SHA | PR | Pattern | +|---|---|---| +| `53a2b19a2` | #2920 | Drops `dataclasses.asdict` (recursive deepcopy) and `torch.save` from request IPC; moves detokenization to the coordinator so it overlaps the next engine step; adds the NVTX ranges that make the host critical path visible. | +| `faced5128` | #3034 | MoE routing replay: records per-token expert choices into a static graph-safe buffer for load-imbalance analysis. | +| `eadbaa618` | #5611 | `/start_profile` and `/stop_profile` endpoints relaying `cudaProfilerStart/Stop` to every engine, to pair with `nsys --capture-range=cudaProfilerApi`. | +| `edd45620b` | #5609 | Prefix-cache hit reporting as `usage.prompt_tokens_details.cached_tokens`. | +| `bcf4c8fb5` | #5607 | Load-aware DP routing: `alpha * match + (1-alpha) * free_capacity`, replacing round-robin. Vectorized numpy because it runs per request. | +| `602fad039` | #5918 | Drops `prompt_tokens` from the wire by default, keeping `prompt_length` for the usage contract. | + +## Test and CI maintenance + +Not optimization patterns, but useful for seeing what breaks when the above lands. + +| SHA | PR | Note | +|---|---|---| +| `369e0eba7` | #3071 | Fixes functional tests broken by #2920 | +| `66ec17eac` | #3357 | Inference functional test fixes | +| `a1165fabc` | #4454 | Fixes GitLab functional tests | + +## Non-performance work + +Included for completeness; these are product features on the same subsystem, not +hot-path optimizations. + +| SHA | PR | Note | +|---|---|---| +| `a79f49d37` | #5276 | Reasoning-token retention delegated to the chat template | +| `82e9dc69c` | #5634 | `nemotron_v3` reasoning parser | + +## Superseded APIs — do not copy from these + +Several commits above contain code that has since been replaced. If you read them +directly, know what no longer exists. + +**Dispatchers.** `InferenceCUDAGraphTokenDispatcher` (#3496) and the manual +`set_inference_cuda_graphed_iteration` / `unset_...` plumbing were replaced in #4258 +by `NCCLAllGatherDispatcher` / `NVLSAllGatherVDispatcher` plus the automatic +`MoELayer.train()` swap. + +**Grouped GEMM backend.** `inference_disable_torch_grouped_mm` (#3496) → +`inference_grouped_gemm_backend: 'auto'|'torch'|'te'` (#3858) → +`'flashinfer'|'torch'` (#4258) → `'flashinfer'|'torch'|'vllm'` with `vllm` as the +default (#4603). + +**Deleted modules.** `megatron/core/inference/moe/pad.py` and the `skip_permute` +branches (added #3858, removed #4258). `smallest_non_decode_cuda_graph_size` +(removed #4587). + +**Symmetric memory globals.** `_GLOBAL_SYMMETRIC_MEMORY_BUFFER_TP/_EP` in +`parallel_state.py` (#3496) → `SymmetricMemoryManager` in +`megatron/core/inference/symmetric_memory.py` (#3851). + +**CUDA graph scope.** The `cuda_graph_scope` list holding +`CudaGraphScope.full_iteration_inference` (#3250, #4440) → `cuda_graph_impl` plus +`inference_cuda_graph_scope` (`none`/`layer`/`block`). `CudaGraphScope` survives +only for checkpoint deserialization; its docstring carries the migration table. + +## Related work on branches, not `main` + +`44334edad` "Fix flashinfer sampling kernels to return int64 instead of int32" +lives on the unmerged `fix_flashinfer_sampling` branch. The same `.long()` casts +reached `main` via `650b7838` (#5791, Keshav Santhanam, co-authored by Siddharth), +which also concluded that FlashInfer sampling must run **eagerly** — its kernel +choice is data-dependent and FlashInfer bakes the philox RNG state into a graph as +a by-value constant. See [host-path.md](host-path.md). diff --git a/skills/optimize-inference-siddharth/references/cuda-graphs.md b/skills/optimize-inference-siddharth/references/cuda-graphs.md new file mode 100644 index 00000000000..32a5ee25554 --- /dev/null +++ b/skills/optimize-inference-siddharth/references/cuda-graphs.md @@ -0,0 +1,302 @@ +# CUDA Graphs for Inference + +Decode is launch-bound. A hybrid or MoE decode step issues hundreds of small +kernels whose combined GPU time is less than the CPU time to launch them, so the +device idles between kernels. CUDA graphs replace per-step launches with one +replay. Getting them to actually engage is most of the work. + +Source commits: `32efeffd` (#3250), `fde3b90a` (#3527), `60a25aa6` (#3525), +`35f76df3` (#4440), `740c16e6` (#5797). + +## Scopes + +Two config fields control this today: + +- `cuda_graph_impl`: `none` | `local` | `transformer_engine` | `full_iteration`. + Inference graphs require **`local`**. +- `inference_cuda_graph_scope`: `none` | `layer` | `block`. + +| Scope | Graphs per step | When to use | +|---|---|---| +| `block` | One, wrapping the whole decoder block | The latency target. Requires the entire block to be graph-safe. | +| `layer` | One per transformer or Mamba layer | Fallback when block capture is not safe. Still cuts launches substantially. | +| `none` | Eager | Baseline / debugging. | + +`local` derives `layer` by default, so **set `block` explicitly** — a config that +merely says `cuda_graph_impl=local` is leaving the biggest win on the table. + +`CudaGraphModule` (`attn`, `mlp`, `moe`, `moe_router`, …) is a separate, +training-oriented axis for capturing sub-layer regions. Not used for inference. + +
+Deprecated API you will see in older commits and docs + +Commits before the refactor use a `cuda_graph_scope` *list* holding +`CudaGraphScope.full_iteration_inference`. The migration guide lives in +[enums.py](megatron/core/transformer/enums.py): + +- `full_iteration` → `cuda_graph_impl="full_iteration"` (training) +- `full_iteration_inference` → `inference_cuda_graph_scope=block` +- everything else → the equivalent `CudaGraphModule` member + +`CudaGraphScope` is retained only for checkpoint deserialization. Do not write +new code against it. +
+ +## Who owns the graph + +Two hooks decide, and they must agree — if both a block and its layers think they +own the graph you get nested capture. + +**`create_mcore_cudagraph_manager(config)`** attaches the manager. Each module +checks whether the scope selects it: + +```python +# megatron/core/models/hybrid/hybrid_model.py +def create_mcore_cudagraph_manager(self, config): + if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: + from megatron.core.transformer.cuda_graphs import CudaGraphManager + self.cudagraph_manager = CudaGraphManager(config) +``` + +**`_should_call_local_cudagraph(*args, **kwargs)`** decides per call whether to +replay or run eager. It checks the scope, that a real inference context is +present, that `attention_mask is None`, and finally +`inference_context.using_cuda_graph_this_step()`. + +Ownership by model family: + +| Model | Owner under `block` scope | +|---|---| +| Hybrid | `HybridModel` — the *model*, so embedding + stack + output layer are one capture | +| GPT | `model.decoder`, the `TransformerBlock` | +| Mamba | The `MambaStack` | + +### Widening the scope is itself an optimization + +Commit `35f76df3` moved hybrid graph ownership up from `HybridStack` to +`HybridModel`, making `HybridStack` a plain `MegatronModule` and adding +`GraphableMegatronModule` to `HybridModel`. The embedding lookup and the final +projection are individually small kernels, but their launch overhead is on the +critical path of every decode step. Folding them into the existing capture cost +nothing and removed those launches. + +**Generalizable:** when you add a new model class, ask what is still outside the +graph. Small kernels bracketing a graphed region are pure launch overhead. + +### What a wide capture costs you + +Widening is the right default, but it is a trade, not a free win, and the things it +forecloses are not obvious until you try them. On Qwen3-30B under full-iteration +inference capture, three separate levers turned out to be unreachable *because* the +capture was wide: + +| Lever | Outcome under wide capture | +|---|---| +| flashinfer sampling backend | Server crash: `RuntimeError: Generator not registered with the capturing graph`. The sampler cannot be inside the capture — see the RNG hazard below. | +| Async scheduling (`async-sched-mode=serial`) | Guarded off for EP (`Async scheduling does not support expert parallelism`); opening the guard ran at batch, but deadlocked on the first single-request decode step. The guard encodes a real limitation. | +| Comm/compute overlap on separate streams | `CUDA_DEVICE_MAX_CONNECTIONS=8` measured flat. Overlap is bounded by the captured graph's structure and data dependencies, not by hardware queue count. Chunked collective pipelining needs concurrent streams under capture and hits the same wall. | + +The async-scheduling result is worth internalizing beyond its own flag: async +overlap cannot hide the post-sampling host chain, because that chain is +**data-dependent on the current step's sampled tokens**. Measured end-to-end gain +from opening the guard was +0.85% — consistent with a genuine serial dependency +rather than a scheduling artifact. + +**Generalizable:** before proposing overlap, streams, or asynchrony as the fix for +GPU idle, check whether the idle is a *data dependency*. If step N+1's input needs +step N's sampled token, no scheduler will overlap them, and the only fix is making +that host chain cheaper. + +## Bucket coverage + +A graph replays at a fixed shape, so every step is padded up to the smallest +captured bucket that fits. Two failure modes: + +- **No bucket fits** → silent fallback to eager. A run with graphs "enabled" can + still be launch-bound. This is the single most common reason a CUDA-graph + change shows no improvement. +- **Nearest bucket is far too large** → wasted compute on padding. + +All the logic is in +[batch_dimensions_utils.py](megatron/core/inference/batch_dimensions_utils.py). +The bucket unit is `InferenceBatchDimensions(token_count, prefill_req_count, +decode_req_count)`. + +### Sizing distributions + +`EXPONENTIAL` (default) halves down from the ceiling. This bounds the graph count +at roughly `log2(max_tokens)` and the relative padding at about 2x: + +```python +sizes = set() +val = cuda_graph_max_tokens +for _ in range(num_cuda_graphs): + rounded = max(rounder, (val // rounder) * rounder) + rounded = math.ceil(rounded / tp_size) * tp_size + sizes.add(rounded) + val //= 2 + if val < 1: + break +``` + +`LINEAR` gives the dense small-batch ladder introduced by `fde3b90a`, which +mirrors vLLM's: + +```python +sizes = ( + [1, 2, 4] + list(range(8, 256, 8)) + list(range(256, cuda_graph_max_tokens + 1, 16)) +) +``` + +The reasoning: decode batches are usually small, so spend granularity at the low +end where relative padding hurts, and be coarse above 256 to bound the graph +count. Note the hygiene steps that follow in both paths — TP-align each entry, +clamp to the ceiling, and force the endpoints to be present. + +`num_cuda_graphs=-1` auto-sizes. Under `EXPONENTIAL` it derives the count from +`log2(cuda_graph_max_tokens)` plus headroom, floored at 4. + +### Two ceilings + +Decode and prefill graphs are sized separately: + +- **Decode** is always capped at `max_requests * (num_speculative_tokens + 1)` — + one token per request, or 1 + speculative. +- **Prefill and mixed** are capped at `cuda_graph_max_tokens`, default **512** + (commit `740c16e6`). Set `cuda_graph_all_prefills` to extend to `max_tokens`, + at the cost of many large graphs. + +At runtime `match_graph_config` picks `min()` over the applicable buckets and +returns `None` to signal eager. + +### Cost of more graphs + +Each bucket is a separate capture: its own activation memory in the graph mempool +plus capture time. `create_cuda_graphs` in +[dynamic_engine.py](megatron/core/inference/engines/dynamic_engine.py) runs a +full forward per bucket and logs elapsed time and memory deltas, so capture cost +is directly observable. `cuda_graph_max_tokens=512` exists to keep this bounded +by default. + +## Making code graph-safe + +### What breaks capture or replay + +1. **Dynamic shapes** — handled by bucketing plus padding. +2. **Reallocated buffers.** A graph records raw pointers. If a buffer is freed and + reallocated between capture and replay, replay writes into freed memory. Any + grow-only buffer must be pre-sized to the worst case *before* capture: + +```python +# megatron/core/inference/engines/dynamic_engine.py +# A forward larger than the capture-time size would reallocate (and free) the buffer +# whose address a captured graph still writes to on replay, corrupting whatever later +# reuses that freed block. +if getattr(model_config, "sequence_parallel", False): + max_ag_numel = self.context.max_tokens * model_config.hidden_size + get_global_memory_buffer().get_tensor((max_ag_numel,), model_config.params_dtype, "mpu") +``` + +3. **Host syncs and data-dependent Python** inside the captured region. +4. **Python-level assignment inside `forward`.** Under `block` scope it executes + only during capture, so replays never see it. Use `copy_()` into a + preallocated buffer instead — this is exactly why + `mtp_decoder_hidden_states` exists. +5. **State baked in by value at capture.** The subtlest failure, because it + produces no error. FlashInfer's sampling kernels bake the philox RNG state into + the graph as a by-value constant, so a captured sampler replays identical + "random" numbers every step. Its kernel choice is also data-dependent. For both + reasons FlashInfer sampling is deliberately left eager — see + [host-path.md](host-path.md). Before widening a graph over anything holding RNG + or counter state, check what gets frozen. + +### The toolkit + +**One contiguous buffer, one memcpy.** +[ContextGPUView](megatron/core/inference/contexts/gpu_view.py) is the only +interface GPU code uses to read context state. Every field is a `view(dtype)` +onto a slice of a single `uint8` buffer, mirroring a pinned CPU buffer with the +identical layout, so publishing a step's bookkeeping is one `cudaMemcpyAsync` +rather than one per field. The convention is worth internalizing: + +``` +context.foo -> CPU (source of truth, used by bookkeeping) +context.gpu_view.foo -> GPU (snapshot, used by forward pass) +``` + +**Cache views instead of reconstructing them.** Slicing and unsqueezing per step +constructs new `TensorImpl`s at 30-60us each: + +```python +# megatron/core/inference/contexts/dynamic_context.py +# Instead of slicing and unsqueezing on every new inference step (constructing +# new TensorImpls at 30-60 us), we fix the underlying storage so views are +# reusable across steps. +self._input_position_views: Dict[int, Tuple[Tensor, Tensor]] = {} +``` + +**Point padding at reserved storage.** Pad tokens get `dummy_block_idx` so the +KV-append kernel writes to a valid throwaway block. Pad rows get routing index +`-1` so they activate no expert. + +**Keep EP sync off the compute stream.** `adjust_batch_dims_for_expert_parallelism` +can do the cross-rank max over ZMQ on the CPU, avoiding both a per-step NCCL +all-reduce on the compute stream and the H2D/D2H pair around it. + +## Expert parallelism: idle ranks still have to run + +Under EP, every rank must issue the same collectives in lockstep, and all ranks +must select the *same* graph bucket. So a rank with no real requests still runs a +forward. Commit `60a25aa6` found that idle rank was building its dummy batch +through the heavyweight graph-capture warmup path, constructing real request +objects and allocating KV blocks every step. + +The fix, `add_dummy_requests_for_expert_parallel_step`, pokes only the +preallocated tensors that `initialize_attention_state` and the forward actually +read: + +```python +N = smallest_cuda_graph_dimensions.decode_req_count +dummy_block_idx = self.block_allocator.dummy_block_idx +self.total_request_count = N +self.active_token_count = N +self.num_prefill_requests = 0 +self.request_query_lengths[0:N].fill_(1) +self.request_kv_length_offsets[0:N].fill_(0) +self.request_to_kv_block_ids[0:N, 0] = dummy_block_idx +self.token_to_block_idx[0:N] = dummy_block_idx +``` + +Hybrid models also need Mamba slots, so `MambaMetadata.batch_allocate_slots(n)` +grabs them in a batch without allocating. + +**Generalizable:** for any lockstep collective, idle participants must still run, +but give them the cheapest valid inputs. Reserve the dummy block and slots once, +then only `fill_` preallocated tensors. Never let the idle path allocate or +construct objects. + +The EP token-count sync all-reduces the max so every rank agrees on the bucket, +and returns `None` (eager everywhere) if any rank is non-decode. Commit `c817dad2` +cut this back to the minimum after it had accumulated dead complexity: + +```python +(max_token_count, max_is_non_decode) = ep_zmq_communicator.sync_all_reduce_max( + local_batch_dims.token_count, int(is_non_decode) +) +``` + +## Expect golden values to change + +Changing bucketization changes which padded shape each real step maps onto, which +changes padding tokens and MoE routing masks, which shifts logits. Commit +`60a25aa6` rewrote a 1448-line golden values file for exactly this reason. That is +expected — regenerate, and sanity-check that the *generated text* is still +coherent rather than assuming a diff means a bug. + +## Validation gates + +Under `block` scope, fp8 requires `--transformer-impl=inference_optimized` and +`--fp8-recipe=mxfp8`. `cuda_graph_impl=local` requires +`inference_dynamic_batching_num_cuda_graphs` to be positive or `-1`. Both are +asserted in [arguments.py](megatron/training/arguments.py). diff --git a/skills/optimize-inference-siddharth/references/decision-gates.md b/skills/optimize-inference-siddharth/references/decision-gates.md new file mode 100644 index 00000000000..c16f6a6052d --- /dev/null +++ b/skills/optimize-inference-siddharth/references/decision-gates.md @@ -0,0 +1,237 @@ +# Decision Gates: Prove the Ceiling Before You Build + +Profiling tells you where time is spent. It does not tell you whether that time is +**recoverable**. Those are different questions, and the gap between them is where +whole sessions get lost — writing a kernel whose best possible outcome was never +worth the week. + +A decision gate is a short, measurement-backed answer to one question: *if this +optimization worked perfectly, how much would the step actually get faster?* If +the answer is small, or if the mechanism you plan to attack is not the mechanism +that is slow, you stop before writing production code. + +Source: the Qwen3-30B-A3B EP4 campaign ledger +(`skills/run-qwen-model/EXPERIMENTS.md`), where three gates each killed a +multi-session effort, and a fourth picked the right candidate out of six. + +## Why share is not headroom + +The instinct is to rank categories by their share of device time and attack the +top one. That instinct is wrong twice over on this stack. + +**Some categories are already at a hardware floor.** A category can be 33% of +device time and have 2% of headroom, because it is moving the bytes it has to +move. Ranking by share puts it first; ranking by headroom puts it last. + +**Device time is not wall time.** Kernels overlap across streams, so the step's +wall clock is the *critical path* through the per-layer dependency chain, not the +sum of kernel durations. A category with a large device-time share that sits on a +side stream, already overlapped with compute, contributes far less to wall time +than its share suggests. Conversely, a tiny kernel on the serial chain costs its +duration *plus* the dispatch gap behind it, every layer. + +So the gate has to produce a **wall-time ceiling**, not a device-time share. + +## The gate, in four steps + +### 1. Establish the floor + +Compute what the operation could not go faster than, from first principles, using +*measured* machine constants rather than datasheet peaks. + +For a memory-bound op, the floor is bytes / achievable bandwidth. Measure the +achievable bandwidth yourself with a streaming-read microbenchmark on the same +device — datasheet HBM numbers overstate it, and being wrong here invalidates the +whole gate. + +For a latency-bound op, decompose the measured time into launch, synchronization, +and transfer, and identify which term dominates. An empty-kernel launch and a +barrier-only kernel are the two calibration points you need. + +### 2. Measure the current cost the same way + +Under **CUDA-graph replay**, not eager, and at the token count decode actually +runs at. Eager per-kernel timings and replay timings differ by more than the +effect sizes you are chasing. + +### 3. Ceiling = current / floor, then subtract what the fix costs + +The gross ceiling is the ratio. The net ceiling subtracts the machinery the fix +introduces: an added zeroing launch, a grid sync, extra atomics, a second +streaming pass. Fusions frequently give back most of the gross win this way, and +occasionally more than all of it. + +Then convert to a fraction of the step, because that is the number that decides +whether to proceed. + +### 4. Write down a verdict with a threshold + +State the gross ceiling, the net ceiling, the fraction of the step, and either +**proceed** or **gated out**, with the reason. Record it in the ledger even when +the answer is "don't build this" — especially then. A recorded negative gate is +what stops the next person re-deriving it. + +## Calibrate the fixed cost of a launch + +Several gates reduce to "is removing this launch worth it," so measure the two +constants once per platform and reuse them. + +On GB200 under full-iteration graph capture, the campaign measured: + +| Constant | Value | How | +|---|---|---| +| Empty-kernel launch | 0.72 µs | Time a no-op kernel under replay | +| Inter-kernel dispatch gap | 0.55 µs | Median gap between consecutive kernels in the trace | +| Floor per launch | 1.27 µs | Sum of the two | +| Graph node cost | ~0.17 µs | `cudaGraphLaunch` 199.1 µs / 1158 nodes | + +The consequence: **a removed launch is worth its own duration plus the dispatch +gap behind it**, so eliminating a 1.5 µs kernel buys ~2.05 µs of wall time, not +1.5. And no kernel can be optimized below 1.27 µs — a 1.33 µs kernel has no +headroom left, however inefficient its body looks. + +## Diagnose the mechanism, not just the magnitude + +A gate that identifies the wrong *cause* is as expensive as no gate. The routing +token-count kernel measured 7.52 µs per layer — clearly anomalous for its byte +count — and the first attempt rewrote its reduction, replacing per-pair +`atomic_add` with a `tl.histogram` variant. It measured **0.96×, a wash**, and +the reason was that atomic contention was never the problem: at `BLOCK_SIZE=1024` +only **2 of the kernel's 152 CTAs received any work**. The cost was launch +overhead and CTA underutilization, which an in-kernel reduction rewrite cannot +touch. + +Before attacking a slow kernel, establish *why* it is slow: bytes, FLOPs, +occupancy, CTA utilization, a dependent load serializing a branch, or pure launch +overhead. The fix follows from the mechanism, and the wrong mechanism produces a +technically correct kernel that changes nothing. + +## Case study 1: the grouped GEMM that could not win + +**The proposal.** MoE expert GEMM was the largest single category at ~33-40% of +decode device time. The obvious plan was a hand-written CUTLASS/CuTe grouped GEMM +to replace the Triton path — a multi-week effort. + +**The gate.** Weight traffic is 302 MB per layer per rank. A measured streaming-read +ceiling of 6.081 TB/s gives a **49.66 µs floor**. Production FC1+FC2 measured +**72.13 µs**. So the entire ceiling for *any* implementation is **1.45×** — and +achieved throughput on valid FLOPs was 63-83 TFLOP/s, about 3% of the device's +BF16 peak, confirming the op is nowhere near compute-bound. + +Padding waste was real but nearly free: 74.3% dead rows at `BLOCK_M=64`, yet +cutting the padding ratio from 3.89× to 1.49× buys only ~1.2×, because the dead +rows re-read weights that are already resident. + +**The verdict.** A hand-written GEMM's whole ceiling was 1.45×, and **1.26× of it +was reachable by retuning Triton tiles alone**. The kernel effort was gated out; +the tile retune shipped instead and delivered **+4.32% end-to-end** for a +day's work. A later profile confirmed the gate independently: expert GEMM ended up +only 229 µs/step — 2.3% of the step — above the weight-bandwidth floor. + +**Generalizable:** roofline the category *before* proposing an implementation +swap. A memory-bound op with a 1.45× ceiling does not deserve a new kernel; it +deserves better tiles. + +## Case study 2: the barrier that was a hardware floor + +**The proposal.** Exposed NVLS EP communication was ~11.7% of device time and +100% exposed (the interval union equalled the sum). Plans on the table: pipeline +the collective in chunks, or fuse ReduceScatter into the FC2 epilogue. + +**The gate.** Decomposing the collectives showed they are **latency-bound, not +bandwidth-bound**: + +``` +AllGather-V 6.57 µs = 0.72 launch + 5.08 barrier + 0.77 transfer +ReduceScatter 7.88 µs = 0.72 launch + 5.03 barrier + 2.13 transfer +``` + +Bytes accounted for only 127 µs of a ~693 µs/step total, and the +ReduceScatter transfer was already at 82% of the NVLink floor — so batching bytes +cannot help. A follow-up microbenchmark isolated the barrier further: a +barrier-only kernel costs 5.75-5.83 µs against a 0.72 µs empty kernel, so the +5.05 µs **is the four-way system-scope flag round trip itself**, not polling +granularity. Rewriting the spin (`atom.cas` → `ld.acquire.sys` poll) measured a +**2.3% regression**, exactly as the gate predicted. + +Of 1060 µs/step, 632 µs was intrinsic and 428 µs was **inter-rank skew** — ranks +waiting on the slowest rank's expert GEMM, which is a routing-balance problem, not +a communication problem. + +**The verdict.** Recoverable critical path: **6.4-7.0% of the step**, and each +candidate gave most of it back. Epilogue fusion keeps the barrier (≈1.4% left); +2-chunk pipelining *adds* a barrier per chunk and needs concurrent streams under +graph capture, which full-iteration capture forbids. CTA count was already at its +optimum. Lever abandoned. + +**Generalizable:** decompose a collective into launch + barrier + transfer before +optimizing it. If the barrier dominates, you are looking at a fabric round trip +and no kernel rewrite will move it. Check whether the residual is really *skew*, +in which case the fix is load balance somewhere else entirely. + +## Case study 3: pricing six fusions before writing one + +**The proposal.** The routing/permute category was 1251 µs/step across 242 +kernels — visibly a launch storm, with several plausible fusions. + +**The gate.** Break the category down per kernel name — launches, device µs, +dispatch gap — then price every candidate against the 1.27 µs per-launch floor: + +| Kernel | µs | × per step | +|---|---:|---:| +| `_moe_sum` | 7.79 | 48 | +| `_count_local_tokens` | 7.52 | 48 | +| `gatherTopK` | 6.05 | 48 | +| `_scatter_token_indices` | 2.66 | 48 | +| router softmax | 1.94 | 48 | +| `_prefix_fill_init` | 1.33 | 48 | + +Candidate ceilings, gross and net: + +| Candidate | Gross | Net after costs | Verdict | +|---|---:|---|---| +| `_moe_sum` → FC2 epilogue | 400 µs (4.03%) | ~3.3%, possibly negative — needs cross-CTA fp32 atomics **plus** a zeroing launch | gated out | +| Cooperative-grid merge of count→fill→scatter | 122 µs (1.23%) | less, after two grid syncs | gated out | +| Fold count + its `zeros` fill into `_prefix_fill_init` | 449 µs (4.52%) | no grid sync, integer-exact | **chosen** | + +**The verdict.** The chosen fusion shipped at **+3.01% end-to-end**. The +`_moe_sum` epilogue candidate — superficially the most attractive, since +`_moe_sum` was the single most expensive routing kernel — had the exact failure +shape of an earlier rejected mega-fusion: cross-CTA atomics plus a zeroing pass +eating the gain. It was never built. + +Notably, `_moe_sum` was later improved anyway, but by a different mechanism the +gate exposed: its per-topk-slot locality test was a uniform scalar branch on a +*dependent* load, serializing the topk walk. Predicating it into the load mask +gave 1.36× and **+1.83%** end-to-end, bit-exact, with no fusion at all. + +**Generalizable:** when a category has several fusion candidates, price all of +them before building any. Rank by *net* ceiling, and prefer candidates that need +no grid sync and no atomics — those two costs are what turn a 4% gross win into a +0% net win. + +## Verdict template + +Record this in the ledger for every gate, including the negative ones: + +``` +GATE — +Question: +Floor: +Current: +Gross ceiling: +Fix costs: +Net ceiling: <% of step after costs> +Mechanism: +Verdict: PROCEED | GATED OUT — +``` + +## When to skip the gate + +Gates cost time too. Skip straight to measurement when the change is cheap and +reversible: flipping an existing flag, retuning tile sizes, swapping a backend +enum. Those are their own experiments — a gate on a config flip costs more than +the flip. + +Gate anything that would take more than about a day to build, anything requiring +a new kernel, and anything whose category you have not yet rooflined. diff --git a/skills/optimize-inference-siddharth/references/host-path.md b/skills/optimize-inference-siddharth/references/host-path.md new file mode 100644 index 00000000000..14b27be5da8 --- /dev/null +++ b/skills/optimize-inference-siddharth/references/host-path.md @@ -0,0 +1,225 @@ +# The Host Path: Per-Step CPU Overhead and IPC + +Once the GPU forward is graphed and fused, the bottleneck moves to the CPU. Every +engine step ends with bookkeeping, detokenization, serialization, and a ZMQ send, +all of which block the next step. This is where several of the largest wins came +from, and it is the part people forget to profile. + +Source commits: `53a2b19a` (#2920), `602fad03` (#5918), `bcf4c8fb` (#5607), +`edd45620` (#5609), `650b7838` (#5791, co-authored). + +## Where per-step host work lives + +Always inspect these four places. + +**The engine step tail** in +[dynamic_engine.py](megatron/core/inference/engines/dynamic_engine.py). The +critical section is bracketed by NVTX ranges so it is directly visible in a trace: +`bookkeeping` → `detokenization` → `coordinator_communication`. Anything added here +runs once per step. + +**The controller**, `text_generation_controller.py`: routing-record bookkeeping, +log-prob detokenization, and building the step-result dict. CPU work here overlaps +already-enqueued GPU kernels, which makes it the right place to *put* unavoidable +work. + +**Serialization** in +[inference_request.py](megatron/core/inference/inference_request.py): +`InferenceRequest.serialize`, `DynamicInferenceRequest.serialize`, +`DynamicInferenceRequestRecord.serialize` and `.merge()`. Runs once per finished +request per step. + +**The coordinator** in +[data_parallel_inference_coordinator/](megatron/core/inference/data_parallel_inference_coordinator/): +per-request routing scoring in `coordinator.py`, pending-count updates in +`handlers.py`. + +## Serialization: the two rules + +### Never `dataclasses.asdict()` on anything holding tensors + +`asdict()` recursively deepcopies every field. With CUDA tensors in the object +that is catastrophic. Use a shallow dict copy and handle fields explicitly: + +```python +# megatron/core/inference/inference_request.py +# Dataclass to dict. +# do not use asdict(self) - it has very high CPU overheads +# and if there are tensors, it will try to deepcopy them +obj = self.__dict__.copy() # shallow dict copy +obj["status"] = self.status.name if self.status else None +obj["sampling_params"] = self.sampling_params.serialize() if self.sampling_params else None +``` + +### Never `torch.save` for IPC + +`torch.save` into a `BytesIO` pickles and writes the full tensor blob. For token +ids, a list is dramatically cheaper: + +```python +def serialize_tensor(tensor: torch.Tensor) -> List: + nvtx_range_push("serialize_tensor") + # simply convert tensor into a list + tensor = tensor.cpu().tolist() + nvtx_range_pop("serialize_tensor") + return tensor +``` + +Note the NVTX range on a function this small. That is deliberate — you cannot +attribute host cost you cannot see. + +## Don't put large tensors on the wire by default + +The engine was serializing and shipping the entire `prompt_tokens` tensor +engine → coordinator → API for every finished request. For long agentic or RL +prompts that dominates wire cost, and the client usually does not want the ids +echoed back. + +Commit `602fad03` made it opt-in: + +```python +# megatron/core/inference/sampling_params.py +# Echo prompt token ids back in the response. When False (default), the engine +# drops prompt_tokens before serializing the finished request, saving the ZMQ +# transmission cost for long prompts. Opt in when the client needs them. +return_prompt_tokens: bool = False +``` + +The API contract still needs `usage.prompt_tokens`, so the technique is to **keep +the scalar and drop the payload**: a `prompt_length: Optional[int]` field is always +populated during serialize even when the tensor is omitted. The drop is wire-only +— `serialize` nulls `self.prompt_tokens` around the `super().serialize()` call and +restores it, so the local object is unchanged. + +Endpoints opt in where they must: `/v1/completions` always echoes +`prompt_token_ids`, so it sets `return_prompt_tokens=True` unconditionally. + +**Generalizable:** for any per-request field on the IPC path, ask what the consumer +actually needs. Usually it is a count or a length, not the tensor. + +## Move CPU work off the engine, don't just make it faster + +The engine was detokenizing the full prompt plus generated sequence for every +finished request, inline in the step loop. Making detokenization faster would have +helped a little; moving it to a different process helped much more, because it then +overlaps the next engine step: + +```python +# Detokenize all finished requests if not using +# the coordinator. Otherwise, the coordinator will +# overlap detokenization with the engine. +if not self.use_coordinator: + nvtx_range_push("detokenization") + ... +``` + +The coordinator gained a tokenizer and a `detokenize()` method; the engine ships +raw token ids. When a coordinator is present the engine does no detokenization at +all. + +**Generalizable:** work that is not needed to compute the *next* step does not +belong in the step loop. Either defer it or move it to a process that runs +concurrently. + +## Load-aware DP routing + +With data-parallel engines, the coordinator decides which rank gets each request. +Round-robin ignores both prefix-cache affinity and current load, so it can send a +request to a busy rank while an idle rank holds exactly the KV blocks that request +needs. + +`bcf4c8fb` replaced round-robin with `LOAD_BALANCED`, which is now the default +policy, and unified everything under one score: + +```python +# megatron/core/inference/data_parallel_inference_coordinator/coordinator.py +if self.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.LOAD_BALANCED: + return self.get_least_loaded_data_parallel_rank() + +# Without prefix caching (or when the request has no hashes to match on) +# fall back to load-balanced routing. +if not self.enable_prefix_caching or not request_hashes: + return self.get_least_loaded_data_parallel_rank() + +match, recency = self._match_vector(request_hashes) +alpha = self.prefix_caching_routing_alpha + +# Vectorized score: alpha * match + (1-alpha) * free_capacity_fraction. +free_slots = np.maximum(0, self.max_requests - self._pending_counts).astype(np.float64) +scores = alpha * match + (1.0 - alpha) * (free_slots / self.max_requests) + +# Tiebreak: highest score, then highest recency, then lowest rank index. +n_ranks = len(self._identities_list) +order = np.lexsort((np.arange(n_ranks), -recency, -scores)) +``` + +`prefix_caching_routing_alpha` (default 0.5) is the tradeoff knob: + +| alpha | Behavior | Risk | +|---|---|---| +| 0 | Pure load balance | Ignores cache; re-prefills work another rank already has | +| 0.5 | Balanced (default) | — | +| 1 | Pure cache affinity | Piles affine requests on one rank, starves the others | + +Two implementation details worth copying. `match` is policy-dependent and +normalized to `[0, 1]` — binary for `first_prefix_block`, normalized prefix depth +for `longest_prefix` — so it is commensurable with the load term. And the whole +computation is vectorized numpy, because it runs per request on the coordinator's +hot path; a Python loop over ranks here would be its own bottleneck. + +Pending counts are maintained incrementally in `handlers.py` (incremented on +dispatch, decremented on finish), and `_remove_engine` rebuilds them when an engine +disconnects mid-flight. + +## Hold the dtype contract at kernel boundaries + +FlashInfer's sampling kernels return sampled token ids as `int32`. The rest of the +pipeline — token buffers, `input_ids`, scatter into KV and embedding gather — is +`int64`. Every return in +[flashinfer_sampling.py](megatron/core/inference/sampling/flashinfer_sampling.py) +is cast with `.long()`. + +A mismatch inserts implicit conversion kernels on the per-step path. Normalize +dtype where the external kernel enters rather than letting it propagate: the cast +is free once, the mismatch costs every step. + +> Provenance note: Siddharth's `44334eda` ("Fix flashinfer sampling kernels to +> return int64 instead of int32") sits on the unmerged `fix_flashinfer_sampling` +> branch. The same casts reached `main` through #5791 (`650b7838`), which he +> co-authored. + +### Corollary: not everything belongs in a graph + +That same PR concluded FlashInfer sampling should run **eagerly**, and the reasons +generalize to any stateful or random kernel you are tempted to capture: + +``` +The sampler runs eagerly. Its kernel choice is data-dependent (it varies with +which filters the batch uses), so it cannot be captured in a CUDA graph; running +eagerly also lets the controller's seeded RNG generator advance its philox offset +normally between steps -- fresh randomness per step, reproducible from the seed. +(FlashInfer bakes the philox state into a graph as a by-value constant at capture, +so a captured sampler replays identical random numbers ...) +``` + +Two distinct hazards: **data-dependent kernel selection** cannot be captured at +all, and **state baked in by value at capture** produces a graph that replays +silently wrong — identical "random" numbers every step, with no error. Before +widening a graph over anything holding RNG or counter state, check what gets frozen. + +Note also how the dispatch flags are read: from the pinned CPU sampling metadata, +so evaluating them costs no GPU sync. Same rule as everywhere else in this skill. + +## Host-side rules, condensed + +1. No `dataclasses.asdict()` on objects holding tensors. +2. No `torch.save` / pickle for IPC; use lists or ndarrays. +3. Drop large tensors from the wire by default; ship the scalar the contract needs. +4. Work not needed for the next step goes off the step loop, ideally to a + concurrent process. +5. Route with load awareness, not blind round-robin; keep the scoring vectorized. +6. Normalize dtypes at external kernel boundaries. +7. NVTX-annotate anything you suspect, including small functions. Unmeasured host + cost is invisible cost. +8. Any per-step recording must use preallocated static buffers and respect + `using_cuda_graph_this_step()`. diff --git a/skills/optimize-inference-siddharth/references/mamba-and-triton.md b/skills/optimize-inference-siddharth/references/mamba-and-triton.md new file mode 100644 index 00000000000..d7f24f63015 --- /dev/null +++ b/skills/optimize-inference-siddharth/references/mamba-and-triton.md @@ -0,0 +1,321 @@ +# Mamba / SSM Inference and Triton Production Rules + +Hybrid Mamba models (Nemotron-H, Nano) push most inference time through hand- +written Triton kernels, so kernel hygiene matters more here than anywhere else in +the stack. The Triton rules in this document are not Mamba-specific — they apply +to any kernel on the inference path. + +Source commits: `ab2b33d5` (#4397), `f29c747f` (#5608), `9b4074b5` (#4764), +`411a5d8b` (#5863), `648bc011` (#5866). + +## The Triton specialization trap + +This is the highest-value rule in the skill relative to how easy it is to get +wrong. Commit `f29c747f` is four lines: + +```diff + def _tensor_get_slice_after_kernel( +- INPUT_BATCH_SIZE: tl.constexpr, +- OUTPUT_BATCH_SIZE: tl.constexpr, ++ INPUT_BATCH_SIZE, ++ OUTPUT_BATCH_SIZE, + ROW_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +``` + +A `tl.constexpr` parameter is **baked into the compiled kernel**, so every +distinct value produces a separate compilation. `INPUT_BATCH_SIZE` is the number +of active requests, which changes essentially every step. The kernel was +JIT-compiling on the hot path, on nearly every step. + +`ROW_SIZE` and `BLOCK_SIZE` correctly stay `constexpr` — they depend only on the +fixed state shape and are needed for `tl.arange` bounds and loop unrolling. + +Current state in +[tensor_ops.py](megatron/core/inference/contexts/attention_context/triton/tensor_ops.py): + +```python +@triton.jit +def _tensor_get_slice_after_kernel( + INPUT_TENSOR, + OUTPUT_TENSOR, + POS_ON_DEVICE, + INPUT_BATCH_SIZE, + OUTPUT_BATCH_SIZE, + ROW_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): +``` + +### The full decision rule + +| Value | How to pass it | +|---|---| +| Block/tile size, `tl.arange` bound, unroll count, `d_conv` | `tl.constexpr` | +| Per-step count, but you want one compilation | plain arg + `do_not_specialize` | +| Per-step count, and the kernel is CUDA-graphed | fixed-address GPU tensor, `tl.load` it | + +The subtle middle case: **even plain `int` arguments get specialized** by +`@triton.jit` on divisibility-by-16 and `== 1`. So demoting from `constexpr` is +not always sufficient. If the value varies per step, be explicit: + +```python +# megatron/core/ssm/ops/mamba_ssm.py +@triton.jit(do_not_specialize=["batch"]) +def _selective_scan_update_kernel( +``` + +The third case is strictly better when applicable: passing the value as a 0-d +tensor and loading it inside the kernel avoids specialization entirely *and* keeps +CUDA graphs valid, since the address is fixed while the value may change between +replays. + +## Never autotune on the hot path + +`@triton.autotune` times every config on first sight of each new key. In +production that is a multi-second stall and a source of nondeterminism. + +Two mitigations, both in use: + +**Keep config lists short.** Commit `9b4074b5` trimmed the varlen causal conv from +four configs to two, keeping one from each regime. The comment explains why both +regimes exist, which is the useful part: + +```python +# megatron/core/ssm/ops/causal_conv1d_varlen.py +# Two block-dim regimes: +# 1. vLLM-style: small BLOCK_T, large BLOCK_C, pipelined. Many small programs maximize +# GPU occupancy, BLOCK_C=256 fully fills HBM transactions, sequence-pure programs +# avoid in-block boundary branching. Usually wins at moderate-to-large conv_dim. +# 2. Large-block fallback: bigger tiles, fewer programs. Can win for small conv_dim +# where vLLM's regime over-parallelizes, or when launch overhead dominates. +``` + +**Route everything through `autotune_configs`.** In +[determinism.py](megatron/core/ssm/ops/determinism.py), this is the single funnel: +in deterministic mode it either enables cached autotuning +(`TRITON_CACHE_AUTOTUNING=1`, Triton >= 3.4.0) or picks the single cheapest config +by `block_product * stages`, tie-broken on warp count — so exactly one +compilation and no timing sweep. + +```python +def autotune_configs(configs): + if not configs or not use_deterministic_mode(): + return configs + if TRITON_HAS_CACHE_RESULTS and os.environ.get("TRITON_CACHE_AUTOTUNING") == "1": + return configs + ... + return [min(configs, key=_estimate_config_cost)] +``` + +Wrap any new autotuned kernel in `autotune_configs([...])`, not a bare +`configs=[...]`. + +## Fused gather-plus-scatter instead of materialized intermediates + +Prefix caching needs to extract Mamba states at block boundaries. The original +implementation was a three-step tensor dance: the chunk scan gathered requested +chunks internally, `mamba_mixer` copied the result into scratch, and conv windows +were extracted with a separate PyTorch gather plus `clamp_` plus `transpose` plus +`copy_`. That materialized an intermediate tensor and made two passes over HBM. + +Commit `648bc011` replaced it with two kernels in +[intermediate_extraction.py](megatron/core/ssm/ops/intermediate_extraction.py). +The module docstring states the design contract: + +``` +These replace the two-step ``states[indices]`` (dense gather) + ``.copy_()`` +(scratch write) pattern with a single kernel that: + +1. Reads a runtime ``real_count`` from a fixed-address GPU tensor. +2. For each slot ``i < real_count``, gathers the source row indexed by the + per-slot index/position and writes it directly into the destination scratch. +3. For each slot ``i >= real_count``, returns immediately (no work, no write). + +This is CUDA-graph safe: the launch grid is sized at capture time to the maximum +possible slot count, but per-program execution is data-conditional on the +runtime ``real_count``, so padded slots cost almost nothing. +``` + +The gating idiom, which generalizes to any graphed kernel with a variable count: + +```python +pid_slot = tl.program_id(0) +pid_col = tl.program_id(1) + +real_count = tl.load(real_count_ptr).to(tl.int32) +if pid_slot >= real_count: + return +``` + +The conv kernel adds two more techniques worth stealing: + +**Fold the transpose into the write address.** The old code did +`.transpose(1, 2).copy_()`; the kernel just indexes the destination differently: + +```python +for j in tl.static_range(D_CONV): + p_raw = abs_pos - D_CONV + j + p = tl.maximum(0, tl.minimum(p_raw, seq_len - 1)) + src = src_ptr + p.to(tl.int64) * src_stride_s + c_idxs.to(tl.int64) * src_stride_c + dst = out_ptr + slot_base + c_idxs.to(tl.int64) * D_CONV + j + val = tl.load(src, mask=c_mask) + tl.store(dst, val, mask=c_mask) +``` + +**`D_CONV` is `constexpr` so the window loop is `tl.static_range`** — fully +unrolled, no dynamic loop. This is the legitimate use of `constexpr`: `d_conv` is +a model constant. + +The scan correspondingly got simpler. `ssd_combined.py` no longer gathers +internally; `return_raw_states` hands the caller the full state tensor and the +caller extracts what it needs. + +### Publishing the count + +```python +# megatron/core/inference/contexts/attention_context/mamba_metadata.py +# Publish real_count to the fixed-address GPU tensor the scatter +# kernels consult. fill_ is async (no host sync) and keeps the tensor +# at the same address captured graphs reference. +self._intermediate_real_count_buffer.fill_(self.intermediate_count) +self.intermediate_real_count = self._intermediate_real_count_buffer +``` + +## Right-size scratch to the real per-step bound + +The extraction scratch was sized `MAX_INTERMEDIATE_OFFSETS_PER_REQUEST (=3) * +max_requests`. But a step processes at most `max_tokens` tokens, and a state can +only be extracted at a block boundary — one per `block_size_tokens`. So there are +two independent bounds and the truth is the tighter one: + +```python +# megatron/core/inference/contexts/dynamic_context.py +# Per-step upper bound on Mamba intermediate-state extractions, shared with +# MambaMetadata and MambaSlotAllocator so scratch/metadata buffers and the +# budget accounting agree. Bounded both by the token budget (one block +# boundary per block_size_tokens) and by the request budget +# (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST per request); +token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) +request_based_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * self.max_requests +self.max_mamba_intermediate_states_per_step = min(token_based_count, request_based_count) +``` + +For low-concurrency long-context configs this is an order of magnitude: 1 request +with 16384 tokens once reserved 65 scratch slots a single request could never +fill. Since scratch is reserved *before* the durable prefix cache, every wasted +slot directly shrinks the cache. + +Two structural points that make this pattern work: + +**One value, shared everywhere.** `max_mamba_intermediate_states_per_step` flows +into `MambaMetadata(max_intermediate_count=...)`, +`MambaSlotAllocator.max_intermediate_count`, and the byte budget. When a bound is +duplicated it drifts; when it is shared, a wrong bound fails loudly. + +**Fail at config time, not at OOM.** If reserving scratch leaves fewer than one +durable slot, the context raises `ValueError` naming the budget and advising which +knob to reduce. + +The same commit fixed a hardcoded `mamba_chunk_size = 128`, now read from config. +States can only be extracted at multiples of the chunk size the kernel actually +runs with, so a hardcoded value silently skipped valid boundaries for any other +chunk size. + +## Bound per-step work by the bucket, not the global max + +Commit `9b4074b5` is essentially this rule applied to Nemotron prefill. The +extraction bookkeeping looped over the *global* `max_intermediate_count`, so every +prefill step did work proportional to the largest possible batch even in a tiny +graph bucket: + +```diff +- max_count = self.max_intermediate_count ++ max_count = padded_prefill_count * MAX_INTERMEDIATE_OFFSETS_PER_REQUEST +... +- self._intermediate_chunk_indices_buffer[real_count:].fill_(0) ++ self._intermediate_chunk_indices_buffer[real_count:max_count].fill_(0) +``` + +Same for the scratch writes: `intermediate_ssm_out.copy_(...)` became +`intermediate_ssm_out[:n].copy_(...)`. Whenever you see a `fill_` or `copy_` over +a whole preallocated buffer on the per-step path, ask what the real bound is. + +## Inner-loop micro-optimizations + +From `ab2b33d5`, in the single-token decode SSM update: + +**Use `exp2` for `exp`.** Hardware has a native `exp2`: + +```python +@triton.jit +def fast_exp(x): + """ + Fast calculation of exponent via exponent of 2. + """ + LOG2E = tl.constexpr(1.4426950408889634) + return tl.math.exp2(LOG2E * x) +``` + +**`torch.empty` when the kernel overwrites every row.** Zero-init is wasted HBM +bandwidth. Applied to the squared-ReLU output in +[activations.py](megatron/core/inference/moe/activations.py). + +**Specialize the launch shape on compute capability.** The tile and warp choice +was derived from `dstate` alone; Blackwell wants a different point: + +```python +# megatron/core/ssm/ops/mamba_ssm.py +is_blackwell = torch.cuda.get_device_capability(x.device)[0] >= 10 +... +else: + # dstate > 64 + if is_blackwell: + # Optimized for B200 with dstate>64 + BLOCK_SIZE_M, num_warps = 32, 8 + elif dstate <= 128: + BLOCK_SIZE_M, num_warps = 4, 4 +``` + +## File map + +| File | Role | +|---|---| +| [mamba_mixer.py](megatron/core/ssm/mamba_mixer.py) | The mixer. Prefill: conv, chunk scan, state update, extraction. Decode: `selective_state_update` + `causal_conv1d_update`. | +| [ssd_combined.py](megatron/core/ssm/ops/ssd_combined.py) | `mamba_chunk_scan_combined_varlen`, the packed varlen chunk scan. `return_raw_states` exposes intermediates. | +| [mamba_ssm.py](megatron/core/ssm/ops/mamba_ssm.py) | `selective_state_update`, the decode-path kernel. | +| [intermediate_extraction.py](megatron/core/ssm/ops/intermediate_extraction.py) | Fused conditional gather+scatter for prefix-cache extraction. | +| [causal_conv1d_varlen.py](megatron/core/ssm/ops/causal_conv1d_varlen.py) | Varlen depthwise causal conv with fused SiLU (prefill). | +| [determinism.py](megatron/core/ssm/ops/determinism.py) | `autotune_configs` — the single autotuning funnel. | +| [mamba_metadata.py](megatron/core/inference/contexts/attention_context/mamba_metadata.py) | Per-step Mamba metadata in fixed-address buffers. | +| [mamba_slot_allocator.py](megatron/core/inference/contexts/mamba_slot_allocator.py) | Durable state cache plus extraction scratch. | + +## How these are tested + +The pattern is consistent and worth imitating: + +**Compare against an obvious reference.** `_ref_conv` is a plain nested-loop +PyTorch implementation; the kernel is checked against it with +`torch.testing.assert_close`. A fast kernel with no slow reference has no test. + +**Assert the gating with a sentinel.** Fill the output with `12345.0`, run with +`real_count < max_count` where the trailing indices are *valid but must not be +gathered*, then assert both that `out[:real_count]` matches the reference and that +`torch.all(out[real_count:] == sentinel)`. The second half proves padded slots +produced no HBM writes — the actual performance claim. + +**Re-derive formulas independently.** `test_max_intermediate_states_per_step_formula` +recomputes `min(token_based, request_based)` in the test, asserts it matches, and +asserts the same value is shared by the allocator, the metadata, and the buffer +shape. It covers both regimes explicitly, including one where the request bound is +strictly tighter. + +**Cover the boundary.** `test_scatter_conv_sub_dconv_clamp` uses `pos=2` with +`d_conv=4` so the window start goes negative, checking all three out-of-range +positions clamp to token 0. + +**Test the error path.** A too-small budget must raise `ValueError`, not silently +OOM later. + +Tests: `tests/unit_tests/ssm/ops/test_ssd_combined.py`, +`tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py`. diff --git a/skills/optimize-inference-siddharth/references/measuring.md b/skills/optimize-inference-siddharth/references/measuring.md new file mode 100644 index 00000000000..2221393fe6d --- /dev/null +++ b/skills/optimize-inference-siddharth/references/measuring.md @@ -0,0 +1,359 @@ +# Measuring Inference Performance + +Measure first, then change one thing, then measure again. Every commit in this +history started from a profile, and the bottleneck was frequently not where it +seemed. This document covers the observability that exists specifically so +inference changes can be attributed. + +Source commits: `eadbaa61` (#5611, profiler endpoints), `edd45620` (#5609, +cache-hit reporting), `faced5128` (#3034, routing instrumentation), `53a2b19a` +(#2920, NVTX ranges). + +## Bracketing an nsys capture with the profile endpoints + +Profiling a server is awkward: you want the steady state, not startup, warmup, or +CUDA-graph capture. The profile endpoints solve this by relaying +`cudaProfilerStart` / `cudaProfilerStop` to every engine, so an outer nsys process +records only the window you ask for. + +From +[profile.py](megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/profile.py): + +``` +POST /start_profile and /stop_profile relay a control signal through the +InferenceClient -> data-parallel coordinator -> every connected EP/DP engine, +which calls cudaProfilerStart()/cudaProfilerStop(). Pair with an outer +`nsys profile --capture-range=cudaProfilerApi` to bracket a capture window. +``` + +Workflow: + +1. Launch the server under nsys with the capture range armed: + +```bash +nsys profile --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o inference_profile \ + +``` + +2. Send warmup load until the engine reaches steady state, including enough steps + that all CUDA graphs are captured. +3. Start recording, drive the load you care about, stop recording: + +```bash +curl -X POST http://localhost:5000/start_profile +# ... drive the workload ... +curl -X POST http://localhost:5000/stop_profile +``` + +Both `/start_profile` and `/v1/start_profile` are routed, likewise for stop. The +signal reaches **every rank**, so a multi-GPU trace covers all of them. + +Then analyze the trace with the `nsight-system-analysis` skill. + +### Flags that deadlock finalization — read before adding host visibility + +On a MoE decode workload under full-iteration CUDA graphs, nsys **hangs in +finalization** and `QdstrmImporter` then rejects the `.qdstrm` if any of the +following is enabled. This was bisected one flag at a time over several sessions +on Qwen3-30B / 4×GB200, and it is independent of `--cuda-graph-trace` level: + +| Trigger | Symptom | +|---|---| +| `osrt` in `--trace` | Finalization deadlock; unusable qdstrm | +| `--sample=process-tree` (CPU / Python sampling) | Same | +| NVTX ranges active under graph capture | Same — capture inflates NVTX event volume enormously | + +The only set that reliably finalizes on this workload: + +```bash +nsys profile --trace=cuda,nvtx --sample=none \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o inference_profile +``` + +This is a real constraint on method, not a nuisance: **the documented way to +attribute host time is exactly the way that breaks.** Two consequences, both with +working substitutes below — attribute host-side idle from the CUDA-API rows in a +clean GPU-only trace (see *Idle is not one thing*), and name host phases with +`perf_counter` instead of NVTX (see *When NVTX is unavailable*). + +### Validate that your instrumentation is not the measurement + +`--cuda-graph-trace=node` resolves individual kernels inside a graph, which is +what makes per-kernel analysis possible, but it instruments every node — 1158 of +them on this model — so it is fair to suspect it inflates the step. + +Check rather than assume. Capturing the same steady-state window both ways gave a +step period of **8882.5 µs (node) against 8945.1 µs (graph)** — 0.71% apart and in +the *wrong direction* for an instrumentation artifact — and host `cudaGraphLaunch` +of 190.0 against 184.7 µs. Node mode also reported 1169 kernels/step, matching the +independent count. Only after that control did a 900 µs/step graph-machinery +attribution become safe to act on. + +Run this control once per workload before trusting node-mode wall-time +attributions. Use `--cuda-graph-trace=graph` when you only need step periods, since +it produces far fewer events. + +## NVTX ranges label the host critical path + +Device timelines show kernels; they do not show why the host fell behind. The +per-step host critical section is annotated, so the CPU side of the trace is +readable without guessing: + +| Range | What it covers | +|---|---| +| `bookkeeping` | Per-step request state updates | +| `detokenization` | Only when no coordinator is present | +| `coordinator_communication` | msgpack serialize + ZMQ send | +| `serialize_tensor` | Individual tensor to list conversion | +| `drain_zmq_socket`, `add_request` | Inbound request handling | + +If you add host work to the step loop, add an NVTX range around it. `nvtx_range_push` +/ `nvtx_range_pop` come from `megatron/core/utils.py`. The precedent is that even +`serialize_tensor` — a two-line function — is annotated, because unmeasured host +cost is invisible cost. + +Two gotchas before you rely on these ranges. The engine's own helpers are **inert +unless NVTX profiling is enabled**, which only the training loop does by default — +a whole 187.9 µs/step host phase (`post_process_requests` via `async_bookkeep`) was +invisible in one host trace for exactly this reason, and had to be found by +elimination instead. And enabling NVTX under graph capture is one of the +finalization-deadlock triggers above, so on this workload the ranges you most want +are the ones you cannot record. + +### When NVTX is unavailable: `perf_counter` phase timing + +Since the range names are the useful part and nsys is the broken part, keep the +names and drop nsys. Monkeypatch the same push/pop symbols with +`time.perf_counter_ns` self-timers over a shared nesting stack — each label then +reports self-time excluding its children — and print a per-step breakdown every N +steps. No profiler, no CUDA-graph interaction, and it names precisely the +between-step Python phases a GPU-only trace can only localize. + +Gate it behind an env var and make `install()` a no-op by default, the same way +every other diagnostic here ships. + +## Built-in counters, no profiler needed + +Two things are exposed through the API and are often enough to confirm or reject a +hypothesis in seconds. + +### Prefix cache hit rate + +`DynamicInferenceRequest.num_cached_tokens` accumulates +`num_matched_blocks * block_size_tokens` across prefill chunks, and surfaces in the +chat-completions response as `usage.prompt_tokens_details.cached_tokens` +(OpenAI-compatible). + +Use it to check that a prefix-caching or routing change actually improved cache +hits, rather than inferring it from end-to-end latency. Particularly relevant when +tuning `prefix_caching_routing_alpha` — see +[host-path.md](host-path.md). + +### MoE expert load imbalance + +Set `moe_enable_routing_replay` (requires `num_moe_experts`). The context then +records, per request, which experts every token was routed to, in a static +CUDA-graph-safe buffer of shape +`[max_tokens, num_moe_layers, moe_router_topk]`, exposed as +`DynamicInferenceRequest.routing_indices` with shape +`[total_tokens, num_layers, topk]`. + +Aggregate a histogram across many requests. A few hot experts with the rest idle +means wasted expert capacity and straggler experts gating every step — a +throughput problem no kernel tuning will fix. + +Implementation: [routing_metadata.py](megatron/core/inference/contexts/routing_metadata.py), +[router_replay.py](megatron/core/transformer/moe/router_replay.py), collected in +`TextGenerationController._router_record_bookkeeping()`. Note that the recording +itself follows the rules in this skill: a preallocated static buffer, toggled by +`using_cuda_graph_this_step()`, sliced to `active_token_count` so padding is +excluded. + +## Benchmark harnesses + +| Tool | Use | +|---|---| +| [tools/run_inference_performance_test.py](tools/run_inference_performance_test.py) | Standalone driver: builds the model and engine, sweeps requests, times steps. Fastest local loop. | +| `tests/performance_tests/shell_test_utils/run_perf_test.sh` | The harness CI perf recipes invoke. | +| `tests/performance_tests/shell_test_utils/compare_to_baseline.py` | Compares a run against checked-in `baseline_values.json`. | +| `tests/test_utils/recipes/h100/*inference*.yaml` | Registered perf and functional recipes. | +| [examples/inference/](examples/inference/) | Server launch and offline inference examples. | + +For cluster runs use the `run-inference-performance-tests` skill (throughput and +latency versus baseline) and `run-inference-functional-tests` (correctness, +including the CUDA-graph paths). Neither is duplicated here. + +## Interpreting the profile + +### Pick a clean window, and measure busy as a union + +Two arithmetic mistakes will make every category number wrong. + +**Do not take a fraction of the capture span.** A span fraction mixes prefill, +graph capture, and ramp into the "steady state." Instead search for the **densest +window** of the length you want — kernel starts per unit time — and verify the +kernel sequence repeats with a stable period inside it. A 1-2 s window is enough. + +**Do not sum per-stream kernel time.** Kernels overlap across streams, so summing +double-counts and can exceed the wall clock. Compute the **union of kernel +intervals** across all streams on the device; that is busy time, and +`window − union` is real idle. + +On the twelve-gate Qwen config a clean 2 s window gave union-busy **81.6%** and +idle **18.4%**, with device-time by category of expert GEMM 587 ms, NVLS comm +247 ms, dense GEMM 236 ms, attention 211 ms, and then routing 69, norm 67, +elementwise 66, `moe_sum` 65, rope/KV 15. Those small buckets being ~65-70 ms each +*and partly overlapped on other streams* is what established that further +small-kernel fusion was worth only ~0.5-1% apiece — a conclusion that a +per-stream sum would have hidden. + +### Idle is not one thing + +Total idle is not an actionable number, because a large fraction of it is not +attackable. Split it by gap size first: + +| Gap size | Share (Qwen, 178 ms idle in 1 s) | What it is | +|---|---:|---| +| < 10 µs | 59.6 ms (33%) | Intra-graph kernel scheduling — unavoidable | +| ≥ 10 µs | 118.8 ms (67%) | Host chain between steps — attackable, median 37 µs | + +Then attribute the large gaps **without a host trace**, using the CUDA-API +(`CUPTI_ACTIVITY_KIND_RUNTIME`) rows that a clean `cuda,nvtx` capture already +contains. For each gap, ask which host API calls fall inside it: + +| Inside the gap | Share of large-gap idle | +|---|---:| +| Nothing — no CUDA call at all (Python/CPU compute) | **73.9%** | +| Token D2H memcpy | 8.9% | +| Graph launch | 7.5% | +| Kernel launch | 6.9% | +| Sync wait | 1.6% | + +So ~74% of attackable idle — about 8-9% of wall time — was **Python on the +critical path**, and only 1.6% was waiting on the GPU. Localize it further by +bracketing each uncovered region with the CUDA APIs immediately before and after +it: 47.6% sat between `cudaMemcpyAsync` and `cuKernelGetName`, another 15.7% +between `cudaLaunchKernel` and `cudaMemcpyAsync`. That pins the cost to +**between-step host orchestration after the sampled-token D2H copy** — post-sampling +bookkeeping, scheduling, attention-state prep — which is where the host-path wins +in [host-path.md](host-path.md) came from. + +This RUNTIME-row method is the recommended one, not a fallback. It works on the +only capture configuration that finalizes. + +### Map the signal to a section + +Map the dominant signal to the section of this skill that addresses it: + +| Signal | Reading | Go to | +|---|---|---| +| GPU idle gaps, host running ahead | Launch-bound: too many launches, or no matching graph bucket | [cuda-graphs.md](cuda-graphs.md) | +| Long host span after the last kernel | Post-processing on the critical path | [host-path.md](host-path.md) | +| NCCL kernels not overlapped | Exposed collective, or EP ranks out of step | [moe-inference.md](moe-inference.md) | +| Multi-ms stalls, or slow early steps | Triton recompilation or autotune | [mamba-and-triton.md](mamba-and-triton.md) | +| Grouped GEMM wide but shallow | Tiles tuned for the wrong batch size | [moe-inference.md](moe-inference.md) | +| Many sub-microsecond kernels in a row | Fusion candidate | [moe-inference.md](moe-inference.md) | + +## Measurement discipline + +- **A/B against a matched config.** Same batch size, sequence length, and + parallelism as the vLLM or baseline number you are comparing to. A latency + comparison across different `max_requests` values means nothing. +- **Exclude capture from the measurement.** Graph capture runs a full forward per + bucket. Warm up until capture is done before you start recording. +- **Change one thing.** Most of these commits are single-mechanism precisely so the + measurement attributes cleanly. +- **Keep the kill switch and measure both ways.** `inference_disable_triton_nvls_kernels`, + `inference_moe_disable_fused_quant_kernels`, and the backend enums exist so a + regression can be bisected by config rather than by revert. +- **Re-measure per model.** The shared-expert CTA cap was a win on one model and a + loss on another (see [moe-inference.md](moe-inference.md)). Overlap and tile + tuning do not transfer. + +## The noise floor is bigger than your win + +Once the easy multi-percent wins are gone, the remaining levers are worth well +under 1% each, and at that point the above discipline is necessary but **not +sufficient**. Across the Qwen campaign, throughput for an *identical* config +measured in different allocations drifted by **−1.61%, −0.57%, and +1.22%** on +three separate occasions — larger than most of the individual wins that followed +(+0.94%, +0.90%, +0.98%, ~+0.5%). Comparing against last session's number will +produce confident, wrong conclusions in both directions. + +Four rules make a sub-1% claim falsifiable. + +**1. Re-baseline in the same session.** Before each lever, run the current best +config in the *same allocation* you will test in, and use that as the reference. +Record it as its own ledger entry. This is the single highest-value habit here. + +**2. Run the arms back to back.** Same allocation, same warm process where +possible, alternating OFF and ON. An OFF arm collected an hour earlier is a +different measurement. + +**3. Judge by distribution separation, not mean delta.** With N timed iterations +per arm, the acceptance criterion is that the **arms do not overlap** — the +slowest ON iteration beats the fastest OFF iteration. Examples that met it: + +``` +QWEN-026: min ON 27,102.9 > max OFF 26,425.7 (fully separated) +QWEN-025: 13 of 15 ON iterations beat all 15 OFF (near-separated) +QWEN-023: slowest ON 25,992.1 > fastest OFF 25,893.9 (fully separated) +``` + +A +1% mean delta with overlapping arms is not a result. A +0.9% delta with +separated arms is, even though it is smaller than the session drift — because +drift is between sessions and separation is measured within one. + +**4. Repeat the pair, and report pairwise deltas.** Two or three OFF/ON pairs, each +delta stated separately (`+0.99% / +1.26% / +0.68%`), tells the reader the spread. +A single averaged number hides it. + +Also discard the first timed iteration or check it explicitly — cold-start outliers +of ~1.3% appeared repeatedly while iterations 2-5 spanned 0.23%. + +## Predict how much of a kernel win will convert + +The most common surprise in this campaign was the ratio between a kernel-level +speedup and the end-to-end result. It went both ways, by large factors, and it is +predictable enough to estimate before you build. + +| Where the work sits | Conversion | Evidence | +|---|---|---| +| Serial per-layer dependency chain | **more than 1:1** | Fused QK-norm: ~1% microbench ceiling → **+2.9% e2e**, ~3×. Removing a launch also removes a graph node and the host dispatch gap behind it, ×48 layers. | +| Overlapped with other work | **well under 1:1** | A 183.9 µs/step host saving converted at only **1/3 to 1/2**, because that phase was already partly overlapped with GPU work. | +| Off the critical path | **~0** | A 1.25× speedup of the MoE activation path delivered **+0.55%, then +0.13%** — a wash. The op was never the wall. | + +The rule: **estimate against the critical path, never against device-time share.** +Before building, establish whether the target is (a) on the serial chain, (b) +already overlapped, or (c) off it — the same determination the +[decision-gates.md](decision-gates.md) ceiling calculation needs. + +The corollary is that a small kernel on the serial chain can be worth more than a +large kernel off it, which inverts the usual ranking. Two ~1.5-4 µs norm fusions +delivered +2.9% and +1.37%, while the largest device-time category in the profile +had 2.3% of headroom left in total. + +## Keep an append-only ledger + +For anything longer than a single session, the ledger *is* the method. Sessions get +preempted, context is lost, and without a durable record the same dead end gets +re-explored — a routing-kernel rewrite that measured a wash was nearly repeated for +exactly this reason. + +Requirements that earned their place: + +- **A fixed protocol table at the top** — cluster, hardware, model, batch, output + length, parallelism, warmup/timed counts, correctness gate, primary metric — plus + the explicit rule that results are not comparable when any of them differ. +- **One row per experiment, including every rejection**, with its root cause. The + negative results are the higher-value half: they are what stops the next person, + and three of them killed multi-week efforts here. +- **Append-only. Never edit a recorded result.** Supersede it with a new entry. +- **A running distance-to-target number** so priorities stay honest. +- **A ranked next-levers list**, re-derived after each profile rather than carried + forward. + +Record the baselines in a fixed order — competitor first, then yours, then the +gap — and only then change code. diff --git a/skills/optimize-inference-siddharth/references/moe-inference.md b/skills/optimize-inference-siddharth/references/moe-inference.md new file mode 100644 index 00000000000..dfdfba041ed --- /dev/null +++ b/skills/optimize-inference-siddharth/references/moe-inference.md @@ -0,0 +1,382 @@ +# MoE Inference: the `inference_optimized` Stack + +The training MoE path is unusable for low-latency decode. Its AlltoAll dispatcher +does host synchronizations (`.item()`, `.tolist()`, `.cpu()`) that make CUDA-graph +capture impossible; the router computes z-loss, aux-loss, token-drop, and +expert-bias updates that are pure training overhead; and TE GroupedGEMM needs a +host-resident `tokens_per_expert`. + +Rather than branch the training code, Siddharth forked a parallel inference +module hierarchy. Understanding that structure is most of what you need. + +Source commits: `7d1c0168` (#3496, foundational), `905c0e38` (#3851), +`589cd9e1` (#3858), `bfd45740` (#4258, architectural rewrite), `442a936a` (#4570), +`20f09364` (#4603), `c817dad2` (#4587), `3a253ac5` (#4922). + +## How the backend is selected + +```mermaid +flowchart TD + A["--transformer-impl inference_optimized"] --> B["gpt_layer_specs / mamba_layer_specs
take the inference branch"] + B --> C["get_inference_optimized_moe_spec()"] + C --> D["MoELayer with MoESubmodules"] + D --> E["InferenceTopKRouter"] + D --> F["InferenceGroupedMLP
(via InferenceSpecProvider)"] + D --> G["SharedExpertMLP"] + H["MoELayer.__init__ -> _setup_inference_mode"] --> I["instantiate dispatcher from
inference_moe_token_dispatcher_type"] + J["MoELayer.train(mode)"] --> K["eval -> inference dispatcher
train -> standard dispatcher"] +``` + +The load-bearing detail: **`MoELayer.train()` performs the swap.** Eval mode means +the inference stack is live. There is no per-iteration flag to set — the original +`set_inference_cuda_graphed_iteration` mechanism from #3496 was removed in #4258 +precisely because manual plumbing was error-prone. + +Relevant files: [backends.py](megatron/core/models/backends.py) +(`InferenceSpecProvider.grouped_mlp_modules`), +[moe_module_specs.py](megatron/core/models/gpt/moe_module_specs.py), +[moe_layer.py](megatron/core/transformer/moe/moe_layer.py). + +## What each replaced component does differently + +### `InferenceTopKRouter` + +In [router.py](megatron/core/transformer/moe/router.py). Strips z-loss, +aux/load-balance loss, token drop, and expert-bias updates. It is +`@torch.compile`d, and it returns **dense** `[num_tokens, topk]` tensors rather +than sparse `[num_tokens, num_experts]`, which is what the fused GEMM kernels +want. + +`moe_router_dtype` must be `fp32` — enforced as a hard config error specifically +to avoid a dtype conversion on every decode step. + +### `InferenceGroupedMLP` + +In [experts.py](megatron/core/transformer/moe/experts.py). Three things matter: + +**Stacked contiguous weights.** TE stores per-expert `weight0..weightN`; grouped +GEMM wants one `[num_experts, out, in]` tensor. Built lazily on first forward and +shared by pointer, never by replacing the `Parameter` — see hard rule 7 in +[SKILL.md](../SKILL.md). + +**GPU-resident expert offsets.** Grouped GEMM needs to know where each expert's +token block starts and ends. The training path gets that from a host-resident +`tokens_per_expert`; the inference path computes it entirely on device. Commit +#3496 used `tokens_per_expert.cumsum(0)`; it is now a Triton prefix-sum kernel in +[permute.py](megatron/core/inference/moe/permute.py) that also folds in the +per-expert alignment rounding: + +```python +h = tl.load(tokens_per_expert_ptr + r, mask=mask, other=0) +# Round up non-zero counts to alignment boundary +if alignment > 1: + h = tl.where(h > 0, ((h + alignment - 1) // alignment) * alignment, h) +inc = tl.cumsum(h, axis=0) +tl.store(exclusive_offsets_ptr + r, inc - h, mask=mask) +tl.store(inclusive_offsets_ptr + r, inc, mask=mask) +``` + +The inclusive offsets are passed straight through as `offs` to `grouped_mm` / +`scaled_grouped_mm`, so no expert boundary ever touches the host. + +**Backend dispatch** on `inference_grouped_gemm_backend`: `vllm` (default), +`flashinfer`, or `torch`. MXFP8 requires `torch`, because +`scaled_grouped_mm` is the only path supporting blockwise 1x32 scaling. Note the +alignment constraint that falls out of combining requirements: + +```python +# scaled_grouped_mm requires each expert's token count aligned to 32, +# but swizzled MXFP8 scales require alignment to 128. Use 128 to +# satisfy both constraints. +expert_alignment = 128 +``` + +For BF16 decode on GB200, all three backends were measured against each other and +**`vllm` won**, so the default is also the fast path — but the two rejections are +informative rather than obvious: + +- **`torch`** (`grouped_mm`, cuBLAS) measured **0.82×**. Reach for it only when + MXFP8 forces it. +- **`flashinfer`** (`cutlass_fused_moe`) measured **90.81 µs against 83.63 µs** for + the retuned Triton path — 8% slower. Getting to that number first required fixing + two independent wiring bugs: an `ActivationType` mis-map that hard-fails, and a + gate/up ordering mismatch that **silently corrupts numerics**. Note the shape of + that second one — a backend can be wired wrongly and still run. +- **`trtllm_bf16_routed_moe`**, the TRT-LLM-Gen kernel vLLM actually wins with, + cannot currently be dropped in at all: `weight_layout=MajorK` is rejected + (`BF16 Moe: weight_layout must be BlockMajorK`) and `BlockMajorK` indexes + `size(3)`, i.e. it requires **4-D pre-shuffled block-major** weights against + mcore's 3-D `[E, 2*ffn, H]`. The blocker is the weight-preparation pipeline, not + the kernel or the EP contract. + +## Dispatchers: NCCL AllGather vs NVLS AllGather-V + +Both live in +[token_dispatcher_inference.py](megatron/core/transformer/moe/token_dispatcher_inference.py) +under `InferenceAllGatherDispatcherBase`. Selected by +`inference_moe_token_dispatcher_type`. + +The progression is worth understanding because it explains why the default is +what it is. AlltoAll was replaced by AllGather/ReduceScatter for graph +friendliness. But a fixed AllGather requires every EP rank to contribute the same +token count, so ranks pad to the global max — wasted compute and communication — +and the original version only worked with decode-only graphs, forcing prefill to +run eager. AllGather-V removes the equal-count requirement. + +| | `nccl` | `nvls` (default) | +|---|---|---| +| Token counts across EP | Must be equal on the graph path | Variable per rank | +| Non-decode CUDA graphs | Disabled | Supported | +| Padding | Pads to max, then compacts | None | +| Requires | Nothing special | Hopper+, NVLink, bf16, 16-byte-aligned | + +Choose `nccl` only when NVLS is unavailable. It is the correctness fallback, not a +tuning option. + +### NVLS eligibility + +Checked centrally so the conditions live in one place: + +```python +# megatron/core/inference/communication/torch_symm_triton/utils.py +def are_tensors_nvls_eligible(*tensors: torch.Tensor) -> bool: + """ + Requirements: + - Hopper+ GPU (SM >= 9) + - All tensor byte sizes are divisible by 16 (128-bit), since NVLS + kernels process data in 128-bit chunks. + """ + if not tensors: + return False + return is_device_nvls_capable(tensors[0].device) and all( + t.element_size() * t.numel() % 16 == 0 for t in tensors + ) +``` + +If you add a collective, call this rather than re-deriving +`dtype == bf16 and device.major >= 9`. Centralizing it was part of #3496. + +### Why symmetric memory and multimem win + +`multimem.st` / `.ld` / `.red` operate on a symmetric-memory buffer mapped across +NVLink. For the small, latency-bound messages of decode: + +- A multicast store *is* the all-gather — one kernel, one end barrier, versus + NCCL's multi-step algorithm. +- Multiple tensors fuse into one kernel and one barrier. The dispatch of + `routing_map`, `probs`, and `hidden` went from 3 collectives to 1. +- Reduce-scatter accumulates in fp32 then casts, so results are *exactly* + reproducible — the unit tests assert `atol=0, rtol=0`. +- Buffers are fixed-address, hence graph-safe. +- CTA count is tunable, so the collective can be throttled to share the GPU with + overlapped compute. + +Buffers come from `SymmetricMemoryManager` in +[symmetric_memory.py](megatron/core/inference/symmetric_memory.py), a lazy +registry keyed by string (`"tp"`, `"ep"`). Lazy specifically so there is no +init-ordering coupling with the context — the eager version broke the RL +integration. + +## Fusing the metadata update + +Per-step metadata (`valid_tokens`, `rank_token_offset`, `ep_max_tokens`) used to +take five kernels plus a NCCL all-gather. `fused_metadata_update` in +[metadata.py](megatron/core/inference/moe/metadata.py) does it in one Triton +kernel that multicast-stores the local count, barriers, then computes all three +in place. From its docstring: + +``` +Replaces the multi-kernel sequence: + dist.all_gather_into_tensor(...) # NCCL + local_tokens_per_rank.sum() # kernel + local_tokens_per_rank[:rank].sum() # kernel + local_tokens_per_rank.max() # kernel + _step_metadata.copy_(...) # kernel +with a single Triton kernel ... +``` + +**Generalizable:** a chain of tiny reductions over a handful of values is all +launch overhead. If you see several sub-microsecond kernels in sequence, they are +one kernel. + +## Buffers are allocated once, at model init + +`allocate_buffers()` is a classmethod called at model init from the dynamic +context, never inside a captured graph. Everything downstream operates on +max-sized buffers and gates work by a single fixed-address scalar +`_valid_tokens_tensor`. Expert output is written *directly* into the +reduce-scatter symmetric buffer via `out=`, avoiding a copy before the collective. + +## Padding must not reach an expert + +CUDA-graph replay pads the token dimension to the captured size. Those rows carry +garbage routing indices, and before `3a253ac5` they were dispatched to real +experts — wasted GEMM work and potentially corrupted reductions. + +The fix publishes a fixed-address GPU scalar `real_token_count` each step, and a +Triton kernel writes `-1` into every topk slot of the padded rows: + +```python +# Mask out CUDA-graph padding rows of the local routing map so the AGV +# propagates -1 into agv_r for those slots; padding tokens then route +# to no expert. +if self.__class__._real_token_count_tensor is not None: + mask_routing_padding( + self.routing_map, self.__class__._real_token_count_tensor, self.sp_rank + ) +``` + +Two subtleties: the count is in the *global* frame, so `sp_rank` shifts local rows +for the comparison; and on capture or dummy steps `real_token_count` is set to 0 +so *all* rows are masked. + +The complementary rule is in the activation and reduction kernels — skip rows +where `permutation_map == -1`, and do not *zero* rows past `valid_tokens`, since +downstream only reads the valid prefix. Commit `20f09364` removed exactly such a +zeroing pass from `moe_sum`. + +## Tune for the typical batch, not the buffer + +The vLLM-derived grouped GEMM originally used `@triton.autotune` over 25 configs. +That meant long compile times and, worse, tile choices made for the wrong batch +size. It was replaced with vLLM's host-side heuristic: + +```python +# megatron/core/inference/moe/vllm_fused_moe.py +def _get_default_config(M: int, E: int, top_k: int) -> dict: + """ + M here is the host-side token-count hint (``num_tokens_hint`` in + ``vllm_fused_moe``), NOT ``hidden_states.size(0)``. The hint is the + expected per-step token count; the worst-case buffer size would over-tune + for prefill on every decode step. + """ + # BLOCK_SIZE_M: shrink at small M to limit per-expert padding waste. + if M <= 32: + block_m = 16 + elif M <= 96: + block_m = 32 + elif M <= 512: + block_m = 64 + else: + block_m = 128 +``` + +The grid is sized from the same hint; on rare prefill spikes each CTA strides over +extra tiles via `tl.range` — correct, with reduced parallelism, which is the right +trade for a rare case. + +**Generalizable:** buffers are sized for the worst case, but kernels should be +*tuned* for the common case. Pass an explicit `num_tokens_hint` rather than +inferring from `tensor.size(0)`. + +### One shared config cannot serve both expert GEMMs + +vLLM's heuristic is the right starting point, not the end of the tuning. It returns +**one** config, but FC1 and FC2 have different output widths — on Qwen3-30B-A3B, +N=768 for FC1 against N=2048 for FC2 — and they want *opposite* `BLOCK_SIZE_N`. A +single shared config cannot express that, so one of the two GEMMs is always +mistiled. + +Splitting into per-GEMM decode-tuned configs (`config_fc1` / `config_fc2`) took +expert-GEMM device time from 77.65 to 61.89 µs (**1.255×**), the whole MoE call from +100.64 to 84.31 µs, and delivered **+4.32% end-to-end** — the largest single win of +that campaign, from tuning alone, after a roofline gate had shown a hand-written +CUTLASS kernel could not beat it (see +[decision-gates.md](decision-gates.md)). + +Two details worth copying: + +**It can be bit-exact.** Changing `BLOCK_SIZE_M` / `BLOCK_SIZE_N` while leaving +`BLOCK_SIZE_K` **unchanged** preserves the fp32 K-reduction order, so the output is +bit-identical (max abs diff 0.0). Retuning K would not be. This makes tile tuning +the cheapest possible optimization to validate — reach for it first. + +**Tuned tiles have a range, so gate them.** The tuned configs measured +1.289× / 1.197× / 1.120× / **0.952×** at 128 / 256 / 384 / 512 tokens — a 5% +*regression* at the top of the intended range. Fall back to the default heuristic +above the token count you actually tuned for (384 here, not 512). Always measure the +upper edge of a tuned range; a win at the decode point can hide a loss just past it. + +### Two built-in fusion flags do not work on this path + +`--moe-router-fusion` and `--moe-permute-fusion` look like free wins and are the +first thing anyone tries. Under `--transformer-impl inference_optimized` both crash: + +``` +AssertionError: hidden_size mismatch: 128 vs 8 +``` + +TE's fused router emits a **dense `num_experts`-wide routing map** (128), while +`InferenceTopKRouter` produces the **dense top-k** contract (8) that the vLLM/NVLS +dispatcher consumes. The built-in fusions are wired to the training MoE path only. A +hand-written fused softmax+top-k that honors the top-k contract *is* worth building +— doing so collapsed four kernels into one, 16.41 → 4.10 µs under replay, for +**+3.89% end-to-end**, with bit-exact probabilities and identical expert sets. + +## Overlapping the shared expert + +The shared expert (a dense FFN over every token) is independent of the routed +path, so it can run concurrently with dispatch, expert GEMMs, and combine. Launch +it on a side stream in `dispatch_preprocess` and join in `combine_postprocess`: + +```python +if self.shared_experts is not None and not self._external_shared_expert_launch: + stream = SharedExpertMLP.stream + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + self._shared_expert_output = apply_module(self.shared_experts)(hidden_states) +``` + +Only enabled for `nvls` plus `use_shared_expert` plus `moe_shared_expert_overlap`. + +**Latent MoEs** (SuperV3, UltraV3) run experts in a projected latent dimension, +but the shared expert needs full hidden dim, so it sits outside the dispatcher's +view. There, `MoELayer.preprocess`/`postprocess` own the launch and the add, and +the dispatcher is told `_external_shared_expert_launch = True`. + +### A cautionary tale on overlap tuning + +Commit `442a936a` capped the AllGather-V to 16 CTAs when overlapping, so the +collective would not starve the shared-expert GEMMs on the side stream. Commit +`9b4074b` **removed that cap** — on Nemotron it measured net worse. + +Overlap tuning is model- and shape-specific. Re-measure per model; do not copy a +CTA cap forward on the assumption it generalizes. + +## Reduce how often you synchronize + +`ep_consensus_interval` (default 20) skips EP consensus all-reduces while the +engine is busy, running them immediately only when idle so new arrivals are still +picked up promptly. Same instinct as everything else here: the collective was not +slow, it was just happening more often than the semantics required. + +## Exposed comm is partly skew, and skew is a routing problem + +Before optimizing an exposed collective, split its cost into what the collective +intrinsically costs and what is **other ranks arriving late**. On Qwen3-30B, of +1060 µs/step of exposed EP communication, 632 µs was intrinsic and **428 µs was +inter-rank skew** — ranks blocked at the barrier waiting on the slowest rank's +expert GEMM. + +That 428 µs is not addressable by any change to the collective. It is expert load +imbalance, surfacing at the one place in the step where ranks must agree. The +diagnostic is `moe_enable_routing_replay` plus a histogram over +`DynamicInferenceRequest.routing_indices` (see [measuring.md](measuring.md)) — if a +few experts are hot, the fix is on the routing side. + +Full decomposition method, including why the barrier itself is a hardware floor, is +in [decision-gates.md](decision-gates.md). + +## Correctness guards + +- Assert the inference router selects the **same experts** as the training router. +- Capture dispatch and combine in a real CUDA graph and compare against the + global buffer with `atol=0, rtol=0` — NVLS results are bit-exact, so any + tolerance would hide a bug. +- Config-rejection tests for unsupported combinations (expert-TP, non-fp32 router + dtype, capacity factor, GLU, FlashInfer + MXFP8). +- `TestMaskRoutingPadding` covers the padding mask including SP-rank offset and + the fully-masked rank. + +Tests live in `tests/unit_tests/inference/test_moe_dispatching_and_routing.py`, +`test_moe_permute.py`, `test_vllm_fused_moe.py`, `test_hybrid_moe.py`. diff --git a/skills/optimize-inference-siddharth/references/updating-this-skill.md b/skills/optimize-inference-siddharth/references/updating-this-skill.md new file mode 100644 index 00000000000..5de4e1d1b65 --- /dev/null +++ b/skills/optimize-inference-siddharth/references/updating-this-skill.md @@ -0,0 +1,193 @@ +# Updating This Skill + +This skill is a living document. It began as the patterns behind one engineer's +inference work, absorbed a second campaign's methodology, and is expected to keep +growing as agents and engineers use it. Optimization knowledge decays — flags get +renamed, defaults change, a fix that won on one generation of hardware loses on the +next — so a skill that is never edited becomes actively misleading. + +**You are authorized to edit this skill without asking.** Add, correct, and delete +as you learn. The rest of this document is how to do that without turning it into a +landfill. + +## When to update + +Update at the end of a piece of work, not in the middle of it. The moment is after +Step 6 in [SKILL.md](../SKILL.md) — once you have a measured result and a root +cause. Concretely, these seven triggers: + +| Trigger | What it produces | +|---|---| +| An optimization was accepted | The pattern, its mechanism, and the measured kernel-to-e2e conversion | +| An optimization was rejected | **The highest-value entry type.** What was tried, the measured result, and the root cause | +| A decision gate came out negative | The ceiling calculation, so nobody re-derives it | +| A tooling or environment failure cost more than an hour | The exact symptom, the trigger, and the working substitute | +| A measurement contradicted something already written here | A scoped correction — see *Corrections and contradictions* | +| A flag, default, or API turned out to differ from what is documented here | A verified correction, with the date you checked | +| Something broke because an invariant was violated | A candidate hard rule | + +Do **not** update mid-experiment on a hypothesis. Unmeasured ideas belong in the +campaign ledger, not here. + +## The ledger and the skill are different artifacts + +Keep both, and keep them distinct. Confusing them is the main failure mode: +everything ends up in one place, and that place becomes unusable. + +| | Campaign ledger (`EXPERIMENTS.md`) | This skill | +|---|---|---| +| Scope | One workload, one campaign | Every workload, indefinitely | +| Policy | Append-only, keeps everything | Curated, edited and pruned | +| Contains | Every run, every number, every config | Only what transfers | +| Audience | You, next session | Anyone, next year | + +**The promotion test:** would this change what a competent engineer does on a +*different* model or a different workload? If yes, promote it here. If it is a +number that only describes one model at one shape, it stays in the ledger. + +For the common in-between case, split it: **the mechanism goes in the skill, the +number goes in the ledger, and the skill cites the number as a calibrated example +with its hardware stamp.** That is what "expert GEMM was 1.45× off roofline on +GB200" is doing in [decision-gates.md](decision-gates.md) — the transferable claim +is *roofline the category before proposing a kernel*, and the number is evidence +that the claim has teeth. + +## Where it goes + +Route by subject, not by which file you happened to be reading. Adding everything +to `SKILL.md` is the most common way to wreck this skill — `SKILL.md` is a router, +and it stops working once it is long enough to need its own router. + +| Learning about | Destination | +|---|---| +| Launch overhead, graph scope, capture safety, bucketing | [cuda-graphs.md](cuda-graphs.md) | +| MoE dispatchers, experts, routing, grouped GEMM, EP comm | [moe-inference.md](moe-inference.md) | +| Per-step CPU work, serialization, IPC, coordinator, DP routing | [host-path.md](host-path.md) | +| Triton rules, kernel hygiene, Mamba/SSM, scratch sizing | [mamba-and-triton.md](mamba-and-triton.md) | +| Profiling tooling, trace analysis, A/B protocol, idle accounting | [measuring.md](measuring.md) | +| Prioritization, ceilings, whether a lever is worth building | [decision-gates.md](decision-gates.md) | +| Comparing against vLLM or another engine | [vllm-differential.md](vllm-differential.md) | +| A new invariant whose violation broke something | `SKILL.md` → Hard rules | +| Flag behavior, defaults, flags that look free and are not | `SKILL.md` → Flags | +| A check that belongs before every PR | [../assets/review-checklist.md](../assets/review-checklist.md) | +| A subject none of the above covers | A new `references/*.md`, linked from `SKILL.md` | + +When you add a new reference file, three things must follow or it is invisible: link +it from the `## References` list in `SKILL.md`, mention it from whichever step it +serves, and extend the frontmatter `description` if it represents a genuinely new +capability — that description is what causes this skill to be selected at all. + +## The bar for an addition + +All four, or it does not go in: + +1. **Measured.** A number, and enough method that someone could reproduce it. "Felt + faster" is not a finding. +2. **Root-caused.** The mechanism, not the symptom. A kernel that was slow because + 2 of 152 CTAs got work teaches something; "the count kernel was slow" teaches + nothing and invites the wrong fix. +3. **Scoped.** State the conditions under which it holds — hardware, model shape, + token count, config. An unscoped claim will be applied where it is false. +4. **Actionable.** It changes a decision. If a reader finishes the paragraph and + does nothing differently, cut it. + +Reject: restatements of content already here, cluster or queue incidents with no +methodological consequence, and single anomalous runs that were never repeated. + +## House style + +Match what is already here, so the file reads as one voice rather than a pile of +contributions: + +- Lead with the finding, then the number. Not the narrative of how you got there. +- Imperative and second person. "Roofline the category before proposing a kernel." +- Quote **exact error strings** for failure modes. `AssertionError: hidden_size + mismatch: 128 vs 8` is findable; "it crashes with a shape error" is not. +- Close a transferable pattern with a **`Generalizable:`** line. That convention is + load-bearing — it is how a reader tells a specific finding from a rule. +- Cite the source: a commit sha for merged work, a ledger entry id for campaign + work. Unattributed claims cannot be re-verified. +- Show real code with a file-path comment as the first line, as the existing + snippets do. +- Tables for enumerable facts, prose for reasoning. Never a table of paragraphs. + +## Corrections and contradictions + +**If something here is wrong, fix it in place.** If it was wrong in a way someone +may have already acted on, say so in one clause rather than silently rewriting. + +**If a new measurement contradicts an existing one, do not overwrite it.** Both are +facts; they differ in scope. Scope them: + +> Capping the AllGather-V to 16 CTAs when overlapping was a win on one model and +> measured net worse on Nemotron. + +That is the correct shape — it preserves both results and tells the reader the +outcome is model-dependent, which is itself the lesson. Overwriting would have +produced a confident claim that is wrong half the time. + +**Never delete a measured negative result** because it is old or inconvenient. Those +entries are the ones preventing repeated work; three of them here each killed a +multi-week effort. If new hardware changes the answer, add the condition, do not +replace the record. + +## What to delete + +Deletion is part of maintenance, not vandalism. Delete: + +- **Content that is now false** — a flag that no longer exists, a default that + changed, a code path that was removed. Verify against the tree before restating + any default, and note the date you checked. +- **Superseded APIs**, once nothing in-tree references them. Until then keep the + old→new mapping in a collapsed `
` block; the deprecated + `CudaGraphScope` block in [cuda-graphs.md](cuda-graphs.md) is the precedent. +- **Duplication.** The same lesson stated in three files means two of them are + drift waiting to happen. Keep the fullest treatment, link to it from the others. +- **Prose that changes no decision**, including your own earlier additions. Prefer + cutting recent additions over the original source-commit material, which carries + provenance you cannot reconstruct. + +When you delete something substantive, note it in the revision log so the removal +is discoverable rather than mysterious. + +## Size budget + +Progressive disclosure only works if the entry point stays navigable. + +| File | Soft target | Action when over | +|---|---|---| +| `SKILL.md` | ~500 lines (currently ~466) | Move detail into a reference; keep the routing line | +| Any reference | ~400 lines | Split by subject, link both from `SKILL.md` | + +These are soft targets, and the line count is a proxy, not the thing that matters. +The real test for `SKILL.md`: **can a reader scan it once and know which reference +answers their question?** If a section explains a mechanism rather than pointing at +one, the mechanism belongs in a reference regardless of the current line count. Do +not trim readable prose to hit a number — that trade makes the skill worse. + +For references the test is subject coherence. If one has grown to cover two +unrelated subjects, that is a split rather than a trim; `mamba-and-triton.md` +covering both Triton hygiene and SSM scratch sizing is already near that line. + +## After any edit + +``` +- [ ] Every intra-skill link still resolves +- [ ] Frontmatter still parses, and `name` still matches the directory +- [ ] `description` extended if a new capability area was added +- [ ] New reference linked from the References list AND from the step it serves +- [ ] Measured claims carry model + hardware + date, or a commit / ledger citation +- [ ] Revision log below appended +``` + +## Revision log + +One line per substantive change, newest last. This exists so a reader can tell what +is original source-commit material and what was added later, and so deletions are +discoverable rather than mysterious. + +| Date | Change | Source | +|---|---|---| +| 2026-07 | Initial skill: the five moves, hard rules 1-9, the CUDA-graph / MoE / Mamba-Triton / host-path / measuring references, commit log, review checklist | Siddharth Singh's 2026 inference work, 29 commits | +| 2026-07-28 | Added `decision-gates.md` (ceiling-before-building, per-launch fixed costs, three gates that each killed a multi-week effort) and `vllm-differential.md` (competitor-trace comparison). Hardened `measuring.md`: nsys flag combinations that deadlock finalization, the node-vs-graph trace control, `perf_counter` phase timing, union-busy idle accounting with gap-size decomposition, the same-session back-to-back A/B protocol, kernel-to-e2e conversion, ledger requirements. Added *What a wide capture costs you* to `cuda-graphs.md`. Added per-GEMM tile tuning, the measured backend comparison, the training-path-only fusion flags, and comm-vs-skew to `moe-inference.md`. Added hard rules 10 (A/B protocol) and 11 (non-bit-exact acceptance) plus the *flags that look like free wins* table | Qwen3-30B-A3B EP4 on 4×GB200, `skills/run-qwen-model/EXPERIMENTS.md` | +| 2026-07-28 | Made the skill self-maintaining: this file, the editing authorization at the top of `SKILL.md`, the Step 6 feedback loop, the checklist's skill-maintenance block, and this log | — | diff --git a/skills/optimize-inference-siddharth/references/vllm-differential.md b/skills/optimize-inference-siddharth/references/vllm-differential.md new file mode 100644 index 00000000000..ba623c0a5f4 --- /dev/null +++ b/skills/optimize-inference-siddharth/references/vllm-differential.md @@ -0,0 +1,140 @@ +# Differential Analysis Against vLLM + +When the goal is stated as "match vLLM," the competitor's trace is not just a +scoreboard — it is a specification. It tells you what the same model on the same +hardware costs when someone else built the decode path, which is a far more useful +target than a roofline, because it is known to be achievable. + +Everything else in this skill looks inward at Megatron. This document is the one +workflow that looks at both. + +Source: the Qwen3-30B-A3B EP4 campaign, where a single kernel-level differential +simultaneously justified the fusion work and killed the grouped-GEMM work. + +## The one question worth asking + +Given matched configs and matched trace windows: + +> **Is any individual kernel slower in mcore, or does mcore just launch more +> kernels to do the same thing?** + +These have opposite fixes, and guessing wrong is expensive. Slower kernels mean +tiling, occupancy, or implementation work. More kernels mean fusion, graph-node +removal, and host-dispatch reduction. A third answer — mcore has *categories the +competitor does not have at all* — means the structure differs and no amount of +kernel work will close it. + +On Qwen3-30B the answer was unambiguous, and it was the second one. One forward +block: **~467 kernels in vLLM against ~1784 in mcore**, roughly 3.8×, for a step +of **4.4 ms against 9.5 ms**. Individual kernels were broadly competitive; mcore +was running three to four times as many of them, most on the serial per-layer +dependency chain. + +## Getting comparable windows + +The comparison is worthless unless both windows cover the same work in the same +regime. Two requirements beyond the usual matched-config discipline: + +**Both must be steady-state decode.** No prefill, no graph capture, no warmup. +Verify by checking that the window's kernel sequence repeats with a stable period. + +**Anchor on a once-per-block kernel, not on wall time.** Pick a kernel that fires +exactly once per transformer block with a stable name — the block's leading norm +is the natural choice — and take the window from one occurrence to the next +occurrence of the *closing* kernel. That gives exactly one forward block on each +side, independent of clock offsets between the two runs. + +For Qwen the anchors were: + +| Engine | Window start | Window end | +|---|---|---| +| vLLM | `triton_red_fused__to_copy_embedding_rms_norm_0` | `triton_red_fused__to_copy_add_mean_moe_forward_mul_pow_rsqrt_0` | +| mcore | `rmsnorm_fwd_tuned_kernel<...>` | `triton_poi_fused_add_copy__0` | + +Extract the ordered kernel list per window from the `.sqlite` export — kernel +rows joined to the string table, filtered to one device and the window, ordered by +start time. Keep start, duration, and name; you need the order, not just the +totals. + +## Decompose into three buckets + +Assign every kernel in both windows to a *role* — norm, QKV GEMM, attention, +router, dispatch, expert GEMM, combine, residual, sampling — and then compare +role by role. Do **not** match on kernel names: Inductor-generated Triton names +differ between engines and encode fusion structure rather than function, which is +exactly the thing you are trying to measure. + +Then split the total gap three ways: + +**Same role, more kernels.** The fusion signal. vLLM runs one +`triton_red_fused__to_copy_add_..._rms_norm` where mcore runs +`triton_poi_fused_add_copy` followed by `rmsnorm_fwd_tuned` — two launches, two +graph nodes, per boundary, per layer. That observation, repeated across the norm +and residual roles, is what produced the fused QK-norm and fused add-norm +optimizations, worth **+2.9%** and **+1.37%** respectively. + +**Same role, slower kernel.** The implementation signal. Rank by +`(mcore_time - competitor_time)` for the role, not by mcore time alone. + +**Role present on one side only.** The structural signal, and usually the largest. +mcore's decode carried **exposed NVLS AllGather-V dispatch and ReduceScatter-V +combine on the per-layer critical path**; vLLM, using a TRT-LLM fused MoE, had +**zero exposed communication** — dispatch, grouped GEMM, activation, and finalize +were one fused unit. No fusion or tile change on mcore's side reaches that; it is +a different decomposition of the same math. + +## Read the competitor's fusion boundaries as the target + +The most directly actionable output is the competitor's *fusion inventory*: which +adjacent operations it has merged that you have not. For Qwen decode: + +| Operation | vLLM | mcore (before) | +|---|---|---| +| Residual add + RMSNorm | one kernel | two kernels | +| Router softmax + top-k | one kernel | four kernels | +| Routing metadata / indirection table | folded in | up to five kernels | +| Dispatch + expert GEMM + activation + finalize | one fused MoE | six-plus kernels with exposed comm between | + +Each row is a candidate with a known-achievable target. Price them with +[decision-gates.md](decision-gates.md) before building — the routing chain +collapsed from five kernels to one across four separate experiments, but the +full MoE mega-fusion was attempted, measured **0.68-0.80×**, and rejected. That +the competitor fuses something does not prove *your* fusion of it will be faster; +it proves the fusion is legal and bounds what it is worth. + +## Do not conclude the competitor's kernel is better without checking why + +Two traps, both hit during this campaign. + +**The competitor may win by weight layout, not by kernel quality.** The TRT-LLM +BF16 fused MoE that vLLM uses was probed directly as a drop-in backend and turned +out to require **4-D pre-shuffled block-major weights**; mcore stores 3-D +`[E, 2*ffn, H]`. `weight_layout=MajorK` is rejected outright and `BlockMajorK` +indexes `size(3)`. The kernel was never the blocker — the weight preparation +pipeline was. + +**An existing backend may be mis-wired rather than slow.** The already-shipped +flashinfer `cutlass_fused_moe` path was documented as "blocked for SwiGLU." It is +not: the kernel supports BF16 gated SwiGLU, and mcore's wiring was broken in two +independent ways — an `ActivationType` mis-map that hard-fails, and a gate/up +ordering mismatch that *silently corrupts numerics*. Both were root-caused and +fixed in a harness. Only then was the comparison meaningful, and the answer was +**90.81 µs against 83.63 µs** for the retuned Triton path, i.e. 8% slower. The +ledger note had been wrong for the wrong reason, and the correct measurement still +rejected the backend. + +**Generalizable:** before adopting a competitor's kernel, verify what it demands +of the weight layout and whether your existing wiring to it is actually correct. +A hard failure and a silent numerics corruption look nothing alike in a profile. + +## What a differential cannot tell you + +It gives you counts and durations, not the critical path. A role with more kernels +in mcore may still cost nothing if those kernels sit on a side stream and overlap +with compute. Convert the differential's findings to wall-time impact with +union-busy analysis and the serial-chain reasoning in +[measuring.md](measuring.md) before ranking them. + +It also will not show host-side differences. Two engines with identical kernel +sequences can differ substantially in Python overhead between steps, which on this +workload was 18.4% of wall time.