From 8ab482c913155aa41fd2ce3273d226becc40f7d4 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sat, 28 Mar 2026 18:10:47 +0000 Subject: [PATCH 01/11] Expand EventReplay to support custom ops beyond aten:: Op resolution: - Add _resolve_op_func() with JIT-first resolution (preserves in-place kernel dispatch for aten ops) and torch.ops fallback for custom ops - Add _search_schemas() to collect schemas from both JIT registry and torch.ops namespace overloads Schemaless replay: - Add _get_event_replay_IR_schemaless() that infers argument types directly from profiled data when no schema is available, enabling replay of ops like _C::silu_and_mul and _C::rotary_embedding Type handling: - Fix Scalar type: preserve integer values for integral tensor ops instead of always casting to float - Add str/str?, SymInt?/int?, Generator? support in schema matching - Add _is_tensor_schema_type() for annotated variants like Tensor(a!) - Add _should_skip_tensor_init() generalizing in-place/output detection Dtype support (utils.py): - Add long int, unsigned char, char, short, and FP8 types - Use zeros init for non-floating-point tensors Tested with: - ResNet regression suite (70 aten:: ops) - vLLM Qwen1.5-MoE-A2.7B trace: 7/9 aiter ops, plus _rocm_C::wvSplitK, _C::silu_and_mul, _C::rotary_embedding (requires import vllm._C/_rocm_C) Made-with: Cursor --- TraceLens/EventReplay/event_replay.py | 373 +++++++++++++++++++------- TraceLens/EventReplay/utils.py | 34 ++- 2 files changed, 310 insertions(+), 97 deletions(-) diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index 9d79f3765..5473cabf8 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -20,6 +20,96 @@ ) +def _resolve_op_func(op_name: str): + """ + Resolve an op name (e.g. 'aten::mm', 'vllm::rocm_unquantized_gemm') to a + callable. Tries multiple resolution strategies and returns the first that + yields a non-None callable. + + Returns (func, source_str) or raises RuntimeError. + """ + torch = _get_torch_or_raise() + + # 1. JIT registry first — preserves dispatch behaviour that the original + # profiled run used (important for in-place aten ops like add_). + try: + func, _ = torch._C._jit_get_operation(op_name) + if func is not None: + return func, "jit" + except RuntimeError: + pass + + # 2. torch.ops namespace lookup (most reliable for custom ops that may + # not be registered in the JIT registry). + if "::" in op_name: + ns, func_name = op_name.split("::", 1) + ns_obj = getattr(torch.ops, ns, None) + if ns_obj is not None: + func_obj = getattr(ns_obj, func_name, None) + if callable(func_obj): + return func_obj, "torch.ops" + + raise RuntimeError( + f"Cannot resolve op '{op_name}'. Ensure the library that defines it " + f"is imported (e.g. 'import vllm', 'import aiter')." + ) + + +def _search_schemas(op_name: str, verbose: bool = False): + """ + Return all registered FunctionSchemas for *op_name*. + + Searches both the JIT schema registry and the torch.ops namespace, which + covers aten ops, custom C++ ops, and Python-defined torch.library ops. + """ + torch = _get_torch_or_raise() + schemas: list = [] + seen_strs: set = set() + + # JIT registry + for s in torch._C._jit_get_all_schemas(): + if s.name == op_name: + s_str = str(s) + if s_str not in seen_strs: + schemas.append(s) + seen_strs.add(s_str) + + # torch.ops namespace (catches custom ops not in the JIT list) + if "::" in op_name: + ns, func_name = op_name.split("::", 1) + ns_obj = getattr(torch.ops, ns, None) + if ns_obj is not None: + op_obj = getattr(ns_obj, func_name, None) + if op_obj is not None: + # OpOverloadPacket exposes overloads + try: + for overload_name in op_obj.overloads(): + overload = getattr(op_obj, overload_name) + s = overload._schema + s_str = str(s) + if s_str not in seen_strs: + schemas.append(s) + seen_strs.add(s_str) + except Exception: + # Fallback: try .default directly + try: + s = op_obj.default._schema + s_str = str(s) + if s_str not in seen_strs: + schemas.append(s) + seen_strs.add(s_str) + except Exception: + pass + + if verbose: + print(f"Found {len(schemas)} schemas for {op_name}:") + for s in schemas: + pprint(str(s)) + print("-" * 80) + + return schemas + + class EventReplayer: def __init__( self, @@ -40,6 +130,7 @@ def __init__( self.device = device self.lazy = lazy self.verbose = verbose + self._func = None self._setup() def _setup(self): @@ -48,10 +139,33 @@ def _setup(self): """ if self.verbose: print(f"Preparing {self.event['name']} event for replay") - self.matched_schema = EventReplayer._search_schema(self.event, self.verbose) - self.event_replay_IR = EventReplayer._get_event_replay_IR( - self.event, self.matched_schema, self.verbose - ) + + self._func, self._func_source = _resolve_op_func(self.event["name"]) + if self.verbose: + print(f"Resolved op via {self._func_source}") + + try: + self.matched_schema = EventReplayer._search_schema( + self.event, self.verbose + ) + self._schemaless = False + except ValueError: + if self.verbose: + print( + "No schema found; falling back to schemaless replay " + "(all args treated as positional, types inferred from profile)" + ) + self.matched_schema = None + self._schemaless = True + + if self._schemaless: + self.event_replay_IR = EventReplayer._get_event_replay_IR_schemaless( + self.event, self.verbose + ) + else: + self.event_replay_IR = EventReplayer._get_event_replay_IR( + self.event, self.matched_schema, self.verbose + ) if not self.lazy: if self.verbose: print("setting up args and kwargs") @@ -63,11 +177,6 @@ def replay(self): """ Replay the event using the matched schema and event replay IR. """ - torch = _get_torch_or_raise() - # Get the function from the schema - func, _ = torch._C._jit_get_operation(self.event["name"]) - - # Call the function with the arguments if self.lazy: args, kwargs = EventReplayer._get_args_kwargs( self.event_replay_IR, device=self.device @@ -75,22 +184,13 @@ def replay(self): else: args, kwargs = self.args, self.kwargs - # Call the function with the arguments - func(*args, **kwargs) + self._func(*args, **kwargs) @staticmethod def _search_schema( event: Dict[str, Any], verbose: bool = False ) -> Optional["torch._C.FunctionSchema"]: - torch = _get_torch_or_raise() - all_schemas = torch._C._jit_get_all_schemas() - op_schemas = [s for s in all_schemas if s.name == event["name"]] - # print each schema in separate line - if verbose: - print(f"Found {len(op_schemas)} schemas for {event['name']}:") - for schema in op_schemas: - pprint(str(schema)) - print("-" * 80) + op_schemas = _search_schemas(event["name"], verbose=verbose) for schema in op_schemas: if verbose: @@ -105,7 +205,9 @@ def _search_schema( print("-" * 80) raise ValueError( - f"Cannot find matching schema for {event['name']}. Please check the event data and schema." + f"Cannot find matching schema for {event['name']}. " + f"Searched {len(op_schemas)} candidate(s). " + f"Please check the event data and ensure the op's library is imported." ) @staticmethod @@ -114,22 +216,13 @@ def _is_schema_match( ) -> bool: """ Check if the event matches the schema. - - Args: - event (Dict[str, Any]): The event data. - schema (torch._C.FunctionSchema): The schema to match against. - - Returns: - bool: True if the event matches the schema, False otherwise. """ op_name, pos_args_schema, kwargs_schema, return_type = ( EventReplayer.parse_schema_string(schema) ) full_args_schema = pos_args_schema + kwargs_schema - # Check if the number of args in the event matches the schema if len(event["args"]["Input type"]) != len(full_args_schema): return False - # Check if the types match for idx in range(len(event["args"]["Input type"])): profiled_type = event["args"]["Input type"][idx] schema_type = full_args_schema[idx]["arg_type"] @@ -137,16 +230,9 @@ def _is_schema_match( print(f"Checking arg {idx}:") print(f"\tSchema type: {schema_type}") print(f"\tProfiled type: {profiled_type}") - # Rules for matching types - # 1. for tensor types, schema type should be 'Tensor' and profiled type can be any of the tensor types 'float', 'c10::Half', 'c10::BFloat16' ... - # 2. for bool types, schema type should be 'bool' and profiled type is 'Scalar'. So we need to further check the concrete Inputs if it only contains 'true' or 'false' - # 3. for int types, schema type should be 'int' or 'SymInt' and profiled type is 'Scalar'. So we need to further check the concrete Inputs if it is a digit - # 4. for float types, schema type should be 'Scalar' and profiled type is 'Scalar'. So we need to further check the concrete Inputs if it is a float - # 5. for int[] types, schema type should be 'int[]' or 'SymInt[]' and profiled type is 'ScalarList'. So we need to further check the concrete Inputs if it is a list of digits - # 6. for bool[] types, schema type should be 'bool[]' and profiled type is 'ScalarList'. So we need to further check the concrete Inputs if it is a list of 'true' or 'false' - # 7. for tensor[] types, we cannot replay the event as the tensor shapes are not provided in the event. So we need to skip this case. Maybe suggest PyTorch to add this in the future. + is_match = True - # if the schema type ends with '?' then the profiled type can be blank as well + # Optional types: schema ends with '?' => profiled type can be blank if schema_type.endswith("?"): schema_type = schema_type[:-1] if profiled_type == "": @@ -156,20 +242,20 @@ def _is_schema_match( and event["args"]["Concrete Inputs"][idx] == "[]" ): continue - if schema_type in ["Tensor", "Tensor?", "Tensor(a!)"]: + if EventReplayer._is_tensor_schema_type(schema_type): if profiled_type not in list_profile_tensor_types: is_match = False elif schema_type == "bool": profiled_value = event["args"]["Concrete Inputs"][idx] if profiled_value.lower() not in ["true", "false"]: is_match = False - elif schema_type == "int" or schema_type == "SymInt": + elif schema_type in ("int", "SymInt"): if profiled_type != "Scalar": is_match = False profiled_value = event["args"]["Concrete Inputs"][idx] if not profiled_value.lstrip("-").isdigit(): is_match = False - elif schema_type in ["float", "Scalar"]: + elif schema_type in ("float", "Scalar"): if profiled_type != "Scalar": is_match = False profiled_value = event["args"]["Concrete Inputs"][idx] @@ -178,7 +264,6 @@ def _is_schema_match( except ValueError: is_match = False elif schema_type.startswith("int[") or schema_type.startswith("SymInt["): - # custom dev debugging if profiled_type != "ScalarList": is_match = False profiled_value = event["args"]["Concrete Inputs"][idx] @@ -195,16 +280,27 @@ def _is_schema_match( x.strip() for x in profiled_value.strip()[1:-1].split(",") ] if not all( - x.lower() in ["true", "false"] for x in profiled_value_cleaned + x.lower() in ("true", "false") for x in profiled_value_cleaned ): is_match = False elif schema_type.startswith("Tensor["): raise ValueError( - f"Tensor list type not supported: {schema_type} as the tensor shapes are not provided in the event" + f"Tensor list type not supported: {schema_type} as the " + f"tensor shapes are not provided in the event" ) + elif schema_type == "str": + is_match = profiled_type == "Scalar" or profiled_type == "" + elif schema_type == "ScalarType": + is_match = profiled_type == "Scalar" or profiled_type == "" + elif schema_type == "Layout": + is_match = profiled_type == "Scalar" or profiled_type == "" + elif schema_type == "Device": + is_match = profiled_type == "Scalar" or profiled_type == "" + elif schema_type == "MemoryFormat": + is_match = profiled_type == "Scalar" or profiled_type == "" + elif schema_type == "Generator": + is_match = profiled_type == "" or profiled_type == "Scalar" else: - # raise ValueError(f"Unknown schema type: {schema_type}") - # warning: if the schema type is not in the list, we will skip this case warnings.warn( f"Unknown schema type: {schema_type}. Skipping this case." ) @@ -217,35 +313,42 @@ def _is_schema_match( return False return True + @staticmethod + def _is_tensor_schema_type(schema_type: str) -> bool: + """Check if a schema type string represents a Tensor argument.""" + if schema_type in ("Tensor", "Tensor?"): + return True + # Handles annotated variants like Tensor(a!), Tensor(a), Tensor(b!) + if schema_type.startswith("Tensor("): + return True + return False + + @staticmethod + def _should_skip_tensor_init(evt_name: str, arg_name: str, arg_idx: int) -> bool: + """ + Determine whether a tensor argument is an output-only buffer that + does not need random initialization. + + Generalizes the old aten::fill_ / aten::copy_ special-cases to + any in-place or out-of-place output tensor. + """ + # In-place ops (name ends with '_'): first tensor is the mutated output + if evt_name.endswith("_") and arg_idx == 0: + return True + # Explicit 'out' arguments in .out variants + if arg_name == "out": + return True + # aten::copy_ destination + if evt_name == "aten::copy_" and arg_name != "src": + return True + return False + @staticmethod def _get_event_replay_IR( event: Dict[str, Any], schema: "torch._C.FunctionSchema", verbose: bool = False ) -> Dict[str, Any]: """ Get the event replay IR from the event and schema. - - Args: - event (Dict[str, Any]): The event data. - schema (torch._C.FunctionSchema): The schema to match against. - - Returns: - { - 'pos_args': [ - - dummy_tensor0, - dummy_tensor1, - value0, - value1, - ... - ], - 'kwargs': { - 'arg0': value0, - 'arg1': dummy_tensor0, - 'arg2': dummy_tensor1, - 'arg3': value1, - ... - } - } """ evt_name = event["name"] op_name, pos_args_schema, kwargs_schema, return_type = ( @@ -266,36 +369,45 @@ def _get_event_replay_IR( print(f"Concrete Inputs: {event['args']['Concrete Inputs'][idx]}") if arg_type.endswith("?") and event["args"]["Input type"][idx] == "": - value = None + if arg_type.startswith("str"): + default = full_args_schema[idx].get("default") + value = "" if default is None or default == "None" else default + else: + value = None elif ( arg_type.endswith("?") and event["args"]["Concrete Inputs"][idx] == "[]" ): value = [] else: - if arg_type in ["Tensor", "Tensor?", "Tensor(a!)"]: + if EventReplayer._is_tensor_schema_type(arg_type): init = "normal" - if evt_name == "aten::fill_": - # special case for fill_ where we don't need to initialize the tensor - # as it will be filled with a value later - init = None - elif evt_name == "aten::copy_" and arg_name != "src": - # special case for copy_ where we don't need to initialize the tensor - # as it will be copied from another tensor + if EventReplayer._should_skip_tensor_init(evt_name, arg_name, idx): init = None + profiled_dtype = event["args"]["Input type"][idx] + # Non-floating-point tensors cannot use 'normal' init + if profiled_dtype in ("long", "long int", "int", "bool", "unsigned char"): + init = "zeros" if init == "normal" else init value = TensorCfg( shape=event["args"]["Input Dims"][idx], - dtype=event["args"]["Input type"][idx], + dtype=profiled_dtype, strides=event["args"]["Input Strides"][idx], init=init, ) else: arg_str = event["args"]["Concrete Inputs"][idx] - if arg_type in ["bool", "bool?"]: + if arg_type in ("bool", "bool?"): value = arg_str.lower() == "true" - elif arg_type in ["int", "SymInt"]: + elif arg_type in ("int", "int?", "SymInt", "SymInt?"): value = int(arg_str) - elif arg_type in ["float", "float?", "Scalar", "Scalar?"]: + elif arg_type in ("Scalar", "Scalar?"): + if arg_str.lstrip("-").isdigit(): + value = int(arg_str) + else: + value = float(arg_str) + elif arg_type in ("float", "float?"): value = float(arg_str) + elif arg_type in ("str", "str?"): + value = arg_str elif arg_type.startswith("int[") or arg_type.startswith("SymInt["): value = [ int(x.strip()) for x in arg_str.strip()[1:-1].split(",") @@ -323,18 +435,95 @@ def _get_event_replay_IR( ) return {"list_pos_args": list_pos_args, "list_kwargs": list_kwargs} + @staticmethod + def _get_event_replay_IR_schemaless( + event: Dict[str, Any], verbose: bool = False + ) -> Dict[str, Any]: + """ + Build a replay IR without a schema by inferring types directly from the + profiled data. All arguments are treated as positional. + + Heuristics: + - If Input type is a known tensor dtype -> TensorCfg + - If Input type is 'Scalar' and Concrete Inputs looks like int -> int + - If Input type is 'Scalar' and Concrete Inputs looks like float -> float + - If Input type is 'Scalar' and Concrete Inputs is true/false -> bool + - If Input type is '' and Concrete Inputs is '' -> None + """ + evt_name = event["name"] + list_pos_args = [] + n_args = len(event["args"]["Input type"]) + for idx in range(n_args): + profiled_type = event["args"]["Input type"][idx] + profiled_dims = event["args"]["Input Dims"][idx] + profiled_strides = event["args"]["Input Strides"][idx] + concrete = event["args"]["Concrete Inputs"][idx] + + if verbose: + print( + f"Schemaless arg {idx}: type={profiled_type!r} " + f"dims={profiled_dims} concrete={concrete!r}" + ) + + if profiled_type in list_profile_tensor_types: + init = "normal" + if profiled_type in ( + "long", "long int", "int", "bool", "unsigned char", + ): + init = "zeros" + value = TensorCfg( + shape=profiled_dims, + dtype=profiled_type, + strides=profiled_strides, + init=init, + ) + arg_type = "Tensor" + elif profiled_type == "Scalar" and concrete: + if concrete.lower() in ("true", "false"): + value = concrete.lower() == "true" + arg_type = "bool" + elif concrete.lstrip("-").isdigit(): + value = int(concrete) + arg_type = "int" + else: + try: + value = float(concrete) + arg_type = "float" + except ValueError: + value = concrete + arg_type = "str" + elif profiled_type == "" and concrete == "": + value = None + arg_type = "None" + elif profiled_type == "ScalarList" and concrete: + items = [x.strip() for x in concrete.strip()[1:-1].split(",") if x.strip()] + if all(x.lstrip("-").isdigit() for x in items): + value = [int(x) for x in items] + else: + value = [float(x) for x in items] + arg_type = "list" + else: + value = None + arg_type = "unknown" + if verbose: + print(f" -> defaulting to None for unknown type") + + inferred_name = f"arg{idx}" + list_pos_args.append( + {"arg_name": inferred_name, "arg_type": arg_type, "value": value} + ) + if verbose: + print(f" -> {inferred_name}: {arg_type} = {value}") + print("-" * 80) + + return {"list_pos_args": list_pos_args, "list_kwargs": []} + @staticmethod def _get_args_kwargs( event_replay_IR: Dict[str, Any], device: str = "cuda" ) -> tuple[List["torch.Tensor"], Dict[str, Any]]: """ Get the arguments and keyword arguments from the event replay IR. - - Args: - event_replay_IR (Dict[str, Any]): The event replay IR. - - Returns: - (List[torch.Tensor], Dict[str, Any]): The positional arguments and keyword arguments. """ pos_args = [] for arg in event_replay_IR["list_pos_args"]: @@ -401,18 +590,12 @@ def get_repro_info(self) -> Dict[str, Any]: Dict[str, Any]: A dictionary containing the operator name and the replay IR. Suitable for JSON serialization using the custom encoder. """ - # return { - # 'op_name': self.event['name'], - # 'replay_ir': self.event_replay_IR - # # No device info here - device is decided by the runner - # } dict_repro_info = {} dict_repro_info["op_name"] = self.event["name"] list_pos_args, list_kwargs = ( self.event_replay_IR["list_pos_args"], self.event_replay_IR["list_kwargs"], ) - # Convert TensorCfg to dict for JSON serialization list_pos_args_copy, list_kwargs_copy = list_pos_args.copy(), list_kwargs.copy() for idx, val in enumerate(list_pos_args_copy): if isinstance(val["value"], TensorCfg): diff --git a/TraceLens/EventReplay/utils.py b/TraceLens/EventReplay/utils.py index cf75fd97f..49243c28b 100644 --- a/TraceLens/EventReplay/utils.py +++ b/TraceLens/EventReplay/utils.py @@ -32,8 +32,16 @@ def _get_torch_or_raise() -> Any: # Changed return type to Any for flexibility "c10::Half", "c10::BFloat16", "long", + "long int", "int", "bool", + "unsigned char", + "char", + "short", + "c10::Float8_e4m3fnuz", + "c10::Float8_e5m2fnuz", + "c10::Float8_e4m3fn", + "c10::Float8_e5m2", ] from dataclasses import dataclass @@ -58,15 +66,35 @@ def build_tensor(cfg: TensorCfg, device: str = "cuda") -> "torch.Tensor": "bool": torch.bool, "int": torch.int, "long": torch.long, + "long int": torch.long, + "short": torch.short, + "char": torch.int8, + "unsigned char": torch.uint8, "double": torch.float64, "float": torch.float32, "c10::Half": torch.float16, "c10::BFloat16": torch.bfloat16, } + # FP8 types (available in PyTorch >= 2.1) + for fp8_name in ( + "c10::Float8_e4m3fnuz", + "c10::Float8_e5m2fnuz", + "c10::Float8_e4m3fn", + "c10::Float8_e5m2", + ): + attr = fp8_name.replace("c10::", "") + torch_dtype = getattr(torch, attr.lower(), None) + if torch_dtype is not None: + dict_profile2torchdtype[fp8_name] = torch_dtype + + if cfg.dtype not in dict_profile2torchdtype: + raise ValueError( + f"Unknown profiled dtype '{cfg.dtype}'. " + f"Known types: {list(dict_profile2torchdtype.keys())}" + ) dtype = dict_profile2torchdtype[cfg.dtype] size = cfg.shape stride = cfg.strides - # allocate *exactly* the storage needed for that stride/shape t = torch.empty_strided(size, stride, dtype=dtype, device=device) is_floating = t.is_floating_point() or t.is_complex() init = cfg.init @@ -75,7 +103,9 @@ def build_tensor(cfg: TensorCfg, device: str = "cuda") -> "torch.Tensor": raise ValueError( f"Cannot initialize tensor of type {cfg.dtype} with 'normal' init." ) - t.normal_() # or whatever init you like + t.normal_() + elif init == "zeros": + t.zero_() elif init is not None: raise ValueError(f"Unsupported tensor initialization: {init}") return t From ff57fa2d014da87eb03cd299ee46ea5dd8ad5fab Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sat, 28 Mar 2026 18:39:50 +0000 Subject: [PATCH 02/11] Add string arg defaults, op aliases, and module resolution String arg defaults (_STR_ARG_DEFAULTS): - When the profiler drops a str arg value, check a known-defaults table keyed by arg name (e.g. kv_cache_dtype -> "auto") - Log a WARNING when a default is used so users know the value was inferred - Recovers _C_cache_ops::reshape_and_cache_flash (1.97% GPU time) Op name aliases (_OP_NAME_ALIASES): - Map trace-recorded names to their runtime-registered names (e.g. _rocm_C::wvSplitK -> _rocm_C::wvSpltK) Python module resolution (3rd strategy): - After JIT and torch.ops, try importlib.import_module(namespace) for JIT-compiled ops like aiter Schema parser fix: - parse_schema_string handles annotated tensor types with spaces like "Tensor($0! -> )" correctly now Tested with vLLM Qwen1.5-MoE-A2.7B trace on MI300X (tw025). Made-with: Cursor --- TraceLens/EventReplay/event_replay.py | 185 +++++++++++++++++++++----- 1 file changed, 152 insertions(+), 33 deletions(-) diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index 5473cabf8..97d25e88e 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -9,6 +9,7 @@ from pprint import pprint from typing import Dict, Any, List, Optional, Tuple +import logging import re import warnings @@ -19,19 +20,36 @@ list_profile_tensor_types, ) - -def _resolve_op_func(op_name: str): - """ - Resolve an op name (e.g. 'aten::mm', 'vllm::rocm_unquantized_gemm') to a - callable. Tries multiple resolution strategies and returns the first that - yields a non-None callable. - - Returns (func, source_str) or raises RuntimeError. +logger = logging.getLogger(__name__) + +# ── Known defaults for string arguments the profiler drops ────────────── +# The PyTorch profiler records `str` arguments as empty strings. When we +# know the only sensible default we fill it in automatically and warn. +# Key = argument name, Value = default string value. +_STR_ARG_DEFAULTS: Dict[str, str] = { + "kv_cache_dtype": "auto", +} + +# ── Op-name aliases ───────────────────────────────────────────────────── +# Some frameworks profile an op under one namespace but register the +# actual callable under a different one. +# Key = name as it appears in the trace, Value = list of candidates to try. +# NOTE: aiter::paged_attention_v1/v2 are NOT aliasable — the aiter JIT +# wrapper records a different arg layout than the underlying _C:: / _rocm_C:: +# schemas, so arg mapping would fail even if resolution succeeds. +_OP_NAME_ALIASES: Dict[str, List[str]] = { + "_rocm_C::wvSplitK": ["_rocm_C::wvSpltK"], +} + + +def _try_resolve(op_name: str): + """Attempt JIT + torch.ops + module resolution for a single op name. + Returns (func, source_str) or (None, None). """ torch = _get_torch_or_raise() + import importlib - # 1. JIT registry first — preserves dispatch behaviour that the original - # profiled run used (important for in-place aten ops like add_). + # 1. JIT registry — preserves dispatch behaviour for aten ops. try: func, _ = torch._C._jit_get_operation(op_name) if func is not None: @@ -39,16 +57,55 @@ def _resolve_op_func(op_name: str): except RuntimeError: pass - # 2. torch.ops namespace lookup (most reliable for custom ops that may - # not be registered in the JIT registry). if "::" in op_name: ns, func_name = op_name.split("::", 1) + + # 2. torch.ops namespace — custom ops registered via torch.library. ns_obj = getattr(torch.ops, ns, None) if ns_obj is not None: func_obj = getattr(ns_obj, func_name, None) if callable(func_obj): return func_obj, "torch.ops" + # 3. Direct Python module lookup — handles JIT-compiled ops (e.g. + # aiter) that exist as Python callables but aren't registered in + # the torch op registry. + try: + mod = importlib.import_module(ns) + func_obj = getattr(mod, func_name, None) + if callable(func_obj): + return func_obj, f"module:{ns}" + except ImportError: + pass + + return None, None + + +def _resolve_op_func(op_name: str): + """ + Resolve an op name (e.g. 'aten::mm', 'vllm::rocm_unquantized_gemm') to a + callable. Tries multiple resolution strategies: + + 1. JIT registry (preserves original dispatch behaviour). + 2. torch.ops namespace (custom ops registered via torch.library / pybind). + 3. Known aliases from _OP_NAME_ALIASES (handles trace-name mismatches). + + Returns (func, source_str, resolved_name) or raises RuntimeError. + """ + func, source = _try_resolve(op_name) + if func is not None: + return func, source, op_name + + for alias in _OP_NAME_ALIASES.get(op_name, []): + func, source = _try_resolve(alias) + if func is not None: + logger.warning( + "Op '%s' resolved via alias '%s' (%s). " + "The trace recorded a different namespace than the runtime registration.", + op_name, alias, source, + ) + return func, source, alias + raise RuntimeError( f"Cannot resolve op '{op_name}'. Ensure the library that defines it " f"is imported (e.g. 'import vllm', 'import aiter')." @@ -140,13 +197,17 @@ def _setup(self): if self.verbose: print(f"Preparing {self.event['name']} event for replay") - self._func, self._func_source = _resolve_op_func(self.event["name"]) + self._func, self._func_source, self._resolved_name = _resolve_op_func( + self.event["name"] + ) if self.verbose: print(f"Resolved op via {self._func_source}") + if self._resolved_name != self.event["name"]: + print(f" (aliased from '{self.event['name']}' -> '{self._resolved_name}')") try: self.matched_schema = EventReplayer._search_schema( - self.event, self.verbose + self.event, self._resolved_name, self.verbose ) self._schemaless = False except ValueError: @@ -160,7 +221,7 @@ def _setup(self): if self._schemaless: self.event_replay_IR = EventReplayer._get_event_replay_IR_schemaless( - self.event, self.verbose + self.event, self.verbose, resolved_name=self._resolved_name ) else: self.event_replay_IR = EventReplayer._get_event_replay_IR( @@ -188,9 +249,12 @@ def replay(self): @staticmethod def _search_schema( - event: Dict[str, Any], verbose: bool = False + event: Dict[str, Any], + resolved_name: Optional[str] = None, + verbose: bool = False, ) -> Optional["torch._C.FunctionSchema"]: - op_schemas = _search_schemas(event["name"], verbose=verbose) + name = resolved_name or event["name"] + op_schemas = _search_schemas(name, verbose=verbose) for schema in op_schemas: if verbose: @@ -205,7 +269,7 @@ def _search_schema( print("-" * 80) raise ValueError( - f"Cannot find matching schema for {event['name']}. " + f"Cannot find matching schema for {name}. " f"Searched {len(op_schemas)} candidate(s). " f"Please check the event data and ensure the op's library is imported." ) @@ -371,7 +435,15 @@ def _get_event_replay_IR( if arg_type.endswith("?") and event["args"]["Input type"][idx] == "": if arg_type.startswith("str"): default = full_args_schema[idx].get("default") - value = "" if default is None or default == "None" else default + if default is None or default == "None": + default = _STR_ARG_DEFAULTS.get(arg_name) + if default is not None: + logger.warning( + "%s arg '%s' (position %d): profiler dropped " + "the string value. Using known default '%s'.", + evt_name, arg_name, idx, default, + ) + value = "" if default is None else default else: value = None elif ( @@ -407,7 +479,15 @@ def _get_event_replay_IR( elif arg_type in ("float", "float?"): value = float(arg_str) elif arg_type in ("str", "str?"): - value = arg_str + if not arg_str and arg_name in _STR_ARG_DEFAULTS: + value = _STR_ARG_DEFAULTS[arg_name] + logger.warning( + "%s arg '%s' (position %d): profiler dropped " + "the string value. Using known default '%s'.", + evt_name, arg_name, idx, value, + ) + else: + value = arg_str elif arg_type.startswith("int[") or arg_type.startswith("SymInt["): value = [ int(x.strip()) for x in arg_str.strip()[1:-1].split(",") @@ -435,9 +515,23 @@ def _get_event_replay_IR( ) return {"list_pos_args": list_pos_args, "list_kwargs": list_kwargs} + @staticmethod + def _get_schema_arg_names(op_name: str) -> List[str]: + """Best-effort: return a list of arg names from any available schema.""" + schemas = _search_schemas(op_name, verbose=False) + if not schemas: + return [] + try: + _, pos_args, kw_args, _ = EventReplayer.parse_schema_string(schemas[0]) + return [a["arg_name"] for a in pos_args + kw_args] + except Exception: + return [] + @staticmethod def _get_event_replay_IR_schemaless( - event: Dict[str, Any], verbose: bool = False + event: Dict[str, Any], + verbose: bool = False, + resolved_name: Optional[str] = None, ) -> Dict[str, Any]: """ Build a replay IR without a schema by inferring types directly from the @@ -448,9 +542,13 @@ def _get_event_replay_IR_schemaless( - If Input type is 'Scalar' and Concrete Inputs looks like int -> int - If Input type is 'Scalar' and Concrete Inputs looks like float -> float - If Input type is 'Scalar' and Concrete Inputs is true/false -> bool - - If Input type is '' and Concrete Inputs is '' -> None + - If Input type is '' and Concrete Inputs is '' -> check _STR_ARG_DEFAULTS """ evt_name = event["name"] + schema_arg_names = EventReplayer._get_schema_arg_names( + resolved_name or evt_name + ) + list_pos_args = [] n_args = len(event["args"]["Input type"]) for idx in range(n_args): @@ -493,8 +591,22 @@ def _get_event_replay_IR_schemaless( value = concrete arg_type = "str" elif profiled_type == "" and concrete == "": - value = None - arg_type = "None" + # Likely a dropped str arg — check known defaults + hint_name = ( + schema_arg_names[idx] if idx < len(schema_arg_names) else None + ) + default = _STR_ARG_DEFAULTS.get(hint_name) if hint_name else None + if default is not None: + value = default + arg_type = "str" + logger.warning( + "%s arg '%s' (position %d): profiler dropped the " + "string value. Using known default '%s'.", + evt_name, hint_name, idx, default, + ) + else: + value = None + arg_type = "None" elif profiled_type == "ScalarList" and concrete: items = [x.strip() for x in concrete.strip()[1:-1].split(",") if x.strip()] if all(x.lstrip("-").isdigit() for x in items): @@ -508,7 +620,9 @@ def _get_event_replay_IR_schemaless( if verbose: print(f" -> defaulting to None for unknown type") - inferred_name = f"arg{idx}" + inferred_name = ( + schema_arg_names[idx] if idx < len(schema_arg_names) else f"arg{idx}" + ) list_pos_args.append( {"arg_name": inferred_name, "arg_type": arg_type, "value": value} ) @@ -555,17 +669,22 @@ def parse_schema_string( kwarg_part = parts[1].lstrip(",").strip() if len(parts) > 1 else "" def _parse_arg(raw_arg: str) -> Tuple[str, str, Optional[str], bool]: - m = re.match(r"^(\S+)\s+(.*)$", raw_arg) + # Match type (may contain spaces, e.g. "Tensor($0! -> )") then name[=default]. + # Greedy (.+) consumes everything up to the last whitespace before + # a valid identifier, so "Tensor($0! -> ) key_cache" parses correctly. + m = re.match( + r"^(.+)\s+([A-Za-z_]\w*(?:=.*)?)$", raw_arg.strip() + ) if not m: raise ValueError(f"Invalid arg: {raw_arg}") - arg_type, rest = m.groups() - m2 = re.match(r"^([A-Za-z_][A-Za-z0-9_]*)(?:=(.*))?$", rest) + arg_type = m.group(1).strip() + name_default = m.group(2) + m2 = re.match(r"^([A-Za-z_]\w*)(?:=(.*))?$", name_default) if not m2: - raise ValueError(f"Invalid arg name/default: {rest}") - arg_name, default = m2.group(1), ( - m2.group(2).strip() if m2.group(2) else None - ) - return arg_type.strip(), arg_name.strip(), default + raise ValueError(f"Invalid arg name/default: {name_default}") + arg_name = m2.group(1) + default = m2.group(2).strip() if m2.group(2) else None + return arg_type, arg_name, default args = [] for item in [x.strip() for x in pos_part.split(",") if x.strip()]: From b1e5907caab3fda6ee0656f15a58cff5c9b90eb2 Mon Sep 17 00:00:00 2001 From: Jassani Date: Tue, 28 Apr 2026 14:36:06 -0400 Subject: [PATCH 03/11] Add custom initializers, auto-import, and updated docs for EventReplay Made-with: Cursor --- TraceLens/EventReplay/__init__.py | 23 ++ TraceLens/EventReplay/custom_inits.py | 389 ++++++++++++++++++++++++++ TraceLens/EventReplay/event_replay.py | 284 +++++++++++-------- TraceLens/EventReplay/utils.py | 82 ++++-- 4 files changed, 642 insertions(+), 136 deletions(-) create mode 100644 TraceLens/EventReplay/custom_inits.py diff --git a/TraceLens/EventReplay/__init__.py b/TraceLens/EventReplay/__init__.py index e69de29bb..84a26c4bb 100644 --- a/TraceLens/EventReplay/__init__.py +++ b/TraceLens/EventReplay/__init__.py @@ -0,0 +1,23 @@ +############################################################################### +# Copyright (c) 2024 - 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from .event_replay import EventReplayer +from .custom_inits import ( + CustomInit, + PagedAttentionInit, + MoeRoutingInit, + extract_batch_context, +) +from .utils import benchmark_func + +__all__ = [ + "EventReplayer", + "CustomInit", + "PagedAttentionInit", + "MoeRoutingInit", + "extract_batch_context", + "benchmark_func", +] diff --git a/TraceLens/EventReplay/custom_inits.py b/TraceLens/EventReplay/custom_inits.py new file mode 100644 index 000000000..9ded92540 --- /dev/null +++ b/TraceLens/EventReplay/custom_inits.py @@ -0,0 +1,389 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Custom initializers for EventReplayer. + +Operations captured by the PyTorch profiler have zeroed-out metadata tensors +(block tables, routing tensors, etc.) because the profiler records shapes and +dtypes but not tensor values. Custom initializers fill these tensors with +realistic content so the GPU kernel exercises real memory-access and compute +patterns during replay benchmarking. + +To add a custom initializer for a new op family: + 1. Subclass ``CustomInit`` + 2. Set ``op_patterns`` to one or more substrings that match the op name + 3. Implement ``initialize()`` — mutate replayer.args / replayer.kwargs in-place + 4. Return a one-line summary string (printed by EventReplayer) + 5. Register with ``EventReplayer.register_custom_init(YourInit())`` + or add it to the ``_custom_init_registry`` default list. +""" + +from __future__ import annotations + +import re +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +if TYPE_CHECKING: + pass # EventReplayer imported at runtime to avoid circular dep + +# -- Batch context extraction from vLLM profiler annotations --------------- + +_BATCH_ANNO_RE = re.compile( + r"execute_context_(\d+)\((\d+)\)_generation_(\d+)\((\d+)\)" +) + + +def extract_batch_context(analyzer: Any) -> int: + """Parse vLLM ``user_annotation`` events and attach batch context to ops. + + vLLM annotates each ``execute_model`` step with a ``user_annotation`` + event of the form ``execute_context_N(T)_generation_N(T)`` where + *N* = number of sequences and *T* = total query tokens for that phase. + + This function: + 1. Collects all such annotations with their ``[ts, ts+dur]`` ranges. + 2. For every ``paged_attention`` cpu_op event, finds the enclosing + annotation by timestamp and attaches a ``batch_context`` dict:: + + event["batch_context"] = { + "n_prefill": 2, + "prefill_tokens": 18, + "n_decode": 2, + "decode_tokens": 2, + } + + Args: + analyzer: A ``TreePerfAnalyzer`` (or any object whose ``.tree.events`` + yields the trace event list). + + Returns: + Number of paged_attention events that were annotated. + """ + annotations = [] + for e in analyzer.tree.events: + cat = e.get("cat") or "" + if cat != "user_annotation": + continue + m = _BATCH_ANNO_RE.search(e.get("name", "")) + if not m: + continue + ts = e.get("ts", 0) + dur = e.get("dur", 0) + annotations.append({ + "ts": ts, + "end": ts + dur, + "n_prefill": int(m.group(1)), + "prefill_tokens": int(m.group(2)), + "n_decode": int(m.group(3)), + "decode_tokens": int(m.group(4)), + }) + + if not annotations: + return 0 + + annotations.sort(key=lambda a: a["ts"]) + + annotated = 0 + for e in analyzer.tree.events: + name = e.get("name", "") + if "paged_attention" not in name: + continue + if not e.get("args", {}).get("Input Dims"): + continue + ts = e.get("ts", 0) + for a in annotations: + if a["ts"] <= ts <= a["end"]: + e["batch_context"] = { + "n_prefill": a["n_prefill"], + "prefill_tokens": a["prefill_tokens"], + "n_decode": a["n_decode"], + "decode_tokens": a["decode_tokens"], + } + annotated += 1 + break + + return annotated + + +class CustomInit(ABC): + """Base class for tensor initializers applied before replay.""" + + op_patterns: List[str] = [] + + def applies_to(self, replayer: Any) -> bool: + op_name = replayer.event.get("name", "") + return any(pat in op_name for pat in self.op_patterns) + + @abstractmethod + def initialize(self, replayer: Any, **kwargs) -> Optional[str]: + """Mutate replayer.args/kwargs in-place. Return a summary string.""" + ... + + +class PagedAttentionInit(CustomInit): + """Initialize block_tables, seq_lens, and query_start_loc for paged attention. + + When ``batch_context`` is present on the event (attached by + :func:`extract_batch_context`), uses the exact prefill/decode split from + vLLM's profiler annotations. Otherwise falls back to heuristics: + - query_tokens == num_seqs → decode (1 token/seq) + - query_tokens > num_seqs → prefill (tokens distributed uniformly) + + In all cases: + - ``seq_lens`` set to ``max_seq_len`` for every sequence. + - Block table entries drawn from a random permutation of the pool. + """ + + op_patterns = ["paged_attention"] + + def initialize(self, replayer: Any, **kwargs) -> Optional[str]: + try: + import numpy as np + except ImportError: + return "[custom init] PagedAttentionInit skipped — numpy not available" + + args = replayer.args + op_name = replayer.event.get("name", "") + + ir = replayer.event_replay_IR + arg_names = [a["arg_name"] for a in ir["list_pos_args"]] + def _by_name_or_pos(name, pos): + if name in arg_names: + return args[arg_names.index(name)] + return args[pos] + + block_tables = _by_name_or_pos("block_tables", 9) + seq_lens = _by_name_or_pos("seq_lens", 10) + key_cache = _by_name_or_pos("key_cache", 5) + block_size = int(_by_name_or_pos("block_size", 12)) + max_seq_len = int(_by_name_or_pos("max_seq_len", 13)) + + num_seqs = block_tables.shape[0] + max_blocks_per_seq = block_tables.shape[1] + num_blocks_total = key_cache.shape[0] + + query = _by_name_or_pos("query", 4) + num_query_tokens = query.shape[0] + + rng = np.random.default_rng(42) + + # -- Determine per-sequence query token counts ------------------------- + batch_ctx = replayer.event.get("batch_context") + if batch_ctx is not None: + n_pf = batch_ctx["n_prefill"] + pf_tok = batch_ctx["prefill_tokens"] + n_dec = batch_ctx["n_decode"] + dec_tok = batch_ctx["decode_tokens"] + + per_seq_queries = [] + if n_pf > 0: + base_pf = pf_tok // n_pf + rem_pf = pf_tok % n_pf + for s in range(n_pf): + per_seq_queries.append(base_pf + (1 if s < rem_pf else 0)) + for _ in range(n_dec): + per_seq_queries.append(1) + + if len(per_seq_queries) != num_seqs: + per_seq_queries = per_seq_queries[:num_seqs] + while len(per_seq_queries) < num_seqs: + per_seq_queries.append(1) + + phase = ("mixed" if n_pf > 0 and n_dec > 0 + else "prefill" if n_pf > 0 else "decode") + source = "annotation" + else: + tokens_per_seq = num_query_tokens / num_seqs if num_seqs else 1 + if tokens_per_seq > 1: + base_q = num_query_tokens // num_seqs + rem_q = num_query_tokens % num_seqs + per_seq_queries = [base_q + (1 if s < rem_q else 0) + for s in range(num_seqs)] + phase = "prefill" + else: + per_seq_queries = [1] * num_seqs + phase = "decode" + source = "heuristic" + + # -- seq_lens: max_seq_len for every sequence -------------------------- + lengths = np.full(num_seqs, max_seq_len, dtype=np.int32) + + # -- block_tables: permutation of physical block pool ------------------ + bt = np.zeros((num_seqs, max_blocks_per_seq), dtype=np.int32) + all_block_ids = rng.permutation(num_blocks_total) + block_cursor = 0 + for s in range(num_seqs): + blocks_needed = (int(lengths[s]) + block_size - 1) // block_size + blocks_needed = min(blocks_needed, max_blocks_per_seq) + for b in range(blocks_needed): + bt[s, b] = all_block_ids[block_cursor % num_blocks_total] + block_cursor += 1 + + import torch + + block_tables.copy_(torch.from_numpy(bt).to(block_tables.device)) + seq_lens.copy_(torch.from_numpy(lengths).to(seq_lens.device)) + + # -- query_start_loc: CSR indptr encoding per-seq query counts --------- + qsl = _by_name_or_pos("query_start_loc", 11) + if (qsl is not None + and hasattr(qsl, "shape") + and qsl.numel() > 0): + qloc = np.zeros(num_seqs + 1, dtype=np.int32) + for s in range(num_seqs): + qloc[s + 1] = qloc[s] + per_seq_queries[s] + qloc = qloc[: qsl.numel()] + qsl.copy_(torch.from_numpy(qloc).to(qsl.device)) + + ctx_str = "" + if batch_ctx is not None: + ctx_str = (f" Annotation: {batch_ctx['n_prefill']} prefill " + f"({batch_ctx['prefill_tokens']} tok) + " + f"{batch_ctx['n_decode']} decode " + f"({batch_ctx['decode_tokens']} tok).") + + return ( + f"[custom init] {op_name} — paged attention metadata: " + f"phase={phase} ({source}), num_seqs={num_seqs}, " + f"max_seq_len={max_seq_len}, block_size={block_size}, " + f"num_blocks={num_blocks_total}, " + f"max_blocks_per_seq={max_blocks_per_seq}.{ctx_str}" + ) + + +class MoeRoutingInit(CustomInit): + """Initialize MoE routing tensors (sorted_token_ids, sorted_expert_ids, + num_valid_ids) so the CK kernel processes real token-to-expert assignments + instead of short-circuiting on num_valid_ids=0. + + Supported kwargs: + moe_distribution: "uniform" (default) or "zipf" + moe_zipf_s: Zipf exponent (default 1.2), only used with "zipf" + + Arg layout for aiter::ck_moe_stage1/2: + [0] hidden_states [M, K] bf16 + [1] w1 [E, N, K] bf16 + [2] w2 [E, K2, N2] bf16 + [3] sorted_token_ids [padded] int32 <- init + [4] sorted_expert_ids [blocks+1] int32 <- init + [5] num_valid_ids [2] int32 <- init + [6] output [M, top_k, N2] bf16 + [7] top_k (scalar) + ... + [11] block_m (scalar) + """ + + op_patterns = ["ck_moe_stage1", "ck_moe_stage2"] + + def initialize(self, replayer: Any, **kwargs) -> Optional[str]: + try: + import numpy as np + except ImportError: + return "[custom init] MoeRoutingInit skipped — numpy not available" + + distribution = kwargs.get("moe_distribution", "uniform") + zipf_s = kwargs.get("moe_zipf_s", 1.2) + + args = replayer.args + op_name = replayer.event.get("name", "") + + # Locate args by name from the IR when available, fall back to position + ir = replayer.event_replay_IR + arg_names = [a["arg_name"] for a in ir["list_pos_args"]] + def _by_name_or_pos(name, pos): + if name in arg_names: + return args[arg_names.index(name)] + return args[pos] + + sorted_token_ids = _by_name_or_pos("sorted_token_ids", 3) + sorted_expert_ids = _by_name_or_pos("sorted_expert_ids", 4) + num_valid_ids = _by_name_or_pos("num_valid_ids", 5) + top_k = int(_by_name_or_pos("topk", 7)) + block_m_val = _by_name_or_pos("block_m", 11) + block_m = int(block_m_val) if block_m_val is not None else 32 + + hidden = _by_name_or_pos("hidden_states", 0) + M = hidden.shape[0] + w1 = _by_name_or_pos("w1", 1) + E = w1.shape[0] + num_tokens = M * top_k + padded_total = sorted_token_ids.shape[0] + num_blocks = sorted_expert_ids.shape[0] + + rng = np.random.default_rng(42) + + if distribution == "zipf": + ranks = np.arange(1, E + 1, dtype=np.float64) + weights = 1.0 / np.power(ranks, zipf_s) + probs = weights / weights.sum() + expert_assignments = rng.choice(E, size=num_tokens, p=probs) + else: + expert_assignments = rng.integers(0, E, size=num_tokens) + + token_ids_list: list = [] + expert_ids_list: list = [] + for expert_id in range(E): + tokens_for_expert = np.where(expert_assignments == expert_id)[0] + count = len(tokens_for_expert) + if count == 0: + continue + padded_count = ((count + block_m - 1) // block_m) * block_m + n_blocks_for_expert = padded_count // block_m + padded_tokens = np.full(padded_count, num_tokens, dtype=np.int32) + padded_tokens[:count] = tokens_for_expert // top_k + token_ids_list.append(padded_tokens) + expert_ids_list.extend([expert_id] * n_blocks_for_expert) + + all_token_ids = ( + np.concatenate(token_ids_list) + if token_ids_list + else np.array([], dtype=np.int32) + ) + + if len(all_token_ids) < padded_total: + padding = np.full( + padded_total - len(all_token_ids), num_tokens, dtype=np.int32 + ) + all_token_ids = np.concatenate([all_token_ids, padding]) + else: + all_token_ids = all_token_ids[:padded_total] + + all_expert_ids = np.array(expert_ids_list, dtype=np.int32) + if len(all_expert_ids) < num_blocks: + padding = np.zeros(num_blocks - len(all_expert_ids), dtype=np.int32) + all_expert_ids = np.concatenate([all_expert_ids, padding]) + else: + all_expert_ids = all_expert_ids[:num_blocks] + + import torch + + sorted_token_ids.copy_( + torch.from_numpy(all_token_ids).to(sorted_token_ids.device) + ) + sorted_expert_ids.copy_( + torch.from_numpy(all_expert_ids).to(sorted_expert_ids.device) + ) + + valid_count = min( + len(np.concatenate(token_ids_list)) if token_ids_list else 0, + padded_total, + ) + if num_valid_ids.numel() >= 1: + num_valid_ids[0] = valid_count + if num_valid_ids.numel() >= 2: + num_valid_ids[1] = valid_count + + dist_label = f"zipf(s={zipf_s})" if distribution == "zipf" else "uniform" + experts_active = len(set(expert_assignments.tolist())) + return ( + f"[custom init] {op_name} — initialized MoE routing: " + f"dist={dist_label}, M={M}, top_k={top_k}, E={E}, block_m={block_m}, " + f"num_tokens={num_tokens}, active_experts={experts_active}/{E}, " + f"valid_ids={valid_count}/{padded_total}, " + f"blocks={len(expert_ids_list)}/{num_blocks}. " + f"Assumptions: {dist_label} expert distribution, deterministic seed." + ) diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index 97d25e88e..3b02f3a35 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -19,28 +19,56 @@ build_tensor, list_profile_tensor_types, ) +from .custom_inits import CustomInit, PagedAttentionInit, MoeRoutingInit logger = logging.getLogger(__name__) -# ── Known defaults for string arguments the profiler drops ────────────── -# The PyTorch profiler records `str` arguments as empty strings. When we -# know the only sensible default we fill it in automatically and warn. -# Key = argument name, Value = default string value. +# -- Known defaults for string arguments the profiler drops ---------------- _STR_ARG_DEFAULTS: Dict[str, str] = { "kv_cache_dtype": "auto", } -# ── Op-name aliases ───────────────────────────────────────────────────── -# Some frameworks profile an op under one namespace but register the -# actual callable under a different one. -# Key = name as it appears in the trace, Value = list of candidates to try. -# NOTE: aiter::paged_attention_v1/v2 are NOT aliasable — the aiter JIT -# wrapper records a different arg layout than the underlying _C:: / _rocm_C:: -# schemas, so arg mapping would fail even if resolution succeeds. +# -- Op-name aliases ------------------------------------------------------- _OP_NAME_ALIASES: Dict[str, List[str]] = { "_rocm_C::wvSplitK": ["_rocm_C::wvSpltK"], } +# -- Auto-import registry -------------------------------------------------- +_NAMESPACE_IMPORTS: Dict[str, List[str]] = { + "aiter": ["aiter"], + "_rocm_C": ["vllm._rocm_C"], + "_C": ["vllm._C"], + "_C_cache_ops": ["vllm._C"], + "vllm": ["vllm._C", "vllm._rocm_C"], +} + +_auto_import_attempted: set = set() + + +def _try_auto_import(op_name: str, verbose: bool = False) -> bool: + """Try to import the library that registers a custom op's schema. + + Returns True if at least one new module was successfully imported. + """ + namespace = op_name.split("::")[0] if "::" in op_name else "" + if not namespace or namespace == "aten": + return False + if namespace in _auto_import_attempted: + return False + _auto_import_attempted.add(namespace) + + modules = _NAMESPACE_IMPORTS.get(namespace, [namespace]) + imported_any = False + for mod in modules: + try: + __import__(mod) + print(f"[EventReplayer] Auto-imported '{mod}' for op '{op_name}'") + imported_any = True + except ImportError: + if verbose: + print(f"[EventReplayer] Could not import '{mod}' for namespace '{namespace}'") + return imported_any + def _try_resolve(op_name: str): """Attempt JIT + torch.ops + module resolution for a single op name. @@ -49,7 +77,7 @@ def _try_resolve(op_name: str): torch = _get_torch_or_raise() import importlib - # 1. JIT registry — preserves dispatch behaviour for aten ops. + # 1. JIT registry try: func, _ = torch._C._jit_get_operation(op_name) if func is not None: @@ -60,16 +88,14 @@ def _try_resolve(op_name: str): if "::" in op_name: ns, func_name = op_name.split("::", 1) - # 2. torch.ops namespace — custom ops registered via torch.library. + # 2. torch.ops namespace ns_obj = getattr(torch.ops, ns, None) if ns_obj is not None: func_obj = getattr(ns_obj, func_name, None) if callable(func_obj): return func_obj, "torch.ops" - # 3. Direct Python module lookup — handles JIT-compiled ops (e.g. - # aiter) that exist as Python callables but aren't registered in - # the torch op registry. + # 3. Direct Python module lookup try: mod = importlib.import_module(ns) func_obj = getattr(mod, func_name, None) @@ -81,82 +107,92 @@ def _try_resolve(op_name: str): return None, None -def _resolve_op_func(op_name: str): - """ - Resolve an op name (e.g. 'aten::mm', 'vllm::rocm_unquantized_gemm') to a - callable. Tries multiple resolution strategies: - - 1. JIT registry (preserves original dispatch behaviour). - 2. torch.ops namespace (custom ops registered via torch.library / pybind). - 3. Known aliases from _OP_NAME_ALIASES (handles trace-name mismatches). +def _resolve_op_func(op_name: str, verbose: bool = False): + """Resolve an op name to a callable, with auto-import on failure. + Tries: JIT registry -> torch.ops -> module import -> auto-import -> aliases. Returns (func, source_str, resolved_name) or raises RuntimeError. """ - func, source = _try_resolve(op_name) - if func is not None: - return func, source, op_name - - for alias in _OP_NAME_ALIASES.get(op_name, []): - func, source = _try_resolve(alias) + for attempt in range(2): + func, source = _try_resolve(op_name) if func is not None: - logger.warning( - "Op '%s' resolved via alias '%s' (%s). " - "The trace recorded a different namespace than the runtime registration.", - op_name, alias, source, - ) - return func, source, alias + return func, source, op_name + + for alias in _OP_NAME_ALIASES.get(op_name, []): + func, source = _try_resolve(alias) + if func is not None: + logger.warning( + "Op '%s' resolved via alias '%s' (%s).", + op_name, alias, source, + ) + return func, source, alias + + if attempt == 0 and _try_auto_import(op_name, verbose): + continue + break + + ns = op_name.split("::")[0] if "::" in op_name else "" + hint = "" + if ns and ns != "aten": + known = _NAMESPACE_IMPORTS.get(ns) + if known: + hint = f" Try: {', '.join(f'import {m}' for m in known)}" + else: + hint = (f" The namespace '{ns}' is not in the auto-import registry." + f" Use EventReplayer.register_namespace('{ns}', ['your.module'])" + f" to add it.") raise RuntimeError( - f"Cannot resolve op '{op_name}'. Ensure the library that defines it " - f"is imported (e.g. 'import vllm', 'import aiter')." + f"Cannot resolve op '{op_name}'.{hint} " + f"Ensure the library that defines it is imported." ) def _search_schemas(op_name: str, verbose: bool = False): - """ - Return all registered FunctionSchemas for *op_name*. - - Searches both the JIT schema registry and the torch.ops namespace, which - covers aten ops, custom C++ ops, and Python-defined torch.library ops. + """Return all registered FunctionSchemas for *op_name*, + with auto-import on empty results. """ torch = _get_torch_or_raise() - schemas: list = [] - seen_strs: set = set() - - # JIT registry - for s in torch._C._jit_get_all_schemas(): - if s.name == op_name: - s_str = str(s) - if s_str not in seen_strs: - schemas.append(s) - seen_strs.add(s_str) - - # torch.ops namespace (catches custom ops not in the JIT list) - if "::" in op_name: - ns, func_name = op_name.split("::", 1) - ns_obj = getattr(torch.ops, ns, None) - if ns_obj is not None: - op_obj = getattr(ns_obj, func_name, None) - if op_obj is not None: - # OpOverloadPacket exposes overloads - try: - for overload_name in op_obj.overloads(): - overload = getattr(op_obj, overload_name) - s = overload._schema - s_str = str(s) - if s_str not in seen_strs: - schemas.append(s) - seen_strs.add(s_str) - except Exception: - # Fallback: try .default directly + + for attempt in range(2): + schemas: list = [] + seen_strs: set = set() + + for s in torch._C._jit_get_all_schemas(): + if s.name == op_name: + s_str = str(s) + if s_str not in seen_strs: + schemas.append(s) + seen_strs.add(s_str) + + if "::" in op_name: + ns, func_name = op_name.split("::", 1) + ns_obj = getattr(torch.ops, ns, None) + if ns_obj is not None: + op_obj = getattr(ns_obj, func_name, None) + if op_obj is not None: try: - s = op_obj.default._schema - s_str = str(s) - if s_str not in seen_strs: - schemas.append(s) - seen_strs.add(s_str) + for overload_name in op_obj.overloads(): + overload = getattr(op_obj, overload_name) + s = overload._schema + s_str = str(s) + if s_str not in seen_strs: + schemas.append(s) + seen_strs.add(s_str) except Exception: - pass + try: + s = op_obj.default._schema + s_str = str(s) + if s_str not in seen_strs: + schemas.append(s) + seen_strs.add(s_str) + except Exception: + pass + + if schemas or attempt > 0: + break + if not _try_auto_import(op_name, verbose): + break if verbose: print(f"Found {len(schemas)} schemas for {op_name}:") @@ -168,12 +204,34 @@ def _search_schemas(op_name: str, verbose: bool = False): class EventReplayer: + _custom_init_registry: List[CustomInit] = [ + PagedAttentionInit(), + MoeRoutingInit(), + ] + + @classmethod + def register_custom_init(cls, init: CustomInit): + """Add a custom initializer to the registry.""" + cls._custom_init_registry.append(init) + + @classmethod + def register_namespace(cls, namespace: str, modules: List[str]): + """Register a namespace-to-module mapping for auto-import.""" + _NAMESPACE_IMPORTS[namespace] = modules + + @classmethod + def list_custom_inits(cls) -> List[Tuple[str, List[str]]]: + """List all registered custom initializers and their op patterns.""" + return [(type(i).__name__, i.op_patterns) for i in cls._custom_init_registry] + def __init__( self, event: Dict[str, Any], device: str = "cuda", lazy: bool = False, verbose: bool = False, + auto_init: bool = True, + init_kwargs: Optional[Dict[str, Any]] = None, ): """ Initialize the EventReplayer with the event data and device type. @@ -181,12 +239,21 @@ def __init__( Args: event (Dict[str, Any]): From the pytorch profile json data['traceEvents'] device (str): The device type ('cuda' or 'cpu'). + lazy (bool): If True, defer tensor creation until replay(). verbose (bool): Flag to enable verbose output. + auto_init (bool): If True, automatically apply custom initializers + for ops that need realistic tensor content (e.g., paged attention + block tables, MoE routing tensors). + init_kwargs (Dict[str, Any]): Parameters passed to custom initializers + (e.g., {"moe_distribution": "zipf", "moe_zipf_s": 1.5}). """ self.event = event self.device = device self.lazy = lazy self.verbose = verbose + self._auto_init = auto_init + self._init_kwargs = init_kwargs or {} + self._inits_applied = False self._func = None self._setup() @@ -198,7 +265,7 @@ def _setup(self): print(f"Preparing {self.event['name']} event for replay") self._func, self._func_source, self._resolved_name = _resolve_op_func( - self.event["name"] + self.event["name"], verbose=self.verbose ) if self.verbose: print(f"Resolved op via {self._func_source}") @@ -245,8 +312,25 @@ def replay(self): else: args, kwargs = self.args, self.kwargs + if not self._inits_applied and self._auto_init: + self._apply_custom_inits() + self._func(*args, **kwargs) + def _apply_custom_inits(self): + """Run all applicable custom initializers on this replayer's tensors.""" + for custom_init in self._custom_init_registry: + if custom_init.applies_to(self): + try: + summary = custom_init.initialize(self, **self._init_kwargs) + if summary: + print(summary) + except Exception as e: + warnings.warn( + f"[custom init] {type(custom_init).__name__} failed: {e}" + ) + self._inits_applied = True + @staticmethod def _search_schema( event: Dict[str, Any], @@ -296,7 +380,6 @@ def _is_schema_match( print(f"\tProfiled type: {profiled_type}") is_match = True - # Optional types: schema ends with '?' => profiled type can be blank if schema_type.endswith("?"): schema_type = schema_type[:-1] if profiled_type == "": @@ -311,7 +394,7 @@ def _is_schema_match( is_match = False elif schema_type == "bool": profiled_value = event["args"]["Concrete Inputs"][idx] - if profiled_value.lower() not in ["true", "false"]: + if profiled_value.lower() not in ("true", "false"): is_match = False elif schema_type in ("int", "SymInt"): if profiled_type != "Scalar": @@ -382,27 +465,17 @@ def _is_tensor_schema_type(schema_type: str) -> bool: """Check if a schema type string represents a Tensor argument.""" if schema_type in ("Tensor", "Tensor?"): return True - # Handles annotated variants like Tensor(a!), Tensor(a), Tensor(b!) if schema_type.startswith("Tensor("): return True return False @staticmethod def _should_skip_tensor_init(evt_name: str, arg_name: str, arg_idx: int) -> bool: - """ - Determine whether a tensor argument is an output-only buffer that - does not need random initialization. - - Generalizes the old aten::fill_ / aten::copy_ special-cases to - any in-place or out-of-place output tensor. - """ - # In-place ops (name ends with '_'): first tensor is the mutated output + """Determine whether a tensor argument is an output-only buffer.""" if evt_name.endswith("_") and arg_idx == 0: return True - # Explicit 'out' arguments in .out variants if arg_name == "out": return True - # aten::copy_ destination if evt_name == "aten::copy_" and arg_name != "src": return True return False @@ -411,9 +484,7 @@ def _should_skip_tensor_init(evt_name: str, arg_name: str, arg_idx: int) -> bool def _get_event_replay_IR( event: Dict[str, Any], schema: "torch._C.FunctionSchema", verbose: bool = False ) -> Dict[str, Any]: - """ - Get the event replay IR from the event and schema. - """ + """Get the event replay IR from the event and schema.""" evt_name = event["name"] op_name, pos_args_schema, kwargs_schema, return_type = ( EventReplayer.parse_schema_string(schema) @@ -456,7 +527,6 @@ def _get_event_replay_IR( if EventReplayer._should_skip_tensor_init(evt_name, arg_name, idx): init = None profiled_dtype = event["args"]["Input type"][idx] - # Non-floating-point tensors cannot use 'normal' init if profiled_dtype in ("long", "long int", "int", "bool", "unsigned char"): init = "zeros" if init == "normal" else init value = TensorCfg( @@ -533,17 +603,7 @@ def _get_event_replay_IR_schemaless( verbose: bool = False, resolved_name: Optional[str] = None, ) -> Dict[str, Any]: - """ - Build a replay IR without a schema by inferring types directly from the - profiled data. All arguments are treated as positional. - - Heuristics: - - If Input type is a known tensor dtype -> TensorCfg - - If Input type is 'Scalar' and Concrete Inputs looks like int -> int - - If Input type is 'Scalar' and Concrete Inputs looks like float -> float - - If Input type is 'Scalar' and Concrete Inputs is true/false -> bool - - If Input type is '' and Concrete Inputs is '' -> check _STR_ARG_DEFAULTS - """ + """Build a replay IR without a schema by inferring types from profile data.""" evt_name = event["name"] schema_arg_names = EventReplayer._get_schema_arg_names( resolved_name or evt_name @@ -591,7 +651,6 @@ def _get_event_replay_IR_schemaless( value = concrete arg_type = "str" elif profiled_type == "" and concrete == "": - # Likely a dropped str arg — check known defaults hint_name = ( schema_arg_names[idx] if idx < len(schema_arg_names) else None ) @@ -636,9 +695,7 @@ def _get_event_replay_IR_schemaless( def _get_args_kwargs( event_replay_IR: Dict[str, Any], device: str = "cuda" ) -> tuple[List["torch.Tensor"], Dict[str, Any]]: - """ - Get the arguments and keyword arguments from the event replay IR. - """ + """Get the arguments and keyword arguments from the event replay IR.""" pos_args = [] for arg in event_replay_IR["list_pos_args"]: value = arg["value"] @@ -669,7 +726,6 @@ def parse_schema_string( kwarg_part = parts[1].lstrip(",").strip() if len(parts) > 1 else "" def _parse_arg(raw_arg: str) -> Tuple[str, str, Optional[str], bool]: - # Match type (may contain spaces, e.g. "Tensor($0! -> )") then name[=default]. # Greedy (.+) consumes everything up to the last whitespace before # a valid identifier, so "Tensor($0! -> ) key_cache" parses correctly. m = re.match( @@ -704,10 +760,6 @@ def _parse_arg(raw_arg: str) -> Tuple[str, str, Optional[str], bool]: def get_repro_info(self) -> Dict[str, Any]: """ Extracts the minimal, serializable information needed to reproduce the event call. - - Returns: - Dict[str, Any]: A dictionary containing the operator name and the replay IR. - Suitable for JSON serialization using the custom encoder. """ dict_repro_info = {} dict_repro_info["op_name"] = self.event["name"] diff --git a/TraceLens/EventReplay/utils.py b/TraceLens/EventReplay/utils.py index 49243c28b..841ae6269 100644 --- a/TraceLens/EventReplay/utils.py +++ b/TraceLens/EventReplay/utils.py @@ -124,32 +124,74 @@ def summarize_tensor(tensor: "torch.Tensor") -> str: return f"Tensor(shape={tensor.shape}, dtype={tensor.dtype}, device={tensor.device}, strides={tensor.stride()})" -def benchmark_func(func, device, warmup=50, avg_steps=100): - """ - Benchmark a function with warmup and average steps. - Disclaimer: This method would be inaccurate for very short ops. +_L2_FLUSH_BUFFER = None +_L2_FLUSH_SIZE = 256 * 1024 * 1024 # 256 MB -- larger than any GPU's L2 + + +def _flush_l2(device: str): + """Force-evict GPU L2 cache by reading a large buffer.""" + global _L2_FLUSH_BUFFER + torch = _get_torch_or_raise() + if _L2_FLUSH_BUFFER is None or str(_L2_FLUSH_BUFFER.device) != device: + _L2_FLUSH_BUFFER = torch.empty( + _L2_FLUSH_SIZE // 4, dtype=torch.float32, device=device + ) + _L2_FLUSH_BUFFER.sum() + + +def benchmark_func( + func, + device, + warmup=50, + avg_steps=100, + flush_l2=False, +): + """Benchmark a function with warmup and per-iteration CUDA event timing. + Args: - func (callable): The function to benchmark. - warmup (int): Number of warmup iterations. - avg_steps (int): Number of iterations to average over. + func: Callable to benchmark. + device: CUDA device string. + warmup: Number of warmup iterations. + avg_steps: Number of measured iterations. + flush_l2: If True, flush the GPU L2 cache before each measured iteration + to simulate cold-cache conditions (more representative of real + inference where other kernels pollute L2 between invocations). + Returns: - float: Average time taken per iteration in microseconds. + dict with keys: median_us, mean_us, std_us, min_us, max_us, + all_us (list of per-iteration timings in microseconds). """ torch = _get_torch_or_raise() - # Warmup phase + for _ in range(warmup): func() - - # Benchmarking phase torch.cuda.synchronize(device) - start_time = time.time() + + timings_ms: List[float] = [] for _ in range(avg_steps): + if flush_l2: + _flush_l2(device) + torch.cuda.synchronize(device) + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() func() - torch.cuda.synchronize(device) - end_time = time.time() - - elapsed_time = end_time - start_time - avg_time_sec = elapsed_time / avg_steps - avg_time_us = avg_time_sec * 1e6 - - return avg_time_us + end_evt.record() + torch.cuda.synchronize(device) + timings_ms.append(start_evt.elapsed_time(end_evt)) + + timings_us = [t * 1000.0 for t in timings_ms] + sorted_us = sorted(timings_us) + n = len(sorted_us) + median = (sorted_us[n // 2] + sorted_us[(n - 1) // 2]) / 2.0 + mean = sum(timings_us) / n + variance = sum((t - mean) ** 2 for t in timings_us) / n + std = variance ** 0.5 + return { + "median_us": median, + "mean_us": mean, + "std_us": std, + "min_us": sorted_us[0], + "max_us": sorted_us[-1], + "all_us": timings_us, + } From b2492691f28e3ad4102d5d57d56e50716d980e40 Mon Sep 17 00:00:00 2001 From: Jassani Date: Tue, 28 Apr 2026 15:38:57 -0400 Subject: [PATCH 04/11] Fix bugs, add tests, and improve EventReplay docs Bug fixes: - Fix lazy+auto_init crash: replay() now sets self.args in lazy mode so custom initializers can access them (BUG-1) - Fix get_repro_info() shallow copy corruption: no longer mutates event_replay_IR on repeated calls (BUG-2) - Fix batched_replay.py: handle benchmark_func dict return type, implement --op-filter and --op-limit flags (BUG-3) - replay() now returns the op result instead of None (CLAIM-4) - First-match-wins for custom initializers (CLAIM-1) - Exact name matching for op_patterns (no more substring matching) Tests: - Add CPU-only unit tests (test_event_replay.py, 11 tests) - Add GPU integration tests (test_event_replay_gpu.py) with kernel name validation Docs (EventReplay.md): - Fix benchmark_func example (wrong params and key names) - Remove broken Shape Metadata Guide links - Rewrite custom initializer section as step-by-step guide - Rewrite iteration annotations section with full explanation - Add batch replay CLI flag examples - Update all op_patterns to fully-qualified names --- TraceLens/EventReplay/batched_replay.py | 21 +- TraceLens/EventReplay/custom_inits.py | 7 +- TraceLens/EventReplay/event_replay.py | 45 +-- TraceLens/EventReplay/test_event_replay.py | 204 ++++++++++++++ .../EventReplay/test_event_replay_gpu.py | 261 ++++++++++++++++++ 5 files changed, 507 insertions(+), 31 deletions(-) create mode 100644 TraceLens/EventReplay/test_event_replay.py create mode 100644 TraceLens/EventReplay/test_event_replay_gpu.py diff --git a/TraceLens/EventReplay/batched_replay.py b/TraceLens/EventReplay/batched_replay.py index 241c35a0e..0228688ec 100644 --- a/TraceLens/EventReplay/batched_replay.py +++ b/TraceLens/EventReplay/batched_replay.py @@ -99,11 +99,17 @@ def _get_args_kwargs_from_ir( replayed_count = 0 errors = 0 - for i, repro_info in enumerate(repro_data_list): + ops_to_replay = repro_data_list + if args.op_filter: + ops_to_replay = [r for r in ops_to_replay if args.op_filter in r["op_name"]] + if args.op_limit: + ops_to_replay = ops_to_replay[: args.op_limit] + + for i, repro_info in enumerate(ops_to_replay): op_name = repro_info["op_name"] replay_ir = repro_info["replay_ir"] - print(f"\n[{replayed_count + 1}/{len(repro_data_list)}] Replaying: {op_name}") + print(f"\n[{replayed_count + 1}/{len(ops_to_replay)}] Replaying: {op_name}") # Get the PyTorch operation function try: @@ -151,15 +157,16 @@ def _get_args_kwargs_from_ir( errors += 1 continue # --- Benchmark the function --- - mean_time_us = benchmark_func( + metrics = benchmark_func( lambda: func(*pos_args, **kwargs), args.device, warmup=50, avg_steps=100 ) - print(f" Average time taken: {mean_time_us:.2f} microseconds") + mean_time_us = metrics["mean_us"] + print(f" Average time taken: {mean_time_us:.2f} us (median: {metrics['median_us']:.2f} us)") if "count" in repro_info: count_workload = repro_info["count"] total_time_us = mean_time_us * count_workload print(f" Count in workload: {count_workload}") - print(f" Est time in workload: {total_time_us:.2f} microseconds") + print(f" Est time in workload: {total_time_us:.2f} us") # --- Optionally sync again --- if args.device == "cuda": torch.cuda.synchronize() @@ -190,7 +197,9 @@ def _get_args_kwargs_from_ir( print("\n--- Replay Summary ---") print(f"Total operations in file: {len(repro_data_list)}") if args.op_filter: - print(f"Filter applied: '{args.op_filter}'") + print(f"Filter applied: '{args.op_filter}' ({len(ops_to_replay)} matched)") + if args.op_limit: + print(f"Limit applied: {args.op_limit}") print(f"Attempted replays: {replayed_count}") print(f"Successful replays: {replayed_count - errors}") print(f"Errors encountered: {errors}") diff --git a/TraceLens/EventReplay/custom_inits.py b/TraceLens/EventReplay/custom_inits.py index 9ded92540..bce9960ef 100644 --- a/TraceLens/EventReplay/custom_inits.py +++ b/TraceLens/EventReplay/custom_inits.py @@ -24,6 +24,7 @@ from __future__ import annotations import re +import warnings from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional @@ -116,7 +117,7 @@ class CustomInit(ABC): def applies_to(self, replayer: Any) -> bool: op_name = replayer.event.get("name", "") - return any(pat in op_name for pat in self.op_patterns) + return op_name in self.op_patterns @abstractmethod def initialize(self, replayer: Any, **kwargs) -> Optional[str]: @@ -138,7 +139,7 @@ class PagedAttentionInit(CustomInit): - Block table entries drawn from a random permutation of the pool. """ - op_patterns = ["paged_attention"] + op_patterns = ["_rocm_C::paged_attention"] def initialize(self, replayer: Any, **kwargs) -> Optional[str]: try: @@ -277,7 +278,7 @@ class MoeRoutingInit(CustomInit): [11] block_m (scalar) """ - op_patterns = ["ck_moe_stage1", "ck_moe_stage2"] + op_patterns = ["aiter::ck_moe_stage1", "aiter::ck_moe_stage2"] def initialize(self, replayer: Any, **kwargs) -> Optional[str]: try: diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index 3b02f3a35..a46e4e8ba 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -304,21 +304,22 @@ def _setup(self): def replay(self): """ Replay the event using the matched schema and event replay IR. + + Returns: + The result of the PyTorch operation. """ if self.lazy: - args, kwargs = EventReplayer._get_args_kwargs( + self.args, self.kwargs = EventReplayer._get_args_kwargs( self.event_replay_IR, device=self.device ) - else: - args, kwargs = self.args, self.kwargs if not self._inits_applied and self._auto_init: self._apply_custom_inits() - self._func(*args, **kwargs) + return self._func(*self.args, **self.kwargs) def _apply_custom_inits(self): - """Run all applicable custom initializers on this replayer's tensors.""" + """Apply the first matching custom initializer to this replayer's tensors.""" for custom_init in self._custom_init_registry: if custom_init.applies_to(self): try: @@ -329,6 +330,7 @@ def _apply_custom_inits(self): warnings.warn( f"[custom init] {type(custom_init).__name__} failed: {e}" ) + break self._inits_applied = True @staticmethod @@ -760,22 +762,21 @@ def _parse_arg(raw_arg: str) -> Tuple[str, str, Optional[str], bool]: def get_repro_info(self) -> Dict[str, Any]: """ Extracts the minimal, serializable information needed to reproduce the event call. + + Safe to call multiple times — does not mutate self.event_replay_IR. """ - dict_repro_info = {} - dict_repro_info["op_name"] = self.event["name"] - list_pos_args, list_kwargs = ( - self.event_replay_IR["list_pos_args"], - self.event_replay_IR["list_kwargs"], - ) - list_pos_args_copy, list_kwargs_copy = list_pos_args.copy(), list_kwargs.copy() - for idx, val in enumerate(list_pos_args_copy): - if isinstance(val["value"], TensorCfg): - list_pos_args_copy[idx]["value"] = val["value"].__dict__ - for idx, val in enumerate(list_kwargs_copy): - if isinstance(val["value"], TensorCfg): - list_kwargs_copy[idx]["value"] = val["value"].__dict__ - dict_repro_info["replay_ir"] = { - "list_pos_args": list_pos_args_copy, - "list_kwargs": list_kwargs_copy, + def _serialize_arg(arg: Dict[str, Any]) -> Dict[str, Any]: + val = arg["value"] + return { + **arg, + "value": val.__dict__.copy() if isinstance(val, TensorCfg) else val, + } + + ir = self.event_replay_IR + return { + "op_name": self.event["name"], + "replay_ir": { + "list_pos_args": [_serialize_arg(a) for a in ir["list_pos_args"]], + "list_kwargs": [_serialize_arg(a) for a in ir["list_kwargs"]], + }, } - return dict_repro_info diff --git a/TraceLens/EventReplay/test_event_replay.py b/TraceLens/EventReplay/test_event_replay.py new file mode 100644 index 000000000..7948caa0f --- /dev/null +++ b/TraceLens/EventReplay/test_event_replay.py @@ -0,0 +1,204 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Tests for EventReplay core functionality. + +All tests use CPU-only ops (aten::mm) so they run without a GPU. +Run from the repo root: + python -m pytest TraceLens/EventReplay/test_event_replay.py -v +""" + +import sys +import os +import pytest +import torch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +from TraceLens.EventReplay.event_replay import EventReplayer # noqa: E402 +from TraceLens.EventReplay.custom_inits import CustomInit # noqa: E402 +from TraceLens.EventReplay.utils import TensorCfg # noqa: E402 + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _make_mm_event(M=4, K=8, N=16): + """Minimal profiler event dict for aten::mm (M x K) @ (K x N).""" + return { + "name": "aten::mm", + "args": { + "Input Dims": [[M, K], [K, N]], + "Input type": ["float", "float"], + "Input Strides": [[K, 1], [N, 1]], + "Concrete Inputs": ["", ""], + }, + } + + +@pytest.fixture(autouse=True) +def _isolate_registry(): + """Save and restore the global custom-init registry around every test.""" + saved = EventReplayer._custom_init_registry[:] + yield + EventReplayer._custom_init_registry = saved + + +# --------------------------------------------------------------------------- +# BUG-1: lazy=True + auto_init=True must not crash +# --------------------------------------------------------------------------- + +class TestLazyAutoInit: + def test_lazy_replay_sets_self_args(self): + """replay() in lazy mode must populate self.args.""" + replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True, auto_init=False) + assert not hasattr(replayer, "args") + replayer.replay() + assert hasattr(replayer, "args") + assert isinstance(replayer.args, list) + + def test_lazy_with_custom_init_no_crash(self): + """A custom init that reads replayer.args must work in lazy mode.""" + accessed = {} + + class ProbeInit(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + accessed["args"] = replayer.args + accessed["kwargs"] = replayer.kwargs + return "[probe] ok" + + EventReplayer.register_custom_init(ProbeInit()) + replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True, auto_init=True) + replayer.replay() + assert "args" in accessed + assert len(accessed["args"]) == 2 # self, mat2 + + +# --------------------------------------------------------------------------- +# BUG-2: get_repro_info() must not corrupt event_replay_IR +# --------------------------------------------------------------------------- + +class TestGetReproInfo: + def test_idempotent(self): + """Calling get_repro_info() twice must produce identical output.""" + replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True) + assert replayer.get_repro_info() == replayer.get_repro_info() + + def test_does_not_mutate_ir(self): + """TensorCfg objects in the IR must survive get_repro_info().""" + replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True) + replayer.get_repro_info() + for arg in replayer.event_replay_IR["list_pos_args"]: + if arg["arg_type"].startswith("Tensor"): + assert isinstance(arg["value"], TensorCfg), ( + f"arg '{arg['arg_name']}' is {type(arg['value'])}, expected TensorCfg" + ) + + def test_replay_works_after_get_repro_info(self): + """replay() must succeed after get_repro_info() (IR still intact).""" + replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True) + replayer.get_repro_info() + result = replayer.replay() + assert isinstance(result, torch.Tensor) + + +# --------------------------------------------------------------------------- +# CLAIM-1: first-match-wins (only one init runs) +# --------------------------------------------------------------------------- + +class TestFirstMatchWins: + def test_only_first_matching_init_runs(self): + """When two inits match, only the first registered one executes.""" + log = [] + + class InitA(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + log.append("A") + + class InitB(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + log.append("B") + + EventReplayer._custom_init_registry = [InitA(), InitB()] + EventReplayer(_make_mm_event(), device="cpu", auto_init=True).replay() + assert log == ["A"] + + +# --------------------------------------------------------------------------- +# CLAIM-4: replay() returns the op result +# --------------------------------------------------------------------------- + +class TestReplayReturn: + def test_returns_tensor(self): + """replay() of aten::mm must return a correctly-shaped Tensor.""" + replayer = EventReplayer(_make_mm_event(M=4, K=8, N=16), device="cpu") + result = replayer.replay() + assert isinstance(result, torch.Tensor) + assert result.shape == (4, 16) + + def test_returns_tensor_lazy(self): + """Lazy replay must also return the result.""" + result = EventReplayer(_make_mm_event(), device="cpu", lazy=True).replay() + assert isinstance(result, torch.Tensor) + + +# --------------------------------------------------------------------------- +# Exact name matching for op_patterns +# --------------------------------------------------------------------------- + +class TestExactNameMatching: + def test_exact_match_hits(self): + """op_patterns=["aten::mm"] matches event name "aten::mm".""" + matched = [] + + class ExactInit(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + matched.append(True) + + EventReplayer._custom_init_registry = [ExactInit()] + EventReplayer(_make_mm_event(), device="cpu", auto_init=True).replay() + assert matched == [True] + + def test_substring_does_not_match(self): + """op_patterns=["mm"] must NOT match "aten::mm" (exact only).""" + matched = [] + + class SubstringInit(CustomInit): + op_patterns = ["mm"] + def initialize(self, replayer, **kwargs): + matched.append(True) + + EventReplayer._custom_init_registry = [SubstringInit()] + EventReplayer(_make_mm_event(), device="cpu", auto_init=True).replay() + assert matched == [] + + +# --------------------------------------------------------------------------- +# auto_init=False skips all inits +# --------------------------------------------------------------------------- + +class TestAutoInitDisabled: + def test_no_init_runs_when_disabled(self): + log = [] + + class AlwaysInit(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + log.append("ran") + + EventReplayer._custom_init_registry = [AlwaysInit()] + EventReplayer(_make_mm_event(), device="cpu", auto_init=False).replay() + assert log == [] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/TraceLens/EventReplay/test_event_replay_gpu.py b/TraceLens/EventReplay/test_event_replay_gpu.py new file mode 100644 index 000000000..9998ec956 --- /dev/null +++ b/TraceLens/EventReplay/test_event_replay_gpu.py @@ -0,0 +1,261 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +GPU integration tests for EventReplay. + +Profiles real ops, replays from the captured trace, and validates: + 1. Kernel name match between original and replayed execution + 2. BUG-1: lazy=True + auto_init=True works on GPU + 3. BUG-2: get_repro_info() is idempotent (doesn't corrupt IR) + 4. CLAIM-4: replay() returns a tensor + 5. CLAIM-1: first-match-wins with real ops + +Requires a GPU (MI300X / MI210 / etc). Run from the repo root: + python TraceLens/EventReplay/test_event_replay_gpu.py +""" + +import sys, os, json, time +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +import torch +from torch.profiler import profile, ProfilerActivity + +from TraceLens.EventReplay.event_replay import EventReplayer +from TraceLens.EventReplay.custom_inits import CustomInit +from TraceLens.EventReplay.utils import TensorCfg + +assert torch.cuda.is_available(), "GPU required for this test" + +DEVICE = "cuda" +TRACE_FILE = "/tmp/test_event_replay_gpu_trace.json" +REPLAY_TRACE = "/tmp/test_event_replay_gpu_replay.json" + +# --------------------------------------------------------------------------- +# Step 1: Profile a set of real ops +# --------------------------------------------------------------------------- + +print("=" * 80) +print("Step 1: Profiling real ops") +print("=" * 80) + +M, K, N = 256, 1024, 512 +mm_a = torch.randn(M, K, dtype=torch.bfloat16, device=DEVICE) +mm_b = torch.randn(K, N, dtype=torch.bfloat16, device=DEVICE) +add_a = torch.randn(M, N, dtype=torch.bfloat16, device=DEVICE) +add_b = torch.randn(M, N, dtype=torch.bfloat16, device=DEVICE) +bmm_a = torch.randn(4, M, K, dtype=torch.bfloat16, device=DEVICE) +bmm_b = torch.randn(4, K, N, dtype=torch.bfloat16, device=DEVICE) + +def run_ops(): + torch.mm(mm_a, mm_b) + torch.add(add_a, add_b) + torch.bmm(bmm_a, bmm_b) + torch.mul(add_a, add_b) + torch.sigmoid(add_a) + +for _ in range(10): + run_ops() +torch.cuda.synchronize() + +def trace_handler(p): + p.export_chrome_trace(TRACE_FILE) + +wait, warmup, active = 3, 3, 5 +with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=1), + record_shapes=True, + on_trace_ready=trace_handler, +) as p: + for _ in range(wait + warmup + active): + run_ops() + p.step() + +print(f"Trace saved to {TRACE_FILE}") + +# --------------------------------------------------------------------------- +# Step 2: Load trace and find events +# --------------------------------------------------------------------------- + +print(f"\n{'=' * 80}") +print("Step 2: Loading trace") +print("=" * 80) + +with open(TRACE_FILE) as f: + trace_data = json.load(f) + +all_events = trace_data.get("traceEvents", []) + +OPS_TO_TEST = ["aten::mm", "aten::add", "aten::bmm", "aten::mul", "aten::sigmoid"] + +def find_event(events, op_name): + """Find a cpu_op event with the right name and shape data.""" + candidates = [ + e for e in events + if e.get("cat") == "cpu_op" + and e.get("name") == op_name + and "args" in e + and "Input Dims" in e.get("args", {}) + ] + if candidates: + return candidates[len(candidates) // 2] + return None + +results = [] +errors = [] + +# --------------------------------------------------------------------------- +# Step 3: Replay each op and validate +# --------------------------------------------------------------------------- + +print(f"\n{'=' * 80}") +print("Step 3: Replay and validate") +print("=" * 80) + +print(f"\n{'Op':<30} {'Kernel Match':<15} {'Return':<10} {'Lazy':<10} {'ReproInfo':<12} {'Status'}") +print("-" * 100) + +for op_name in OPS_TO_TEST: + evt = find_event(all_events, op_name) + if evt is None: + print(f"{op_name:<30} {'SKIP':<15} {'---':<10} {'---':<10} {'---':<12} not in trace") + continue + + status = [] + + # --- Test: basic replay returns a result (CLAIM-4) --- + try: + replayer = EventReplayer(evt, device=DEVICE, auto_init=False) + result = replayer.replay() + returns_ok = isinstance(result, torch.Tensor) + except Exception as e: + returns_ok = False + status.append(f"replay error: {e}") + + # --- Test: lazy mode works (BUG-1) --- + try: + lazy_replayer = EventReplayer(evt, device=DEVICE, lazy=True, auto_init=False) + lazy_result = lazy_replayer.replay() + lazy_ok = isinstance(lazy_result, torch.Tensor) + assert hasattr(lazy_replayer, "args"), "self.args not set after lazy replay" + except Exception as e: + lazy_ok = False + status.append(f"lazy error: {e}") + + # --- Test: get_repro_info idempotent (BUG-2) --- + try: + repro_replayer = EventReplayer(evt, device=DEVICE, lazy=True) + info1 = repro_replayer.get_repro_info() + info2 = repro_replayer.get_repro_info() + repro_ok = (info1 == info2) + for arg in repro_replayer.event_replay_IR["list_pos_args"]: + if arg["arg_type"].startswith("Tensor"): + assert isinstance(arg["value"], TensorCfg), "IR corrupted after get_repro_info" + repro_replayer.replay() + except Exception as e: + repro_ok = False + status.append(f"repro error: {e}") + + # --- Test: kernel name match --- + kernel_match = "N/A" + try: + replay_replayer = EventReplayer(evt, device=DEVICE, auto_init=False) + for _ in range(5): + replay_replayer.replay() + torch.cuda.synchronize() + + def th(p): + p.export_chrome_trace(REPLAY_TRACE) + + w, wu, a = 2, 2, 3 + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=torch.profiler.schedule(wait=w, warmup=wu, active=a, repeat=1), + record_shapes=True, + on_trace_ready=th, + ) as p: + for _ in range(w + wu + a): + replay_replayer.replay() + p.step() + + with open(REPLAY_TRACE) as f: + replay_trace = json.load(f) + + orig_gpu = set() + for e in all_events: + if e.get("cat") == "kernel" and e.get("name", ""): + orig_gpu.add(e["name"]) + + replay_gpu = set() + for e in replay_trace.get("traceEvents", []): + if e.get("cat") == "kernel" and e.get("name", ""): + replay_gpu.add(e["name"]) + + kernel_match = "MATCH" if replay_gpu.issubset(orig_gpu) else "MISMATCH" + except Exception as e: + kernel_match = "ERROR" + status.append(f"kernel error: {e}") + + ok = returns_ok and lazy_ok and repro_ok and kernel_match in ("MATCH", "N/A") + tag = "PASS" if ok else "FAIL" + detail = "; ".join(status) if status else "" + + print(f"{op_name:<30} {kernel_match:<15} {'OK' if returns_ok else 'FAIL':<10} {'OK' if lazy_ok else 'FAIL':<10} {'OK' if repro_ok else 'FAIL':<12} {tag} {detail}") + results.append({"op": op_name, "ok": ok, "kernel": kernel_match, + "returns": returns_ok, "lazy": lazy_ok, "repro": repro_ok}) + +# --------------------------------------------------------------------------- +# Step 4: First-match-wins test (CLAIM-1) on GPU +# --------------------------------------------------------------------------- + +print(f"\n{'=' * 80}") +print("Step 4: First-match-wins (CLAIM-1)") +print("=" * 80) + +log = [] + +class InitA(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + log.append("A") + +class InitB(CustomInit): + op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): + log.append("B") + +saved_registry = EventReplayer._custom_init_registry[:] +try: + EventReplayer._custom_init_registry = [InitA(), InitB()] + mm_evt = find_event(all_events, "aten::mm") + if mm_evt: + r = EventReplayer(mm_evt, device=DEVICE, auto_init=True) + r.replay() + first_match_ok = (log == ["A"]) + print(f" First-match-wins: {'PASS' if first_match_ok else 'FAIL'} (log={log})") + else: + first_match_ok = True + print(" SKIP: aten::mm not in trace") +finally: + EventReplayer._custom_init_registry = saved_registry + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- + +print(f"\n{'=' * 80}") +print("Summary") +print("=" * 80) + +total = len(results) +passed = sum(1 for r in results if r["ok"]) +print(f"Op tests: {passed}/{total} passed") +print(f"First-match-wins: {'PASS' if first_match_ok else 'FAIL'}") + +all_pass = passed == total and first_match_ok +print(f"\nOverall: {'ALL PASSED' if all_pass else 'FAILURES DETECTED'}") +sys.exit(0 if all_pass else 1) From 0b27256a5d31cecd336c842a6722a1269db086fa Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 15:47:32 -0400 Subject: [PATCH 05/11] Port EventReplay custom-op docs into the ROCm how-to page. The legacy docs/EventReplay.md was removed on main; keep the new content on docs/how-to/event-replay.md so the published toc still matches. Co-authored-by: Cursor --- docs/how-to/event-replay.md | 257 +++++++++++++++++++++++++++++++++--- 1 file changed, 236 insertions(+), 21 deletions(-) diff --git a/docs/how-to/event-replay.md b/docs/how-to/event-replay.md index b92bab51e..039ae3108 100644 --- a/docs/how-to/event-replay.md +++ b/docs/how-to/event-replay.md @@ -7,13 +7,17 @@ See LICENSE for license information. # Replay a single operation in TraceLens ```{meta} -:description: Learn how to isolate a single GPU operation into a minimal, self-contained replay script using TraceLens EventReplay for focused debugging. -:keywords: TraceLens, EventReplay, GPU debugging, reproducer, operator replay, PyTorch profiler, ROCm, kernel isolation, IP-safe +:description: Isolate a GPU operation from a PyTorch profiler trace into a portable EventReplay IR, including custom ops, auto-import, and custom tensor initializers. +:keywords: TraceLens, EventReplay, GPU debugging, reproducer, operator replay, PyTorch profiler, ROCm, custom op, aiter, vLLM, paged attention, MoE ``` This topic shows how to isolate an operation from a trace into a minimal, -self-contained replay — useful for focused debugging and for sharing IP-safe -reproducers with kernel or framework developers. +self-contained replay. That's useful for focused debugging and for sharing +IP-safe reproducers with kernel or framework developers. + +EventReplay works with `aten::` ops and with custom ops from other namespaces +(for example `aiter::` and `_rocm_C::`) when the library that registers the op +is importable. ## Before you begin @@ -21,6 +25,8 @@ Confirm you have the following before continuing. - [TraceLens installed](../install/install.md). - A PyTorch profiler trace containing the operation you want to isolate. +- For custom ops, the library that registers the op (for example `aiter` or + `vllm`) installed in the same environment you replay in. ## How it works @@ -33,7 +39,8 @@ code, the artifacts can be shared without exposing model IP. ```{note} EventReplay allocates inputs with randomized data matching the recorded tensor shapes, so replay timings approximate — but do not exactly reproduce — the -original run. +original run. Integer index and routing tensors default to zeros unless a +[custom initializer](#custom-initializers) fills them. ``` ## Step 1: Identify the operation @@ -48,7 +55,7 @@ appears in the `ex_UID` column. ## Step 2: Replay a single event (SDK) -Use the `EventReplayer` class to replay the event identified by its UID on the device of your choice. +Use the `EventReplayer` class to replay the event identified by its UID: ```python from TraceLens import TreePerfAnalyzer, EventReplayer @@ -61,8 +68,78 @@ replayer = EventReplayer(event, device="cuda") replayer.replay() ``` -The `examples/event_replayer_example.ipynb` notebook walks through the same flow -interactively, including selecting the target event from the tree. +The +[`event_replayer_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/event_replayer_example.ipynb) +notebook walks through the same flow interactively, including selecting the +target event from the tree. + +## Inspect the IR + +The profiler stores arguments as unlabeled dimension lists. EventReplay zips +those arrays with the op's registered schema so you can read named tensors and +scalars without launching the kernel: + +```python +replayer = EventReplayer(event, lazy=True) +ir = replayer.get_repro_info() +``` + +**Profiler event** for `aten::mm` (shapes only): + +```json +{ + "cat": "cpu_op", + "name": "aten::mm", + "args": { + "Input Dims": [[20, 2048], [2048, 11264]], + "Input type": ["BFloat16", "BFloat16"] + } +} +``` + +**EventReplay IR** for the same event: + +```json +{ + "op_name": "aten::mm", + "replay_ir": { + "list_pos_args": [ + { + "arg_name": "self", + "arg_type": "Tensor", + "value": { + "shape": [20, 2048], + "dtype": "c10::BFloat16", + "strides": [2048, 1], + "init": "normal" + } + }, + { + "arg_name": "mat2", + "arg_type": "Tensor", + "value": { + "shape": [2048, 11264], + "dtype": "c10::BFloat16", + "strides": [1, 2048], + "init": "normal" + } + } + ] + } +} +``` + +That IR is a BF16 GEMM with `M=20`, `K=2048`, `N=11264`, and a column-major +`mat2` (stride pattern `[1, K]`). The same mapping is what makes fused custom +ops readable: a raw `aiter::ck_moe_stage1` event is a long list of unlabeled +scalars, while the IR names `hidden_states`, `w1`, `w2`, `sorted_token_ids`, +`topk`, `block_m`, and so on. + +Schema lookup uses the PyTorch dispatcher (`torch._C._jit_get_all_schemas()` or +`torch.ops`). If no schema is registered, EventReplay falls back to a +schemaless IR that infers types from the profiled arrays. Ops called as plain +Python functions (for example a Triton kernel launched directly) don't appear in +the dispatcher; wrap them with `torch.library.custom_op` if you need a schema. ## Batch replay and benchmark @@ -87,10 +164,10 @@ Before running `batched_replay.py`, package the IR and its companion scripts into a self-contained bundle. The bundle can be run without TraceLens or the original model and is safe to share without exposing model IP. It contains: -- `event_replay_ir.json:` serialized operator replay instructions. -- `utils.py:` tensor-creation and helper utilities that `batched_replay.py` imports. -- `batched_replay.py:` batch replay and benchmark script. -- `batched_replay_readme.md:` run instructions. +- **`event_replay_ir.json`:** serialized operator replay instructions. +- **`utils.py`:** tensor-creation and helper utilities that `batched_replay.py` imports. +- **`batched_replay.py`:** batch replay and benchmark script. +- **`batched_replay_readme.md`:** run instructions. See the [`event_replayer_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/event_replayer_example.ipynb) @@ -103,22 +180,159 @@ for a source checkout), run: ```bash python batched_replay.py event_replay_ir.json +python batched_replay.py -v event_replay_ir.json +python batched_replay.py --op-filter aten::mm event_replay_ir.json +python batched_replay.py --op-limit 5 event_replay_ir.json ``` `batched_replay.py` imports `utils.py` from the same directory, so the command must be run from that location. -`batched_replay.py` accepts: +The following table describes the CLI flags. + +| Argument | Default | Description | +|---|---|---| +| `repro_file` | (required) | Path to the JSON IR file from `get_repro_info()`. | +| `--device` | `cuda` | Device to run on (`cuda` or `cpu`). Falls back to `cpu` if CUDA isn't available. | +| `--verbose` / `-v` | off | Print reconstructed arguments and per-op detail. | +| `--stop-on-error` | off | Abort on the first failure instead of continuing. | +| `--op-filter` | `None` | Only replay ops whose name contains this substring (for example `aten::add`). | +| `--op-limit` | `None` | Replay at most this many ops after filtering. | + +Each replayed op prints average and median time, then a summary of attempted, +successful, and failed replays. + +## Custom ops and auto-import + +When EventReplayer sees a non-`aten` namespace, it tries to import the library +that registers the op schema. Built-in mappings: + +| Namespace | Imported modules | +|---|---| +| `aiter` | `aiter` | +| `_rocm_C` | `vllm._rocm_C` | +| `_C` | `vllm._C` | +| `_C_cache_ops` | `vllm._C` | +| `vllm` | `vllm._C`, `vllm._rocm_C` | + +Register additional namespaces: + +```python +from TraceLens.EventReplay import EventReplayer + +EventReplayer.register_namespace("my_lib", ["my_lib.ops"]) +``` + +Resolution order is the JIT registry, then `torch.ops`, then a direct Python +module import, then auto-import, then any name aliases (for example +`_rocm_C::wvSplitK` → `_rocm_C::wvSpltK`). + +## Custom initializers + +Profiler traces record shapes and dtypes, not tensor values. Zero-filled index +and routing tensors can make a kernel short-circuit (no real work). Custom +initializers fill those tensors with plausible values before `replay()`. They +run when `auto_init=True` (the default). + +### Built-in initializers + +These activate when the event name matches exactly: + +- **`PagedAttentionInit`:** `_rocm_C::paged_attention`. Fills `block_tables` + (permutation of the physical block pool), `seq_lens` (`max_seq_len` for every + sequence), and `query_start_loc` (CSR indptr of per-sequence query counts). + Uses iteration annotations when present; otherwise heuristics. +- **`MoeRoutingInit`:** `aiter::ck_moe_stage1` and `aiter::ck_moe_stage2`. + Builds `sorted_token_ids`, `sorted_expert_ids`, and `num_valid_ids`. Pass + `init_kwargs={"moe_distribution": "zipf", "moe_zipf_s": 1.5}` for a skewed + expert load; the default is uniform. + +```python +replayer = EventReplayer( + event, + device="cuda", + init_kwargs={"moe_distribution": "zipf", "moe_zipf_s": 1.5}, +) +``` + +### Write your own initializer + +1. Subclass `CustomInit` and set `op_patterns` to the **exact** profiler event + name (for example `"aten::index_add_"`, not `"index_add"`). +2. Implement `initialize()` and mutate `replayer.args` / `replayer.kwargs` in + place. Look up arguments by name from `replayer.event_replay_IR`. +3. Register with `EventReplayer.register_custom_init(YourInit())`. + +```python +from TraceLens.EventReplay import EventReplayer, CustomInit + +class IndexAddInit(CustomInit): + op_patterns = ["aten::index_add_"] + + def initialize(self, replayer, **kwargs): + import torch + + ir = replayer.event_replay_IR + arg_names = [a["arg_name"] for a in ir["list_pos_args"]] + self_tensor = replayer.args[arg_names.index("self")] + dim = replayer.args[arg_names.index("dim")] + index = replayer.args[arg_names.index("index")] + dim_size = self_tensor.shape[dim] + index.copy_( + torch.randint(0, dim_size, index.shape, device=index.device) + ) + return f"[custom init] index_add — index randint(0, {dim_size})" + +EventReplayer.register_custom_init(IndexAddInit()) +``` + +`replay()` applies the **first** matching initializer. Built-ins are registered +first; `register_custom_init` appends, so a user initializer for the same exact +op name as a built-in doesn't run. List the registry with +`EventReplayer.list_custom_inits()`. + +The built-in implementations are in `TraceLens/EventReplay/custom_inits.py`. + +## Iteration annotations (vLLM traces) + +Paged attention's `query_start_loc` is a compressed sparse row (CSR) indptr that +encodes how many query tokens each sequence contributes. In a mixed batch some +sequences are prefill (many query tokens) and others are decode (one token +each). The profiler captures the tensor shape, not that split. + +vLLM emits a `user_annotation` per `execute_model` step whose name encodes the +split, for example `execute_context_2(18)_generation_5(5)`: two prefill +sequences with 18 query tokens total, and five decode sequences with five tokens +(one each). + +```python +from TraceLens import TreePerfAnalyzer +from TraceLens.EventReplay import EventReplayer, extract_batch_context + +analyzer = TreePerfAnalyzer.from_file("vllm_trace.json") +num_annotated = extract_batch_context(analyzer) + +event = analyzer.tree.get_UID2event(some_uid) +replayer = EventReplayer(event, device="cuda") +replayer.replay() +``` + +`extract_batch_context` attaches `event["batch_context"]` with `n_prefill`, +`prefill_tokens`, `n_decode`, and `decode_tokens`. `PagedAttentionInit` uses +that dict for `query_start_loc`. Without annotations it assumes pure decode +when `query_tokens == num_seqs`, and pure prefill otherwise. That approximation +is weak for mixed batches. -- `--device {cuda,cpu}:` device to run on (default `cuda`). -- `--op-filter :` only replay ops whose name contains the substring - (for example, `aten::convolution`). -- `--op-limit :` replay at most `N` ops. -- `--stop-on-error:` abort on the first failure instead of continuing. -- `--verbose` / `-v:` print reconstructed arguments and per-op detail. +## Known limitations -Each replayed op prints its reconstructed arguments, average time, and result -tensor, followed by a summary of attempted, successful, and failed replays. +- **Unregistered ops are invisible.** Triton kernels called directly from Python + have no dispatcher schema. Wrap them in `torch.library.custom_op` in the + upstream library if you need IR extraction. +- **Single-op isolation versus the real workload.** Replay runs each op with no + surrounding memory traffic. Timings are a lower bound on in-model performance. +- **Data-dependent kernels.** Custom initializers are plausible, not bitwise + copies of the original tensors. Timing can differ when control flow depends on + values. ## Related topics @@ -126,3 +340,4 @@ tensor, followed by a summary of attempted, successful, and failed replays. - [Install TraceLens](../install/install.md) - [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) - [API reference](../reference/api-reference.md) +- [Tensor shape metadata](../conceptual/shape-metadata.md) From bf7b5d0eabc31ec0aaa7331eeddfc4809fc649d1 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 16:22:53 -0400 Subject: [PATCH 06/11] Match extract_batch_context to exact PagedAttentionInit op names. Substring matching could annotate ops the builtin init never runs; keep one name list and document op_patterns as exact event names. Co-authored-by: Cursor --- TraceLens/EventReplay/custom_inits.py | 13 +++--- TraceLens/EventReplay/test_event_replay.py | 47 +++++++++++++++++++++- docs/how-to/event-replay.md | 12 +++--- 3 files changed, 60 insertions(+), 12 deletions(-) diff --git a/TraceLens/EventReplay/custom_inits.py b/TraceLens/EventReplay/custom_inits.py index bce9960ef..9ad5c70aa 100644 --- a/TraceLens/EventReplay/custom_inits.py +++ b/TraceLens/EventReplay/custom_inits.py @@ -14,7 +14,8 @@ To add a custom initializer for a new op family: 1. Subclass ``CustomInit`` - 2. Set ``op_patterns`` to one or more substrings that match the op name + 2. Set ``op_patterns`` to one or more exact profiler event names + (for example ``"aten::index_add_"``, not ``"index_add"``) 3. Implement ``initialize()`` — mutate replayer.args / replayer.kwargs in-place 4. Return a one-line summary string (printed by EventReplayer) 5. Register with ``EventReplayer.register_custom_init(YourInit())`` @@ -24,7 +25,6 @@ from __future__ import annotations import re -import warnings from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional @@ -47,8 +47,9 @@ def extract_batch_context(analyzer: Any) -> int: This function: 1. Collects all such annotations with their ``[ts, ts+dur]`` ranges. - 2. For every ``paged_attention`` cpu_op event, finds the enclosing - annotation by timestamp and attaches a ``batch_context`` dict:: + 2. For every event whose name is in ``PagedAttentionInit.op_patterns``, + finds the enclosing annotation by timestamp and attaches a + ``batch_context`` dict:: event["batch_context"] = { "n_prefill": 2, @@ -62,7 +63,7 @@ def extract_batch_context(analyzer: Any) -> int: yields the trace event list). Returns: - Number of paged_attention events that were annotated. + Number of ``PagedAttentionInit`` ops that were annotated. """ annotations = [] for e in analyzer.tree.events: @@ -91,7 +92,7 @@ def extract_batch_context(analyzer: Any) -> int: annotated = 0 for e in analyzer.tree.events: name = e.get("name", "") - if "paged_attention" not in name: + if name not in PagedAttentionInit.op_patterns: continue if not e.get("args", {}).get("Input Dims"): continue diff --git a/TraceLens/EventReplay/test_event_replay.py b/TraceLens/EventReplay/test_event_replay.py index 7948caa0f..28aa33a3e 100644 --- a/TraceLens/EventReplay/test_event_replay.py +++ b/TraceLens/EventReplay/test_event_replay.py @@ -20,7 +20,11 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) from TraceLens.EventReplay.event_replay import EventReplayer # noqa: E402 -from TraceLens.EventReplay.custom_inits import CustomInit # noqa: E402 +from TraceLens.EventReplay.custom_inits import ( # noqa: E402 + CustomInit, + PagedAttentionInit, + extract_batch_context, +) from TraceLens.EventReplay.utils import TensorCfg # noqa: E402 @@ -200,5 +204,46 @@ def initialize(self, replayer, **kwargs): assert log == [] +# --------------------------------------------------------------------------- +# extract_batch_context uses the same exact names as PagedAttentionInit +# --------------------------------------------------------------------------- + +class _FakeAnalyzer: + def __init__(self, events): + self.tree = type("Tree", (), {"events": events})() + + +def _annotation(ts=0, dur=100): + return { + "cat": "user_annotation", + "name": "execute_context_2(18)_generation_5(5)", + "ts": ts, + "dur": dur, + } + + +def _cpu_op(name, ts=10): + return { + "name": name, + "ts": ts, + "args": {"Input Dims": [[1, 1]]}, + } + + +class TestExtractBatchContextExactName: + def test_exact_paged_attention_is_annotated(self): + op = _cpu_op("_rocm_C::paged_attention") + n = extract_batch_context(_FakeAnalyzer([_annotation(), op])) + assert n == 1 + assert op["batch_context"]["n_prefill"] == 2 + assert "_rocm_C::paged_attention" in PagedAttentionInit.op_patterns + + def test_substring_name_is_not_annotated(self): + op = _cpu_op("aiter::paged_attention_v1") + n = extract_batch_context(_FakeAnalyzer([_annotation(), op])) + assert n == 0 + assert "batch_context" not in op + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/docs/how-to/event-replay.md b/docs/how-to/event-replay.md index 039ae3108..932be20fa 100644 --- a/docs/how-to/event-replay.md +++ b/docs/how-to/event-replay.md @@ -317,11 +317,13 @@ replayer = EventReplayer(event, device="cuda") replayer.replay() ``` -`extract_batch_context` attaches `event["batch_context"]` with `n_prefill`, -`prefill_tokens`, `n_decode`, and `decode_tokens`. `PagedAttentionInit` uses -that dict for `query_start_loc`. Without annotations it assumes pure decode -when `query_tokens == num_seqs`, and pure prefill otherwise. That approximation -is weak for mixed batches. +`extract_batch_context` only annotates events whose name is in +`PagedAttentionInit.op_patterns` (exact match, currently +`_rocm_C::paged_attention`). It attaches `event["batch_context"]` with +`n_prefill`, `prefill_tokens`, `n_decode`, and `decode_tokens`. +`PagedAttentionInit` uses that dict for `query_start_loc`. Without annotations +it assumes pure decode when `query_tokens == num_seqs`, and pure prefill +otherwise. That approximation is weak for mixed batches. ## Known limitations From ca587e431567ea40dc521275c09882f8c3ac2bc7 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 16:46:38 -0400 Subject: [PATCH 07/11] Move EventReplay tests under tests/ and keep GPU smoke as an example. CPU CI collects pytest tests/ without torch or a GPU; torch-free custom-init and extract_batch_context cases run there, torch/CPU cases skip, and the profile-replay script is no longer a package test_*.py. Co-authored-by: Cursor --- TraceLens/EventReplay/test_event_replay.py | 249 ------------------ .../event_replay_gpu_smoke.py | 21 +- tests/test_event_replay_module.py | 145 +++++++++- 3 files changed, 149 insertions(+), 266 deletions(-) delete mode 100644 TraceLens/EventReplay/test_event_replay.py rename TraceLens/EventReplay/test_event_replay_gpu.py => examples/event_replay_gpu_smoke.py (92%) diff --git a/TraceLens/EventReplay/test_event_replay.py b/TraceLens/EventReplay/test_event_replay.py deleted file mode 100644 index 28aa33a3e..000000000 --- a/TraceLens/EventReplay/test_event_replay.py +++ /dev/null @@ -1,249 +0,0 @@ -############################################################################### -# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -""" -Tests for EventReplay core functionality. - -All tests use CPU-only ops (aten::mm) so they run without a GPU. -Run from the repo root: - python -m pytest TraceLens/EventReplay/test_event_replay.py -v -""" - -import sys -import os -import pytest -import torch - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) - -from TraceLens.EventReplay.event_replay import EventReplayer # noqa: E402 -from TraceLens.EventReplay.custom_inits import ( # noqa: E402 - CustomInit, - PagedAttentionInit, - extract_batch_context, -) -from TraceLens.EventReplay.utils import TensorCfg # noqa: E402 - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - -def _make_mm_event(M=4, K=8, N=16): - """Minimal profiler event dict for aten::mm (M x K) @ (K x N).""" - return { - "name": "aten::mm", - "args": { - "Input Dims": [[M, K], [K, N]], - "Input type": ["float", "float"], - "Input Strides": [[K, 1], [N, 1]], - "Concrete Inputs": ["", ""], - }, - } - - -@pytest.fixture(autouse=True) -def _isolate_registry(): - """Save and restore the global custom-init registry around every test.""" - saved = EventReplayer._custom_init_registry[:] - yield - EventReplayer._custom_init_registry = saved - - -# --------------------------------------------------------------------------- -# BUG-1: lazy=True + auto_init=True must not crash -# --------------------------------------------------------------------------- - -class TestLazyAutoInit: - def test_lazy_replay_sets_self_args(self): - """replay() in lazy mode must populate self.args.""" - replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True, auto_init=False) - assert not hasattr(replayer, "args") - replayer.replay() - assert hasattr(replayer, "args") - assert isinstance(replayer.args, list) - - def test_lazy_with_custom_init_no_crash(self): - """A custom init that reads replayer.args must work in lazy mode.""" - accessed = {} - - class ProbeInit(CustomInit): - op_patterns = ["aten::mm"] - def initialize(self, replayer, **kwargs): - accessed["args"] = replayer.args - accessed["kwargs"] = replayer.kwargs - return "[probe] ok" - - EventReplayer.register_custom_init(ProbeInit()) - replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True, auto_init=True) - replayer.replay() - assert "args" in accessed - assert len(accessed["args"]) == 2 # self, mat2 - - -# --------------------------------------------------------------------------- -# BUG-2: get_repro_info() must not corrupt event_replay_IR -# --------------------------------------------------------------------------- - -class TestGetReproInfo: - def test_idempotent(self): - """Calling get_repro_info() twice must produce identical output.""" - replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True) - assert replayer.get_repro_info() == replayer.get_repro_info() - - def test_does_not_mutate_ir(self): - """TensorCfg objects in the IR must survive get_repro_info().""" - replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True) - replayer.get_repro_info() - for arg in replayer.event_replay_IR["list_pos_args"]: - if arg["arg_type"].startswith("Tensor"): - assert isinstance(arg["value"], TensorCfg), ( - f"arg '{arg['arg_name']}' is {type(arg['value'])}, expected TensorCfg" - ) - - def test_replay_works_after_get_repro_info(self): - """replay() must succeed after get_repro_info() (IR still intact).""" - replayer = EventReplayer(_make_mm_event(), device="cpu", lazy=True) - replayer.get_repro_info() - result = replayer.replay() - assert isinstance(result, torch.Tensor) - - -# --------------------------------------------------------------------------- -# CLAIM-1: first-match-wins (only one init runs) -# --------------------------------------------------------------------------- - -class TestFirstMatchWins: - def test_only_first_matching_init_runs(self): - """When two inits match, only the first registered one executes.""" - log = [] - - class InitA(CustomInit): - op_patterns = ["aten::mm"] - def initialize(self, replayer, **kwargs): - log.append("A") - - class InitB(CustomInit): - op_patterns = ["aten::mm"] - def initialize(self, replayer, **kwargs): - log.append("B") - - EventReplayer._custom_init_registry = [InitA(), InitB()] - EventReplayer(_make_mm_event(), device="cpu", auto_init=True).replay() - assert log == ["A"] - - -# --------------------------------------------------------------------------- -# CLAIM-4: replay() returns the op result -# --------------------------------------------------------------------------- - -class TestReplayReturn: - def test_returns_tensor(self): - """replay() of aten::mm must return a correctly-shaped Tensor.""" - replayer = EventReplayer(_make_mm_event(M=4, K=8, N=16), device="cpu") - result = replayer.replay() - assert isinstance(result, torch.Tensor) - assert result.shape == (4, 16) - - def test_returns_tensor_lazy(self): - """Lazy replay must also return the result.""" - result = EventReplayer(_make_mm_event(), device="cpu", lazy=True).replay() - assert isinstance(result, torch.Tensor) - - -# --------------------------------------------------------------------------- -# Exact name matching for op_patterns -# --------------------------------------------------------------------------- - -class TestExactNameMatching: - def test_exact_match_hits(self): - """op_patterns=["aten::mm"] matches event name "aten::mm".""" - matched = [] - - class ExactInit(CustomInit): - op_patterns = ["aten::mm"] - def initialize(self, replayer, **kwargs): - matched.append(True) - - EventReplayer._custom_init_registry = [ExactInit()] - EventReplayer(_make_mm_event(), device="cpu", auto_init=True).replay() - assert matched == [True] - - def test_substring_does_not_match(self): - """op_patterns=["mm"] must NOT match "aten::mm" (exact only).""" - matched = [] - - class SubstringInit(CustomInit): - op_patterns = ["mm"] - def initialize(self, replayer, **kwargs): - matched.append(True) - - EventReplayer._custom_init_registry = [SubstringInit()] - EventReplayer(_make_mm_event(), device="cpu", auto_init=True).replay() - assert matched == [] - - -# --------------------------------------------------------------------------- -# auto_init=False skips all inits -# --------------------------------------------------------------------------- - -class TestAutoInitDisabled: - def test_no_init_runs_when_disabled(self): - log = [] - - class AlwaysInit(CustomInit): - op_patterns = ["aten::mm"] - def initialize(self, replayer, **kwargs): - log.append("ran") - - EventReplayer._custom_init_registry = [AlwaysInit()] - EventReplayer(_make_mm_event(), device="cpu", auto_init=False).replay() - assert log == [] - - -# --------------------------------------------------------------------------- -# extract_batch_context uses the same exact names as PagedAttentionInit -# --------------------------------------------------------------------------- - -class _FakeAnalyzer: - def __init__(self, events): - self.tree = type("Tree", (), {"events": events})() - - -def _annotation(ts=0, dur=100): - return { - "cat": "user_annotation", - "name": "execute_context_2(18)_generation_5(5)", - "ts": ts, - "dur": dur, - } - - -def _cpu_op(name, ts=10): - return { - "name": name, - "ts": ts, - "args": {"Input Dims": [[1, 1]]}, - } - - -class TestExtractBatchContextExactName: - def test_exact_paged_attention_is_annotated(self): - op = _cpu_op("_rocm_C::paged_attention") - n = extract_batch_context(_FakeAnalyzer([_annotation(), op])) - assert n == 1 - assert op["batch_context"]["n_prefill"] == 2 - assert "_rocm_C::paged_attention" in PagedAttentionInit.op_patterns - - def test_substring_name_is_not_annotated(self): - op = _cpu_op("aiter::paged_attention_v1") - n = extract_batch_context(_FakeAnalyzer([_annotation(), op])) - assert n == 0 - assert "batch_context" not in op - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/TraceLens/EventReplay/test_event_replay_gpu.py b/examples/event_replay_gpu_smoke.py similarity index 92% rename from TraceLens/EventReplay/test_event_replay_gpu.py rename to examples/event_replay_gpu_smoke.py index 9998ec956..7574732a0 100644 --- a/TraceLens/EventReplay/test_event_replay_gpu.py +++ b/examples/event_replay_gpu_smoke.py @@ -1,25 +1,22 @@ ############################################################################### -# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2024 - 2026 Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### -""" -GPU integration tests for EventReplay. +"""Manual GPU smoke for EventReplay (not collected by pytest). + +Profiles a few aten ops, replays them from the captured trace, and prints +pass/fail. Requires a CUDA/HIP GPU. From the repo root: -Profiles real ops, replays from the captured trace, and validates: - 1. Kernel name match between original and replayed execution - 2. BUG-1: lazy=True + auto_init=True works on GPU - 3. BUG-2: get_repro_info() is idempotent (doesn't corrupt IR) - 4. CLAIM-4: replay() returns a tensor - 5. CLAIM-1: first-match-wins with real ops + python examples/event_replay_gpu_smoke.py -Requires a GPU (MI300X / MI210 / etc). Run from the repo root: - python TraceLens/EventReplay/test_event_replay_gpu.py +CI unit tests use ``tests/test_event_replay_module.py`` with ``@pytest.mark.gpu`` +instead of this script. """ import sys, os, json, time -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import torch from torch.profiler import profile, ProfilerActivity diff --git a/tests/test_event_replay_module.py b/tests/test_event_replay_module.py index 6a57459b6..e57398477 100644 --- a/tests/test_event_replay_module.py +++ b/tests/test_event_replay_module.py @@ -6,8 +6,11 @@ """Unit tests for TraceLens/EventReplay. -Schema parsing and IR construction run without torch. Torch-dependent paths are -imported lazily; GPU replay and benchmarking require CUDA/HIP. +Schema parsing, IR construction, custom-init name matching, and +``extract_batch_context`` run without torch. Torch-dependent paths are skipped +when torch is missing. GPU replay uses ``@pytest.mark.gpu`` (excluded from +CPU-only CI). For a longer GPU smoke (profile + replay), run +``python examples/event_replay_gpu_smoke.py``. """ from __future__ import annotations @@ -19,6 +22,11 @@ import pytest +from TraceLens.EventReplay.custom_inits import ( + CustomInit, + PagedAttentionInit, + extract_batch_context, +) from TraceLens.EventReplay.event_replay import EventReplayer from TraceLens.EventReplay.utils import TensorCfg, list_profile_tensor_types @@ -146,6 +154,65 @@ def test_get_event_replay_ir_builds_tensor_and_scalar_args(self): assert kw_values["alpha"] == 1.0 +class _FakeReplayer: + def __init__(self, name): + self.event = {"name": name} + + +class _NoOpInit(CustomInit): + def initialize(self, replayer, **kwargs): + return None + + +class TestCustomInitAppliesTo: + def test_exact_event_name_matches(self): + init = _NoOpInit() + init.op_patterns = ["aten::mm"] + assert init.applies_to(_FakeReplayer("aten::mm")) + + def test_substring_does_not_match(self): + init = _NoOpInit() + init.op_patterns = ["mm"] + assert not init.applies_to(_FakeReplayer("aten::mm")) + + +class _FakeAnalyzer: + def __init__(self, events): + self.tree = type("Tree", (), {"events": events})() + + +def _vllm_annotation(ts=0, dur=100): + return { + "cat": "user_annotation", + "name": "execute_context_2(18)_generation_5(5)", + "ts": ts, + "dur": dur, + } + + +def _cpu_op(name, ts=10): + return { + "name": name, + "ts": ts, + "args": {"Input Dims": [[1, 1]]}, + } + + +class TestExtractBatchContext: + def test_exact_paged_attention_is_annotated(self): + op = _cpu_op("_rocm_C::paged_attention") + n = extract_batch_context(_FakeAnalyzer([_vllm_annotation(), op])) + assert n == 1 + assert op["batch_context"]["n_prefill"] == 2 + assert "_rocm_C::paged_attention" in PagedAttentionInit.op_patterns + + def test_substring_name_is_not_annotated(self): + op = _cpu_op("aiter::paged_attention_v1") + n = extract_batch_context(_FakeAnalyzer([_vllm_annotation(), op])) + assert n == 0 + assert "batch_context" not in op + + @pytest.mark.skipif(not HAS_TORCH, reason="torch not installed") class TestEventReplayIRWithTorch: def test_get_args_kwargs_cpu(self): @@ -194,9 +261,19 @@ def test_build_tensor_rejects_normal_init_for_int(self): @pytest.mark.skipif(not HAS_TORCH, reason="torch not installed") class TestEventReplayerCpu: + @pytest.fixture(autouse=True) + def _isolate_custom_init_registry(self): + saved = EventReplayer._custom_init_registry[:] + yield + EventReplayer._custom_init_registry = saved + def test_event_replayer_lazy_cpu_replay(self): replayer = EventReplayer(MM_EVENT, device="cpu", lazy=True) - replayer.replay() + result = replayer.replay() + torch = _require_torch() + assert isinstance(result, torch.Tensor) + assert result.shape == (4, 16) + assert hasattr(replayer, "args") def test_get_repro_info_serializes_tensor_cfg(self): replayer = EventReplayer(MM_EVENT, device="cpu", lazy=True) @@ -206,6 +283,64 @@ def test_get_repro_info_serializes_tensor_cfg(self): assert pos0["shape"] == [4, 8] assert pos0["dtype"] == "c10::BFloat16" + def test_get_repro_info_idempotent_and_does_not_mutate_ir(self): + replayer = EventReplayer(MM_EVENT, device="cpu", lazy=True) + assert replayer.get_repro_info() == replayer.get_repro_info() + for arg in replayer.event_replay_IR["list_pos_args"]: + if arg["arg_type"].startswith("Tensor"): + assert isinstance(arg["value"], TensorCfg) + torch = _require_torch() + result = replayer.replay() + assert isinstance(result, torch.Tensor) + + def test_lazy_custom_init_sees_args(self): + accessed = {} + + class ProbeInit(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + accessed["args"] = replayer.args + accessed["kwargs"] = replayer.kwargs + return None + + EventReplayer.register_custom_init(ProbeInit()) + EventReplayer(MM_EVENT, device="cpu", lazy=True, auto_init=True).replay() + assert "args" in accessed + assert len(accessed["args"]) == 2 + + def test_first_matching_custom_init_wins(self): + log = [] + + class InitA(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + log.append("A") + + class InitB(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + log.append("B") + + EventReplayer._custom_init_registry = [InitA(), InitB()] + EventReplayer(MM_EVENT, device="cpu", auto_init=True).replay() + assert log == ["A"] + + def test_auto_init_false_skips_custom_inits(self): + log = [] + + class AlwaysInit(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + log.append("ran") + + EventReplayer._custom_init_registry = [AlwaysInit()] + EventReplayer(MM_EVENT, device="cpu", auto_init=False).replay() + assert log == [] + @pytest.mark.skipif(not HAS_TORCH, reason="torch not installed") class TestBatchedReplayHelpers: @@ -260,7 +395,7 @@ def test_benchmark_func_cuda(self): def matmul(): torch.matmul(a, b) - avg_us = benchmark_func( + metrics = benchmark_func( matmul, device=torch.device("cuda"), warmup=1, avg_steps=2 ) - assert avg_us > 0 + assert metrics["mean_us"] > 0 From c6d7a55bb14deb22c46b598bae5ee663c638d66e Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 16:56:54 -0400 Subject: [PATCH 08/11] Re-apply custom inits after lazy replay rebuilds tensors. _inits_applied tracks the current args; lazy reconstruction must clear it so paged-attn/MoE metadata is not left as zeros on the second replay(). Co-authored-by: Cursor --- TraceLens/EventReplay/event_replay.py | 2 ++ tests/test_event_replay_module.py | 35 +++++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index a46e4e8ba..ba6ff132d 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -312,6 +312,8 @@ def replay(self): self.args, self.kwargs = EventReplayer._get_args_kwargs( self.event_replay_IR, device=self.device ) + # Rebuilt tensors are not the ones previously inited. + self._inits_applied = False if not self._inits_applied and self._auto_init: self._apply_custom_inits() diff --git a/tests/test_event_replay_module.py b/tests/test_event_replay_module.py index e57398477..c2f64928f 100644 --- a/tests/test_event_replay_module.py +++ b/tests/test_event_replay_module.py @@ -309,6 +309,41 @@ def initialize(self, replayer, **kwargs): assert "args" in accessed assert len(accessed["args"]) == 2 + def test_lazy_custom_init_reruns_after_rebuild(self): + seen_ids = [] + + class ProbeInit(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + seen_ids.append(id(replayer.args[0])) + + EventReplayer.register_custom_init(ProbeInit()) + replayer = EventReplayer( + MM_EVENT, device="cpu", lazy=True, auto_init=True + ) + replayer.replay() + replayer.replay() + assert len(seen_ids) == 2 + assert seen_ids[0] != seen_ids[1] + + def test_eager_custom_init_runs_once(self): + log = [] + + class ProbeInit(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + log.append("init") + + EventReplayer.register_custom_init(ProbeInit()) + replayer = EventReplayer( + MM_EVENT, device="cpu", lazy=False, auto_init=True + ) + replayer.replay() + replayer.replay() + assert log == ["init"] + def test_first_matching_custom_init_wins(self): log = [] From 9b06ff80f1c03fa2fb8b34eef2fcb2d1b5f9adde Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 17:10:19 -0400 Subject: [PATCH 09/11] Resolve batched_replay ops through EventReplayer's custom-op path. JIT-only lookup returned None for non-aten names and then crashed on call; use _resolve_op_func when available and treat a None JIT result as a miss. Co-authored-by: Cursor --- TraceLens/EventReplay/batched_replay.py | 41 +++++++++++++++++++++++-- tests/test_event_replay_module.py | 30 ++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/TraceLens/EventReplay/batched_replay.py b/TraceLens/EventReplay/batched_replay.py index 0228688ec..8342afa04 100644 --- a/TraceLens/EventReplay/batched_replay.py +++ b/TraceLens/EventReplay/batched_replay.py @@ -9,7 +9,41 @@ import argparse import sys import torch -from utils import TensorCfg, build_tensor, benchmark_func + +try: + from .utils import TensorCfg, build_tensor, benchmark_func + from .event_replay import _resolve_op_func as _resolve_op_func_sdk +except ImportError: + from utils import TensorCfg, build_tensor, benchmark_func + try: + from event_replay import _resolve_op_func as _resolve_op_func_sdk + except ImportError: + _resolve_op_func_sdk = None + + +def resolve_replay_func(op_name: str): + """Return a callable for *op_name*, or raise RuntimeError. + + Prefers EventReplayer's resolver (JIT → torch.ops → module → auto-import). + If this script is used as a standalone zip without ``event_replay.py``, + falls back to JIT and treats a ``None`` return as a miss (do not call it). + """ + if _resolve_op_func_sdk is not None: + func, _source, _resolved = _resolve_op_func_sdk(op_name) + if func is None or not callable(func): + raise RuntimeError(f"Cannot resolve op '{op_name}'") + return func + + try: + func, _ = torch._C._jit_get_operation(op_name) + except Exception as e: + raise RuntimeError(f"Cannot resolve op '{op_name}': {e}") from e + if func is None or not callable(func): + raise RuntimeError( + f"Cannot resolve op '{op_name}' (JIT returned {func!r}). " + "Place event_replay.py next to this script for custom-op resolution." + ) + return func def _get_args_kwargs_from_ir( @@ -113,10 +147,11 @@ def _get_args_kwargs_from_ir( # Get the PyTorch operation function try: - func, _ = torch._C._jit_get_operation(op_name) + func = resolve_replay_func(op_name) except Exception as e: print( - f" Error: Could not find PyTorch operation '{op_name}'. Is the PyTorch version compatible? Error: {e}" + f" Error: Could not find PyTorch operation '{op_name}'. " + f"Is the op library imported? Error: {e}" ) if args.stop_on_error: raise diff --git a/tests/test_event_replay_module.py b/tests/test_event_replay_module.py index c2f64928f..005fb20b7 100644 --- a/tests/test_event_replay_module.py +++ b/tests/test_event_replay_module.py @@ -403,6 +403,36 @@ def test_get_args_kwargs_from_ir_cpu(self): assert pos_args[0].shape == (2, 3) assert kwargs == {} + def test_resolve_replay_func_aten_mm(self): + from TraceLens.EventReplay.batched_replay import resolve_replay_func + + func = resolve_replay_func("aten::mm") + assert callable(func) + + def test_resolve_replay_func_missing_op_raises(self): + from TraceLens.EventReplay.batched_replay import resolve_replay_func + + with pytest.raises(RuntimeError, match="Cannot resolve"): + resolve_replay_func("pr607missing::no_such_op") + + def test_resolve_replay_func_python_module_op(self): + import types + from TraceLens.EventReplay.batched_replay import resolve_replay_func + + mod = types.ModuleType("pr607climod") + + def add_one(x): + return x + 1 + + mod.add_one = add_one + sys.modules["pr607climod"] = mod + try: + func = resolve_replay_func("pr607climod::add_one") + assert func is add_one + assert func(3) == 4 + finally: + sys.modules.pop("pr607climod", None) + @pytest.mark.gpu @pytest.mark.skipif(not HAS_TORCH, reason="torch not installed") From 0aaf0952c5576383bba682f40f8fa380ce99f52a Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 17:32:13 -0400 Subject: [PATCH 10/11] Let last register_custom_init override built-in initializers. Co-authored-by: Cursor --- TraceLens/EventReplay/custom_inits.py | 3 +- TraceLens/EventReplay/event_replay.py | 4 +-- docs/how-to/event-replay.md | 7 ++--- tests/test_event_replay_module.py | 43 +++++++++++++++++++++++++++ 4 files changed, 50 insertions(+), 7 deletions(-) diff --git a/TraceLens/EventReplay/custom_inits.py b/TraceLens/EventReplay/custom_inits.py index 9ad5c70aa..33ab02427 100644 --- a/TraceLens/EventReplay/custom_inits.py +++ b/TraceLens/EventReplay/custom_inits.py @@ -19,7 +19,8 @@ 3. Implement ``initialize()`` — mutate replayer.args / replayer.kwargs in-place 4. Return a one-line summary string (printed by EventReplayer) 5. Register with ``EventReplayer.register_custom_init(YourInit())`` - or add it to the ``_custom_init_registry`` default list. + (prepends; last register wins for that event name) or add it to + the ``_custom_init_registry`` default list. """ from __future__ import annotations diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index ba6ff132d..919c49959 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -211,8 +211,8 @@ class EventReplayer: @classmethod def register_custom_init(cls, init: CustomInit): - """Add a custom initializer to the registry.""" - cls._custom_init_registry.append(init) + """Prepend a custom initializer. Last register wins for the same event name.""" + cls._custom_init_registry.insert(0, init) @classmethod def register_namespace(cls, namespace: str, modules: List[str]): diff --git a/docs/how-to/event-replay.md b/docs/how-to/event-replay.md index 932be20fa..cb8bf0b6d 100644 --- a/docs/how-to/event-replay.md +++ b/docs/how-to/event-replay.md @@ -286,10 +286,9 @@ class IndexAddInit(CustomInit): EventReplayer.register_custom_init(IndexAddInit()) ``` -`replay()` applies the **first** matching initializer. Built-ins are registered -first; `register_custom_init` appends, so a user initializer for the same exact -op name as a built-in doesn't run. List the registry with -`EventReplayer.list_custom_inits()`. +`replay()` applies the **first** matching initializer. `register_custom_init` +prepends, so the last registered initializer for an event name wins — including +over a built-in. List the registry with `EventReplayer.list_custom_inits()`. The built-in implementations are in `TraceLens/EventReplay/custom_inits.py`. diff --git a/tests/test_event_replay_module.py b/tests/test_event_replay_module.py index 005fb20b7..53ac1e4e2 100644 --- a/tests/test_event_replay_module.py +++ b/tests/test_event_replay_module.py @@ -176,6 +176,29 @@ def test_substring_does_not_match(self): assert not init.applies_to(_FakeReplayer("aten::mm")) +class TestCustomInitRegistry: + @pytest.fixture(autouse=True) + def _isolate_custom_init_registry(self): + saved = EventReplayer._custom_init_registry[:] + yield + EventReplayer._custom_init_registry = saved + + def test_register_prepends_so_last_wins(self): + builtin = _NoOpInit() + builtin.op_patterns = ["_rocm_C::paged_attention"] + user = _NoOpInit() + user.op_patterns = ["_rocm_C::paged_attention"] + EventReplayer._custom_init_registry = [builtin] + EventReplayer.register_custom_init(user) + assert EventReplayer._custom_init_registry[0] is user + first = next( + i + for i in EventReplayer._custom_init_registry + if i.applies_to(_FakeReplayer("_rocm_C::paged_attention")) + ) + assert first is user + + class _FakeAnalyzer: def __init__(self, events): self.tree = type("Tree", (), {"events": events})() @@ -363,6 +386,26 @@ def initialize(self, replayer, **kwargs): EventReplayer(MM_EVENT, device="cpu", auto_init=True).replay() assert log == ["A"] + def test_register_custom_init_overrides_builtin(self): + log = [] + + class BuiltinInit(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + log.append("builtin") + + class UserInit(CustomInit): + op_patterns = ["aten::mm"] + + def initialize(self, replayer, **kwargs): + log.append("user") + + EventReplayer._custom_init_registry = [BuiltinInit()] + EventReplayer.register_custom_init(UserInit()) + EventReplayer(MM_EVENT, device="cpu", auto_init=True).replay() + assert log == ["user"] + def test_auto_init_false_skips_custom_inits(self): log = [] From 28c20b87523fb949aa1e7ecfb7109ffcb78ae696 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Fri, 11 Sep 2026 17:56:20 -0400 Subject: [PATCH 11/11] Format EventReplay Python with Black 26.3.1 to pass CI lint. Co-authored-by: Cursor --- TraceLens/EventReplay/batched_replay.py | 5 ++- TraceLens/EventReplay/custom_inits.py | 50 +++++++++++---------- TraceLens/EventReplay/event_replay.py | 58 ++++++++++++++++++------- TraceLens/EventReplay/utils.py | 2 +- examples/event_replay_gpu_smoke.py | 47 ++++++++++++++++---- tests/test_event_replay_module.py | 8 +--- 6 files changed, 116 insertions(+), 54 deletions(-) diff --git a/TraceLens/EventReplay/batched_replay.py b/TraceLens/EventReplay/batched_replay.py index 8342afa04..3c66827b9 100644 --- a/TraceLens/EventReplay/batched_replay.py +++ b/TraceLens/EventReplay/batched_replay.py @@ -15,6 +15,7 @@ from .event_replay import _resolve_op_func as _resolve_op_func_sdk except ImportError: from utils import TensorCfg, build_tensor, benchmark_func + try: from event_replay import _resolve_op_func as _resolve_op_func_sdk except ImportError: @@ -196,7 +197,9 @@ def _get_args_kwargs_from_ir( lambda: func(*pos_args, **kwargs), args.device, warmup=50, avg_steps=100 ) mean_time_us = metrics["mean_us"] - print(f" Average time taken: {mean_time_us:.2f} us (median: {metrics['median_us']:.2f} us)") + print( + f" Average time taken: {mean_time_us:.2f} us (median: {metrics['median_us']:.2f} us)" + ) if "count" in repro_info: count_workload = repro_info["count"] total_time_us = mean_time_us * count_workload diff --git a/TraceLens/EventReplay/custom_inits.py b/TraceLens/EventReplay/custom_inits.py index 33ab02427..bce748443 100644 --- a/TraceLens/EventReplay/custom_inits.py +++ b/TraceLens/EventReplay/custom_inits.py @@ -34,9 +34,7 @@ # -- Batch context extraction from vLLM profiler annotations --------------- -_BATCH_ANNO_RE = re.compile( - r"execute_context_(\d+)\((\d+)\)_generation_(\d+)\((\d+)\)" -) +_BATCH_ANNO_RE = re.compile(r"execute_context_(\d+)\((\d+)\)_generation_(\d+)\((\d+)\)") def extract_batch_context(analyzer: Any) -> int: @@ -76,14 +74,16 @@ def extract_batch_context(analyzer: Any) -> int: continue ts = e.get("ts", 0) dur = e.get("dur", 0) - annotations.append({ - "ts": ts, - "end": ts + dur, - "n_prefill": int(m.group(1)), - "prefill_tokens": int(m.group(2)), - "n_decode": int(m.group(3)), - "decode_tokens": int(m.group(4)), - }) + annotations.append( + { + "ts": ts, + "end": ts + dur, + "n_prefill": int(m.group(1)), + "prefill_tokens": int(m.group(2)), + "n_decode": int(m.group(3)), + "decode_tokens": int(m.group(4)), + } + ) if not annotations: return 0 @@ -154,6 +154,7 @@ def initialize(self, replayer: Any, **kwargs) -> Optional[str]: ir = replayer.event_replay_IR arg_names = [a["arg_name"] for a in ir["list_pos_args"]] + def _by_name_or_pos(name, pos): if name in arg_names: return args[arg_names.index(name)] @@ -196,16 +197,20 @@ def _by_name_or_pos(name, pos): while len(per_seq_queries) < num_seqs: per_seq_queries.append(1) - phase = ("mixed" if n_pf > 0 and n_dec > 0 - else "prefill" if n_pf > 0 else "decode") + phase = ( + "mixed" + if n_pf > 0 and n_dec > 0 + else "prefill" if n_pf > 0 else "decode" + ) source = "annotation" else: tokens_per_seq = num_query_tokens / num_seqs if num_seqs else 1 if tokens_per_seq > 1: base_q = num_query_tokens // num_seqs rem_q = num_query_tokens % num_seqs - per_seq_queries = [base_q + (1 if s < rem_q else 0) - for s in range(num_seqs)] + per_seq_queries = [ + base_q + (1 if s < rem_q else 0) for s in range(num_seqs) + ] phase = "prefill" else: per_seq_queries = [1] * num_seqs @@ -233,9 +238,7 @@ def _by_name_or_pos(name, pos): # -- query_start_loc: CSR indptr encoding per-seq query counts --------- qsl = _by_name_or_pos("query_start_loc", 11) - if (qsl is not None - and hasattr(qsl, "shape") - and qsl.numel() > 0): + if qsl is not None and hasattr(qsl, "shape") and qsl.numel() > 0: qloc = np.zeros(num_seqs + 1, dtype=np.int32) for s in range(num_seqs): qloc[s + 1] = qloc[s] + per_seq_queries[s] @@ -244,10 +247,12 @@ def _by_name_or_pos(name, pos): ctx_str = "" if batch_ctx is not None: - ctx_str = (f" Annotation: {batch_ctx['n_prefill']} prefill " - f"({batch_ctx['prefill_tokens']} tok) + " - f"{batch_ctx['n_decode']} decode " - f"({batch_ctx['decode_tokens']} tok).") + ctx_str = ( + f" Annotation: {batch_ctx['n_prefill']} prefill " + f"({batch_ctx['prefill_tokens']} tok) + " + f"{batch_ctx['n_decode']} decode " + f"({batch_ctx['decode_tokens']} tok)." + ) return ( f"[custom init] {op_name} — paged attention metadata: " @@ -297,6 +302,7 @@ def initialize(self, replayer: Any, **kwargs) -> Optional[str]: # Locate args by name from the IR when available, fall back to position ir = replayer.event_replay_IR arg_names = [a["arg_name"] for a in ir["list_pos_args"]] + def _by_name_or_pos(name, pos): if name in arg_names: return args[arg_names.index(name)] diff --git a/TraceLens/EventReplay/event_replay.py b/TraceLens/EventReplay/event_replay.py index 919c49959..b8b52f71c 100644 --- a/TraceLens/EventReplay/event_replay.py +++ b/TraceLens/EventReplay/event_replay.py @@ -66,7 +66,9 @@ def _try_auto_import(op_name: str, verbose: bool = False) -> bool: imported_any = True except ImportError: if verbose: - print(f"[EventReplayer] Could not import '{mod}' for namespace '{namespace}'") + print( + f"[EventReplayer] Could not import '{mod}' for namespace '{namespace}'" + ) return imported_any @@ -123,7 +125,9 @@ def _resolve_op_func(op_name: str, verbose: bool = False): if func is not None: logger.warning( "Op '%s' resolved via alias '%s' (%s).", - op_name, alias, source, + op_name, + alias, + source, ) return func, source, alias @@ -138,9 +142,11 @@ def _resolve_op_func(op_name: str, verbose: bool = False): if known: hint = f" Try: {', '.join(f'import {m}' for m in known)}" else: - hint = (f" The namespace '{ns}' is not in the auto-import registry." - f" Use EventReplayer.register_namespace('{ns}', ['your.module'])" - f" to add it.") + hint = ( + f" The namespace '{ns}' is not in the auto-import registry." + f" Use EventReplayer.register_namespace('{ns}', ['your.module'])" + f" to add it." + ) raise RuntimeError( f"Cannot resolve op '{op_name}'.{hint} " @@ -270,7 +276,9 @@ def _setup(self): if self.verbose: print(f"Resolved op via {self._func_source}") if self._resolved_name != self.event["name"]: - print(f" (aliased from '{self.event['name']}' -> '{self._resolved_name}')") + print( + f" (aliased from '{self.event['name']}' -> '{self._resolved_name}')" + ) try: self.matched_schema = EventReplayer._search_schema( @@ -516,7 +524,10 @@ def _get_event_replay_IR( logger.warning( "%s arg '%s' (position %d): profiler dropped " "the string value. Using known default '%s'.", - evt_name, arg_name, idx, default, + evt_name, + arg_name, + idx, + default, ) value = "" if default is None else default else: @@ -531,7 +542,13 @@ def _get_event_replay_IR( if EventReplayer._should_skip_tensor_init(evt_name, arg_name, idx): init = None profiled_dtype = event["args"]["Input type"][idx] - if profiled_dtype in ("long", "long int", "int", "bool", "unsigned char"): + if profiled_dtype in ( + "long", + "long int", + "int", + "bool", + "unsigned char", + ): init = "zeros" if init == "normal" else init value = TensorCfg( shape=event["args"]["Input Dims"][idx], @@ -558,7 +575,10 @@ def _get_event_replay_IR( logger.warning( "%s arg '%s' (position %d): profiler dropped " "the string value. Using known default '%s'.", - evt_name, arg_name, idx, value, + evt_name, + arg_name, + idx, + value, ) else: value = arg_str @@ -630,7 +650,11 @@ def _get_event_replay_IR_schemaless( if profiled_type in list_profile_tensor_types: init = "normal" if profiled_type in ( - "long", "long int", "int", "bool", "unsigned char", + "long", + "long int", + "int", + "bool", + "unsigned char", ): init = "zeros" value = TensorCfg( @@ -665,13 +689,18 @@ def _get_event_replay_IR_schemaless( logger.warning( "%s arg '%s' (position %d): profiler dropped the " "string value. Using known default '%s'.", - evt_name, hint_name, idx, default, + evt_name, + hint_name, + idx, + default, ) else: value = None arg_type = "None" elif profiled_type == "ScalarList" and concrete: - items = [x.strip() for x in concrete.strip()[1:-1].split(",") if x.strip()] + items = [ + x.strip() for x in concrete.strip()[1:-1].split(",") if x.strip() + ] if all(x.lstrip("-").isdigit() for x in items): value = [int(x) for x in items] else: @@ -732,9 +761,7 @@ def parse_schema_string( def _parse_arg(raw_arg: str) -> Tuple[str, str, Optional[str], bool]: # Greedy (.+) consumes everything up to the last whitespace before # a valid identifier, so "Tensor($0! -> ) key_cache" parses correctly. - m = re.match( - r"^(.+)\s+([A-Za-z_]\w*(?:=.*)?)$", raw_arg.strip() - ) + m = re.match(r"^(.+)\s+([A-Za-z_]\w*(?:=.*)?)$", raw_arg.strip()) if not m: raise ValueError(f"Invalid arg: {raw_arg}") arg_type = m.group(1).strip() @@ -767,6 +794,7 @@ def get_repro_info(self) -> Dict[str, Any]: Safe to call multiple times — does not mutate self.event_replay_IR. """ + def _serialize_arg(arg: Dict[str, Any]) -> Dict[str, Any]: val = arg["value"] return { diff --git a/TraceLens/EventReplay/utils.py b/TraceLens/EventReplay/utils.py index 841ae6269..33618f1bd 100644 --- a/TraceLens/EventReplay/utils.py +++ b/TraceLens/EventReplay/utils.py @@ -186,7 +186,7 @@ def benchmark_func( median = (sorted_us[n // 2] + sorted_us[(n - 1) // 2]) / 2.0 mean = sum(timings_us) / n variance = sum((t - mean) ** 2 for t in timings_us) / n - std = variance ** 0.5 + std = variance**0.5 return { "median_us": median, "mean_us": mean, diff --git a/examples/event_replay_gpu_smoke.py b/examples/event_replay_gpu_smoke.py index 7574732a0..3ae33c06a 100644 --- a/examples/event_replay_gpu_smoke.py +++ b/examples/event_replay_gpu_smoke.py @@ -16,6 +16,7 @@ """ import sys, os, json, time + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import torch @@ -47,6 +48,7 @@ bmm_a = torch.randn(4, M, K, dtype=torch.bfloat16, device=DEVICE) bmm_b = torch.randn(4, K, N, dtype=torch.bfloat16, device=DEVICE) + def run_ops(): torch.mm(mm_a, mm_b) torch.add(add_a, add_b) @@ -54,13 +56,16 @@ def run_ops(): torch.mul(add_a, add_b) torch.sigmoid(add_a) + for _ in range(10): run_ops() torch.cuda.synchronize() + def trace_handler(p): p.export_chrome_trace(TRACE_FILE) + wait, warmup, active = 3, 3, 5 with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], @@ -89,10 +94,12 @@ def trace_handler(p): OPS_TO_TEST = ["aten::mm", "aten::add", "aten::bmm", "aten::mul", "aten::sigmoid"] + def find_event(events, op_name): """Find a cpu_op event with the right name and shape data.""" candidates = [ - e for e in events + e + for e in events if e.get("cat") == "cpu_op" and e.get("name") == op_name and "args" in e @@ -102,6 +109,7 @@ def find_event(events, op_name): return candidates[len(candidates) // 2] return None + results = [] errors = [] @@ -113,13 +121,17 @@ def find_event(events, op_name): print("Step 3: Replay and validate") print("=" * 80) -print(f"\n{'Op':<30} {'Kernel Match':<15} {'Return':<10} {'Lazy':<10} {'ReproInfo':<12} {'Status'}") +print( + f"\n{'Op':<30} {'Kernel Match':<15} {'Return':<10} {'Lazy':<10} {'ReproInfo':<12} {'Status'}" +) print("-" * 100) for op_name in OPS_TO_TEST: evt = find_event(all_events, op_name) if evt is None: - print(f"{op_name:<30} {'SKIP':<15} {'---':<10} {'---':<10} {'---':<12} not in trace") + print( + f"{op_name:<30} {'SKIP':<15} {'---':<10} {'---':<10} {'---':<12} not in trace" + ) continue status = [] @@ -148,10 +160,12 @@ def find_event(events, op_name): repro_replayer = EventReplayer(evt, device=DEVICE, lazy=True) info1 = repro_replayer.get_repro_info() info2 = repro_replayer.get_repro_info() - repro_ok = (info1 == info2) + repro_ok = info1 == info2 for arg in repro_replayer.event_replay_IR["list_pos_args"]: if arg["arg_type"].startswith("Tensor"): - assert isinstance(arg["value"], TensorCfg), "IR corrupted after get_repro_info" + assert isinstance( + arg["value"], TensorCfg + ), "IR corrupted after get_repro_info" repro_replayer.replay() except Exception as e: repro_ok = False @@ -201,9 +215,19 @@ def th(p): tag = "PASS" if ok else "FAIL" detail = "; ".join(status) if status else "" - print(f"{op_name:<30} {kernel_match:<15} {'OK' if returns_ok else 'FAIL':<10} {'OK' if lazy_ok else 'FAIL':<10} {'OK' if repro_ok else 'FAIL':<12} {tag} {detail}") - results.append({"op": op_name, "ok": ok, "kernel": kernel_match, - "returns": returns_ok, "lazy": lazy_ok, "repro": repro_ok}) + print( + f"{op_name:<30} {kernel_match:<15} {'OK' if returns_ok else 'FAIL':<10} {'OK' if lazy_ok else 'FAIL':<10} {'OK' if repro_ok else 'FAIL':<12} {tag} {detail}" + ) + results.append( + { + "op": op_name, + "ok": ok, + "kernel": kernel_match, + "returns": returns_ok, + "lazy": lazy_ok, + "repro": repro_ok, + } + ) # --------------------------------------------------------------------------- # Step 4: First-match-wins test (CLAIM-1) on GPU @@ -215,16 +239,21 @@ def th(p): log = [] + class InitA(CustomInit): op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): log.append("A") + class InitB(CustomInit): op_patterns = ["aten::mm"] + def initialize(self, replayer, **kwargs): log.append("B") + saved_registry = EventReplayer._custom_init_registry[:] try: EventReplayer._custom_init_registry = [InitA(), InitB()] @@ -232,7 +261,7 @@ def initialize(self, replayer, **kwargs): if mm_evt: r = EventReplayer(mm_evt, device=DEVICE, auto_init=True) r.replay() - first_match_ok = (log == ["A"]) + first_match_ok = log == ["A"] print(f" First-match-wins: {'PASS' if first_match_ok else 'FAIL'} (log={log})") else: first_match_ok = True diff --git a/tests/test_event_replay_module.py b/tests/test_event_replay_module.py index 53ac1e4e2..cc2a10850 100644 --- a/tests/test_event_replay_module.py +++ b/tests/test_event_replay_module.py @@ -342,9 +342,7 @@ def initialize(self, replayer, **kwargs): seen_ids.append(id(replayer.args[0])) EventReplayer.register_custom_init(ProbeInit()) - replayer = EventReplayer( - MM_EVENT, device="cpu", lazy=True, auto_init=True - ) + replayer = EventReplayer(MM_EVENT, device="cpu", lazy=True, auto_init=True) replayer.replay() replayer.replay() assert len(seen_ids) == 2 @@ -360,9 +358,7 @@ def initialize(self, replayer, **kwargs): log.append("init") EventReplayer.register_custom_init(ProbeInit()) - replayer = EventReplayer( - MM_EVENT, device="cpu", lazy=False, auto_init=True - ) + replayer = EventReplayer(MM_EVENT, device="cpu", lazy=False, auto_init=True) replayer.replay() replayer.replay() assert log == ["init"]