From b8f1f4a2008b5ff182126be33136417b08787273 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Thu, 10 Sep 2026 18:32:45 -0400 Subject: [PATCH 1/3] Inference split: keep host ops between step annotations. vLLM compute_logits and similar post-step work sit after the execute_* envelope; bounding the CPU window by the next matching annotation keeps lm_head GEMMs without leaking the following step. Co-authored-by: Cursor --- .../split_inference/trace_extraction.py | 67 +++++++++++++++++-- .../test_split_inference_trace_annotation.py | 62 +++++++++++++++++ 2 files changed, 124 insertions(+), 5 deletions(-) diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index a035ea7b6..0f42efcee 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -23,6 +23,52 @@ iteration_details, ) + +def _annotation_pattern(name: str): + """Return the first iteration-root pattern that matches ``name``, else None.""" + for pattern in ITERATION_PATTERNS + ITERATION_BACKUP_PATTERNS: + if pattern.match(name or ""): + return pattern + return None + + +def _next_same_pattern_ts(root: dict, events: list[dict]) -> float | None: + """Timestamp of the next same-pattern iteration annotation on this thread. + + Serving runtimes (vLLM ``compute_logits``, SGLang ``_compute_lm_head``) often + launch work *after* the step annotation's duration ends and *before* the + next step annotation starts. Bounding the CPU window by the annotation + duration drops that work. Bounding by the next matching annotation keeps it + without leaking the following step's nested ops. + """ + pattern = _annotation_pattern(root.get("name", "")) + if pattern is None: + return None + iter_ts = root.get("ts", 0) + iter_tid = root.get("tid") + iter_pid = root.get("pid") + next_ts = None + for e in events: + ts = e.get("ts") + if ts is None or ts <= iter_ts: + continue + if e.get("tid") != iter_tid or e.get("pid") != iter_pid: + continue + if pattern.match(e.get("name") or ""): + if next_ts is None or ts < next_ts: + next_ts = ts + return next_ts + + +def _cpu_window_end(root: dict, next_sibling_ts: float | None, events: list[dict]): + """Exclusive end of the CPU window for one iteration root.""" + if next_sibling_ts is not None: + return next_sibling_ts + next_ts = _next_same_pattern_ts(root, events) + if next_ts is not None: + return next_ts + return root.get("ts", 0) + root.get("dur", 0) + GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"] @@ -73,7 +119,12 @@ def extract_iteration( flow_corr_map: dict, meta_events: list[dict], ) -> dict: - """Extract a single iteration trace.""" + """Extract CPU/GPU events for one or more consecutive iteration roots. + + CPU ops are kept from each root's start until the next same-pattern + iteration annotation (not merely ``ts + dur``). GPU kernels still attach + via correlation, so they may sit outside that window. + """ filtered_events = [] gpu_dur = 0 @@ -86,10 +137,16 @@ def extract_iteration( # Compute the global time window for all iteration roots if not iteration_roots: return trace_json.copy(), [], 0, 0, 0 + roots_by_ts = sorted(iteration_roots, key=lambda r: r.get("ts", 0)) + sibling_end = [] + for i, root in enumerate(roots_by_ts): + next_sibling = ( + roots_by_ts[i + 1].get("ts") if i + 1 < len(roots_by_ts) else None + ) + sibling_end.append(_cpu_window_end(root, next_sibling, events)) + root_end = {id(root): end for root, end in zip(roots_by_ts, sibling_end)} min_iter_ts = min(root.get("ts", 0) for root in iteration_roots) - max_iter_end = max( - root.get("ts", 0) + root.get("dur", 0) for root in iteration_roots - ) + max_iter_end = max(sibling_end) # Collect all relevant tid/pid pairs tid_pid_set = {(root.get("tid"), root.get("pid")) for root in iteration_roots} @@ -117,7 +174,7 @@ def extract_iteration( iter_tid = iteration_root.get("tid") iter_pid = iteration_root.get("pid") iter_ts = iteration_root.get("ts", 0) - iter_end = iter_ts + iteration_root.get("dur", 0) + iter_end = root_end[id(iteration_root)] correlation_ids: set[int] = set() diff --git a/tests/test_split_inference_trace_annotation.py b/tests/test_split_inference_trace_annotation.py index 26a12cd1a..6753079ca 100644 --- a/tests/test_split_inference_trace_annotation.py +++ b/tests/test_split_inference_trace_annotation.py @@ -300,6 +300,68 @@ def test_preprocess_trace_collects_flow_and_gpu_maps(): assert len(meta) == 1 +def test_extract_iteration_keeps_cpu_ops_between_step_annotations(): + """Keep host ops that sit after a step annotation and before the next one. + + vLLM ``compute_logits`` / lm_head GEMMs are launched in that gap; bounding + the split window by ``annotation.ts + dur`` drops them. + """ + names = [VLLM_PRIMARY_ANNOTATION.format(i=i) for i in range(3)] + trace = make_trace(names) + events = trace["traceEvents"] + events.append( + { + "name": "compute_logits", + "cat": "cpu_op", + "ph": "X", + "ts": 1200, + "dur": 40, + "tid": 10, + "pid": 1, + "args": {"correlation": 9001}, + } + ) + events.append( + { + "name": "vllm::rocm_unquantized_gemm", + "cat": "kernel", + "ph": "X", + "ts": 1250, + "dur": 15, + "tid": 99, + "pid": 1, + "args": {"correlation": 9001}, + } + ) + gpu_map, flow_map, meta = split.preprocess_trace(events) + roots = split.find_iteration_roots(events) + assert roots is not None and len(roots) == 3 + + out, _, num_gpu, _, _ = split.extract_iteration( + roots, events, trace, gpu_map, flow_map, meta + ) + out_names = {e["name"] for e in out["traceEvents"]} + assert "compute_logits" in out_names + assert "vllm::rocm_unquantized_gemm" in out_names + # Original per-root kernels (2 each) plus the inter-step GEMM. + assert num_gpu == 7 + + out0, _, _, _, _ = split.extract_iteration( + [roots[0]], events, trace, gpu_map, flow_map, meta + ) + names0 = {e["name"] for e in out0["traceEvents"]} + assert "compute_logits" in names0 + assert "vllm::rocm_unquantized_gemm" in names0 + assert "cpu_op_1_0" not in names0 + + out1, _, _, _, _ = split.extract_iteration( + [roots[1]], events, trace, gpu_map, flow_map, meta + ) + names1 = {e["name"] for e in out1["traceEvents"]} + assert "compute_logits" not in names1 + assert "cpu_op_0_0" not in names1 + + def test_extract_iteration_empty_roots(): trace = make_trace([VLLM_PRIMARY_ANNOTATION.format(i=0)]) events = trace["traceEvents"] From de202a6c1b2aeefec2996c4db1c034cd56693391 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Thu, 10 Sep 2026 22:57:29 -0400 Subject: [PATCH 2/3] Use TraceKeys for split CPU-window field names. Co-authored-by: Cursor --- .../split_inference/trace_extraction.py | 35 +++++++++++-------- .../test_split_inference_trace_annotation.py | 15 ++++++++ 2 files changed, 35 insertions(+), 15 deletions(-) diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index 0f42efcee..8e50166b3 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -13,6 +13,7 @@ from tqdm import tqdm +from ...util import TraceEventUtils from ..annotation_utils import ( ITERATION_BACKUP_PATTERNS, ITERATION_PATTERNS, @@ -23,6 +24,8 @@ iteration_details, ) +_K = TraceEventUtils.TraceKeys + def _annotation_pattern(name: str): """Return the first iteration-root pattern that matches ``name``, else None.""" @@ -41,20 +44,20 @@ def _next_same_pattern_ts(root: dict, events: list[dict]) -> float | None: duration drops that work. Bounding by the next matching annotation keeps it without leaking the following step's nested ops. """ - pattern = _annotation_pattern(root.get("name", "")) + pattern = _annotation_pattern(root.get(_K.Name, "")) if pattern is None: return None - iter_ts = root.get("ts", 0) - iter_tid = root.get("tid") - iter_pid = root.get("pid") + iter_ts = root.get(_K.TimeStamp, 0) + iter_tid = root.get(_K.TID) + iter_pid = root.get(_K.PID) next_ts = None for e in events: - ts = e.get("ts") + ts = e.get(_K.TimeStamp) if ts is None or ts <= iter_ts: continue - if e.get("tid") != iter_tid or e.get("pid") != iter_pid: + if e.get(_K.TID) != iter_tid or e.get(_K.PID) != iter_pid: continue - if pattern.match(e.get("name") or ""): + if pattern.match(e.get(_K.Name) or ""): if next_ts is None or ts < next_ts: next_ts = ts return next_ts @@ -67,7 +70,7 @@ def _cpu_window_end(root: dict, next_sibling_ts: float | None, events: list[dict next_ts = _next_same_pattern_ts(root, events) if next_ts is not None: return next_ts - return root.get("ts", 0) + root.get("dur", 0) + return root.get(_K.TimeStamp, 0) + root.get(_K.Duration, 0) GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"] @@ -137,18 +140,20 @@ def extract_iteration( # Compute the global time window for all iteration roots if not iteration_roots: return trace_json.copy(), [], 0, 0, 0 - roots_by_ts = sorted(iteration_roots, key=lambda r: r.get("ts", 0)) + roots_by_ts = sorted(iteration_roots, key=lambda r: r.get(_K.TimeStamp, 0)) sibling_end = [] for i, root in enumerate(roots_by_ts): next_sibling = ( - roots_by_ts[i + 1].get("ts") if i + 1 < len(roots_by_ts) else None + roots_by_ts[i + 1].get(_K.TimeStamp) + if i + 1 < len(roots_by_ts) + else None ) sibling_end.append(_cpu_window_end(root, next_sibling, events)) root_end = {id(root): end for root, end in zip(roots_by_ts, sibling_end)} - min_iter_ts = min(root.get("ts", 0) for root in iteration_roots) + min_iter_ts = min(root.get(_K.TimeStamp, 0) for root in iteration_roots) max_iter_end = max(sibling_end) # Collect all relevant tid/pid pairs - tid_pid_set = {(root.get("tid"), root.get("pid")) for root in iteration_roots} + tid_pid_set = {(root.get(_K.TID), root.get(_K.PID)) for root in iteration_roots} # Pre-filter all CPU events in the global window and by tid/pid cpu_events = [] @@ -171,9 +176,9 @@ def extract_iteration( for iteration_root in tqdm(iteration_roots): start_time = [] end_time = [] - iter_tid = iteration_root.get("tid") - iter_pid = iteration_root.get("pid") - iter_ts = iteration_root.get("ts", 0) + iter_tid = iteration_root.get(_K.TID) + iter_pid = iteration_root.get(_K.PID) + iter_ts = iteration_root.get(_K.TimeStamp, 0) iter_end = root_end[id(iteration_root)] correlation_ids: set[int] = set() diff --git a/tests/test_split_inference_trace_annotation.py b/tests/test_split_inference_trace_annotation.py index 6753079ca..6bfe3f550 100644 --- a/tests/test_split_inference_trace_annotation.py +++ b/tests/test_split_inference_trace_annotation.py @@ -362,6 +362,21 @@ def test_extract_iteration_keeps_cpu_ops_between_step_annotations(): assert "cpu_op_0_0" not in names1 +def test_cpu_window_end_falls_back_to_annotation_duration(): + from TraceLens.TraceUtils.split_inference import trace_extraction as te + + assert te._annotation_pattern("not_an_iteration") is None + lone = { + "name": VLLM_PRIMARY_ANNOTATION.format(i=0), + "ts": 1000, + "dur": 100, + "tid": 10, + "pid": 1, + } + assert te._cpu_window_end(lone, None, []) == 1100 + assert te._next_same_pattern_ts({"name": "not_an_iteration", "ts": 0}, []) is None + + def test_extract_iteration_empty_roots(): trace = make_trace([VLLM_PRIMARY_ANNOTATION.format(i=0)]) events = trace["traceEvents"] From df78358d6279035fd704daf6929d460e7700b27d Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Thu, 10 Sep 2026 23:27:05 -0400 Subject: [PATCH 3/3] Black-format split CPU-window helpers. Co-authored-by: Cursor --- TraceLens/TraceUtils/split_inference/trace_extraction.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index 8e50166b3..907033e70 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -72,6 +72,7 @@ def _cpu_window_end(root: dict, next_sibling_ts: float | None, events: list[dict return next_ts return root.get(_K.TimeStamp, 0) + root.get(_K.Duration, 0) + GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"] @@ -144,9 +145,7 @@ def extract_iteration( sibling_end = [] for i, root in enumerate(roots_by_ts): next_sibling = ( - roots_by_ts[i + 1].get(_K.TimeStamp) - if i + 1 < len(roots_by_ts) - else None + roots_by_ts[i + 1].get(_K.TimeStamp) if i + 1 < len(roots_by_ts) else None ) sibling_end.append(_cpu_window_end(root, next_sibling, events)) root_end = {id(root): end for root, end in zip(roots_by_ts, sibling_end)}