EventReplay: extend beyond aten ops with auto-import and custom initializers by ajassani · Pull Request #607 · AMD-AGI/TraceLens

ajassani · 2026-04-28T19:41:28Z

Summary

EventReplay previously only worked with aten:: ops. This PR extends it to support custom ops from any namespace (vLLM, aiter, etc.) and adds custom initializers so data-dependent ops produce realistic replay behavior.

Extending beyond aten

Auto-import for custom op namespaces: when EventReplay encounters a non-aten op (e.g., _rocm_C::paged_attention, aiter::ck_moe_stage1), it automatically imports the library that registers the op's schema. Supports aiter, _rocm_C, _C, vllm out of the box, and users can register additional namespaces.
Schemaless fallback: ops that lack a registered schema can still be replayed with heuristic type inference from the profiler data.

Custom initializers for data-dependent ops

PagedAttentionInit — fills block_tables, seq_lens, and query_start_loc with realistic values so the attention kernel does real work instead of short-circuiting on zeros.
MoeRoutingInit — constructs a complete token-to-expert routing table (sorted_token_ids, sorted_expert_ids, num_valid_ids) with configurable distribution (uniform or Zipf).
User-extensible: subclass CustomInit, set op_patterns to the exact op name, implement initialize(), and register. First-match-wins, exact name matching.

Iteration annotations (vLLM)

extract_batch_context parses vLLM's per-iteration user_annotation events to get the exact prefill/decode split, so PagedAttentionInit builds an accurate query_start_loc for mixed batches.

Bug fixes

replay() with lazy=True + auto_init=True crashed (AttributeError) — fixed
get_repro_info() corrupted event_replay_IR via shallow copy — fixed
batched_replay.py: benchmark_func dict return type crash + dead --op-filter/--op-limit flags — fixed
replay() now returns the op result
Custom init matching changed from substring to exact name

Tests and docs

11 CPU-only unit tests + GPU integration test with kernel name validation
Docs rewritten: IR interpretability examples, step-by-step custom init guide, iteration annotations explained

Test plan

CPU unit tests pass (11/11, ~3s)
GPU integration tests pass (5/5 ops, MI300X)
- Kernel name match: all MATCH
- Lazy mode, get_repro_info idempotency, return value, first-match-wins: all PASS

ajassani

Review of #607 (inline comments on the specific lines).

Not merge-ready yet: rebase onto current main first (this branch is ~202 commits behind; docs now live at docs/how-to/event-replay.md, CI is pytest tests/ -m "not gpu" --cov-fail-under=95, and tests/test_event_replay.py already exists). The custom-op resolution + custom-init design is the right shape; the inlined notes are the correctness/CI gaps.

ajassani · 2026-09-11T17:24:47Z

        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

-        # Call the function with the arguments
-        func(*args, **kwargs)
+        if not self._inits_applied and self._auto_init:
+            self._apply_custom_inits()


Blocking — lazy + custom init is wrong on the 2nd replay().

Lazy rebuilds self.args from the IR every call, but _inits_applied stays True, so paged-attn / MoE tensors go back to zeros after the first iteration.

Re-apply inits whenever lazy rebuilds args (or reset _inits_applied when args are reconstructed).

ajassani · 2026-09-11T17:24:47Z

+    @classmethod
+    def register_custom_init(cls, init: CustomInit):
+        """Add a custom initializer to the registry."""
+        cls._custom_init_registry.append(init)


Blocking — user inits can never override builtins.

register_custom_init appends, and _apply_custom_inits is first-match-wins. A user init for _rocm_C::paged_attention will never run.

Insert at the front, or skip the builtin when a later init matches the same exact name.

ajassani · 2026-09-11T17:24:47Z

+_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)


Sticky even on failed import. This set is process-global and is populated before the import succeeds. After a miss, register_namespace() never retries.

Only record successful attempts, or drop the namespace from the set inside register_namespace.

ajassani · 2026-09-11T17:24:47Z

+        try:
+            self.matched_schema = EventReplayer._search_schema(
+                self.event, self._resolved_name, 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


Schemaless fallback swallows the ValueError and proceeds with heuristic types. Unless verbose=True, a wrong-signature custom-op call looks like success.

Please logger.warning here (op name + how many schemas were searched).

ajassani · 2026-09-11T17:24:47Z

                        init = None
+                    profiled_dtype = event["args"]["Input type"][idx]
+                    if profiled_dtype in ("long", "long int", "int", "bool", "unsigned char"):
+                        init = "zeros" if init == "normal" else init


char / short were added to list_profile_tensor_types but not this zeros-init branch. build_tensor(..., init="normal") then raises on those integer dtypes.

Same gap in _get_event_replay_IR_schemaless (~631).

ajassani · 2026-09-11T17:24:48Z

+    annotated = 0
+    for e in analyzer.tree.events:
+        name = e.get("name", "")
+        if "paged_attention" not in name:


This still uses substring ("paged_attention" in name) while CustomInit.applies_to is exact-name. Easy to annotate a similarly named op that the init will then ignore — or miss a renamed exact op.

ajassani · 2026-09-11T17:24:48Z

+        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))


