diff --git a/.coveragerc b/.coveragerc index c4bb7d51..0ae5f3a6 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/TraceLens/TraceUtils/kernel_shape_tool/README.md b/TraceLens/TraceUtils/kernel_shape_tool/README.md new file mode 100644 index 00000000..06c93b76 --- /dev/null +++ b/TraceLens/TraceUtils/kernel_shape_tool/README.md @@ -0,0 +1,60 @@ + + +# Kernel shape profiler + +Adds `Input Dims` / `Input type` / `Input Strides` to PyTorch profiler traces +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** — 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: + +``` +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` | 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 + +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 | +| `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. | + +> 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/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py new file mode 100644 index 00000000..f772d70d --- /dev/null +++ b/TraceLens/TraceUtils/kernel_shape_tool/kernel_shape_profiler.py @@ -0,0 +1,761 @@ +############################################################################### +# 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 +``cpu_op`` events carrying ``Input Dims`` / ``Input type``. ``sitecustomize.py`` +drives ``enable()`` / ``disable()`` from the profiler window. See README.md. +""" + +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__) + + +def _active_default_device_override(): + """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: + 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 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() + try: + yield + finally: + try: + torch.set_default_device(saved_device) + except Exception: + pass + try: + torch.set_default_dtype(saved_dtype) + except Exception: + pass + + +_lock = threading.Lock() +_enabled = False +# 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 +_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: + """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: (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"), + ( + "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 (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"), + ("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 kernels looked up from the module __dict__ on every call + ("sglang.srt.layers.quantization.fp8_utils", "w8a8_block_fp8_matmul_triton"), + ("sglang.srt.layers.quantization.fp8_utils", "gemm_a8w8_blockscale"), + # ── 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 ── + ("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+ 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"), + ("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); left to + # auto-discovery rather than pinned to its long current name here. +] + +# 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.kernels.ops.", + "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", +} + +# Same, for PEP 563 string annotations. +_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 + + # String annotations (PEP 563) + if isinstance(annotation, str): + if annotation in _STRING_TYPE_MAP: + return _STRING_TYPE_MAP[annotation] + 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 + + # Real type annotations: direct match, then Optional[X] / Union[X, None] + if annotation in _TYPE_MAP: + return _TYPE_MAP[annotation] + 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 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] = [] + 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 and return values. +_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) + + +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, or None on failure.""" + 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 stash the real return 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): + # A leaked reference must be a no-op when profiling is inactive. + 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 + + # 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) + + _stash_non_tensor_args(op_name, nt_vals) + try: + torch_op(*tensor_args) + return _pop_return_value(op_name) + except Exception: + # 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) + + # Lets enable() recognise its own wrappers and never re-wrap one. + 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 + + +def _resolve_target(module_path: str, attr_name: str): + """Resolve *attr_name* (``"func"`` or ``"Class.method"``) in *module_path*. + + 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): + """Rebind every ``sys.modules`` reference to *original_fn* to *wrapper_fn*. + + 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()): + 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 + + +def _make_record_function_wrapper( + qualified_name: str, + original_fn: Callable, +) -> Callable: + """Fallback wrapper: emit a ``record_function`` event with shapes in the name. + + Used when a ``torch.library`` schema can't be built (no annotations, ``*args``). + """ + + @functools.wraps(original_fn) + def wrapper(*args, **kwargs): + # 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* likely launches a GPU kernel. + + 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(): + 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* into ``sys.modules``. + + *prefix* is a 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 + + # 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() + + def _restore_defaults(): + 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: not entry points, and some + # set the default device at import. + 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 ``sys.modules`` under ``_AUTO_DISCOVER_PREFIXES`` for likely kernel launchers. + + Returns a list of ``(module_path, function_name)`` pairs. + """ + # Force-import submodules first so deeper kernels appear 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 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 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) + + try: + sig = inspect.signature(obj) + except (ValueError, TypeError): + continue + if not sig.parameters: + continue + + # 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 + + 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 + + +def enable(): + """Patch registered kernel entry points to appear as cpu_op.""" + global _enabled + with _lock: + if _enabled: + return + + # _lib / _op_counter are process-persistent; only per-cycle patches rebuild. + _patches.clear() + _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(): + 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 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 + + 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 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: + 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 → 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: + # No annotations / registration failed → record_function + 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 leaked wrapper short-circuits to the original. + _enabled = False + for container, name, original_fn in reversed(_patches): + try: + setattr(container, name, original_fn) + except Exception: + pass + _patches.clear() + # Keep _lib / _op_counter / _built_wrappers alive (see _lib docs). + logger.info("kernel_shape_profiler disabled: all patches restored") + + +def is_enabled() -> bool: + return _enabled diff --git a/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py b/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py new file mode 100644 index 00000000..72648efc --- /dev/null +++ b/TraceLens/TraceUtils/kernel_shape_tool/sitecustomize.py @@ -0,0 +1,247 @@ +############################################################################### +# 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 +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 +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 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: + # 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) + 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() 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 + 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__ 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 + 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 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, +} + + +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: + # 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 + _try_pending() + if _PENDING_PATCHES: + _install_import_hook() + + +_bootstrap() diff --git a/tests/test_kernel_shape_profiler.py b/tests/test_kernel_shape_profiler.py new file mode 100644 index 00000000..47113723 --- /dev/null +++ b/tests/test_kernel_shape_profiler.py @@ -0,0 +1,636 @@ +############################################################################### +# 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 + +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. +_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 +# the tests must not run concurrently in the same interpreter. +pytestmark = pytest.mark.xdist_group("kernel_shape_tool") + + +@pytest.fixture(scope="module") +def ksp(): + return _KSP + + +@pytest.fixture(autouse=True) +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() + + +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. + _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. + _REAL_FORCE_IMPORT("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)) + + _REAL_FORCE_IMPORT("fakewalkpkg") + # A second pass finds every submodule already imported and skips it. + _REAL_FORCE_IMPORT("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) + # 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() + + 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 00000000..077e0802 --- /dev/null +++ b/tests/test_kernel_shape_sitecustomize.py @@ -0,0 +1,353 @@ +############################################################################### +# 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 + +torch = pytest.importorskip("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 + + +@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 +# --------------------------------------------------------------------------- + + +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, hermetic_engine): + ksp = hermetic_engine + + # 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, hermetic_engine + ): + import torch.cuda.profiler as tcp + + 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. + 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, hermetic_engine): + ksp = hermetic_engine + + 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()