From 93691f456f61fb2b8a3253fface8f6bfa9ff0e04 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Thu, 3 Sep 2026 20:21:23 +0000 Subject: [PATCH 01/10] Add launcher-wrapping kernel shape profiler (registry approach) Restores the registry-based kernel_shape_profiler as a standalone, non-invasive tool: wraps kernel *launcher* functions as torch.library custom ops so they appear as cpu_op events with Input Dims / Input type, covering any backend behind the launcher (Triton, ASM, CK, aiter C++ bindings, collectives) rather than only JIT-dispatched kernels. - kernel_shape_profiler.py: explicit registry plus filtered auto-discovery; builds a schema from each launcher's signature, passes non-tensor args and the return value through a thread-local side channel, and rebinds every module-level reference so 'from X import Y' callers are intercepted too. Recovered from the retired SGLang patch and adapted to import standalone. - sitecustomize.py: auto-loaded via PYTHONPATH, drives enable()/disable() from the torch profiler window so nothing is wrapped outside a run. - README.md: approach, safety properties, and a comparison against the JIT-hook tracer on the jit-shape-tracer branch. Two fixes over the recovered original: - Registry refreshed for current layouts. Only 5 of the original 24 entries still resolved against sglang 0.5.18 + aiter (sglang moved Triton kernels to sglang.kernels.ops.*, aiter regrouped its Triton ops). Current paths added alongside the legacy ones, which are kept since unresolvable entries are skipped silently; 21 of 40 now resolve. - Guard against double wrapping. Frameworks keep compat re-exports, so after the first entry is wrapped _patch_all_references rebinds the second path to our own wrapper, which was then wrapped again and produced two nested annotations for one call. Wrappers now carry a _kernel_shape_wrapper marker that enable() skips. Validated on MI300X: synthetic launcher and the real aiter gemm_a8w8_blockscale both annotated with correct fp8/bf16 per-operand dtypes, exactly once per call, with references restored on disable(). --- .../kernel_shape_tool/README.md | 223 ++++ .../kernel_shape_profiler.py | 952 ++++++++++++++++++ .../kernel_shape_tool/sitecustomize.py | 282 ++++++ 3 files changed, 1457 insertions(+) create mode 100644 examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md create mode 100644 examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py create mode 100644 examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md b/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md new file mode 100644 index 000000000..429b51c5f --- /dev/null +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md @@ -0,0 +1,223 @@ +# Kernel shape profiler (launcher-wrapping approach) + +Adds `Input Dims` / `Input type` / `Input Strides` to PyTorch profiler traces +for GPU kernels that would otherwise appear with no operand metadata, so they +can be rooflined. + +It works by wrapping **kernel launcher functions** — the Python entry points +like `gemm_a8w8_blockscale`, `invoke_fused_moe_kernel` or `rmsnorm` — as +`torch.library` custom ops. Each wrapped launcher then shows up in the trace as +a `cpu_op` named after itself, carrying its tensor operands. + +> This is the registry-based approach. The alternative, on the +> `jit-shape-tracer` branch, hooks the Triton/FlyDSL JIT launch boundary +> instead. See [Comparison](#comparison-with-the-jit-hook-approach) for which +> one to use. + +## Files + +| File | Role | +|------|------| +| `kernel_shape_profiler.py` | The core. Wraps launcher functions as `sglang_profiler::__` custom ops via an explicit registry plus filtered auto-discovery, and rebinds every module-level reference to them. | +| `sitecustomize.py` | Auto-loaded shim. Drives `enable()` / `disable()` from the torch-profiler window so nothing is wrapped outside a profiling run. | + +## Activation + +No serving-framework source is patched. Put this directory on `PYTHONPATH` and +set the flag: + +```bash +export PYTHONPATH=/path/to/kernel_shape_tool:$PYTHONPATH +export TRACELENS_SHAPE_DISCOVERY=1 +``` + +CPython auto-imports `sitecustomize` at interpreter startup for every process, +so the server and all TP workers pick it up. When +`TRACELENS_SHAPE_DISCOVERY` is unset or `0`, every hook short-circuits, so it is +safe to leave the directory on `PYTHONPATH` permanently. + +### Optional knobs + +| Env var | Default | Meaning | +|---------|---------|---------| +| `TRACELENS_SHAPE_DISCOVERY` | `0` | Master switch. Must be truthy to do anything. | +| `TRACELENS_SHAPE_FORCE_RECORD_SHAPES` | `1` | Force `record_shapes=True` on the profiler (`Input Dims` only surface when shapes are recorded). Set `0` to respect the server's own setting. | + +## What lands in the trace + +For a launch of `gemm_a8w8_blockscale(x, w, x_scale, w_scale)`: + +``` +sglang_profiler::fp8_utils_gemm_a8w8_blockscale_12 <- cpu_op, named after the launcher + Input Dims: [[1025, 7168], [7168, 2112], [1025, 56], [56, 17]] + Input Strides: [[7168, 1], [1, 7168], [56, 1], [1, 56]] + Input type: ['c10::Float8_e4m3fnuz', 'c10::Float8_e4m3fnuz', 'float', 'float'] + _gemm_a8w8_blockscale_kernel_GROUP_K_128_... <- the real GPU kernel +``` + +The op is named `__`. The counter is +monotonic per process and never reset, so op names from earlier +`enable()`/`disable()` cycles stay valid. + +## How it works + +### 1. Choosing what to wrap + +Two sources are merged in `enable()`: + +**An explicit registry** (`_KERNEL_ENTRY_POINTS`) of `(module_path, function)` +pairs covering SGLang Triton attention, fused MoE, layernorm, FP8 quantization, +LoRA, aiter ops and FlashInfer MoE. + +**Filtered auto-discovery** (`_discover_kernel_entry_points`) which force-imports +everything under `sglang.srt.`, `aiter.ops.` and `flashinfer.`, then keeps only +functions that look like kernel launchers — either they have a `Tensor` +annotation, or (when unannotated) their source contains a launch pattern such as +`[grid`, `torch.ops.` or `sgl_kernel.` (`_is_likely_kernel_launcher`). + +Two things are deliberately excluded: +- `@triton.jit` objects (`JITFunction` / `Autotuner`). Replacing them in module + globals breaks Triton's global resolution for device-side calls between JIT + kernels. Only plain Python functions are wrapped. +- Test / benchmark / autotune modules, which are not entry points and some of + which call `torch.set_default_device("cuda")` at import time. + +### 2. Building the op + +`_build_schema_from_sig` derives a `torch.library` schema from the function +signature, mapping annotated parameters to schema types (`_infer_schema_type` +handles both real annotations and PEP 563 string annotations). Tensor +parameters become schema arguments; everything else is passed through a +thread-local side channel (`_stash_non_tensor_args`), because a schema of +`-> ()` cannot carry them. The return value comes back the same way +(`_stash_return_value`). + +If no schema can be built (no annotations, `*args`/`**kwargs`), the function +falls back to `_make_record_function_wrapper`, which emits a `record_function` +event with the shapes embedded in the *event name* instead of as structured +args. + +### 3. Applying the wrapper + +`_patch_all_references` scans all of `sys.modules` and replaces **every** +module-level attribute pointing at the original function. This is required +because of the `from X import Y` pattern: patching only the defining module +would miss callers that already captured a local binding. + +Launchers whose reference was captured as an *instance attribute* before +`enable()` (e.g. `self.fn = dispatch_w8a8_block_fp8_linear()`) cannot be +intercepted this way. For those, the registry targets the **inner** kernel that +the wrapper looks up from module globals on every call — which is why +`gemm_a8w8_blockscale` is registered rather than the outer dispatch wrapper. + +## Safety properties + +The implementation carries a few hard-won guards worth preserving: + +**The `Library` is never torn down.** Dropping it destroys the registered ops. +A wrapper reference can outlive `disable()` — a module imported lazily *during* +a profiling window may have captured the wrapper via `from X import Y` and is +not in `_patches` to be restored. If the backing op were freed, that leaked +wrapper would dispatch into freed memory and segfault. The `Library`, the +monotonic op counter and the wrapper cache are all kept alive for the process +lifetime, and the `if not _enabled` guard in each wrapper routes leaked calls +straight to the original function. + +**Global torch state is restored around imports.** `enable()` imports a large +number of modules, some of which mutate process-global torch state at import +time — most notably `torch.set_default_device("cuda")`. That mutation is not +undone by `disable()`. A leaked default device corrupts downstream CPU tensor +creation: a buffer such as `seq_lens_cpu` is suddenly allocated on CUDA, which +surfaces as `Buffer seq_lens_cpu has different device than before` and later as +out-of-bounds GPU memory faults in index kernels. +`_preserve_global_torch_state` and the per-import restore in +`_force_import_submodules` keep that from escaping into the serving path. + +**Every wrapper falls back to the original.** Unbindable signatures, all-`None` +tensor arguments, and dispatch failures all call the original function +directly, so wrapping can never change behaviour. + +**Wrappers are never wrapped again.** Frameworks often keep a compat +re-export, so the same function is reachable by two module paths. Once the +first is wrapped, `_patch_all_references` rebinds *every* reference to it — +including the second path — so the second registry entry would otherwise +resolve to our own wrapper and wrap it a second time, nesting two annotations +around a single call and double-counting it. Wrappers carry a +`_kernel_shape_wrapper` marker and `enable()` skips them. + +## Registry staleness + +The explicit registry is coupled to framework internals, and it drifts. Measured +against SGLang 0.5.18 + matching aiter, only **5 of the original 24 entries +still resolved**: SGLang had moved its Triton kernels out of +`sglang.srt.layers.*` into a separate `sglang.kernels.ops.*` package, and aiter +had regrouped its Triton ops into subpackages. + +Current-layout paths have been added alongside the legacy ones (**21 of 40** +resolve now; the 19 that do not are the legacy paths, kept deliberately). +`_resolve_target` returns `None` for a path that does not exist and the entry is +skipped silently, so one registry serves several framework versions. + +Two entries cannot be fixed by a path change: aiter's batched blockscale GEMM +was *renamed* (now +`batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant`), and +`flashinfer` is not installed in this image. Auto-discovery covers the former. + +The practical lesson: **auto-discovery is doing most of the work**, and the +registry is best treated as a hint list for launchers the heuristics miss. Check +what actually resolved before trusting the registry on a new stack. + +## Validation + +Verified with `validate_kernel_shape_profiler.py`: + +- A synthetic launcher is registered, produces a `cpu_op` with + `Input Dims: [[4, 8], [8, 16]]` and per-tensor dtypes, returns the correct + value, and passes its non-tensor argument through the side channel. +- Both the definition-site binding and a `from X import Y` binding are rebound, + and both are restored on `disable()`. +- A wrapper captured while enabled still returns correct results after + `disable()` (the leaked-reference case the persistent `Library` protects). +- The real aiter `gemm_a8w8_blockscale` is annotated with + `Input Dims: [[1025, 7168], [2112, 7168], [1025, 56], [17, 56], []]` and + `['c10::Float8_e4m3fnuz', 'c10::Float8_e4m3fnuz', 'float', 'float', '']`, + exactly once per call. +- `enable()` took **18–19 s** with auto-discovery scoped to `aiter.ops.` alone, + which is the cost of the package walk. Budget for more with `sglang.srt.` + included. + +### What the operands mean at launcher level + +Note the difference from a JIT-boundary tracer. Here the recorded operands are +the launcher's **call arguments**, so for the same GEMM you see `w` as +`[2112, 7168]` (before the launcher transposes it) and the optional `y` output +as an empty slot because it was not passed. A JIT-level tracer sees the +**kernel's** operands: `[7168, 2112]` for the transposed view and the +materialized `[1025, 2112]` output. Neither is wrong; they describe different +boundaries. For roofline work the kernel-level view is usually the one that +matches the GPU kernel's actual traffic. + +## Comparison with the JIT-hook approach + +| | This branch (launcher wrapping) | `jit-shape-tracer` branch (JIT hooks) | +|---|---|---| +| Interception point | the Python launcher function | `JITFunction.run` / `Autotuner.run` / FlyDSL `__call__` | +| Backend coverage | **any** backend — Triton, ASM, CK, aiter C++ bindings, RCCL collectives, torch ops | Triton and FlyDSL only | +| Needs to know the framework | **yes** — module paths in the registry are coupled to SGLang / aiter versions | no | +| Finds unlisted kernels | only via auto-discovery heuristics | automatically, all of them | +| `enable()` cost | high — walks and imports whole package trees, rebinds `sys.modules` | negligible — installs a few method wrappers | +| Risk | rebinding module globals and import side effects (see Safety properties) | contained to the JIT classes | +| Naming | op named after the launcher, with a numeric suffix | op name *equals* the kernel name, plus a launcher-level op | + +Use this branch when you need shapes for non-Triton work (ASM GEMMs, CK MoE, +aiter C++ kernels, collectives). Use the JIT-hook branch when you want +framework independence, complete Triton coverage without a registry, and a much +smaller blast radius. + +## Limitation: CUDA graphs + +Kernels that only run inside replayed CUDA graphs execute no Python, so no +wrapper fires during replay. Their shapes are recorded when the graph is +**captured**. If the capture happens inside a profiler window, the capture-time +trace holds those shapes — so a decode-path analysis needs the graph-capture +trace, not just the serving-window trace. diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py new file mode 100644 index 000000000..fead57869 --- /dev/null +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py @@ -0,0 +1,952 @@ +""" +Automatic tensor shape metadata for Triton / FlashInfer / aiter kernels +in PyTorch profiler traces. + +When enabled, targeted kernel entry-point functions are registered as +torch custom ops via torch.library so they appear as ``cpu_op`` events +with ``Input Dims`` and ``Input type`` in profiler traces. + +Usage: + from kernel_shape_profiler import enable, disable + enable() # before profiling starts + disable() # after profiling stops + +In practice you do not call these by hand: the co-located ``sitecustomize.py`` +is auto-loaded via ``PYTHONPATH`` and drives them from the torch-profiler +window, so no serving-framework source needs patching. See README.md. + +Design: + We maintain an explicit registry of kernel entry points. For each one + we create a ``torch.library`` custom-op wrapper and then replace **every + module-level reference** to the original function across all of + ``sys.modules``. This handles the common ``from X import Y`` pattern + where patching only the definition module would miss callers that + already captured a local binding. + + Functions whose references were captured as *instance attributes* + before ``enable()`` (e.g. ``self.fn = dispatch()``) cannot be + intercepted directly. For those cases the registry should target the + **inner kernel** that the wrapper calls via module-global lookup at + call time (e.g. ``gemm_a8w8_blockscale`` inside + ``aiter_w8a8_block_fp8_linear``). + +Coverage: + Because interception happens at the *launcher* function rather than at a + JIT boundary, the backend behind the launcher is irrelevant: an ASM GEMM, + a CK MoE kernel, an aiter C++ binding and a Triton kernel are all + annotated the same way, as long as the launcher is a Python function that + can be reached by name. The cost is that the launcher must be *known* + (registry) or *guessed* (auto-discovery heuristics), and the framework + module paths in ``_KERNEL_ENTRY_POINTS`` make this file version-coupled to + SGLang / aiter internals. +""" + +import contextlib +import functools +import importlib +import inspect +import logging +import pkgutil +import sys +import threading +from typing import Any, Callable, List, Optional, Tuple + +import torch +from torch.library import Library + +logger = logging.getLogger(__name__) + + +@contextlib.contextmanager +def _preserve_global_torch_state(): + """Snapshot and restore process-global torch defaults. + + ``enable()`` imports a large number of modules (both the explicit + registry via ``_resolve_target`` and the auto-discovery + ``_force_import_submodules``). Some of those modules mutate + process-global torch state at import time — most notably + ``torch.set_default_device("cuda")`` — and that mutation is NOT undone + by ``disable()`` (which only restores patched *function references*). + + A leaked default device corrupts downstream CPU tensor creation: e.g. + a buffer such as ``seq_lens_cpu`` is suddenly allocated on cuda, which + surfaces as ``Buffer seq_lens_cpu has different device than before`` in + the input-buffer pool and, later, as out-of-bounds GPU memory access + faults in index kernels (e.g. ``write_req_to_token_pool_triton``). + + Restoring the default device and dtype around the import-heavy regions + keeps these side effects from escaping into the serving path. In the + healthy case (nothing mutates the defaults) this is a no-op. + """ + get_default_device = getattr(torch, "get_default_device", None) + saved_device = get_default_device() if get_default_device is not None else None + saved_dtype = torch.get_default_dtype() + try: + yield + finally: + if saved_device is not None: + try: + torch.set_default_device(saved_device) + except Exception: + pass + try: + torch.set_default_dtype(saved_dtype) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# State +# --------------------------------------------------------------------------- +_lock = threading.Lock() +_enabled = False +# The torch.library.Library is created once and kept alive for the whole +# process lifetime. It is intentionally NEVER torn down: dropping it destroys +# the registered custom ops (freeing their OperatorName / schema). A wrapper +# reference can outlive a disable() — e.g. a module imported lazily *during* a +# profiling window captured the wrapper via ``from X import Y`` and is not in +# ``_patches`` to be restored. If the backing op were freed, that leaked +# wrapper would later dispatch into freed memory and segfault. Keeping the +# Library (and the monotonic op counter) alive makes such a stale dispatch +# safe; the ``if not _enabled`` guard in each wrapper then routes it straight +# to the original function. +_lib: Optional[Library] = None +# Monotonic — NEVER reset, so op names from earlier enable() cycles stay valid. +_op_counter = 0 +# Each entry: (module_obj, attr_name, original_fn) +_patches: List[Tuple[Any, str, Callable]] = [] +# Persistent cache of built wrappers keyed by qualified function name. +# Value: (wrapper_fn, original_fn). Reused across enable()/disable() cycles so +# each op is defined exactly once and a given function maps to a stable wrapper +# object (so references that leaked across cycles still point at a live op). +_built_wrappers: dict = {} + + +def _get_or_create_lib() -> Library: + """Return the process-wide custom-op Library, creating it on first use.""" + global _lib + if _lib is None: + _lib = Library("sglang_profiler", "FRAGMENT") + return _lib + + +# --------------------------------------------------------------------------- +# Registry of kernel entry points to wrap. +# +# Each entry is (module_path, function_name). +# +# **Guidelines for choosing what to register:** +# +# 1. Prefer functions that are called via *module-global name lookup* +# at call time. These are always patchable because Python resolves +# the name in the module's ``__dict__`` on every call. +# +# 2. Avoid outer "dispatch" wrappers whose references get captured as +# instance attributes (e.g. ``self.w8a8_block_fp8_linear = +# dispatch_w8a8_block_fp8_linear()``). Instead register the *inner* +# kernel they call. +# +# 3. For functions imported via ``from X import Y`` into multiple +# modules, the ``_patch_all_references()`` helper will find and +# replace them everywhere in ``sys.modules``. +# --------------------------------------------------------------------------- +_KERNEL_ENTRY_POINTS = [ + # ── Triton attention ── + ("sglang.srt.layers.attention.triton_ops.decode_attention", "decode_attention_fwd"), + ( + "sglang.srt.layers.attention.triton_ops.decode_attention", + "decode_attention_fwd_normal", + ), + ( + "sglang.srt.layers.attention.triton_ops.decode_attention", + "decode_attention_fwd_grouped", + ), + ("sglang.srt.layers.attention.triton_ops.extend_attention", "extend_attention_fwd"), + ( + "sglang.srt.layers.attention.triton_ops.prefill_attention", + "context_attention_fwd", + ), + # ── Fused MoE ── + ("sglang.srt.layers.moe.fused_moe_triton.fused_moe", "invoke_fused_moe_kernel"), + ("sglang.srt.layers.moe.fused_moe_triton.fused_moe", "moe_align_block_size"), + ( + "sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels", + "fused_append_shared_experts", + ), + # ── MoE TopK ── + ("sglang.srt.layers.moe.topk", "biased_grouped_topk_gpu"), + # ── Layer norm ── + # The actual module-level names are rmsnorm / fused_add_rmsnorm + # (imported from sgl_kernel on CUDA, or aiter on HIP). + ("sglang.srt.layers.layernorm", "rmsnorm"), + ("sglang.srt.layers.layernorm", "fused_add_rmsnorm"), + ("sglang.srt.layers.layernorm", "gemma_rmsnorm"), + ("sglang.srt.layers.layernorm", "gemma_fused_add_rmsnorm"), + # ── FP8 quantization ── + ("sglang.srt.layers.quantization.fp8_utils", "per_token_group_quant_fp8"), + ("sglang.srt.layers.quantization.fp8_utils", "scaled_fp8_quant"), + # Inner kernel called by triton_w8a8_block_fp8_linear via global lookup: + ("sglang.srt.layers.quantization.fp8_utils", "w8a8_block_fp8_matmul_triton"), + # Inner kernel called by aiter_w8a8_block_fp8_linear via global lookup + # (the outer wrapper is captured by reference at model init, but this + # inner kernel is looked up from the module __dict__ on every call): + ("sglang.srt.layers.quantization.fp8_utils", "gemm_a8w8_blockscale"), + # ── LoRA Triton (inner kernel functions, no *args/**kwargs) ── + ("sglang.srt.lora.triton_ops.sgemm_lora_a", "sgemm_lora_a_fwd"), + ("sglang.srt.lora.triton_ops.sgemm_lora_b", "sgemm_lora_b_fwd"), + # ── aiter (AMD) ops — definition-site patching ── + ("aiter.ops.triton.gemm_a8w8_blockscale", "gemm_a8w8_blockscale"), + ("aiter.ops.triton.batched_gemm_a8w8_blockscale", "batched_gemm_a8w8_blockscale"), + ("aiter.ops.norm", "rms_norm"), + ("aiter.ops.norm", "fused_add_rms_norm"), + # ── FlashInfer MoE (cutedsl) ── + ("flashinfer.moe", "moe_gemm_fp8_nt_groupwise"), + # ───────────────────────────────────────────────────────────────────── + # Current-layout paths (SGLang 0.5.18+ / matching aiter). + # + # SGLang moved its Triton kernels out of ``sglang.srt.layers.*`` into a + # separate ``sglang.kernels.ops.*`` package, and aiter regrouped its + # Triton ops into subpackages. Measured against + # sglang 0.5.18 + aiter, only 5 of the 24 legacy entries above still + # resolve, so the current paths are listed here as well. + # + # Both sets are kept on purpose: ``_resolve_target`` returns None for a + # path that does not exist and the entry is skipped silently, so one + # registry works across framework versions. Duplicates are harmless -- + # ``_wrapped_ids`` in enable() prevents double-wrapping the same object. + # ───────────────────────────────────────────────────────────────────── + # ── SGLang Triton attention ── + ("sglang.kernels.ops.attention.decode_attention", "decode_attention_fwd"), + ("sglang.kernels.ops.attention.decode_attention", "decode_attention_fwd_normal"), + ("sglang.kernels.ops.attention.decode_attention", "decode_attention_fwd_grouped"), + ("sglang.kernels.ops.attention.extend_attention", "extend_attention_fwd"), + ("sglang.kernels.ops.attention.prefill_attention", "context_attention_fwd"), + # ── SGLang fused MoE ── + ("sglang.kernels.ops.moe.fused_moe_triton_kernels", "invoke_fused_moe_kernel"), + ("sglang.kernels.ops.moe.fused_moe_triton_kernels", "fused_append_shared_experts"), + ("sglang.kernels.ops.moe.moe_align", "moe_align_block_size"), + # ── SGLang layer norm ── + ("sglang.kernels.ops.layernorm.norm", "rmsnorm"), + ("sglang.kernels.ops.layernorm.norm", "fused_add_rmsnorm"), + ("sglang.kernels.ops.layernorm.minimax_m3_rmsnorm", "gemma_rmsnorm"), + ("sglang.kernels.ops.layernorm", "gemma_fused_add_rmsnorm"), + # ── SGLang LoRA Triton ── + ("sglang.kernels.ops.gemm.sgemm_lora_a", "sgemm_lora_a_fwd"), + ("sglang.kernels.ops.gemm.sgemm_lora_b", "sgemm_lora_b_fwd"), + # ── aiter regrouped Triton ops ── + ("aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale", "gemm_a8w8_blockscale"), + ("aiter.ops.triton.normalization.rmsnorm", "rms_norm"), + # aiter's batched blockscale GEMM was renamed, not just moved (it is now + # batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant + # under aiter.ops.triton.gemm.batched). Auto-discovery picks it up, so it + # is deliberately not pinned to a name here. +] + +# --------------------------------------------------------------------------- +# Auto-discovery prefixes. +# +# In addition to the explicit registry above, ``enable()`` scans every +# already-loaded module whose name starts with one of these prefixes and +# wraps only functions likely to launch kernels. Discovery is filtered by +# signature/source heuristics to avoid wrapping unrelated utility code. +# --------------------------------------------------------------------------- +_AUTO_DISCOVER_PREFIXES: Tuple[str, ...] = ( + "flashinfer.", + "sglang.srt.", + "aiter.ops.", +) + + +# --------------------------------------------------------------------------- +# Schema building — works with or without type annotations +# --------------------------------------------------------------------------- + +# Python type → torch schema type +_TYPE_MAP = { + torch.Tensor: "Tensor", + Optional[torch.Tensor]: "Tensor?", + int: "int", + float: "float", + bool: "bool", + str: "str", + torch.dtype: "ScalarType", +} + +# String annotation variants produced by ``from __future__ import annotations`` +# (PEP 563) – annotations are stored as literal strings in the source code. +_STRING_TYPE_MAP = { + "torch.Tensor": "Tensor", + "Tensor": "Tensor", + "Optional[torch.Tensor]": "Tensor?", + "Optional[Tensor]": "Tensor?", + "int": "int", + "float": "float", + "bool": "bool", + "str": "str", + "torch.dtype": "ScalarType", +} + + +def _infer_schema_type(param: inspect.Parameter) -> Optional[str]: + """Map a parameter's annotation to a torch schema type string.""" + annotation = param.annotation + if annotation is inspect._empty: + return None + + # ── Handle string annotations (PEP 563) ── + if isinstance(annotation, str): + if annotation in _STRING_TYPE_MAP: + return _STRING_TYPE_MAP[annotation] + # Check "Optional[X]" pattern in string form + if annotation.startswith("Optional[") and annotation.endswith("]"): + inner = annotation[len("Optional[") : -1] + base = _STRING_TYPE_MAP.get(inner) + if base is not None: + return base if base.endswith("?") else base + "?" + return None + + # ── Handle real type annotations ── + # Check direct match + if annotation in _TYPE_MAP: + return _TYPE_MAP[annotation] + # Check Optional[X] (Union[X, None]) + origin = getattr(annotation, "__origin__", None) + if origin is type(None): + return None + args = getattr(annotation, "__args__", ()) + if args and type(None) in args: + for a in args: + if a is not type(None) and a in _TYPE_MAP: + return _TYPE_MAP[a] + "?" + return None + + +def _build_schema_from_sig( + sig: inspect.Signature, + skip_self: bool = False, +) -> Optional[Tuple[str, List[str], List[str]]]: + """ + Build schema string from signature annotations. + Returns (schema_str, tensor_param_names, non_tensor_param_names) or None + if there are no tensor params or the signature can't be mapped. + """ + tensor_params: List[str] = [] + non_tensor_params: List[str] = [] + schema_parts: List[str] = [] + + for name, param in sig.parameters.items(): + if skip_self and name == "self": + continue + if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): + return None + + stype = _infer_schema_type(param) + if stype is not None and "Tensor" in stype: + tensor_params.append(name) + schema_parts.append(f"{stype} {name}") + else: + non_tensor_params.append(name) + + if not tensor_params: + return None + + schema_str = f"({', '.join(schema_parts)}) -> ()" + return schema_str, tensor_params, non_tensor_params + + +# --------------------------------------------------------------------------- +# Thread-local side channel for non-tensor args +# --------------------------------------------------------------------------- +_tls = threading.local() + + +def _stash_non_tensor_args(op_name: str, values: dict): + if not hasattr(_tls, "stash"): + _tls.stash = {} + _tls.stash[op_name] = values + + +def _pop_non_tensor_args(op_name: str) -> dict: + if not hasattr(_tls, "stash"): + return {} + return _tls.stash.pop(op_name, {}) + + +def _stash_return_value(op_name: str, value: Any): + if not hasattr(_tls, "returns"): + _tls.returns = {} + _tls.returns[op_name] = value + + +def _pop_return_value(op_name: str) -> Any: + if not hasattr(_tls, "returns"): + return None + return _tls.returns.pop(op_name, None) + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + + +def _next_op_name(base: str) -> str: + global _op_counter + sanitized = base.replace(".", "_").replace("::", "_").replace("-", "_") + name = f"{sanitized}_{_op_counter}" + _op_counter += 1 + return name + + +def _register_op( + op_name: str, + schema_str: str, + original_fn: Callable, + tensor_param_names: List[str], + non_tensor_param_names: List[str], + sig: inspect.Signature, + skip_self: bool = False, +) -> Optional[Callable]: + """ + Register a function as a torch custom op and return a dispatch wrapper. + Returns None if registration fails. + """ + try: + lib = _get_or_create_lib() + lib.define(op_name + schema_str) + + def impl(*tensor_args): + nt_args = _pop_non_tensor_args(op_name) + full_kwargs = {} + t_idx = 0 + for pname, param in sig.parameters.items(): + if skip_self and pname == "self": + continue + if pname in tensor_param_names: + full_kwargs[pname] = tensor_args[t_idx] + t_idx += 1 + elif pname in non_tensor_param_names: + if pname in nt_args: + full_kwargs[pname] = nt_args[pname] + elif param.default is not inspect._empty: + full_kwargs[pname] = param.default + result = original_fn(**full_kwargs) + # Schema is -> () so we can't return the actual value + # through the dispatcher. Stash it for the caller. + _stash_return_value(op_name, result) + + lib.impl(op_name, impl, dispatch_key="CompositeExplicitAutograd") + + torch_op = getattr(torch.ops.sglang_profiler, op_name) + + @functools.wraps(original_fn) + def dispatch_wrapper(*args, **kwargs): + # When profiling is not active, never route through the custom op. + # A wrapper reference may outlive disable() (captured by a module + # imported lazily while profiling was on, so _patch_all_references + # could not restore it). Falling back to the original keeps such + # leaked bindings correct and crash-free. + if not _enabled: + return original_fn(*args, **kwargs) + try: + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + except TypeError: + return original_fn(*args, **kwargs) + + tensor_args = [] + nt_vals = {} + for pname, val in bound.arguments.items(): + if skip_self and pname == "self": + continue + if pname in tensor_param_names: + tensor_args.append(val) + elif pname in non_tensor_param_names: + nt_vals[pname] = val + + # If every tensor arg is None (all Optional[Tensor] and not + # provided), torch dispatch will fail with "no tensor arguments". + # Fall back to calling the original function directly. + if not any(isinstance(t, torch.Tensor) for t in tensor_args): + return original_fn(*args, **kwargs) + + _stash_non_tensor_args(op_name, nt_vals) + try: + torch_op(*tensor_args) + return _pop_return_value(op_name) + except Exception: + # Dispatch failed (e.g. device/type mismatch, None for + # non-optional Tensor, schema arity error). Clean up + # thread-local stash and fall back to the original call. + _pop_non_tensor_args(op_name) + _pop_return_value(op_name) + return original_fn(*args, **kwargs) + + # Mark so enable() can recognise its own wrappers and never wrap one + # again (see the guard in enable()). + dispatch_wrapper._kernel_shape_wrapper = True + return dispatch_wrapper + + except Exception as e: + logger.debug("Failed to register %s: %s", op_name, e) + return None + + +# --------------------------------------------------------------------------- +# Module + attribute resolution +# --------------------------------------------------------------------------- + + +def _resolve_target(module_path: str, attr_name: str): + """ + Resolve a target function from *module_path* and *attr_name*. + *attr_name* can be ``"func_name"`` or ``"ClassName.method_name"``. + + Returns ``(container, attr_name, original_fn, is_method)`` or ``None``. + """ + try: + mod = importlib.import_module(module_path) + except ImportError: + return None + + if "." in attr_name: + cls_name, method_name = attr_name.split(".", 1) + cls = getattr(mod, cls_name, None) + if cls is None: + return None + fn = getattr(cls, method_name, None) + if fn is None: + return None + return cls, method_name, fn, True + else: + fn = getattr(mod, attr_name, None) + if fn is None: + return None + return mod, attr_name, fn, False + + +def _patch_all_references(original_fn: Callable, wrapper_fn: Callable): + """ + Scan ``sys.modules`` and replace **every** module-level attribute that + points to *original_fn* with *wrapper_fn*. + + This handles the common ``from X import Y`` pattern: if module A + defines ``Y`` and module B does ``from A import Y``, both A and B + will have their binding replaced. + + Returns a list of ``(module, attr_name, original_fn)`` for later + restoration. + """ + patches = [] + for _mod_name, mod in list(sys.modules.items()): + if mod is None: + continue + try: + mod_dict = vars(mod) + except TypeError: + continue + for attr_name in list(mod_dict.keys()): + if attr_name.startswith("__"): + continue + try: + if mod_dict[attr_name] is original_fn: + setattr(mod, attr_name, wrapper_fn) + patches.append((mod, attr_name, original_fn)) + except Exception: + pass + return patches + + +# --------------------------------------------------------------------------- +# Lightweight wrappers for functions that can't use torch.library +# --------------------------------------------------------------------------- + + +def _make_record_function_wrapper( + qualified_name: str, + original_fn: Callable, +) -> Callable: + """ + Create a wrapper that uses ``torch.profiler.record_function`` to emit + a ``cpu_op`` event with tensor shapes embedded in the event name. + + Used for functions where we can't build a ``torch.library`` schema + (e.g. no type annotations, ``*args``/``**kwargs``, etc.). + """ + + @functools.wraps(original_fn) + def wrapper(*args, **kwargs): + # See dispatch_wrapper: a leaked reference must be a no-op when + # profiling is inactive. + if not _enabled: + return original_fn(*args, **kwargs) + shape_parts: List[str] = [] + for i, arg in enumerate(args): + if isinstance(arg, torch.Tensor): + shape_parts.append(f"arg{i}:{list(arg.shape)}") + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + shape_parts.append(f"{k}:{list(v.shape)}") + if shape_parts: + event_name = f"{qualified_name}({', '.join(shape_parts)})" + else: + event_name = qualified_name + with torch.profiler.record_function(event_name): + return original_fn(*args, **kwargs) + + wrapper._kernel_shape_wrapper = True + return wrapper + + +# --------------------------------------------------------------------------- +# Kernel-launch detection heuristics +# --------------------------------------------------------------------------- + +# Substrings that strongly indicate a function launches a GPU kernel. +_KERNEL_SOURCE_INDICATORS = ( + "[grid", # Triton launch pattern: kernel[grid](...) + "torch.ops.", # Custom C++/CUDA op dispatch + "sgl_kernel.", # sgl-kernel extension entry points +) + + +def _source_launches_kernel(fn: Callable) -> bool: + """Return True if *fn* source contains known kernel-launch patterns.""" + try: + source = inspect.getsource(fn) + except (OSError, TypeError): + return False + return any(marker in source for marker in _KERNEL_SOURCE_INDICATORS) + + +def _is_likely_kernel_launcher(fn: Callable, sig: inspect.Signature) -> bool: + """ + Decide whether *fn* is likely to launch a GPU kernel. + + Priority: + 1) Tensor annotation exists -> include. + 2) Non-tensor annotations only -> exclude. + 3) No annotations -> fallback to source pattern matching. + """ + has_any_annotation = False + for param in sig.parameters.values(): + if param.annotation is inspect._empty: + continue + has_any_annotation = True + stype = _infer_schema_type(param) + if stype is not None and "Tensor" in stype: + return True + + if has_any_annotation: + return False + + return _source_launches_kernel(fn) + + +def _force_import_submodules(prefix: str) -> None: + """ + Recursively import submodules under *prefix* so they appear in + ``sys.modules`` before auto-discovery runs. + + *prefix* should be a top-level package name without a trailing dot + (e.g. ``"sglang.srt"``). + """ + try: + pkg = importlib.import_module(prefix) + except ImportError: + return + + pkg_path = getattr(pkg, "__path__", None) + if pkg_path is None: + return + + # Snapshot the global defaults once; restore them after *every* import so + # a module that calls torch.set_default_device("cuda") at import time + # cannot taint subsequently imported modules during discovery (nor leak + # into the serving path). The leaf-name skip list below catches the known + # offenders, but this restore makes the discovery robust to any other + # module with the same import-time side effect. + get_default_device = getattr(torch, "get_default_device", None) + saved_device = get_default_device() if get_default_device is not None else None + saved_dtype = torch.get_default_dtype() + + def _restore_defaults(): + if saved_device is not None: + try: + torch.set_default_device(saved_device) + except Exception: + pass + try: + torch.set_default_dtype(saved_dtype) + except Exception: + pass + + for _importer, mod_name, _is_pkg in pkgutil.walk_packages( + pkg_path, prefix=prefix + "." + ): + if mod_name in sys.modules: + continue + # Skip test / benchmark / autotune modules. These are not kernel + # entry points and several of them (e.g. aiter.ops.flydsl.test_*) + # call ``torch.set_default_device("cuda")`` at import time, which + # leaks a CUDA default-device mode into the importing process and + # corrupts CPU tensor creation downstream (e.g. seq_lens_cpu). + leaf = mod_name.rsplit(".", 1)[-1] + if ( + "test" in leaf + or leaf.startswith("test_") + or leaf.startswith("bench_") + or leaf.endswith("_test") + or leaf.endswith("_tune") + ): + continue + try: + importlib.import_module(mod_name) + except Exception: + # Optional modules may fail to import depending on environment. + pass + finally: + _restore_defaults() + + +def _discover_kernel_entry_points() -> List[Tuple[str, str]]: + """ + Scan already-loaded ``sys.modules`` for modules whose name starts + with one of ``_AUTO_DISCOVER_PREFIXES`` and collect only functions + likely to launch GPU kernels. + + Returns a list of ``(module_path, function_name)`` pairs. + """ + # Force-import submodules under each discovery prefix first so deeper + # kernels become visible in sys.modules. + for prefix in _AUTO_DISCOVER_PREFIXES: + _force_import_submodules(prefix.rstrip(".")) + + results: List[Tuple[str, str]] = [] + seen_ids: set = set() + + for mod_name, mod in list(sys.modules.items()): + if mod is None: + continue + if not any(mod_name.startswith(p) for p in _AUTO_DISCOVER_PREFIXES): + continue + try: + mod_dict = vars(mod) + except TypeError: + continue + for attr_name in list(mod_dict.keys()): + if attr_name.startswith("__"): + continue + obj = mod_dict[attr_name] + + # Only wrap regular Python functions. + # @triton.jit objects (JITFunction / Autotuner) must NOT be + # wrapped — replacing them in module globals breaks the + # Triton compiler's global resolution for device-side calls + # between JIT kernels (e.g. remap_xcd, _rmsmorm_op, etc.). + if not inspect.isfunction(obj): + continue + + # Only include functions *defined* within a target namespace + fn_module = getattr(obj, "__module__", "") or "" + if not any(fn_module.startswith(p) for p in _AUTO_DISCOVER_PREFIXES): + continue + + obj_id = id(obj) + if obj_id in seen_ids: + continue + seen_ids.add(obj_id) + + # Require at least one parameter + try: + sig = inspect.signature(obj) + except (ValueError, TypeError): + continue + if not sig.parameters: + continue + + # Skip signatures we cannot map to a torch schema. + if any( + p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD) + for p in sig.parameters.values() + ): + continue + + # Core filter: only keep likely kernel launchers. + if not _is_likely_kernel_launcher(obj, sig): + continue + + results.append((mod_name, attr_name)) + + logger.debug( + "Auto-discovered %d kernel candidates from %s", + len(results), + ", ".join(_AUTO_DISCOVER_PREFIXES), + ) + return results + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def enable(): + """Patch registered kernel entry points to appear as cpu_op.""" + global _enabled + with _lock: + if _enabled: + return + + # NOTE: the Library and op counter are process-persistent (see the + # _lib / _op_counter docs). We do NOT recreate or reset them here so + # ops registered in earlier cycles stay valid. Only the per-cycle + # reference patches are rebuilt. + _patches.clear() + _wrapped_ids: set = set() # track function ids to avoid double-wrapping + + # Backstop guard around the import-heavy region. Both the explicit + # registry (_resolve_target -> importlib.import_module) and the + # auto-discovery (_discover_kernel_entry_points -> _force_import_*) + # import modules that may mutate process-global torch defaults at + # import time; restore them on exit so the leak cannot escape into + # the serving path. (No-op when nothing mutates the defaults.) + with _preserve_global_torch_state(): + # Merge explicit registry with filtered auto-discovered functions. + # Duplicates are harmless — _wrapped_ids prevents double-wrapping. + all_entry_points = ( + list(_KERNEL_ENTRY_POINTS) + _discover_kernel_entry_points() + ) + + for module_path, attr_name in all_entry_points: + resolved = _resolve_target(module_path, attr_name) + if resolved is None: + continue + + container, name, original_fn, is_method = resolved + is_plain_function = not is_method + + # Already one of our own wrappers. This happens when a + # framework keeps a compat re-export and both the old and new + # module paths are in the registry: the first entry wraps the + # function and _patch_all_references rebinds every reference to + # it, so the second entry now resolves to the wrapper. Wrapping + # again would nest two annotations around a single call and + # double-count it. + if getattr(original_fn, "_kernel_shape_wrapper", False): + logger.debug( + "Skipping already-wrapped %s.%s", module_path, attr_name + ) + continue + + # Skip if this exact function object was already wrapped + fn_id = id(original_fn) + if fn_id in _wrapped_ids: + logger.debug("Skipping duplicate %s.%s", module_path, attr_name) + continue + _wrapped_ids.add(fn_id) + + qualified_name = f"{module_path}.{name}" + + # Reuse a wrapper built in a previous cycle if the underlying + # function object is unchanged. This keeps each op defined + # exactly once and ensures a given function maps to a stable + # wrapper, so a reference that leaked across cycles still + # targets a live op. + wrapper = None + cached = _built_wrappers.get(qualified_name) + if cached is not None and cached[1] is original_fn: + wrapper = cached[0] + + if wrapper is None: + # ── Regular functions ── + try: + sig = inspect.signature(original_fn) + except (ValueError, TypeError): + logger.debug( + "Cannot inspect signature of %s — skipping", + qualified_name, + ) + continue + + base = f"{module_path.split('.')[-1]}_{name}" + schema_info = _build_schema_from_sig(sig, skip_self=is_method) + + if schema_info is not None: + # Full tensor annotations → use torch.library custom op + schema_str, t_names, nt_names = schema_info + op_name = _next_op_name(base) + wrapper = _register_op( + op_name, + schema_str, + original_fn, + t_names, + nt_names, + sig, + skip_self=is_method, + ) + if wrapper is not None: + logger.debug( + "Registered %s as custom op %s", qualified_name, op_name + ) + + if wrapper is None: + # Fallback: no annotations or schema registration failed + # → use record_function wrapper (shapes in event name) + wrapper = _make_record_function_wrapper( + qualified_name, + original_fn, + ) + logger.debug( + "Registered %s via record_function", qualified_name + ) + + _built_wrappers[qualified_name] = (wrapper, original_fn) + + # --- Apply patches --- + if is_plain_function: + ref_patches = _patch_all_references(original_fn, wrapper) + _patches.extend(ref_patches) + if not ref_patches: + setattr(container, name, wrapper) + _patches.append((container, name, original_fn)) + else: + setattr(container, name, wrapper) + _patches.append((container, name, original_fn)) + + n_discovered = len(all_entry_points) - len(_KERNEL_ENTRY_POINTS) + + _enabled = True + logger.info( + "kernel_shape_profiler enabled: %d references patched across " + "%d entry points (%d explicit + %d auto-discovered)", + len(_patches), + len(_KERNEL_ENTRY_POINTS) + n_discovered, + len(_KERNEL_ENTRY_POINTS), + n_discovered, + ) + + +def disable(): + """Restore all patched functions to originals.""" + global _enabled + with _lock: + if not _enabled: + return + # Flip the flag first so any wrapper invoked concurrently — or one that + # leaked past restoration — short-circuits to the original instead of + # dispatching into a custom op. + _enabled = False + for container, name, original_fn in reversed(_patches): + try: + setattr(container, name, original_fn) + except Exception: + pass + _patches.clear() + # Intentionally keep _lib, _op_counter and _built_wrappers alive: the + # registered ops must outlive any wrapper reference that may have + # leaked (see module-level _lib docs). + logger.info("kernel_shape_profiler disabled: all patches restored") + + +def is_enabled() -> bool: + return _enabled diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py b/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py new file mode 100644 index 000000000..c5617f521 --- /dev/null +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py @@ -0,0 +1,282 @@ +""" +Auto-loaded shim that turns on kernel-shape annotation without patching the +inference server (SGLang / vLLM / any torch workload). + +CPython imports a top-level module named ``sitecustomize`` automatically at +interpreter startup for every process, as long as its directory is on +``sys.path`` (i.e. on ``PYTHONPATH``). This shim uses that hook to drive +``kernel_shape_profiler.enable()`` / ``disable()`` from the torch-profiler +window, so the launcher wrapping costs nothing outside a profiling run. + +Activation (both required): + export PYTHONPATH=/path/to/kernel_shape_tool:$PYTHONPATH + export TRACELENS_SHAPE_DISCOVERY=1 + +Behaviour when ``TRACELENS_SHAPE_DISCOVERY`` is unset/false: this shim installs +nothing observable -- every hook short-circuits, so it is safe to leave the +directory on ``PYTHONPATH`` permanently. + +torch is usually not imported yet when ``sitecustomize`` runs, so patches are +registered as *pending* and applied by an ``__import__`` hook the moment the +target module appears in ``sys.modules``. All hooks are idempotent (guarded by a +sentinel attribute) and crash-proof (wrapped in ``try/except``) so they can +never break the workload. + +Note on timing: unlike a JIT-hook tracer, ``enable()`` here is *expensive* -- +it walks ``sglang.srt`` / ``aiter.ops`` / ``flashinfer`` with +``pkgutil.walk_packages``, then rebinds module-level references across all of +``sys.modules``. It therefore runs once, at the first profiler start, and the +first ``start()`` call is measurably slower than subsequent ones. +""" + +import builtins +import os +import sys +import threading + +_ENV_FLAG = "TRACELENS_SHAPE_DISCOVERY" +_FORCE_RECORD_SHAPES_ENV = "TRACELENS_SHAPE_FORCE_RECORD_SHAPES" + + +def _flag_on(name: str, default: str = "0") -> bool: + val = os.environ.get(name, default) + return val.strip().lower() not in ("", "0", "false", "no", "off") + + +def _shape_discovery_on() -> bool: + return _flag_on(_ENV_FLAG, "0") + + +# --------------------------------------------------------------------------- +# Lazy handle to the co-located profiler. +# +# We intentionally do NOT import kernel_shape_profiler (which imports torch) at +# sitecustomize time -- that would force a heavy torch import at interpreter +# startup and run torch's import-time side effects too early. Instead we import +# it lazily, once a hook actually needs it and torch is already loaded. +# --------------------------------------------------------------------------- +_profiler = None + + +def _get_profiler(): + global _profiler + if _profiler is None: + # Make sure this file's own directory is importable even if only the + # parent ended up on sys.path. + here = os.path.dirname(os.path.abspath(__file__)) + if here not in sys.path: + sys.path.insert(0, here) + import kernel_shape_profiler as _ksp # noqa: E402 + + _profiler = _ksp + return _profiler + + +def _enable_profiler() -> None: + try: + if _shape_discovery_on(): + _get_profiler().enable() + except Exception: + pass + + +def _disable_profiler() -> None: + try: + profiler = _get_profiler() + if profiler.is_enabled(): + profiler.disable() + except Exception: + pass + + +# Number of torch-profiler windows currently open. Nested/overlapping windows +# must not disable the profiler until the last one closes. +_profiler_active = [0] + + +# --------------------------------------------------------------------------- +# Patch: torch.profiler.profile.start / stop +# enable() runs just before the profiler starts recording so the first +# captured launch already carries shape metadata; stop() restores the +# original function references. +# --------------------------------------------------------------------------- +def _patch_torch_profiler_profile() -> bool: + try: + import torch.profiler as tp + except Exception: + return False + + cls = getattr(tp, "profile", None) + if cls is None or not isinstance(cls, type): + return False + if getattr(cls, "_tracelens_shape_patched", False): + return True + + orig_start = cls.start + orig_stop = cls.stop + + def _start(self, *a, **kw): + if _shape_discovery_on(): + _enable_profiler() + _profiler_active[0] += 1 + self._tracelens_incremented = True + return orig_start(self, *a, **kw) + + def _stop(self, *a, **kw): + try: + return orig_stop(self, *a, **kw) + finally: + if getattr(self, "_tracelens_incremented", False): + self._tracelens_incremented = False + _profiler_active[0] = max(0, _profiler_active[0] - 1) + if _profiler_active[0] == 0: + _disable_profiler() + + cls.start = _start + cls.stop = _stop + cls._tracelens_shape_patched = True + return True + + +# --------------------------------------------------------------------------- +# Patch: _KinetoProfile.__init__ -> force record_shapes=True +# The wrapped launchers only surface "Input Dims" / "Input type" in the trace +# when the profiler was constructed with record_shapes=True. Force it on (opt +# out with TRACELENS_SHAPE_FORCE_RECORD_SHAPES=0) so shape annotation "just +# works" regardless of how the server configured its profiler. +# --------------------------------------------------------------------------- +def _patch_kineto_record_shapes() -> bool: + try: + from torch.profiler.profiler import _KinetoProfile + except Exception: + return False + + if getattr(_KinetoProfile, "_tracelens_record_shapes_patched", False): + return True + + orig_init = _KinetoProfile.__init__ + + def _patched_init(self, *args, **kwargs): + if _shape_discovery_on() and _flag_on(_FORCE_RECORD_SHAPES_ENV, "1"): + if kwargs.get("record_shapes") is False: + sys.stderr.write( + "[tracelens-shape] forcing record_shapes=True " + "(needed for kernel shape annotation)\n" + ) + kwargs["record_shapes"] = True + orig_init(self, *args, **kwargs) + + _KinetoProfile.__init__ = _patched_init + _KinetoProfile._tracelens_record_shapes_patched = True + return True + + +def _patch_torch_profiler_both() -> bool: + a = _patch_torch_profiler_profile() + b = _patch_kineto_record_shapes() + return a and b + + +# --------------------------------------------------------------------------- +# Patch: torch.cuda.profiler.start / stop (legacy profiling API) +# --------------------------------------------------------------------------- +def _patch_torch_cuda_profiler() -> bool: + try: + import torch.cuda.profiler as tcp + except Exception: + return False + if getattr(tcp, "_tracelens_shape_patched", False): + return True + if not (hasattr(tcp, "start") and hasattr(tcp, "stop")): + return False + + orig_start = tcp.start + orig_stop = tcp.stop + + def _start(*a, **kw): + if _shape_discovery_on(): + _enable_profiler() + _profiler_active[0] += 1 + return orig_start(*a, **kw) + + def _stop(*a, **kw): + try: + return orig_stop(*a, **kw) + finally: + if _profiler_active[0] > 0: + _profiler_active[0] = max(0, _profiler_active[0] - 1) + if _profiler_active[0] == 0: + _disable_profiler() + + tcp.start = _start + tcp.stop = _stop + tcp._tracelens_shape_patched = True + return True + + +# --------------------------------------------------------------------------- +# Pending-patch registry + import hook. +# +# Only the profiler entry points need patching here: the launcher wrapping +# itself is done by kernel_shape_profiler.enable(), which resolves its targets +# by importing them on demand. Kernels launched during CUDA-graph replay need +# no special handling: their shapes are recorded when the graph is *captured*, +# and capture that happens inside a profiler window is already covered. +# --------------------------------------------------------------------------- +_PENDING_PATCHES = { + "torch.profiler": _patch_torch_profiler_both, + "torch.cuda.profiler": _patch_torch_cuda_profiler, +} + + +def _try_pending() -> None: + for mod_name in list(_PENDING_PATCHES.keys()): + if mod_name in sys.modules: + fn = _PENDING_PATCHES.get(mod_name) + if fn is None: + continue + try: + if fn(): + _PENDING_PATCHES.pop(mod_name, None) + except Exception: + _PENDING_PATCHES.pop(mod_name, None) + + +def _install_import_hook() -> None: + if getattr(sys, "_tracelens_shape_import_hook", False): + return + sys._tracelens_shape_import_hook = True + + orig_import = builtins.__import__ + tls = threading.local() + + def _wrapped(name, globals=None, locals=None, fromlist=(), level=0): + module = orig_import(name, globals, locals, fromlist, level) + if not _PENDING_PATCHES: + return module + if getattr(tls, "in_hook", False): + return module + tls.in_hook = True + try: + _try_pending() + finally: + tls.in_hook = False + return module + + builtins.__import__ = _wrapped + + +def _bootstrap() -> None: + # Even when the flag is off we still install the (cheap) profiler patches: + # they all short-circuit via _shape_discovery_on(), and installing + # unconditionally keeps behaviour stable if the flag is toggled between + # fork/exec boundaries. + if getattr(sys, "_tracelens_shape_bootstrapped", False): + return + sys._tracelens_shape_bootstrapped = True + _try_pending() + if _PENDING_PATCHES: + _install_import_hook() + + +_bootstrap() From 7aacf6321bf55fcd3ea4f487fddf306a1d350bf6 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Thu, 3 Sep 2026 20:30:18 +0000 Subject: [PATCH 02/10] Add sglang.kernels.ops. to auto-discovery prefixes SGLang 0.5.18 moved its Triton kernels out of sglang.srt.layers.* into a separate sglang.kernels.ops.* package, which auto-discovery was not scanning. Of the launchers that actually appear in DeepSeek-R1 traces, discovery now finds 6 of 7 instead of 1 of 7: concat_and_cast_mha_k_triton, set_mla_kv_buffer_triton, vocab_parallel_embedding, compute_position_triton and clamp_position_cuda are all gained. The seventh, create_flashinfer_kv_indices_triton, is a @triton.jit kernel rather than a Python launcher, so it is never wrapped by design. Candidate count goes from 1186 to 2012. Requiring a source-level launch pattern in addition to a Tensor annotation would cut that to 413, but it also drops real launchers (compute_position_triton, clamp_position_cuda, write_cache_indices), so the looser filter is kept and the trade-off is documented in the README. --- .../kernel_shape_tool/README.md | 27 ++++++++++++++++--- .../kernel_shape_profiler.py | 5 ++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md b/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md index 429b51c5f..8251bc3a6 100644 --- a/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md @@ -70,10 +70,29 @@ pairs covering SGLang Triton attention, fused MoE, layernorm, FP8 quantization, LoRA, aiter ops and FlashInfer MoE. **Filtered auto-discovery** (`_discover_kernel_entry_points`) which force-imports -everything under `sglang.srt.`, `aiter.ops.` and `flashinfer.`, then keeps only -functions that look like kernel launchers — either they have a `Tensor` -annotation, or (when unannotated) their source contains a launch pattern such as -`[grid`, `torch.ops.` or `sgl_kernel.` (`_is_likely_kernel_launcher`). +everything under `sglang.kernels.ops.`, `sglang.srt.`, `aiter.ops.` and +`flashinfer.`, then keeps only functions that look like kernel launchers — +either they have a `Tensor` annotation, or (when unannotated) their source +contains a launch pattern such as `[grid`, `torch.ops.` or `sgl_kernel.` +(`_is_likely_kernel_launcher`). + +`sglang.kernels.ops.` matters a lot on 0.5.18: that is where the Triton +launchers moved to. Measured against the launchers that actually appear in +DeepSeek-R1 traces, discovery finds 6 of 7 with the prefix and only 1 of 7 +without it (`concat_and_cast_mha_k_triton`, `set_mla_kv_buffer_triton`, +`vocab_parallel_embedding`, `compute_position_triton` and `clamp_position_cuda` +are all gained). The seventh, +`create_flashinfer_kv_indices_triton`, is a `@triton.jit` kernel rather than a +Python launcher, so it is deliberately never wrapped (see below); it is reached +through its calling attention-backend method instead. + +The cost is breadth: candidates go from 1,186 to 2,012, and in a fully +annotated package the "has a Tensor annotation" rule admits plain helpers such +as `_assert_contiguous` and `_affine_supported`, which then show up in the trace +as ops. Tightening the rule to *also* require a source launch pattern cuts +candidates to 413 but drops real launchers (`compute_position_triton`, +`clamp_position_cuda`, `write_cache_indices` among them), so the looser rule is +kept on purpose. Expect trace noise and filter by name when analysing. Two things are deliberately excluded: - `@triton.jit` objects (`JITFunction` / `Autotuner`). Replacing them in module diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py index fead57869..dd4e771d0 100644 --- a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py @@ -252,6 +252,11 @@ def _get_or_create_lib() -> Library: # --------------------------------------------------------------------------- _AUTO_DISCOVER_PREFIXES: Tuple[str, ...] = ( "flashinfer.", + # SGLang 0.5.18+ moved its Triton kernels out of sglang.srt.layers.* into + # this package, which is where most launchers now live. sglang.srt. is kept + # because quantization / MoE-runner launchers still sit there (and for + # older versions). + "sglang.kernels.ops.", "sglang.srt.", "aiter.ops.", ) From 9a1f47ef5a21447bfee95a30f1da8efc48923eff Mon Sep 17 00:00:00 2001 From: mohammad abdul basit Date: Fri, 4 Sep 2026 00:55:09 +0000 Subject: [PATCH 03/10] kernel_shape_tool: use active default-device override for state preservation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port the active-device fix from main's kernel_shape_profiler patch (#980): read torch.utils._device.CURRENT_DEVICE (None when unset) instead of torch.get_default_device() (always concrete). Restoring the raw override — including None — truly clears any default-device mode a module leaks at import time, whereas restoring concrete 'cpu' would install a device mode. Co-authored-by: Cursor --- .../kernel_shape_profiler.py | 44 +++++++++++++------ 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py index dd4e771d0..b6980b985 100644 --- a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py @@ -57,6 +57,26 @@ logger = logging.getLogger(__name__) +def _active_default_device_override(): + """Return the active ``torch.set_default_device`` override, or ``None`` if unset. + + Unlike ``torch.get_default_device()`` — which always resolves to a concrete + device (``cpu`` when no override is installed) — this reads the raw + ``torch.utils._device.CURRENT_DEVICE`` sentinel, which is ``None`` when no + override is active. That distinction matters for restoration: passing the + concrete ``cpu`` back into ``torch.set_default_device`` *installs* a device + mode, whereas passing ``None`` truly clears any mode a module leaked at + import time. + """ + device_mod = sys.modules.get("torch.utils._device") + if device_mod is None: + try: + import torch.utils._device as device_mod + except Exception: + return None + return getattr(device_mod, "CURRENT_DEVICE", None) + + @contextlib.contextmanager def _preserve_global_torch_state(): """Snapshot and restore process-global torch defaults. @@ -78,17 +98,15 @@ def _preserve_global_torch_state(): keeps these side effects from escaping into the serving path. In the healthy case (nothing mutates the defaults) this is a no-op. """ - get_default_device = getattr(torch, "get_default_device", None) - saved_device = get_default_device() if get_default_device is not None else None + saved_device = _active_default_device_override() saved_dtype = torch.get_default_dtype() try: yield finally: - if saved_device is not None: - try: - torch.set_default_device(saved_device) - except Exception: - pass + try: + torch.set_default_device(saved_device) + except Exception: + pass try: torch.set_default_dtype(saved_dtype) except Exception: @@ -670,16 +688,14 @@ def _force_import_submodules(prefix: str) -> None: # into the serving path). The leaf-name skip list below catches the known # offenders, but this restore makes the discovery robust to any other # module with the same import-time side effect. - get_default_device = getattr(torch, "get_default_device", None) - saved_device = get_default_device() if get_default_device is not None else None + saved_device = _active_default_device_override() saved_dtype = torch.get_default_dtype() def _restore_defaults(): - if saved_device is not None: - try: - torch.set_default_device(saved_device) - except Exception: - pass + try: + torch.set_default_device(saved_device) + except Exception: + pass try: torch.set_default_dtype(saved_dtype) except Exception: From 6c2a761094d3ec8236beab4598e72a4fb596292b Mon Sep 17 00:00:00 2001 From: mohammad abdul basit Date: Wed, 9 Sep 2026 15:21:05 +0000 Subject: [PATCH 04/10] patchless shape discovery --- .../kernel_shape_tool/README.md | 242 ++--------- .../kernel_shape_profiler.py | 388 ++++-------------- .../kernel_shape_tool/sitecustomize.py | 95 ++--- 3 files changed, 138 insertions(+), 587 deletions(-) diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md b/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md index 8251bc3a6..854e9ce1a 100644 --- a/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md @@ -1,30 +1,35 @@ -# Kernel shape profiler (launcher-wrapping approach) +# Kernel shape profiler Adds `Input Dims` / `Input type` / `Input Strides` to PyTorch profiler traces -for GPU kernels that would otherwise appear with no operand metadata, so they -can be rooflined. +for GPU kernels (Triton / aiter / FlashInfer) that would otherwise appear with +no operand metadata, so they can be rooflined. -It works by wrapping **kernel launcher functions** — the Python entry points -like `gemm_a8w8_blockscale`, `invoke_fused_moe_kernel` or `rmsnorm` — as +It works by wrapping **kernel launcher functions** — Python entry points like +`gemm_a8w8_blockscale`, `invoke_fused_moe_kernel` or `rmsnorm` — as `torch.library` custom ops. Each wrapped launcher then shows up in the trace as -a `cpu_op` named after itself, carrying its tensor operands. +a `cpu_op` named after itself, carrying its tensor operands: -> This is the registry-based approach. The alternative, on the -> `jit-shape-tracer` branch, hooks the Triton/FlyDSL JIT launch boundary -> instead. See [Comparison](#comparison-with-the-jit-hook-approach) for which -> one to use. +``` +sglang_profiler::fp8_utils_gemm_a8w8_blockscale_12 <- cpu_op, named after the launcher + Input Dims: [[1025, 7168], [7168, 2112], [1025, 56], [56, 17]] + Input Strides: [[7168, 1], [1, 7168], [56, 1], [1, 56]] + Input type: ['c10::Float8_e4m3fnuz', 'c10::Float8_e4m3fnuz', 'float', 'float'] + _gemm_a8w8_blockscale_kernel_GROUP_K_128_... <- the real GPU kernel +``` + +Launchers come from an explicit registry plus auto-discovery, so no +serving-framework source is patched. ## Files | File | Role | |------|------| -| `kernel_shape_profiler.py` | The core. Wraps launcher functions as `sglang_profiler::__` custom ops via an explicit registry plus filtered auto-discovery, and rebinds every module-level reference to them. | -| `sitecustomize.py` | Auto-loaded shim. Drives `enable()` / `disable()` from the torch-profiler window so nothing is wrapped outside a profiling run. | +| `kernel_shape_profiler.py` | Wraps launcher functions as custom ops and rebinds every module-level reference to them. | +| `sitecustomize.py` | Auto-loaded shim that drives `enable()` / `disable()` from the torch-profiler window so nothing is wrapped outside a profiling run. | ## Activation -No serving-framework source is patched. Put this directory on `PYTHONPATH` and -set the flag: +Put this directory on `PYTHONPATH` and set the flag: ```bash export PYTHONPATH=/path/to/kernel_shape_tool:$PYTHONPATH @@ -32,211 +37,18 @@ export TRACELENS_SHAPE_DISCOVERY=1 ``` CPython auto-imports `sitecustomize` at interpreter startup for every process, -so the server and all TP workers pick it up. When -`TRACELENS_SHAPE_DISCOVERY` is unset or `0`, every hook short-circuits, so it is -safe to leave the directory on `PYTHONPATH` permanently. +so the server and all TP workers pick it up. When `TRACELENS_SHAPE_DISCOVERY` +is unset or `0`, every hook short-circuits, so it is safe to leave the directory +on `PYTHONPATH` permanently. ### Optional knobs | Env var | Default | Meaning | |---------|---------|---------| -| `TRACELENS_SHAPE_DISCOVERY` | `0` | Master switch. Must be truthy to do anything. | +| `TRACELENS_SHAPE_DISCOVERY` | `0` | Master switch | | `TRACELENS_SHAPE_FORCE_RECORD_SHAPES` | `1` | Force `record_shapes=True` on the profiler (`Input Dims` only surface when shapes are recorded). Set `0` to respect the server's own setting. | -## What lands in the trace - -For a launch of `gemm_a8w8_blockscale(x, w, x_scale, w_scale)`: - -``` -sglang_profiler::fp8_utils_gemm_a8w8_blockscale_12 <- cpu_op, named after the launcher - Input Dims: [[1025, 7168], [7168, 2112], [1025, 56], [56, 17]] - Input Strides: [[7168, 1], [1, 7168], [56, 1], [1, 56]] - Input type: ['c10::Float8_e4m3fnuz', 'c10::Float8_e4m3fnuz', 'float', 'float'] - _gemm_a8w8_blockscale_kernel_GROUP_K_128_... <- the real GPU kernel -``` - -The op is named `__`. The counter is -monotonic per process and never reset, so op names from earlier -`enable()`/`disable()` cycles stay valid. - -## How it works - -### 1. Choosing what to wrap - -Two sources are merged in `enable()`: - -**An explicit registry** (`_KERNEL_ENTRY_POINTS`) of `(module_path, function)` -pairs covering SGLang Triton attention, fused MoE, layernorm, FP8 quantization, -LoRA, aiter ops and FlashInfer MoE. - -**Filtered auto-discovery** (`_discover_kernel_entry_points`) which force-imports -everything under `sglang.kernels.ops.`, `sglang.srt.`, `aiter.ops.` and -`flashinfer.`, then keeps only functions that look like kernel launchers — -either they have a `Tensor` annotation, or (when unannotated) their source -contains a launch pattern such as `[grid`, `torch.ops.` or `sgl_kernel.` -(`_is_likely_kernel_launcher`). - -`sglang.kernels.ops.` matters a lot on 0.5.18: that is where the Triton -launchers moved to. Measured against the launchers that actually appear in -DeepSeek-R1 traces, discovery finds 6 of 7 with the prefix and only 1 of 7 -without it (`concat_and_cast_mha_k_triton`, `set_mla_kv_buffer_triton`, -`vocab_parallel_embedding`, `compute_position_triton` and `clamp_position_cuda` -are all gained). The seventh, -`create_flashinfer_kv_indices_triton`, is a `@triton.jit` kernel rather than a -Python launcher, so it is deliberately never wrapped (see below); it is reached -through its calling attention-backend method instead. - -The cost is breadth: candidates go from 1,186 to 2,012, and in a fully -annotated package the "has a Tensor annotation" rule admits plain helpers such -as `_assert_contiguous` and `_affine_supported`, which then show up in the trace -as ops. Tightening the rule to *also* require a source launch pattern cuts -candidates to 413 but drops real launchers (`compute_position_triton`, -`clamp_position_cuda`, `write_cache_indices` among them), so the looser rule is -kept on purpose. Expect trace noise and filter by name when analysing. - -Two things are deliberately excluded: -- `@triton.jit` objects (`JITFunction` / `Autotuner`). Replacing them in module - globals breaks Triton's global resolution for device-side calls between JIT - kernels. Only plain Python functions are wrapped. -- Test / benchmark / autotune modules, which are not entry points and some of - which call `torch.set_default_device("cuda")` at import time. - -### 2. Building the op - -`_build_schema_from_sig` derives a `torch.library` schema from the function -signature, mapping annotated parameters to schema types (`_infer_schema_type` -handles both real annotations and PEP 563 string annotations). Tensor -parameters become schema arguments; everything else is passed through a -thread-local side channel (`_stash_non_tensor_args`), because a schema of -`-> ()` cannot carry them. The return value comes back the same way -(`_stash_return_value`). - -If no schema can be built (no annotations, `*args`/`**kwargs`), the function -falls back to `_make_record_function_wrapper`, which emits a `record_function` -event with the shapes embedded in the *event name* instead of as structured -args. - -### 3. Applying the wrapper - -`_patch_all_references` scans all of `sys.modules` and replaces **every** -module-level attribute pointing at the original function. This is required -because of the `from X import Y` pattern: patching only the defining module -would miss callers that already captured a local binding. - -Launchers whose reference was captured as an *instance attribute* before -`enable()` (e.g. `self.fn = dispatch_w8a8_block_fp8_linear()`) cannot be -intercepted this way. For those, the registry targets the **inner** kernel that -the wrapper looks up from module globals on every call — which is why -`gemm_a8w8_blockscale` is registered rather than the outer dispatch wrapper. - -## Safety properties - -The implementation carries a few hard-won guards worth preserving: - -**The `Library` is never torn down.** Dropping it destroys the registered ops. -A wrapper reference can outlive `disable()` — a module imported lazily *during* -a profiling window may have captured the wrapper via `from X import Y` and is -not in `_patches` to be restored. If the backing op were freed, that leaked -wrapper would dispatch into freed memory and segfault. The `Library`, the -monotonic op counter and the wrapper cache are all kept alive for the process -lifetime, and the `if not _enabled` guard in each wrapper routes leaked calls -straight to the original function. - -**Global torch state is restored around imports.** `enable()` imports a large -number of modules, some of which mutate process-global torch state at import -time — most notably `torch.set_default_device("cuda")`. That mutation is not -undone by `disable()`. A leaked default device corrupts downstream CPU tensor -creation: a buffer such as `seq_lens_cpu` is suddenly allocated on CUDA, which -surfaces as `Buffer seq_lens_cpu has different device than before` and later as -out-of-bounds GPU memory faults in index kernels. -`_preserve_global_torch_state` and the per-import restore in -`_force_import_submodules` keep that from escaping into the serving path. - -**Every wrapper falls back to the original.** Unbindable signatures, all-`None` -tensor arguments, and dispatch failures all call the original function -directly, so wrapping can never change behaviour. - -**Wrappers are never wrapped again.** Frameworks often keep a compat -re-export, so the same function is reachable by two module paths. Once the -first is wrapped, `_patch_all_references` rebinds *every* reference to it — -including the second path — so the second registry entry would otherwise -resolve to our own wrapper and wrap it a second time, nesting two annotations -around a single call and double-counting it. Wrappers carry a -`_kernel_shape_wrapper` marker and `enable()` skips them. - -## Registry staleness - -The explicit registry is coupled to framework internals, and it drifts. Measured -against SGLang 0.5.18 + matching aiter, only **5 of the original 24 entries -still resolved**: SGLang had moved its Triton kernels out of -`sglang.srt.layers.*` into a separate `sglang.kernels.ops.*` package, and aiter -had regrouped its Triton ops into subpackages. - -Current-layout paths have been added alongside the legacy ones (**21 of 40** -resolve now; the 19 that do not are the legacy paths, kept deliberately). -`_resolve_target` returns `None` for a path that does not exist and the entry is -skipped silently, so one registry serves several framework versions. - -Two entries cannot be fixed by a path change: aiter's batched blockscale GEMM -was *renamed* (now -`batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant`), and -`flashinfer` is not installed in this image. Auto-discovery covers the former. - -The practical lesson: **auto-discovery is doing most of the work**, and the -registry is best treated as a hint list for launchers the heuristics miss. Check -what actually resolved before trusting the registry on a new stack. - -## Validation - -Verified with `validate_kernel_shape_profiler.py`: - -- A synthetic launcher is registered, produces a `cpu_op` with - `Input Dims: [[4, 8], [8, 16]]` and per-tensor dtypes, returns the correct - value, and passes its non-tensor argument through the side channel. -- Both the definition-site binding and a `from X import Y` binding are rebound, - and both are restored on `disable()`. -- A wrapper captured while enabled still returns correct results after - `disable()` (the leaked-reference case the persistent `Library` protects). -- The real aiter `gemm_a8w8_blockscale` is annotated with - `Input Dims: [[1025, 7168], [2112, 7168], [1025, 56], [17, 56], []]` and - `['c10::Float8_e4m3fnuz', 'c10::Float8_e4m3fnuz', 'float', 'float', '']`, - exactly once per call. -- `enable()` took **18–19 s** with auto-discovery scoped to `aiter.ops.` alone, - which is the cost of the package walk. Budget for more with `sglang.srt.` - included. - -### What the operands mean at launcher level - -Note the difference from a JIT-boundary tracer. Here the recorded operands are -the launcher's **call arguments**, so for the same GEMM you see `w` as -`[2112, 7168]` (before the launcher transposes it) and the optional `y` output -as an empty slot because it was not passed. A JIT-level tracer sees the -**kernel's** operands: `[7168, 2112]` for the transposed view and the -materialized `[1025, 2112]` output. Neither is wrong; they describe different -boundaries. For roofline work the kernel-level view is usually the one that -matches the GPU kernel's actual traffic. - -## Comparison with the JIT-hook approach - -| | This branch (launcher wrapping) | `jit-shape-tracer` branch (JIT hooks) | -|---|---|---| -| Interception point | the Python launcher function | `JITFunction.run` / `Autotuner.run` / FlyDSL `__call__` | -| Backend coverage | **any** backend — Triton, ASM, CK, aiter C++ bindings, RCCL collectives, torch ops | Triton and FlyDSL only | -| Needs to know the framework | **yes** — module paths in the registry are coupled to SGLang / aiter versions | no | -| Finds unlisted kernels | only via auto-discovery heuristics | automatically, all of them | -| `enable()` cost | high — walks and imports whole package trees, rebinds `sys.modules` | negligible — installs a few method wrappers | -| Risk | rebinding module globals and import side effects (see Safety properties) | contained to the JIT classes | -| Naming | op named after the launcher, with a numeric suffix | op name *equals* the kernel name, plus a launcher-level op | - -Use this branch when you need shapes for non-Triton work (ASM GEMMs, CK MoE, -aiter C++ kernels, collectives). Use the JIT-hook branch when you want -framework independence, complete Triton coverage without a registry, and a much -smaller blast radius. - -## Limitation: CUDA graphs - -Kernels that only run inside replayed CUDA graphs execute no Python, so no -wrapper fires during replay. Their shapes are recorded when the graph is -**captured**. If the capture happens inside a profiler window, the capture-time -trace holds those shapes — so a decode-path analysis needs the graph-capture -trace, not just the serving-window trace. +> Note on CUDA graphs: kernels that only run inside replayed graphs execute no +> Python, so their shapes are recorded at graph **capture** time. A decode-path +> analysis therefore needs the graph-capture trace, not just the serving-window +> trace. diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py index b6980b985..0aace2f2e 100644 --- a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py @@ -1,44 +1,8 @@ -""" -Automatic tensor shape metadata for Triton / FlashInfer / aiter kernels -in PyTorch profiler traces. - -When enabled, targeted kernel entry-point functions are registered as -torch custom ops via torch.library so they appear as ``cpu_op`` events -with ``Input Dims`` and ``Input type`` in profiler traces. - -Usage: - from kernel_shape_profiler import enable, disable - enable() # before profiling starts - disable() # after profiling stops - -In practice you do not call these by hand: the co-located ``sitecustomize.py`` -is auto-loaded via ``PYTHONPATH`` and drives them from the torch-profiler -window, so no serving-framework source needs patching. See README.md. - -Design: - We maintain an explicit registry of kernel entry points. For each one - we create a ``torch.library`` custom-op wrapper and then replace **every - module-level reference** to the original function across all of - ``sys.modules``. This handles the common ``from X import Y`` pattern - where patching only the definition module would miss callers that - already captured a local binding. - - Functions whose references were captured as *instance attributes* - before ``enable()`` (e.g. ``self.fn = dispatch()``) cannot be - intercepted directly. For those cases the registry should target the - **inner kernel** that the wrapper calls via module-global lookup at - call time (e.g. ``gemm_a8w8_blockscale`` inside - ``aiter_w8a8_block_fp8_linear``). - -Coverage: - Because interception happens at the *launcher* function rather than at a - JIT boundary, the backend behind the launcher is irrelevant: an ASM GEMM, - a CK MoE kernel, an aiter C++ binding and a Triton kernel are all - annotated the same way, as long as the launcher is a Python function that - can be reached by name. The cost is that the launcher must be *known* - (registry) or *guessed* (auto-discovery heuristics), and the framework - module paths in ``_KERNEL_ENTRY_POINTS`` make this file version-coupled to - SGLang / aiter internals. +"""Tensor shape metadata for Triton / FlashInfer / aiter kernels in profiler traces. + +Registered kernel launchers are wrapped as torch custom ops so they appear as +``cpu_op`` events carrying ``Input Dims`` / ``Input type``. ``sitecustomize.py`` +drives ``enable()`` / ``disable()`` from the profiler window. See README.md. """ import contextlib @@ -58,15 +22,11 @@ def _active_default_device_override(): - """Return the active ``torch.set_default_device`` override, or ``None`` if unset. - - Unlike ``torch.get_default_device()`` — which always resolves to a concrete - device (``cpu`` when no override is installed) — this reads the raw - ``torch.utils._device.CURRENT_DEVICE`` sentinel, which is ``None`` when no - override is active. That distinction matters for restoration: passing the - concrete ``cpu`` back into ``torch.set_default_device`` *installs* a device - mode, whereas passing ``None`` truly clears any mode a module leaked at - import time. + """Active ``set_default_device`` override, or ``None`` if unset. + + Reads the raw ``CURRENT_DEVICE`` sentinel, not ``get_default_device()`` + (always concrete): restoring ``None`` clears a device mode a module leaked + at import time, whereas restoring a concrete ``cpu`` would *install* one. """ device_mod = sys.modules.get("torch.utils._device") if device_mod is None: @@ -79,24 +39,11 @@ def _active_default_device_override(): @contextlib.contextmanager def _preserve_global_torch_state(): - """Snapshot and restore process-global torch defaults. - - ``enable()`` imports a large number of modules (both the explicit - registry via ``_resolve_target`` and the auto-discovery - ``_force_import_submodules``). Some of those modules mutate - process-global torch state at import time — most notably - ``torch.set_default_device("cuda")`` — and that mutation is NOT undone - by ``disable()`` (which only restores patched *function references*). - - A leaked default device corrupts downstream CPU tensor creation: e.g. - a buffer such as ``seq_lens_cpu`` is suddenly allocated on cuda, which - surfaces as ``Buffer seq_lens_cpu has different device than before`` in - the input-buffer pool and, later, as out-of-bounds GPU memory access - faults in index kernels (e.g. ``write_req_to_token_pool_triton``). - - Restoring the default device and dtype around the import-heavy regions - keeps these side effects from escaping into the serving path. In the - healthy case (nothing mutates the defaults) this is a no-op. + """Snapshot and restore global torch default device & dtype. + + ``enable()`` imports modules that may mutate global state (e.g. + ``set_default_device("cuda")``); a leaked default device corrupts later CPU + tensor creation (``Buffer seq_lens_cpu has different device than before``). """ saved_device = _active_default_device_override() saved_dtype = torch.get_default_dtype() @@ -113,31 +60,15 @@ def _preserve_global_torch_state(): pass -# --------------------------------------------------------------------------- -# State -# --------------------------------------------------------------------------- _lock = threading.Lock() _enabled = False -# The torch.library.Library is created once and kept alive for the whole -# process lifetime. It is intentionally NEVER torn down: dropping it destroys -# the registered custom ops (freeing their OperatorName / schema). A wrapper -# reference can outlive a disable() — e.g. a module imported lazily *during* a -# profiling window captured the wrapper via ``from X import Y`` and is not in -# ``_patches`` to be restored. If the backing op were freed, that leaked -# wrapper would later dispatch into freed memory and segfault. Keeping the -# Library (and the monotonic op counter) alive makes such a stale dispatch -# safe; the ``if not _enabled`` guard in each wrapper then routes it straight -# to the original function. +# Created once and NEVER torn down: dropping it frees the registered ops, so a +# wrapper reference that leaked past disable() would dispatch into freed memory. +# The ``if not _enabled`` guard in each wrapper routes leaked calls to the original. _lib: Optional[Library] = None -# Monotonic — NEVER reset, so op names from earlier enable() cycles stay valid. -_op_counter = 0 -# Each entry: (module_obj, attr_name, original_fn) -_patches: List[Tuple[Any, str, Callable]] = [] -# Persistent cache of built wrappers keyed by qualified function name. -# Value: (wrapper_fn, original_fn). Reused across enable()/disable() cycles so -# each op is defined exactly once and a given function maps to a stable wrapper -# object (so references that leaked across cycles still point at a live op). -_built_wrappers: dict = {} +_op_counter = 0 # monotonic, never reset, so old op names stay valid +_patches: List[Tuple[Any, str, Callable]] = [] # (module, attr, original_fn) +_built_wrappers: dict = {} # {qualified_name: (wrapper, original_fn)}, reused across cycles def _get_or_create_lib() -> Library: @@ -148,26 +79,8 @@ def _get_or_create_lib() -> Library: return _lib -# --------------------------------------------------------------------------- -# Registry of kernel entry points to wrap. -# -# Each entry is (module_path, function_name). -# -# **Guidelines for choosing what to register:** -# -# 1. Prefer functions that are called via *module-global name lookup* -# at call time. These are always patchable because Python resolves -# the name in the module's ``__dict__`` on every call. -# -# 2. Avoid outer "dispatch" wrappers whose references get captured as -# instance attributes (e.g. ``self.w8a8_block_fp8_linear = -# dispatch_w8a8_block_fp8_linear()``). Instead register the *inner* -# kernel they call. -# -# 3. For functions imported via ``from X import Y`` into multiple -# modules, the ``_patch_all_references()`` helper will find and -# replace them everywhere in ``sys.modules``. -# --------------------------------------------------------------------------- +# Registry of kernel entry points to wrap: (module_path, function_name). +# Register the inner kernel, not dispatch wrappers captured as instance attrs. _KERNEL_ENTRY_POINTS = [ # ── Triton attention ── ("sglang.srt.layers.attention.triton_ops.decode_attention", "decode_attention_fwd"), @@ -193,9 +106,7 @@ def _get_or_create_lib() -> Library: ), # ── MoE TopK ── ("sglang.srt.layers.moe.topk", "biased_grouped_topk_gpu"), - # ── Layer norm ── - # The actual module-level names are rmsnorm / fused_add_rmsnorm - # (imported from sgl_kernel on CUDA, or aiter on HIP). + # ── Layer norm (rmsnorm / fused_add_rmsnorm from sgl_kernel or aiter) ── ("sglang.srt.layers.layernorm", "rmsnorm"), ("sglang.srt.layers.layernorm", "fused_add_rmsnorm"), ("sglang.srt.layers.layernorm", "gemma_rmsnorm"), @@ -203,36 +114,21 @@ def _get_or_create_lib() -> Library: # ── FP8 quantization ── ("sglang.srt.layers.quantization.fp8_utils", "per_token_group_quant_fp8"), ("sglang.srt.layers.quantization.fp8_utils", "scaled_fp8_quant"), - # Inner kernel called by triton_w8a8_block_fp8_linear via global lookup: + # inner kernels looked up from the module __dict__ on every call ("sglang.srt.layers.quantization.fp8_utils", "w8a8_block_fp8_matmul_triton"), - # Inner kernel called by aiter_w8a8_block_fp8_linear via global lookup - # (the outer wrapper is captured by reference at model init, but this - # inner kernel is looked up from the module __dict__ on every call): ("sglang.srt.layers.quantization.fp8_utils", "gemm_a8w8_blockscale"), - # ── LoRA Triton (inner kernel functions, no *args/**kwargs) ── + # ── LoRA Triton ── ("sglang.srt.lora.triton_ops.sgemm_lora_a", "sgemm_lora_a_fwd"), ("sglang.srt.lora.triton_ops.sgemm_lora_b", "sgemm_lora_b_fwd"), - # ── aiter (AMD) ops — definition-site patching ── + # ── aiter (AMD) ops ── ("aiter.ops.triton.gemm_a8w8_blockscale", "gemm_a8w8_blockscale"), ("aiter.ops.triton.batched_gemm_a8w8_blockscale", "batched_gemm_a8w8_blockscale"), ("aiter.ops.norm", "rms_norm"), ("aiter.ops.norm", "fused_add_rms_norm"), # ── FlashInfer MoE (cutedsl) ── ("flashinfer.moe", "moe_gemm_fp8_nt_groupwise"), - # ───────────────────────────────────────────────────────────────────── - # Current-layout paths (SGLang 0.5.18+ / matching aiter). - # - # SGLang moved its Triton kernels out of ``sglang.srt.layers.*`` into a - # separate ``sglang.kernels.ops.*`` package, and aiter regrouped its - # Triton ops into subpackages. Measured against - # sglang 0.5.18 + aiter, only 5 of the 24 legacy entries above still - # resolve, so the current paths are listed here as well. - # - # Both sets are kept on purpose: ``_resolve_target`` returns None for a - # path that does not exist and the entry is skipped silently, so one - # registry works across framework versions. Duplicates are harmless -- - # ``_wrapped_ids`` in enable() prevents double-wrapping the same object. - # ───────────────────────────────────────────────────────────────────── + # Current-layout paths (SGLang 0.5.18+ moved kernels to sglang.kernels.ops.*; + # aiter regrouped its ops). Unresolved legacy paths above are skipped silently. # ── SGLang Triton attention ── ("sglang.kernels.ops.attention.decode_attention", "decode_attention_fwd"), ("sglang.kernels.ops.attention.decode_attention", "decode_attention_fwd_normal"), @@ -254,37 +150,22 @@ def _get_or_create_lib() -> Library: # ── aiter regrouped Triton ops ── ("aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale", "gemm_a8w8_blockscale"), ("aiter.ops.triton.normalization.rmsnorm", "rms_norm"), - # aiter's batched blockscale GEMM was renamed, not just moved (it is now - # batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant - # under aiter.ops.triton.gemm.batched). Auto-discovery picks it up, so it - # is deliberately not pinned to a name here. + # aiter's batched blockscale GEMM was renamed (not just moved); left to + # auto-discovery rather than pinned to its long current name here. ] -# --------------------------------------------------------------------------- -# Auto-discovery prefixes. -# -# In addition to the explicit registry above, ``enable()`` scans every -# already-loaded module whose name starts with one of these prefixes and -# wraps only functions likely to launch kernels. Discovery is filtered by -# signature/source heuristics to avoid wrapping unrelated utility code. -# --------------------------------------------------------------------------- +# Auto-discovery prefixes: enable() scans loaded modules under these and wraps +# functions that look like kernel launchers (signature/source heuristics). _AUTO_DISCOVER_PREFIXES: Tuple[str, ...] = ( "flashinfer.", - # SGLang 0.5.18+ moved its Triton kernels out of sglang.srt.layers.* into - # this package, which is where most launchers now live. sglang.srt. is kept - # because quantization / MoE-runner launchers still sit there (and for - # older versions). "sglang.kernels.ops.", "sglang.srt.", "aiter.ops.", ) -# --------------------------------------------------------------------------- -# Schema building — works with or without type annotations -# --------------------------------------------------------------------------- - -# Python type → torch schema type +# Schema building — works with or without type annotations. +# Python type → torch schema type. _TYPE_MAP = { torch.Tensor: "Tensor", Optional[torch.Tensor]: "Tensor?", @@ -295,8 +176,7 @@ def _get_or_create_lib() -> Library: torch.dtype: "ScalarType", } -# String annotation variants produced by ``from __future__ import annotations`` -# (PEP 563) – annotations are stored as literal strings in the source code. +# Same, for PEP 563 string annotations. _STRING_TYPE_MAP = { "torch.Tensor": "Tensor", "Tensor": "Tensor", @@ -316,11 +196,10 @@ def _infer_schema_type(param: inspect.Parameter) -> Optional[str]: if annotation is inspect._empty: return None - # ── Handle string annotations (PEP 563) ── + # String annotations (PEP 563) if isinstance(annotation, str): if annotation in _STRING_TYPE_MAP: return _STRING_TYPE_MAP[annotation] - # Check "Optional[X]" pattern in string form if annotation.startswith("Optional[") and annotation.endswith("]"): inner = annotation[len("Optional[") : -1] base = _STRING_TYPE_MAP.get(inner) @@ -328,11 +207,9 @@ def _infer_schema_type(param: inspect.Parameter) -> Optional[str]: return base if base.endswith("?") else base + "?" return None - # ── Handle real type annotations ── - # Check direct match + # Real type annotations: direct match, then Optional[X] / Union[X, None] if annotation in _TYPE_MAP: return _TYPE_MAP[annotation] - # Check Optional[X] (Union[X, None]) origin = getattr(annotation, "__origin__", None) if origin is type(None): return None @@ -348,10 +225,10 @@ def _build_schema_from_sig( sig: inspect.Signature, skip_self: bool = False, ) -> Optional[Tuple[str, List[str], List[str]]]: - """ - Build schema string from signature annotations. - Returns (schema_str, tensor_param_names, non_tensor_param_names) or None - if there are no tensor params or the signature can't be mapped. + """Build a schema from signature annotations. + + Returns ``(schema_str, tensor_params, non_tensor_params)``, or ``None`` if + there are no tensor params or the signature can't be mapped. """ tensor_params: List[str] = [] non_tensor_params: List[str] = [] @@ -377,9 +254,7 @@ def _build_schema_from_sig( return schema_str, tensor_params, non_tensor_params -# --------------------------------------------------------------------------- -# Thread-local side channel for non-tensor args -# --------------------------------------------------------------------------- +# Thread-local side channel for non-tensor args and return values. _tls = threading.local() @@ -407,11 +282,6 @@ def _pop_return_value(op_name: str) -> Any: return _tls.returns.pop(op_name, None) -# --------------------------------------------------------------------------- -# Registration -# --------------------------------------------------------------------------- - - def _next_op_name(base: str) -> str: global _op_counter sanitized = base.replace(".", "_").replace("::", "_").replace("-", "_") @@ -429,10 +299,7 @@ def _register_op( sig: inspect.Signature, skip_self: bool = False, ) -> Optional[Callable]: - """ - Register a function as a torch custom op and return a dispatch wrapper. - Returns None if registration fails. - """ + """Register a function as a torch custom op and return a dispatch wrapper, or None on failure.""" try: lib = _get_or_create_lib() lib.define(op_name + schema_str) @@ -453,8 +320,7 @@ def impl(*tensor_args): elif param.default is not inspect._empty: full_kwargs[pname] = param.default result = original_fn(**full_kwargs) - # Schema is -> () so we can't return the actual value - # through the dispatcher. Stash it for the caller. + # Schema is -> () so stash the real return for the caller. _stash_return_value(op_name, result) lib.impl(op_name, impl, dispatch_key="CompositeExplicitAutograd") @@ -463,11 +329,7 @@ def impl(*tensor_args): @functools.wraps(original_fn) def dispatch_wrapper(*args, **kwargs): - # When profiling is not active, never route through the custom op. - # A wrapper reference may outlive disable() (captured by a module - # imported lazily while profiling was on, so _patch_all_references - # could not restore it). Falling back to the original keeps such - # leaked bindings correct and crash-free. + # A leaked reference must be a no-op when profiling is inactive. if not _enabled: return original_fn(*args, **kwargs) try: @@ -486,9 +348,7 @@ def dispatch_wrapper(*args, **kwargs): elif pname in non_tensor_param_names: nt_vals[pname] = val - # If every tensor arg is None (all Optional[Tensor] and not - # provided), torch dispatch will fail with "no tensor arguments". - # Fall back to calling the original function directly. + # torch dispatch fails with "no tensor arguments" if all are None. if not any(isinstance(t, torch.Tensor) for t in tensor_args): return original_fn(*args, **kwargs) @@ -497,15 +357,12 @@ def dispatch_wrapper(*args, **kwargs): torch_op(*tensor_args) return _pop_return_value(op_name) except Exception: - # Dispatch failed (e.g. device/type mismatch, None for - # non-optional Tensor, schema arity error). Clean up - # thread-local stash and fall back to the original call. + # Dispatch failed: clear the stash and fall back to the original. _pop_non_tensor_args(op_name) _pop_return_value(op_name) return original_fn(*args, **kwargs) - # Mark so enable() can recognise its own wrappers and never wrap one - # again (see the guard in enable()). + # Lets enable() recognise its own wrappers and never re-wrap one. dispatch_wrapper._kernel_shape_wrapper = True return dispatch_wrapper @@ -514,15 +371,8 @@ def dispatch_wrapper(*args, **kwargs): return None -# --------------------------------------------------------------------------- -# Module + attribute resolution -# --------------------------------------------------------------------------- - - def _resolve_target(module_path: str, attr_name: str): - """ - Resolve a target function from *module_path* and *attr_name*. - *attr_name* can be ``"func_name"`` or ``"ClassName.method_name"``. + """Resolve *attr_name* (``"func"`` or ``"Class.method"``) in *module_path*. Returns ``(container, attr_name, original_fn, is_method)`` or ``None``. """ @@ -548,16 +398,10 @@ def _resolve_target(module_path: str, attr_name: str): def _patch_all_references(original_fn: Callable, wrapper_fn: Callable): - """ - Scan ``sys.modules`` and replace **every** module-level attribute that - points to *original_fn* with *wrapper_fn*. + """Rebind every ``sys.modules`` reference to *original_fn* to *wrapper_fn*. - This handles the common ``from X import Y`` pattern: if module A - defines ``Y`` and module B does ``from A import Y``, both A and B - will have their binding replaced. - - Returns a list of ``(module, attr_name, original_fn)`` for later - restoration. + Handles the ``from X import Y`` pattern. Returns ``(module, attr, original_fn)`` + tuples for later restoration. """ patches = [] for _mod_name, mod in list(sys.modules.items()): @@ -579,27 +423,18 @@ def _patch_all_references(original_fn: Callable, wrapper_fn: Callable): return patches -# --------------------------------------------------------------------------- -# Lightweight wrappers for functions that can't use torch.library -# --------------------------------------------------------------------------- - - def _make_record_function_wrapper( qualified_name: str, original_fn: Callable, ) -> Callable: - """ - Create a wrapper that uses ``torch.profiler.record_function`` to emit - a ``cpu_op`` event with tensor shapes embedded in the event name. + """Fallback wrapper: emit a ``record_function`` event with shapes in the name. - Used for functions where we can't build a ``torch.library`` schema - (e.g. no type annotations, ``*args``/``**kwargs``, etc.). + Used when a ``torch.library`` schema can't be built (no annotations, ``*args``). """ @functools.wraps(original_fn) def wrapper(*args, **kwargs): - # See dispatch_wrapper: a leaked reference must be a no-op when - # profiling is inactive. + # A leaked reference must be a no-op when profiling is inactive. if not _enabled: return original_fn(*args, **kwargs) shape_parts: List[str] = [] @@ -620,10 +455,7 @@ def wrapper(*args, **kwargs): return wrapper -# --------------------------------------------------------------------------- -# Kernel-launch detection heuristics -# --------------------------------------------------------------------------- - +# Kernel-launch detection heuristics. # Substrings that strongly indicate a function launches a GPU kernel. _KERNEL_SOURCE_INDICATORS = ( "[grid", # Triton launch pattern: kernel[grid](...) @@ -642,13 +474,10 @@ def _source_launches_kernel(fn: Callable) -> bool: def _is_likely_kernel_launcher(fn: Callable, sig: inspect.Signature) -> bool: - """ - Decide whether *fn* is likely to launch a GPU kernel. + """Decide whether *fn* likely launches a GPU kernel. - Priority: - 1) Tensor annotation exists -> include. - 2) Non-tensor annotations only -> exclude. - 3) No annotations -> fallback to source pattern matching. + A Tensor annotation includes it; non-tensor-only annotations exclude it; no + annotations falls back to source pattern matching. """ has_any_annotation = False for param in sig.parameters.values(): @@ -666,12 +495,9 @@ def _is_likely_kernel_launcher(fn: Callable, sig: inspect.Signature) -> bool: def _force_import_submodules(prefix: str) -> None: - """ - Recursively import submodules under *prefix* so they appear in - ``sys.modules`` before auto-discovery runs. + """Recursively import submodules under *prefix* into ``sys.modules``. - *prefix* should be a top-level package name without a trailing dot - (e.g. ``"sglang.srt"``). + *prefix* is a package name without a trailing dot (e.g. ``"sglang.srt"``). """ try: pkg = importlib.import_module(prefix) @@ -682,12 +508,8 @@ def _force_import_submodules(prefix: str) -> None: if pkg_path is None: return - # Snapshot the global defaults once; restore them after *every* import so - # a module that calls torch.set_default_device("cuda") at import time - # cannot taint subsequently imported modules during discovery (nor leak - # into the serving path). The leaf-name skip list below catches the known - # offenders, but this restore makes the discovery robust to any other - # module with the same import-time side effect. + # Restore global defaults after every import so a module that mutates them + # can't taint later imports or the serving path. saved_device = _active_default_device_override() saved_dtype = torch.get_default_dtype() @@ -706,11 +528,8 @@ def _restore_defaults(): ): if mod_name in sys.modules: continue - # Skip test / benchmark / autotune modules. These are not kernel - # entry points and several of them (e.g. aiter.ops.flydsl.test_*) - # call ``torch.set_default_device("cuda")`` at import time, which - # leaks a CUDA default-device mode into the importing process and - # corrupts CPU tensor creation downstream (e.g. seq_lens_cpu). + # Skip test / benchmark / autotune modules: not entry points, and some + # set the default device at import. leaf = mod_name.rsplit(".", 1)[-1] if ( "test" in leaf @@ -730,15 +549,11 @@ def _restore_defaults(): def _discover_kernel_entry_points() -> List[Tuple[str, str]]: - """ - Scan already-loaded ``sys.modules`` for modules whose name starts - with one of ``_AUTO_DISCOVER_PREFIXES`` and collect only functions - likely to launch GPU kernels. + """Scan ``sys.modules`` under ``_AUTO_DISCOVER_PREFIXES`` for likely kernel launchers. Returns a list of ``(module_path, function_name)`` pairs. """ - # Force-import submodules under each discovery prefix first so deeper - # kernels become visible in sys.modules. + # Force-import submodules first so deeper kernels appear in sys.modules. for prefix in _AUTO_DISCOVER_PREFIXES: _force_import_submodules(prefix.rstrip(".")) @@ -759,15 +574,12 @@ def _discover_kernel_entry_points() -> List[Tuple[str, str]]: continue obj = mod_dict[attr_name] - # Only wrap regular Python functions. - # @triton.jit objects (JITFunction / Autotuner) must NOT be - # wrapped — replacing them in module globals breaks the - # Triton compiler's global resolution for device-side calls - # between JIT kernels (e.g. remap_xcd, _rmsmorm_op, etc.). + # Only plain Python functions. @triton.jit objects must NOT be wrapped: + # rebinding them breaks Triton's device-side global resolution. if not inspect.isfunction(obj): continue - # Only include functions *defined* within a target namespace + # Only functions defined within a target namespace. fn_module = getattr(obj, "__module__", "") or "" if not any(fn_module.startswith(p) for p in _AUTO_DISCOVER_PREFIXES): continue @@ -777,7 +589,6 @@ def _discover_kernel_entry_points() -> List[Tuple[str, str]]: continue seen_ids.add(obj_id) - # Require at least one parameter try: sig = inspect.signature(obj) except (ValueError, TypeError): @@ -785,14 +596,13 @@ def _discover_kernel_entry_points() -> List[Tuple[str, str]]: if not sig.parameters: continue - # Skip signatures we cannot map to a torch schema. + # Skip signatures we can't map to a torch schema. if any( p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD) for p in sig.parameters.values() ): continue - # Core filter: only keep likely kernel launchers. if not _is_likely_kernel_launcher(obj, sig): continue @@ -806,11 +616,6 @@ def _discover_kernel_entry_points() -> List[Tuple[str, str]]: return results -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - - def enable(): """Patch registered kernel entry points to appear as cpu_op.""" global _enabled @@ -818,22 +623,12 @@ def enable(): if _enabled: return - # NOTE: the Library and op counter are process-persistent (see the - # _lib / _op_counter docs). We do NOT recreate or reset them here so - # ops registered in earlier cycles stay valid. Only the per-cycle - # reference patches are rebuilt. + # _lib / _op_counter are process-persistent; only per-cycle patches rebuild. _patches.clear() - _wrapped_ids: set = set() # track function ids to avoid double-wrapping - - # Backstop guard around the import-heavy region. Both the explicit - # registry (_resolve_target -> importlib.import_module) and the - # auto-discovery (_discover_kernel_entry_points -> _force_import_*) - # import modules that may mutate process-global torch defaults at - # import time; restore them on exit so the leak cannot escape into - # the serving path. (No-op when nothing mutates the defaults.) + _wrapped_ids: set = set() # function ids, to avoid double-wrapping + + # Import-heavy region may mutate global torch defaults; restore on exit. with _preserve_global_torch_state(): - # Merge explicit registry with filtered auto-discovered functions. - # Duplicates are harmless — _wrapped_ids prevents double-wrapping. all_entry_points = ( list(_KERNEL_ENTRY_POINTS) + _discover_kernel_entry_points() ) @@ -846,20 +641,14 @@ def enable(): container, name, original_fn, is_method = resolved is_plain_function = not is_method - # Already one of our own wrappers. This happens when a - # framework keeps a compat re-export and both the old and new - # module paths are in the registry: the first entry wraps the - # function and _patch_all_references rebinds every reference to - # it, so the second entry now resolves to the wrapper. Wrapping - # again would nest two annotations around a single call and - # double-count it. + # Already our own wrapper (a compat re-export resolved to it). + # Re-wrapping would nest annotations and double-count the call. if getattr(original_fn, "_kernel_shape_wrapper", False): logger.debug( "Skipping already-wrapped %s.%s", module_path, attr_name ) continue - # Skip if this exact function object was already wrapped fn_id = id(original_fn) if fn_id in _wrapped_ids: logger.debug("Skipping duplicate %s.%s", module_path, attr_name) @@ -868,18 +657,14 @@ def enable(): qualified_name = f"{module_path}.{name}" - # Reuse a wrapper built in a previous cycle if the underlying - # function object is unchanged. This keeps each op defined - # exactly once and ensures a given function maps to a stable - # wrapper, so a reference that leaked across cycles still - # targets a live op. + # Reuse a prior-cycle wrapper if the function is unchanged, so + # each op is defined once and leaked refs stay live. wrapper = None cached = _built_wrappers.get(qualified_name) if cached is not None and cached[1] is original_fn: wrapper = cached[0] if wrapper is None: - # ── Regular functions ── try: sig = inspect.signature(original_fn) except (ValueError, TypeError): @@ -893,7 +678,7 @@ def enable(): schema_info = _build_schema_from_sig(sig, skip_self=is_method) if schema_info is not None: - # Full tensor annotations → use torch.library custom op + # Full tensor annotations → torch.library custom op schema_str, t_names, nt_names = schema_info op_name = _next_op_name(base) wrapper = _register_op( @@ -911,8 +696,7 @@ def enable(): ) if wrapper is None: - # Fallback: no annotations or schema registration failed - # → use record_function wrapper (shapes in event name) + # No annotations / registration failed → record_function wrapper = _make_record_function_wrapper( qualified_name, original_fn, @@ -953,9 +737,7 @@ def disable(): with _lock: if not _enabled: return - # Flip the flag first so any wrapper invoked concurrently — or one that - # leaked past restoration — short-circuits to the original instead of - # dispatching into a custom op. + # Flip the flag first so any leaked wrapper short-circuits to the original. _enabled = False for container, name, original_fn in reversed(_patches): try: @@ -963,9 +745,7 @@ def disable(): except Exception: pass _patches.clear() - # Intentionally keep _lib, _op_counter and _built_wrappers alive: the - # registered ops must outlive any wrapper reference that may have - # leaked (see module-level _lib docs). + # Keep _lib / _op_counter / _built_wrappers alive (see _lib docs). logger.info("kernel_shape_profiler disabled: all patches restored") diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py b/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py index c5617f521..7813d0fd3 100644 --- a/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py +++ b/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py @@ -1,32 +1,15 @@ -""" -Auto-loaded shim that turns on kernel-shape annotation without patching the -inference server (SGLang / vLLM / any torch workload). - -CPython imports a top-level module named ``sitecustomize`` automatically at -interpreter startup for every process, as long as its directory is on -``sys.path`` (i.e. on ``PYTHONPATH``). This shim uses that hook to drive -``kernel_shape_profiler.enable()`` / ``disable()`` from the torch-profiler -window, so the launcher wrapping costs nothing outside a profiling run. - -Activation (both required): - export PYTHONPATH=/path/to/kernel_shape_tool:$PYTHONPATH - export TRACELENS_SHAPE_DISCOVERY=1 - -Behaviour when ``TRACELENS_SHAPE_DISCOVERY`` is unset/false: this shim installs -nothing observable -- every hook short-circuits, so it is safe to leave the -directory on ``PYTHONPATH`` permanently. - -torch is usually not imported yet when ``sitecustomize`` runs, so patches are -registered as *pending* and applied by an ``__import__`` hook the moment the -target module appears in ``sys.modules``. All hooks are idempotent (guarded by a -sentinel attribute) and crash-proof (wrapped in ``try/except``) so they can -never break the workload. - -Note on timing: unlike a JIT-hook tracer, ``enable()`` here is *expensive* -- -it walks ``sglang.srt`` / ``aiter.ops`` / ``flashinfer`` with -``pkgutil.walk_packages``, then rebinds module-level references across all of -``sys.modules``. It therefore runs once, at the first profiler start, and the -first ``start()`` call is measurably slower than subsequent ones. +"""Auto-loaded shim that drives kernel-shape annotation without patching the server. + +CPython auto-imports ``sitecustomize`` at interpreter startup for any process +whose ``sys.path`` (``PYTHONPATH``) includes this directory. The shim wraps +``torch.profiler`` start/stop to call ``kernel_shape_profiler.enable()`` / +``disable()`` around each profiling window, so launcher wrapping costs nothing +outside a profiling run. Gated on ``TRACELENS_SHAPE_DISCOVERY``; when unset, +every hook short-circuits, so it is safe to leave on ``PYTHONPATH`` permanently. + +torch is usually not imported yet when this runs, so profiler patches are +registered as *pending* and applied by an ``__import__`` hook once the target +module loads. All hooks are idempotent and wrapped in ``try/except``. """ import builtins @@ -47,22 +30,17 @@ def _shape_discovery_on() -> bool: return _flag_on(_ENV_FLAG, "0") -# --------------------------------------------------------------------------- -# Lazy handle to the co-located profiler. -# -# We intentionally do NOT import kernel_shape_profiler (which imports torch) at -# sitecustomize time -- that would force a heavy torch import at interpreter -# startup and run torch's import-time side effects too early. Instead we import -# it lazily, once a hook actually needs it and torch is already loaded. -# --------------------------------------------------------------------------- +# Lazy handle to the co-located profiler. We do NOT import kernel_shape_profiler +# (which imports torch) at sitecustomize time -- that would force a heavy torch +# import at startup. Import it lazily once a hook needs it and torch is loaded. _profiler = None def _get_profiler(): global _profiler if _profiler is None: - # Make sure this file's own directory is importable even if only the - # parent ended up on sys.path. + # Ensure this file's directory is importable even if only the parent + # ended up on sys.path. here = os.path.dirname(os.path.abspath(__file__)) if here not in sys.path: sys.path.insert(0, here) @@ -94,12 +72,8 @@ def _disable_profiler() -> None: _profiler_active = [0] -# --------------------------------------------------------------------------- -# Patch: torch.profiler.profile.start / stop -# enable() runs just before the profiler starts recording so the first -# captured launch already carries shape metadata; stop() restores the -# original function references. -# --------------------------------------------------------------------------- +# Patch torch.profiler.profile.start/stop: enable() just before recording starts +# (so the first captured launch carries shapes), disable() on the last stop(). def _patch_torch_profiler_profile() -> bool: try: import torch.profiler as tp @@ -138,13 +112,8 @@ def _stop(self, *a, **kw): return True -# --------------------------------------------------------------------------- -# Patch: _KinetoProfile.__init__ -> force record_shapes=True -# The wrapped launchers only surface "Input Dims" / "Input type" in the trace -# when the profiler was constructed with record_shapes=True. Force it on (opt -# out with TRACELENS_SHAPE_FORCE_RECORD_SHAPES=0) so shape annotation "just -# works" regardless of how the server configured its profiler. -# --------------------------------------------------------------------------- +# Patch _KinetoProfile.__init__ to force record_shapes=True -- "Input Dims" only +# surface when shapes are recorded (opt out with TRACELENS_SHAPE_FORCE_RECORD_SHAPES=0). def _patch_kineto_record_shapes() -> bool: try: from torch.profiler.profiler import _KinetoProfile @@ -177,9 +146,7 @@ def _patch_torch_profiler_both() -> bool: return a and b -# --------------------------------------------------------------------------- -# Patch: torch.cuda.profiler.start / stop (legacy profiling API) -# --------------------------------------------------------------------------- +# Patch torch.cuda.profiler.start/stop (legacy profiling API). def _patch_torch_cuda_profiler() -> bool: try: import torch.cuda.profiler as tcp @@ -214,15 +181,8 @@ def _stop(*a, **kw): return True -# --------------------------------------------------------------------------- -# Pending-patch registry + import hook. -# -# Only the profiler entry points need patching here: the launcher wrapping -# itself is done by kernel_shape_profiler.enable(), which resolves its targets -# by importing them on demand. Kernels launched during CUDA-graph replay need -# no special handling: their shapes are recorded when the graph is *captured*, -# and capture that happens inside a profiler window is already covered. -# --------------------------------------------------------------------------- +# Pending-patch registry + import hook. Only profiler entry points are patched +# here; launcher wrapping is done by kernel_shape_profiler.enable() on demand. _PENDING_PATCHES = { "torch.profiler": _patch_torch_profiler_both, "torch.cuda.profiler": _patch_torch_cuda_profiler, @@ -267,10 +227,9 @@ def _wrapped(name, globals=None, locals=None, fromlist=(), level=0): def _bootstrap() -> None: - # Even when the flag is off we still install the (cheap) profiler patches: - # they all short-circuit via _shape_discovery_on(), and installing - # unconditionally keeps behaviour stable if the flag is toggled between - # fork/exec boundaries. + # Install the cheap profiler patches even when the flag is off: they all + # short-circuit via _shape_discovery_on(), keeping behaviour stable if the + # flag is toggled across fork/exec boundaries. if getattr(sys, "_tracelens_shape_bootstrapped", False): return sys._tracelens_shape_bootstrapped = True From 05127cb7fa82be3563ea298d828cf1566f652da8 Mon Sep 17 00:00:00 2001 From: mohammad abdul basit Date: Thu, 10 Sep 2026 15:26:15 +0000 Subject: [PATCH 05/10] move location --- .../TraceUtils}/kernel_shape_tool/README.md | 0 .../TraceUtils}/kernel_shape_tool/kernel_shape_profiler.py | 0 .../TraceUtils}/kernel_shape_tool/sitecustomize.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename {examples/custom_workflows/inference_analysis => TraceLens/TraceUtils}/kernel_shape_tool/README.md (100%) rename {examples/custom_workflows/inference_analysis => TraceLens/TraceUtils}/kernel_shape_tool/kernel_shape_profiler.py (100%) rename {examples/custom_workflows/inference_analysis => TraceLens/TraceUtils}/kernel_shape_tool/sitecustomize.py (100%) diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md b/TraceLens/TraceUtils/kernel_shape_tool/README.md similarity index 100% rename from examples/custom_workflows/inference_analysis/kernel_shape_tool/README.md rename to TraceLens/TraceUtils/kernel_shape_tool/README.md diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py similarity index 100% rename from examples/custom_workflows/inference_analysis/kernel_shape_tool/kernel_shape_profiler.py rename to TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py diff --git a/examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py b/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py similarity index 100% rename from examples/custom_workflows/inference_analysis/kernel_shape_tool/sitecustomize.py rename to TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py From 5c5b136a5614114d1cb14a56d3c2fb0848e3558e Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 12:02:57 -0500 Subject: [PATCH 06/10] Fix black formatting in kernel_shape_profiler Co-authored-by: Cursor --- .../TraceUtils/kernel_shape_tool/kernel_shape_profiler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py index 0aace2f2e..f7a697495 100644 --- a/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py +++ b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py @@ -68,7 +68,8 @@ def _preserve_global_torch_state(): _lib: Optional[Library] = None _op_counter = 0 # monotonic, never reset, so old op names stay valid _patches: List[Tuple[Any, str, Callable]] = [] # (module, attr, original_fn) -_built_wrappers: dict = {} # {qualified_name: (wrapper, original_fn)}, reused across cycles +# {qualified_name: (wrapper, original_fn)}, reused across cycles +_built_wrappers: dict = {} def _get_or_create_lib() -> Library: From 8db997376767fa56fc9827cdbcbe086b36b413e8 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 12:08:16 -0500 Subject: [PATCH 07/10] Use black suggestion for _built_wrappers formatting Co-authored-by: Cursor --- .../TraceUtils/kernel_shape_tool/kernel_shape_profiler.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py index f7a697495..01ad65620 100644 --- a/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py +++ b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py @@ -68,8 +68,9 @@ def _preserve_global_torch_state(): _lib: Optional[Library] = None _op_counter = 0 # monotonic, never reset, so old op names stay valid _patches: List[Tuple[Any, str, Callable]] = [] # (module, attr, original_fn) -# {qualified_name: (wrapper, original_fn)}, reused across cycles -_built_wrappers: dict = {} +_built_wrappers: dict = ( + {} +) # {qualified_name: (wrapper, original_fn)}, reused across cycles def _get_or_create_lib() -> Library: From 0b89eb0b46f80971fc2f0558bc1ff4168a02e24b Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 14:50:04 -0500 Subject: [PATCH 08/10] added tests --- .../TraceUtils/kernel_shape_tool/README.md | 6 + .../kernel_shape_profiler.py | 6 + .../kernel_shape_tool/sitecustomize.py | 6 + tests/test_kernel_shape_profiler.py | 627 ++++++++++++++++++ tests/test_kernel_shape_sitecustomize.py | 340 ++++++++++ 5 files changed, 985 insertions(+) create mode 100644 tests/test_kernel_shape_profiler.py create mode 100644 tests/test_kernel_shape_sitecustomize.py diff --git a/TraceLens/TraceUtils/kernel_shape_tool/README.md b/TraceLens/TraceUtils/kernel_shape_tool/README.md index 854e9ce1a..06c93b768 100644 --- a/TraceLens/TraceUtils/kernel_shape_tool/README.md +++ b/TraceLens/TraceUtils/kernel_shape_tool/README.md @@ -1,3 +1,9 @@ + + # Kernel shape profiler Adds `Input Dims` / `Input type` / `Input Strides` to PyTorch profiler traces diff --git a/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py index 01ad65620..f772d70df 100644 --- a/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py +++ b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + """Tensor shape metadata for Triton / FlashInfer / aiter kernels in profiler traces. Registered kernel launchers are wrapped as torch custom ops so they appear as diff --git a/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py b/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py index 7813d0fd3..72648efcf 100644 --- a/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py +++ b/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py @@ -1,3 +1,9 @@ +############################################################################### +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + """Auto-loaded shim that drives kernel-shape annotation without patching the server. CPython auto-imports ``sitecustomize`` at interpreter startup for any process diff --git a/tests/test_kernel_shape_profiler.py b/tests/test_kernel_shape_profiler.py new file mode 100644 index 000000000..60eea1612 --- /dev/null +++ b/tests/test_kernel_shape_profiler.py @@ -0,0 +1,627 @@ +############################################################################### +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for the no-patch kernel shape profiler. + +Covers the launcher-wrapping engine in +``TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py``: schema +inference, custom-op registration and dispatch, reference patching, and the +``enable()`` / ``disable()`` lifecycle. All tensors are CPU-only so the suite +runs in CPU CI (no GPU / sglang / aiter install required). +""" + +import inspect +import sys +import types +from pathlib import Path +from typing import Optional, Union + +import pytest +import torch + +# The tool is delivered on PYTHONPATH (see sitecustomize.py), so it is imported +# as a top-level module rather than through the TraceLens package. +_TOOL_DIR = ( + Path(__file__).parent.parent / "TraceLens" / "TraceUtils" / "kernel_shape_tool" +) + +# Serialise every test in this file onto a single xdist worker: the profiler +# keeps process-global state (a persistent Library, the ``_enabled`` flag), so +# the tests must not run concurrently in the same interpreter. +pytestmark = pytest.mark.xdist_group("kernel_shape_tool") + + +@pytest.fixture(scope="module") +def ksp(): + if str(_TOOL_DIR) not in sys.path: + sys.path.insert(0, str(_TOOL_DIR)) + import kernel_shape_profiler as _ksp + + return _ksp + + +@pytest.fixture(autouse=True) +def _disabled_after_each(ksp): + """Guarantee the global profiler is disabled between tests.""" + yield + if ksp.is_enabled(): + ksp.disable() + + +def _param(annotation): + return inspect.Parameter( + "p", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=annotation + ) + + +def _make_kernel_module(name): + """A fake module exposing kernel-launcher-like callables to wrap.""" + mod = types.ModuleType(name) + + def my_kernel(x: torch.Tensor, weight: torch.Tensor, alpha: float = 1.0): + return x + weight * alpha + + my_kernel.__module__ = name + + def my_norm(a, b): # no annotations -> record_function fallback path + return a + b + + my_norm.__module__ = name + + class MyLayer: + def forward(self, x: torch.Tensor) -> None: # method -> is_method path + return x * 2 + + MyLayer.__module__ = name + + mod.my_kernel = my_kernel + mod.my_norm = my_norm + mod.MyLayer = MyLayer + return mod + + +# --------------------------------------------------------------------------- +# Schema-type inference +# --------------------------------------------------------------------------- + + +class TestInferSchemaType: + def test_no_annotation_returns_none(self, ksp): + assert ksp._infer_schema_type(_param(inspect._empty)) is None + + def test_real_types(self, ksp): + assert ksp._infer_schema_type(_param(torch.Tensor)) == "Tensor" + assert ksp._infer_schema_type(_param(Optional[torch.Tensor])) == "Tensor?" + assert ksp._infer_schema_type(_param(int)) == "int" + assert ksp._infer_schema_type(_param(float)) == "float" + assert ksp._infer_schema_type(_param(bool)) == "bool" + assert ksp._infer_schema_type(_param(str)) == "str" + assert ksp._infer_schema_type(_param(torch.dtype)) == "ScalarType" + + def test_union_with_none(self, ksp): + assert ksp._infer_schema_type(_param(Union[int, None])) == "int?" + + def test_unknown_real_type_returns_none(self, ksp): + assert ksp._infer_schema_type(_param(list)) is None + + def test_string_annotations(self, ksp): + assert ksp._infer_schema_type(_param("torch.Tensor")) == "Tensor" + assert ksp._infer_schema_type(_param("Tensor")) == "Tensor" + assert ksp._infer_schema_type(_param("Optional[torch.Tensor]")) == "Tensor?" + assert ksp._infer_schema_type(_param("Optional[Tensor]")) == "Tensor?" + assert ksp._infer_schema_type(_param("int")) == "int" + assert ksp._infer_schema_type(_param("Optional[int]")) == "int?" + + def test_unknown_string_annotation_returns_none(self, ksp): + assert ksp._infer_schema_type(_param("SomeCustomType")) is None + assert ksp._infer_schema_type(_param("Optional[SomeCustomType]")) is None + + +# --------------------------------------------------------------------------- +# Schema building from a signature +# --------------------------------------------------------------------------- + + +class TestBuildSchemaFromSig: + def test_tensor_and_non_tensor_params(self, ksp): + def fn(x: torch.Tensor, weight: torch.Tensor, alpha: float): + return None + + schema, tensor_params, non_tensor = ksp._build_schema_from_sig( + inspect.signature(fn) + ) + # Only tensor params appear in the op schema; non-tensor args travel + # through the thread-local side channel instead. + assert schema == "(Tensor x, Tensor weight) -> ()" + assert tensor_params == ["x", "weight"] + assert non_tensor == ["alpha"] + + def test_no_tensor_params_returns_none(self, ksp): + def fn(a: int, b: float): + return None + + assert ksp._build_schema_from_sig(inspect.signature(fn)) is None + + def test_var_args_returns_none(self, ksp): + def fn(x: torch.Tensor, *args): + return None + + assert ksp._build_schema_from_sig(inspect.signature(fn)) is None + + def test_skip_self_for_methods(self, ksp): + def method(self, x: torch.Tensor): + return None + + schema, tensor_params, non_tensor = ksp._build_schema_from_sig( + inspect.signature(method), skip_self=True + ) + assert schema == "(Tensor x) -> ()" + assert tensor_params == ["x"] + assert non_tensor == [] + + +# --------------------------------------------------------------------------- +# Small helpers +# --------------------------------------------------------------------------- + + +class TestHelpers: + def test_next_op_name_sanitises_and_increments(self, ksp): + a = ksp._next_op_name("mod.sub::fn-name") + b = ksp._next_op_name("mod.sub::fn-name") + assert a.startswith("mod_sub_fn_name_") + assert a != b # monotonic counter + + def test_non_tensor_arg_stash_roundtrip(self, ksp): + ksp._stash_non_tensor_args("op1", {"alpha": 2.0}) + assert ksp._pop_non_tensor_args("op1") == {"alpha": 2.0} + # A second pop returns the empty default. + assert ksp._pop_non_tensor_args("op1") == {} + + def test_return_value_stash_roundtrip(self, ksp): + sentinel = object() + ksp._stash_return_value("op2", sentinel) + assert ksp._pop_return_value("op2") is sentinel + assert ksp._pop_return_value("op2") is None + + def test_pop_helpers_empty_in_fresh_thread(self, ksp): + import threading + + results = {} + + def worker(): + # A thread that never stashed sees empty thread-local defaults. + results["nt"] = ksp._pop_non_tensor_args("never") + results["ret"] = ksp._pop_return_value("never") + + thread = threading.Thread(target=worker) + thread.start() + thread.join() + assert results["nt"] == {} + assert results["ret"] is None + + def test_get_or_create_lib_is_cached(self, ksp): + lib = ksp._get_or_create_lib() + assert lib is ksp._get_or_create_lib() + + def test_register_op_returns_none_on_bad_schema(self, ksp): + def fn(x: torch.Tensor): + return x + + # A malformed schema string makes ``lib.define`` raise -> None. + result = ksp._register_op( + "bad_op", " -> nonsense (", fn, ["x"], [], inspect.signature(fn) + ) + assert result is None + + def test_active_default_device_override_default_none(self, ksp): + assert ksp._active_default_device_override() is None + + def test_active_default_device_override_reimports(self, ksp, monkeypatch): + # Force the ``sys.modules`` miss so the helper re-imports the device mod. + monkeypatch.delitem(sys.modules, "torch.utils._device", raising=False) + assert ksp._active_default_device_override() is None + + def test_preserve_global_torch_state_restores_dtype(self, ksp): + original = torch.get_default_dtype() + try: + with ksp._preserve_global_torch_state(): + torch.set_default_dtype(torch.float64) + assert torch.get_default_dtype() == torch.float64 + assert torch.get_default_dtype() == original + finally: + torch.set_default_dtype(original) + + def test_preserve_global_torch_state_swallows_restore_errors( + self, ksp, monkeypatch + ): + def boom(*args, **kwargs): + raise RuntimeError("cannot restore") + + monkeypatch.setattr(torch, "set_default_device", boom) + monkeypatch.setattr(torch, "set_default_dtype", boom) + # Restore failures on exit must be swallowed, not propagated. + with ksp._preserve_global_torch_state(): + pass + + +# --------------------------------------------------------------------------- +# Kernel-launcher heuristics +# --------------------------------------------------------------------------- + + +class TestLauncherHeuristics: + def test_source_launches_kernel_detects_marker(self, ksp): + def launcher(x): + return torch.ops.aten.relu(x) # 'torch.ops.' marker + + assert ksp._source_launches_kernel(launcher) is True + + def test_source_launches_kernel_no_marker(self, ksp): + def plain(x): + return x + 1 + + assert ksp._source_launches_kernel(plain) is False + + def test_source_launches_kernel_no_source(self, ksp): + # Builtins have no retrievable source. + assert ksp._source_launches_kernel(len) is False + + def test_tensor_annotation_is_launcher(self, ksp): + def fn(x: torch.Tensor): + return x + + assert ksp._is_likely_kernel_launcher(fn, inspect.signature(fn)) is True + + def test_non_tensor_annotation_excluded(self, ksp): + def fn(a: int, b: float): + return a + + assert ksp._is_likely_kernel_launcher(fn, inspect.signature(fn)) is False + + def test_unannotated_falls_back_to_source(self, ksp): + def fn(a, b): + return torch.ops.aten.add(a, b) + + assert ksp._is_likely_kernel_launcher(fn, inspect.signature(fn)) is True + + def plain(a, b): + return a + b + + assert ksp._is_likely_kernel_launcher(plain, inspect.signature(plain)) is False + + +# --------------------------------------------------------------------------- +# Target resolution & reference patching +# --------------------------------------------------------------------------- + + +class TestResolveTarget: + def test_resolves_module_function(self, ksp): + import math + + result = ksp._resolve_target("math", "sqrt") + assert result == (math, "sqrt", math.sqrt, False) + + def test_resolves_class_method(self, ksp, monkeypatch): + name = "fake_resolve_mod" + mod = _make_kernel_module(name) + monkeypatch.setitem(sys.modules, name, mod) + + container, attr, fn, is_method = ksp._resolve_target(name, "MyLayer.forward") + assert container is mod.MyLayer + assert attr == "forward" + assert fn is mod.MyLayer.forward + assert is_method is True + + def test_missing_module_returns_none(self, ksp): + assert ksp._resolve_target("no_such_module_xyz", "foo") is None + + def test_missing_attr_returns_none(self, ksp): + assert ksp._resolve_target("math", "definitely_not_here") is None + + def test_missing_class_returns_none(self, ksp): + assert ksp._resolve_target("math", "Nope.method") is None + + def test_missing_method_returns_none(self, ksp, monkeypatch): + name = "fake_resolve_mod2" + mod = _make_kernel_module(name) + monkeypatch.setitem(sys.modules, name, mod) + assert ksp._resolve_target(name, "MyLayer.no_such_method") is None + + +class TestPatchAllReferences: + def test_rebinds_every_reference(self, ksp, monkeypatch): + def original(): + return "orig" + + def wrapper(): + return "wrapped" + + mod_a = types.ModuleType("fake_ref_a") + mod_b = types.ModuleType("fake_ref_b") + mod_a.f = original + mod_b.g = original + monkeypatch.setitem(sys.modules, "fake_ref_a", mod_a) + monkeypatch.setitem(sys.modules, "fake_ref_b", mod_b) + + patches = ksp._patch_all_references(original, wrapper) + + assert mod_a.f is wrapper + assert mod_b.g is wrapper + # Restoration data is returned for later undo. + restored = {(m, a) for (m, a, _orig) in patches} + assert (mod_a, "f") in restored + assert (mod_b, "g") in restored + + def test_skips_none_and_non_dict_modules(self, ksp, monkeypatch): + def original(): + return "orig" + + def wrapper(): + return "wrapped" + + # ``None`` placeholders and non-module objects can live in sys.modules. + monkeypatch.setitem(sys.modules, "fake_none_mod", None) + monkeypatch.setitem(sys.modules, "fake_int_mod", 42) + # Must scan without raising despite the odd entries. + ksp._patch_all_references(original, wrapper) + + +# --------------------------------------------------------------------------- +# record_function fallback wrapper +# --------------------------------------------------------------------------- + + +class TestRecordFunctionWrapper: + def test_passthrough_when_disabled(self, ksp): + def original(x): + return x + 1 + + wrapper = ksp._make_record_function_wrapper("mod.fn", original) + assert getattr(wrapper, "_kernel_shape_wrapper", False) is True + # Profiler is disabled -> wrapper is a transparent passthrough. + assert wrapper(torch.zeros(2)).tolist() == [1.0, 1.0] + + def test_emits_event_when_enabled(self, ksp, monkeypatch): + calls = {} + + def original(x, scale): + calls["ran"] = True + return x + + wrapper = ksp._make_record_function_wrapper("mod.fn", original) + monkeypatch.setattr(ksp, "_enabled", True) + out = wrapper(torch.ones(3), scale=torch.ones(1)) + assert calls["ran"] is True + assert torch.allclose(out, torch.ones(3)) + + def test_event_without_tensor_args_when_enabled(self, ksp, monkeypatch): + def original(flag): + return flag + + wrapper = ksp._make_record_function_wrapper("mod.noargs", original) + monkeypatch.setattr(ksp, "_enabled", True) + assert wrapper(True) is True + + +# --------------------------------------------------------------------------- +# enable() / disable() lifecycle +# --------------------------------------------------------------------------- + + +class TestEnableDisable: + def _install(self, ksp, monkeypatch, entry_points): + name = "fake_kernel_mod" + mod = _make_kernel_module(name) + monkeypatch.setitem(sys.modules, name, mod) + monkeypatch.setattr(ksp, "_KERNEL_ENTRY_POINTS", entry_points(name)) + return name, mod + + def test_wrap_and_restore(self, ksp, monkeypatch): + name, mod = self._install( + ksp, + monkeypatch, + lambda n: [(n, "my_kernel"), (n, "my_norm"), (n, "MyLayer.forward")], + ) + orig_kernel = mod.my_kernel + orig_norm = mod.my_norm + orig_method = mod.MyLayer.forward + + ksp.enable() + assert ksp.is_enabled() is True + assert getattr(mod.my_kernel, "_kernel_shape_wrapper", False) is True + assert getattr(mod.my_norm, "_kernel_shape_wrapper", False) is True + assert getattr(mod.MyLayer.forward, "_kernel_shape_wrapper", False) is True + + x = torch.ones(2, 3) + w = torch.full((2, 3), 2.0) + + # Custom-op dispatch path returns the same value as the original. + assert torch.allclose(mod.my_kernel(x, w, alpha=1.5), x + w * 1.5) + # record_function fallback path. + assert torch.allclose(mod.my_norm(x, w), x + w) + # Wrapped bound method still returns the correct result. + assert torch.allclose(mod.MyLayer().forward(x), x * 2) + + ksp.disable() + assert ksp.is_enabled() is False + assert mod.my_kernel is orig_kernel + assert mod.my_norm is orig_norm + assert mod.MyLayer.forward is orig_method + + def test_dispatch_falls_back_on_bind_error(self, ksp, monkeypatch): + name, mod = self._install(ksp, monkeypatch, lambda n: [(n, "my_kernel")]) + ksp.enable() + wrapper = mod.my_kernel + x = torch.ones(2) + w = torch.ones(2) + # An unbindable call routes through the original, which then raises. + with pytest.raises(TypeError): + wrapper(x, w, not_a_real_kwarg=1) + + def test_dispatch_falls_back_when_no_tensor_args(self, ksp, monkeypatch): + name, mod = self._install(ksp, monkeypatch, lambda n: [(n, "my_kernel")]) + ksp.enable() + wrapper = mod.my_kernel + # All-None inputs cannot dispatch; the original is called instead. + with pytest.raises(TypeError): + wrapper(None, None) + + def test_leaked_wrapper_is_noop_after_disable(self, ksp, monkeypatch): + name, mod = self._install(ksp, monkeypatch, lambda n: [(n, "my_kernel")]) + ksp.enable() + leaked = mod.my_kernel # capture the live wrapper + ksp.disable() + + x = torch.ones(2) + w = torch.full((2,), 3.0) + # A reference that escaped disable() must fall straight through. + assert torch.allclose(leaked(x, w, alpha=1.0), x + w) + + def test_wrapper_reused_across_cycles(self, ksp, monkeypatch): + name, mod = self._install(ksp, monkeypatch, lambda n: [(n, "my_kernel")]) + ksp.enable() + first = mod.my_kernel + ksp.disable() + ksp.enable() + second = mod.my_kernel + # Same underlying function => cached wrapper is reused. + assert first is second + + def test_duplicate_entry_points_wrap_once(self, ksp, monkeypatch): + name, mod = self._install( + ksp, monkeypatch, lambda n: [(n, "my_kernel"), (n, "my_kernel")] + ) + ksp.enable() + assert getattr(mod.my_kernel, "_kernel_shape_wrapper", False) is True + + def test_enable_and_disable_are_idempotent(self, ksp, monkeypatch): + name, mod = self._install(ksp, monkeypatch, lambda n: [(n, "my_kernel")]) + ksp.enable() + ksp.enable() # second call is a no-op + assert ksp.is_enabled() is True + ksp.disable() + ksp.disable() # second call is a no-op + assert ksp.is_enabled() is False + + def test_unresolvable_entry_points_are_skipped(self, ksp, monkeypatch): + monkeypatch.setattr( + ksp, + "_KERNEL_ENTRY_POINTS", + [("module_that_does_not_exist", "nope")], + ) + ksp.enable() + assert ksp.is_enabled() is True # enables cleanly with nothing patched + + def test_already_wrapped_entry_is_skipped(self, ksp, monkeypatch): + name = "fake_prewrapped_mod" + mod = types.ModuleType(name) + + def already(x: torch.Tensor): + return x + + already._kernel_shape_wrapper = True + already.__module__ = name + mod.already = already + monkeypatch.setitem(sys.modules, name, mod) + monkeypatch.setattr(ksp, "_KERNEL_ENTRY_POINTS", [(name, "already")]) + + ksp.enable() + assert ksp.is_enabled() is True + # A target that is already our wrapper is left untouched (no re-wrap). + assert mod.already is already + + def test_enable_setattr_fallback_without_references(self, ksp, monkeypatch): + name, mod = self._install(ksp, monkeypatch, lambda n: [(n, "my_kernel")]) + orig = mod.my_kernel + # Simulate a launcher with no discoverable ``from x import y`` refs so + # enable() falls back to a direct attribute rebind. + monkeypatch.setattr(ksp, "_patch_all_references", lambda _o, _w: []) + + ksp.enable() + assert getattr(mod.my_kernel, "_kernel_shape_wrapper", False) is True + ksp.disable() + assert mod.my_kernel is orig + + def test_uninspectable_entry_is_skipped(self, ksp, monkeypatch): + name = "fake_uninspectable_mod" + mod = types.ModuleType(name) + mod.weird = object() # has no introspectable signature + monkeypatch.setitem(sys.modules, name, mod) + monkeypatch.setattr(ksp, "_KERNEL_ENTRY_POINTS", [(name, "weird")]) + + ksp.enable() + assert ksp.is_enabled() is True + + +# --------------------------------------------------------------------------- +# Auto-discovery of kernel launchers under target namespaces +# --------------------------------------------------------------------------- + + +class TestAutoDiscovery: + def test_force_import_submodules_missing_pkg(self, ksp): + # Non-existent package: returns without raising. + ksp._force_import_submodules("no_such_pkg_abcxyz") + + def test_force_import_submodules_non_package(self, ksp): + # A plain module has no ``__path__`` to walk: returns without raising. + ksp._force_import_submodules("math") + + def test_force_import_submodules_walks_and_skips_tests( + self, ksp, tmp_path, monkeypatch + ): + pkg = tmp_path / "fakewalkpkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "kernel_ops.py").write_text("def go():\n return 1\n") + (pkg / "test_skip.py").write_text("raise RuntimeError('must not import')\n") + # A non-test module that fails to import must be swallowed. + (pkg / "badmod.py").write_text("raise ImportError('boom')\n") + monkeypatch.syspath_prepend(str(tmp_path)) + + ksp._force_import_submodules("fakewalkpkg") + # A second pass finds every submodule already imported and skips it. + ksp._force_import_submodules("fakewalkpkg") + + assert "fakewalkpkg.kernel_ops" in sys.modules + # ``test_``-prefixed leaves are filtered out before import. + assert "fakewalkpkg.test_skip" not in sys.modules + # Import failures are swallowed, leaving the module unloaded. + assert "fakewalkpkg.badmod" not in sys.modules + + def test_discover_finds_launchers_under_prefix(self, ksp, tmp_path, monkeypatch): + # Build a fake ``aiter.ops`` package (matches _AUTO_DISCOVER_PREFIXES), + # exercising every candidate-filtering branch of the discovery scan. + root = tmp_path / "disc" + ops = root / "aiter" / "ops" + ops.mkdir(parents=True) + (root / "aiter" / "__init__.py").write_text("") + (ops / "__init__.py").write_text("") + (ops / "mykernels.py").write_text( + "from json import dumps\n" # function from another module -> skipped + "X = 5\n" # non-function attribute -> skipped + "def my_launch(a, b):\n" + " return torch.ops.aten.add(a, b)\n" # source marker -> launcher + "alias = my_launch\n" # duplicate object id -> skipped + "def noparams():\n return 1\n" # no params -> skipped + "def varargs(*a):\n return a\n" # *args -> skipped + "def helper(a: int):\n return a\n" # non-tensor annotation -> skipped + ) + monkeypatch.syspath_prepend(str(root)) + for mod_name in ("aiter", "aiter.ops", "aiter.ops.mykernels"): + monkeypatch.delitem(sys.modules, mod_name, raising=False) + # Odd entries under the scanned prefix must be tolerated. + monkeypatch.setitem(sys.modules, "aiter.ops.none_entry", None) + monkeypatch.setitem(sys.modules, "aiter.ops.weird_entry", 42) + + discovered = ksp._discover_kernel_entry_points() + + assert ("aiter.ops.mykernels", "my_launch") in discovered + # The filtered-out candidates must not appear. + names = {attr for _mod, attr in discovered} + assert names.isdisjoint({"dumps", "noparams", "varargs", "helper", "alias"}) diff --git a/tests/test_kernel_shape_sitecustomize.py b/tests/test_kernel_shape_sitecustomize.py new file mode 100644 index 000000000..f007ab759 --- /dev/null +++ b/tests/test_kernel_shape_sitecustomize.py @@ -0,0 +1,340 @@ +############################################################################### +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for the auto-loaded kernel-shape ``sitecustomize`` shim. + +Covers ``TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py``: env-flag +gating, lazy profiler resolution, the ``torch.profiler`` start/stop wrapping and +``record_shapes`` forcing, and the pending-patch import hook. The module is +loaded under a private name (never the real ``sitecustomize``) and every global +it patches (``builtins.__import__``, ``torch.profiler`` internals) is captured +and restored so the shim cannot leak into the rest of the test session. +""" + +import builtins +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch + +_TOOL_DIR = ( + Path(__file__).parent.parent / "TraceLens" / "TraceUtils" / "kernel_shape_tool" +) +_SITE_PATH = _TOOL_DIR / "sitecustomize.py" + +# The shim mutates process-global torch state; keep its tests on one worker. +pytestmark = pytest.mark.xdist_group("kernel_shape_tool") + + +@pytest.fixture(scope="module") +def site(): + """Load the shim under a private name, restoring all globals afterwards.""" + if str(_TOOL_DIR) not in sys.path: + sys.path.insert(0, str(_TOOL_DIR)) + + import torch.cuda.profiler as tcp + import torch.profiler as tp + from torch.profiler.profiler import _KinetoProfile + + saved = { + "import": builtins.__import__, + "p_start": tp.profile.start, + "p_stop": tp.profile.stop, + "k_init": _KinetoProfile.__init__, + "c_start": getattr(tcp, "start", None), + "c_stop": getattr(tcp, "stop", None), + "bootstrapped": getattr(sys, "_tracelens_shape_bootstrapped", None), + "hook_flag": getattr(sys, "_tracelens_shape_import_hook", None), + } + + spec = importlib.util.spec_from_file_location( + "tracelens_site_under_test", _SITE_PATH + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) # runs _bootstrap() + + yield module + + # Restore every global the shim may have patched. + builtins.__import__ = saved["import"] + tp.profile.start = saved["p_start"] + tp.profile.stop = saved["p_stop"] + _KinetoProfile.__init__ = saved["k_init"] + if saved["c_start"] is not None: + tcp.start = saved["c_start"] + if saved["c_stop"] is not None: + tcp.stop = saved["c_stop"] + for obj, attr in ( + (tp.profile, "_tracelens_shape_patched"), + (_KinetoProfile, "_tracelens_record_shapes_patched"), + (tcp, "_tracelens_shape_patched"), + ): + if hasattr(obj, attr): + try: + delattr(obj, attr) + except (AttributeError, TypeError): + pass + + +@pytest.fixture(autouse=True) +def _clean_flag(monkeypatch): + """Default every test to shape-discovery disabled unless it opts in.""" + monkeypatch.delenv("TRACELENS_SHAPE_DISCOVERY", raising=False) + yield + + +# --------------------------------------------------------------------------- +# Env-flag parsing +# --------------------------------------------------------------------------- + + +class TestFlagParsing: + def test_flag_on_truthy_values(self, site, monkeypatch): + for value in ("1", "true", "TRUE", "yes", "on", " 1 "): + monkeypatch.setenv("SOME_FLAG", value) + assert site._flag_on("SOME_FLAG") is True + + def test_flag_on_falsy_values(self, site, monkeypatch): + for value in ("0", "false", "no", "off", ""): + monkeypatch.setenv("SOME_FLAG", value) + assert site._flag_on("SOME_FLAG") is False + + def test_flag_on_default(self, site, monkeypatch): + monkeypatch.delenv("SOME_FLAG", raising=False) + assert site._flag_on("SOME_FLAG", "0") is False + assert site._flag_on("SOME_FLAG", "1") is True + + def test_shape_discovery_flag(self, site, monkeypatch): + assert site._shape_discovery_on() is False + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + assert site._shape_discovery_on() is True + + +# --------------------------------------------------------------------------- +# Lazy profiler resolution and gating +# --------------------------------------------------------------------------- + + +class TestProfilerGating: + def test_get_profiler_returns_engine(self, site): + import kernel_shape_profiler as ksp + + assert site._get_profiler() is ksp + + def test_get_profiler_inserts_own_dir(self, site, monkeypatch): + # Force a cold resolve with the tool dir absent from sys.path so the + # lazy import re-inserts it. + monkeypatch.setattr(site, "_profiler", None) + tool_dir = str(_TOOL_DIR) + monkeypatch.setattr(sys, "path", [p for p in sys.path if p != tool_dir]) + prof = site._get_profiler() + assert hasattr(prof, "enable") + assert tool_dir in sys.path + + def test_enable_disable_gated_on_flag(self, site, monkeypatch): + import kernel_shape_profiler as ksp + + # Flag off -> enable is a no-op. + site._enable_profiler() + assert ksp.is_enabled() is False + + # Flag on -> enable/disable drive the real engine. + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + try: + site._enable_profiler() + assert ksp.is_enabled() is True + site._disable_profiler() + assert ksp.is_enabled() is False + finally: + if ksp.is_enabled(): + ksp.disable() + + def test_enable_swallows_exceptions(self, site, monkeypatch): + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + + def boom(): + raise RuntimeError("profiler import failed") + + monkeypatch.setattr(site, "_get_profiler", boom) + # Must not propagate. + site._enable_profiler() + site._disable_profiler() + + +# --------------------------------------------------------------------------- +# torch.profiler patching +# --------------------------------------------------------------------------- + + +class TestProfilerPatching: + def test_patches_are_idempotent(self, site): + # Two consecutive calls: the second exercises the already-patched guard. + assert site._patch_torch_profiler_profile() is True + assert site._patch_torch_profiler_profile() is True + assert site._patch_kineto_record_shapes() is True + assert site._patch_kineto_record_shapes() is True + assert site._patch_torch_profiler_both() is True + assert site._patch_torch_cuda_profiler() is True + + def test_record_shapes_forced_when_enabled(self, site, monkeypatch): + from torch.profiler.profiler import _KinetoProfile + + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + kp = _KinetoProfile(record_shapes=False) + assert kp.record_shapes is True + + def test_record_shapes_not_forced_when_disabled(self, site): + from torch.profiler.profiler import _KinetoProfile + + kp = _KinetoProfile(record_shapes=False) + assert kp.record_shapes is False + + def test_record_shapes_opt_out(self, site, monkeypatch): + from torch.profiler.profiler import _KinetoProfile + + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + monkeypatch.setenv("TRACELENS_SHAPE_FORCE_RECORD_SHAPES", "0") + kp = _KinetoProfile(record_shapes=False) + assert kp.record_shapes is False + + def test_cuda_profiler_wrappers_toggle_engine(self, site, monkeypatch): + import torch.cuda.profiler as tcp + + import kernel_shape_profiler as ksp + + # Substitute innocuous start/stop (the real ones need a CUDA device) and + # force a fresh patch over them so the wrappers can be exercised on CPU. + started, stopped = [], [] + monkeypatch.setattr(tcp, "start", lambda *a, **k: started.append(1)) + monkeypatch.setattr(tcp, "stop", lambda *a, **k: stopped.append(1)) + monkeypatch.setattr(tcp, "_tracelens_shape_patched", False, raising=False) + + assert site._patch_torch_cuda_profiler() is True + # A second call short-circuits via the already-patched guard. + assert site._patch_torch_cuda_profiler() is True + + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + try: + tcp.start() + assert started == [1] + assert ksp.is_enabled() is True + tcp.stop() + assert stopped == [1] + assert ksp.is_enabled() is False + finally: + if ksp.is_enabled(): + ksp.disable() + + def test_profiler_window_toggles_engine(self, site, monkeypatch): + import kernel_shape_profiler as ksp + + monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") + assert ksp.is_enabled() is False + try: + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU] + ): + # start() enabled the launcher-wrapping engine... + assert ksp.is_enabled() is True + # ...and stop() disabled it once the last window closed. + assert ksp.is_enabled() is False + finally: + if ksp.is_enabled(): + ksp.disable() + + +# --------------------------------------------------------------------------- +# Pending-patch registry and import hook +# --------------------------------------------------------------------------- + + +class TestImportHook: + def test_try_pending_applies_and_pops(self, site, monkeypatch): + calls = [] + monkeypatch.setitem( + site._PENDING_PATCHES, "os", lambda: (calls.append(1), True)[1] + ) + site._try_pending() + assert calls == [1] + assert "os" not in site._PENDING_PATCHES + + def test_try_pending_pops_on_exception(self, site, monkeypatch): + def boom(): + raise RuntimeError("patch failed") + + monkeypatch.setitem(site._PENDING_PATCHES, "sys", boom) + site._try_pending() + assert "sys" not in site._PENDING_PATCHES + + def test_try_pending_skips_none_fn(self, site, monkeypatch): + monkeypatch.setitem(site._PENDING_PATCHES, "os", None) + site._try_pending() + # A ``None`` patch fn is skipped without being popped. + assert "os" in site._PENDING_PATCHES + + def test_install_import_hook_processes_pending(self, site, monkeypatch): + orig_import = builtins.__import__ + monkeypatch.setattr(sys, "_tracelens_shape_import_hook", False, raising=False) + try: + calls = [] + monkeypatch.setitem( + site._PENDING_PATCHES, "os", lambda: (calls.append(1), True)[1] + ) + site._install_import_hook() + assert builtins.__import__ is not orig_import + # A subsequent import runs the wrapped __import__ -> _try_pending. + builtins.__import__("math") + assert calls == [1] + assert "os" not in site._PENDING_PATCHES + # A further import now short-circuits (no pending patches left). + builtins.__import__("math") + assert calls == [1] + # Re-installing is a no-op while the flag is set. + hooked = builtins.__import__ + site._install_import_hook() + assert builtins.__import__ is hooked + finally: + builtins.__import__ = orig_import + + def test_import_hook_is_reentrancy_safe(self, site, monkeypatch): + orig_import = builtins.__import__ + monkeypatch.setattr(sys, "_tracelens_shape_import_hook", False, raising=False) + try: + + def nested_patch(): + # Importing from inside the hook must be short-circuited by the + # reentrancy guard rather than recursing. + builtins.__import__("math") + return True + + monkeypatch.setitem(site._PENDING_PATCHES, "os", nested_patch) + site._install_import_hook() + builtins.__import__("math") + assert "os" not in site._PENDING_PATCHES + finally: + builtins.__import__ = orig_import + + def test_bootstrap_installs_import_hook_when_pending(self, site, monkeypatch): + orig_import = builtins.__import__ + monkeypatch.setattr(sys, "_tracelens_shape_bootstrapped", False, raising=False) + monkeypatch.setattr(sys, "_tracelens_shape_import_hook", False, raising=False) + try: + # A pending patch whose module is not yet imported keeps the registry + # non-empty, so _bootstrap installs the import hook. + monkeypatch.setitem( + site._PENDING_PATCHES, "module.not.imported.yet", lambda: True + ) + site._bootstrap() + assert builtins.__import__ is not orig_import + finally: + builtins.__import__ = orig_import + + def test_bootstrap_is_idempotent(self, site): + # The fixture already bootstrapped; a second call must return early. + assert getattr(sys, "_tracelens_shape_bootstrapped", False) is True + site._bootstrap() From ff879d11f957aab68ec3a738619f172a73756526 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 15:30:53 -0500 Subject: [PATCH 09/10] fix tests --- tests/test_kernel_shape_profiler.py | 30 +++++++++++++++--------- tests/test_kernel_shape_sitecustomize.py | 24 ++++++++++++++----- 2 files changed, 37 insertions(+), 17 deletions(-) diff --git a/tests/test_kernel_shape_profiler.py b/tests/test_kernel_shape_profiler.py index 60eea1612..d2c2623da 100644 --- a/tests/test_kernel_shape_profiler.py +++ b/tests/test_kernel_shape_profiler.py @@ -27,6 +27,14 @@ _TOOL_DIR = ( Path(__file__).parent.parent / "TraceLens" / "TraceUtils" / "kernel_shape_tool" ) +if str(_TOOL_DIR) not in sys.path: + sys.path.insert(0, str(_TOOL_DIR)) +import kernel_shape_profiler as _KSP # noqa: E402 + +# Real submodule walker; the autouse fixture stubs the module global so +# enable() never imports real kernel packages in CI. Tests needing the real +# walker call this captured reference. +_REAL_FORCE_IMPORT = _KSP._force_import_submodules # Serialise every test in this file onto a single xdist worker: the profiler # keeps process-global state (a persistent Library, the ``_enabled`` flag), so @@ -36,16 +44,13 @@ @pytest.fixture(scope="module") def ksp(): - if str(_TOOL_DIR) not in sys.path: - sys.path.insert(0, str(_TOOL_DIR)) - import kernel_shape_profiler as _ksp - - return _ksp + return _KSP @pytest.fixture(autouse=True) -def _disabled_after_each(ksp): - """Guarantee the global profiler is disabled between tests.""" +def _hermetic_and_disabled(ksp, monkeypatch): + """Stub the package walker so enable() imports nothing; disable afterwards.""" + monkeypatch.setattr(ksp, "_force_import_submodules", lambda _prefix: None) yield if ksp.is_enabled(): ksp.disable() @@ -566,11 +571,11 @@ def test_uninspectable_entry_is_skipped(self, ksp, monkeypatch): class TestAutoDiscovery: def test_force_import_submodules_missing_pkg(self, ksp): # Non-existent package: returns without raising. - ksp._force_import_submodules("no_such_pkg_abcxyz") + _REAL_FORCE_IMPORT("no_such_pkg_abcxyz") def test_force_import_submodules_non_package(self, ksp): # A plain module has no ``__path__`` to walk: returns without raising. - ksp._force_import_submodules("math") + _REAL_FORCE_IMPORT("math") def test_force_import_submodules_walks_and_skips_tests( self, ksp, tmp_path, monkeypatch @@ -584,9 +589,9 @@ def test_force_import_submodules_walks_and_skips_tests( (pkg / "badmod.py").write_text("raise ImportError('boom')\n") monkeypatch.syspath_prepend(str(tmp_path)) - ksp._force_import_submodules("fakewalkpkg") + _REAL_FORCE_IMPORT("fakewalkpkg") # A second pass finds every submodule already imported and skips it. - ksp._force_import_submodules("fakewalkpkg") + _REAL_FORCE_IMPORT("fakewalkpkg") assert "fakewalkpkg.kernel_ops" in sys.modules # ``test_``-prefixed leaves are filtered out before import. @@ -618,6 +623,9 @@ def test_discover_finds_launchers_under_prefix(self, ksp, tmp_path, monkeypatch) # Odd entries under the scanned prefix must be tolerated. monkeypatch.setitem(sys.modules, "aiter.ops.none_entry", None) monkeypatch.setitem(sys.modules, "aiter.ops.weird_entry", 42) + # Real walker, restricted to our fake package only. + monkeypatch.setattr(ksp, "_force_import_submodules", _REAL_FORCE_IMPORT) + monkeypatch.setattr(ksp, "_AUTO_DISCOVER_PREFIXES", ("aiter.ops",)) discovered = ksp._discover_kernel_entry_points() diff --git a/tests/test_kernel_shape_sitecustomize.py b/tests/test_kernel_shape_sitecustomize.py index f007ab759..46d414c58 100644 --- a/tests/test_kernel_shape_sitecustomize.py +++ b/tests/test_kernel_shape_sitecustomize.py @@ -88,6 +88,16 @@ def _clean_flag(monkeypatch): yield +@pytest.fixture +def hermetic_engine(monkeypatch): + """Neuter the engine so enable() is a pure state flip (imports nothing).""" + import kernel_shape_profiler as ksp + + monkeypatch.setattr(ksp, "_force_import_submodules", lambda _prefix: None) + monkeypatch.setattr(ksp, "_KERNEL_ENTRY_POINTS", []) + return ksp + + # --------------------------------------------------------------------------- # Env-flag parsing # --------------------------------------------------------------------------- @@ -136,8 +146,8 @@ def test_get_profiler_inserts_own_dir(self, site, monkeypatch): assert hasattr(prof, "enable") assert tool_dir in sys.path - def test_enable_disable_gated_on_flag(self, site, monkeypatch): - import kernel_shape_profiler as ksp + def test_enable_disable_gated_on_flag(self, site, monkeypatch, hermetic_engine): + ksp = hermetic_engine # Flag off -> enable is a no-op. site._enable_profiler() @@ -202,10 +212,12 @@ def test_record_shapes_opt_out(self, site, monkeypatch): kp = _KinetoProfile(record_shapes=False) assert kp.record_shapes is False - def test_cuda_profiler_wrappers_toggle_engine(self, site, monkeypatch): + def test_cuda_profiler_wrappers_toggle_engine( + self, site, monkeypatch, hermetic_engine + ): import torch.cuda.profiler as tcp - import kernel_shape_profiler as ksp + ksp = hermetic_engine # Substitute innocuous start/stop (the real ones need a CUDA device) and # force a fresh patch over them so the wrappers can be exercised on CPU. @@ -230,8 +242,8 @@ def test_cuda_profiler_wrappers_toggle_engine(self, site, monkeypatch): if ksp.is_enabled(): ksp.disable() - def test_profiler_window_toggles_engine(self, site, monkeypatch): - import kernel_shape_profiler as ksp + def test_profiler_window_toggles_engine(self, site, monkeypatch, hermetic_engine): + ksp = hermetic_engine monkeypatch.setenv("TRACELENS_SHAPE_DISCOVERY", "1") assert ksp.is_enabled() is False From d5fdb7de85d84d12f443c2a3178336cbe09205d5 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Tue, 15 Sep 2026 11:07:14 -0500 Subject: [PATCH 10/10] skip torch tests --- .coveragerc | 1 + tests/test_kernel_shape_profiler.py | 3 ++- tests/test_kernel_shape_sitecustomize.py | 3 ++- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.coveragerc b/.coveragerc index c4bb7d51c..0ae5f3a60 100644 --- a/.coveragerc +++ b/.coveragerc @@ -6,6 +6,7 @@ omit = TraceLens/PerfModel/benchmarking/* TraceLens/PerfModel/origami_helper.py TraceLens/PerfModel/run_perf_model.py + TraceLens/TraceUtils/kernel_shape_tool/* [paths] source = diff --git a/tests/test_kernel_shape_profiler.py b/tests/test_kernel_shape_profiler.py index d2c2623da..47113723e 100644 --- a/tests/test_kernel_shape_profiler.py +++ b/tests/test_kernel_shape_profiler.py @@ -20,7 +20,8 @@ from typing import Optional, Union import pytest -import torch + +torch = pytest.importorskip("torch") # The tool is delivered on PYTHONPATH (see sitecustomize.py), so it is imported # as a top-level module rather than through the TraceLens package. diff --git a/tests/test_kernel_shape_sitecustomize.py b/tests/test_kernel_shape_sitecustomize.py index 46d414c58..077e0802a 100644 --- a/tests/test_kernel_shape_sitecustomize.py +++ b/tests/test_kernel_shape_sitecustomize.py @@ -20,7 +20,8 @@ from pathlib import Path import pytest -import torch + +torch = pytest.importorskip("torch") _TOOL_DIR = ( Path(__file__).parent.parent / "TraceLens" / "TraceUtils" / "kernel_shape_tool"