Hardcoded positional fallback will silently init the wrong tensor if the schema renames/reorders args. Same pattern in MoeRoutingInit (~298).

Fail if the name is missing instead of guessing a slot.

ajassani · 2026-09-11T17:24:48Z

+"""
+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
+"""


Blocking for CI. Current unit workflow is pytest tests/ -m "not gpu" --cov-fail-under=95. These 11 tests are not collected, and the new custom_inits.py / expanded event_replay.py will drop coverage.

Move this file under tests/. Also: tests/test_event_replay.py already exists on current main (GPU ResNet integration) — rebase will collide.

ajassani · 2026-09-11T17:24:48Z

+from TraceLens.EventReplay.custom_inits import CustomInit
+from TraceLens.EventReplay.utils import TensorCfg
+
+assert torch.cuda.is_available(), "GPU required for this test"


Module-level assert means this file cannot be collected as a pytest test. Fine as a manual script — please don't name it test_*.py under a collected path after the rebase (or mark/skip at collection time).

ajassani · 2026-09-11T17:24:48Z

+        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"


orig_gpu is every kernel in a 5-op trace. replay_gpu.issubset(orig_gpu) is almost always true even if this op launched the wrong kernel.

Compare against kernels from this cpu_op (or this op's original vs replayed kernel set).

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

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

Made-with: Cursor

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

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 <cursoragent@cursor.com>

codecov-commenter · 2026-09-11T20:21:15Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

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 <cursoragent@cursor.com>

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 <cursoragent@cursor.com>

_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 <cursoragent@cursor.com>

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 <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

ajassani · 2026-09-11T21:39:09Z

@gabeweisz @kyle-hoffmeyer this is ready for another look. Rebased onto current main, and the blocking review items (docs/CI layout, lazy custom inits, CLI op resolve, user init override) are addressed.

Co-authored-by: Cursor <cursoragent@cursor.com>

ajassani changed the title ~~EventReplay: custom initializers, bug fixes, tests, and docs~~ EventReplay: extend beyond aten ops with auto-import, custom initializers, and iteration annotations Apr 28, 2026

ajassani changed the title ~~EventReplay: extend beyond aten ops with auto-import, custom initializers, and iteration annotations~~ EventReplay: extend beyond aten ops with auto-import and custom initializers Apr 28, 2026

ajassani force-pushed the feature/event-replay-custom-ops branch from 94dd686 to 95406f3 Compare April 28, 2026 19:49

ajassani mentioned this pull request May 5, 2026

Extend Replay to any general dispatch functions beyond aten:: #312

Open

ajassani commented Sep 11, 2026

View reviewed changes

ajassani and others added 5 commits September 11, 2026 15:34

Add custom initializers, auto-import, and updated docs for EventReplay

b1e5907

Made-with: Cursor

Port EventReplay custom-op docs into the ROCm how-to page.

0b27256

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 <cursoragent@cursor.com>

ajassani force-pushed the feature/event-replay-custom-ops branch from 95406f3 to 0b27256 Compare September 11, 2026 19:47

ajassani requested review from devalshahamd, gabeweisz and tsrikris as code owners September 11, 2026 19:47

ajassani and others added 5 commits September 11, 2026 16:22

Match extract_batch_context to exact PagedAttentionInit op names.

bf7b5d0

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 <cursoragent@cursor.com>

Re-apply custom inits after lazy replay rebuilds tensors.

c6d7a55

_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 <cursoragent@cursor.com>

Resolve batched_replay ops through EventReplayer's custom-op path.

9b06ff8

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 <cursoragent@cursor.com>

Let last register_custom_init override built-in initializers.

0aaf095

Co-authored-by: Cursor <cursoragent@cursor.com>

ajassani requested a review from kyle-hoffmeyer September 11, 2026 21:39

Format EventReplay Python with Black 26.3.1 to pass CI lint.

28c20b8

Co-authored-by: Cursor <cursoragent@cursor.com>

Navigation Menu

Uh oh!

EventReplay: extend beyond aten ops with auto-import and custom initializers - #607

EventReplay: extend beyond aten ops with auto-import and custom initializers#607
ajassani wants to merge 11 commits into
mainfrom
feature/event-replay-custom-ops

ajassani commented Apr 28, 2026 •

edited

Loading

Uh oh!

ajassani left a comment

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

ajassani Sep 11, 2026

Uh oh!

codecov-commenter commented Sep 11, 2026

Uh oh!

ajassani commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ajassani commented Apr 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Extending beyond aten

Custom initializers for data-dependent ops

Iteration annotations (vLLM)

Bug fixes

Tests and docs

Test plan

Uh oh!

ajassani left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Sep 11, 2026

Codecov Report

Uh oh!

ajassani commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ajassani commented Apr 28, 2026 •

edited

Loading