Inference split: keep host ops between step annotations

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,55 @@
     iteration_details,
 )
 
+_K = TraceEventUtils.TraceKeys
+
+
+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(_K.Name, ""))
+    if pattern is None:
+        return None
+    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(_K.TimeStamp)
+        if ts is None or ts <= iter_ts:
+            continue
+        if e.get(_K.TID) != iter_tid or e.get(_K.PID) != iter_pid:
+            continue
+        if pattern.match(e.get(_K.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(_K.TimeStamp, 0) + root.get(_K.Duration, 0)
+
+
 GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"]
 
 
@@ -73,7 +123,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,12 +141,18 @@ def extract_iteration(
     # Compute the global time window for all iteration roots
     if not iteration_roots:
         return trace_json.copy(), [], 0, 0, 0
-    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
-    )
+    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(_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(_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 = []
@@ -114,10 +175,10 @@ 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_end = iter_ts + iteration_root.get("dur", 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()
 

tests/test_split_inference_trace_annotation.py

@@ -300,6 +300,83 @@ 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_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"]