Place launch-less device work in its step by device time

Muse Spark's prefill all-gathers (layernorm sequence parallelism) run as
64 MiB peer copies issued through driver calls Kineto does not record:
2046 per rank per prefill window on B200 (run 38948461605), all
unattributed and unstepped. Such work now joins the step whose device
span on its rank contains it (step_inferred: device_time); all 2046 land
in a step.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: hbarclay

inferencex-e2e/benchmarks/profiling/extract.py

@@ -1028,6 +1028,7 @@ def extract_replay(trace, captured, rank, window, step_log, copies, clocks, rout
         key = step_key(row.get("marks") or {})
         row["step"] = list(key) if key else None
         by_step[key].append(row)
+    place_by_device_time(by_step)
 
     rank_dir = os.path.join(out_dir, f"window{window}", str(rank))
     os.makedirs(rank_dir, exist_ok=True)
@@ -1073,6 +1074,31 @@ def extract_replay(trace, captured, rank, window, step_log, copies, clocks, rout
     return entry, index
 
 
+def place_by_device_time(by_step):
+    """Device work with no CPU launch (copies a library issues through driver calls Kineto
+    does not record, e.g. Muse Spark's copy-engine all-gathers under layernorm sequence
+    parallelism: 64 MiB peer copies) goes in the step whose device span on this rank
+    contains it, marked step_inferred; outside every step it stays unstepped."""
+    loose = by_step.get(None) or []
+    spans = sorted((min(r["ts_us"] for r in rows), max(r["ts_us"] + r["dur_us"] for r in rows), key)
+                   for key, rows in by_step.items() if key is not None and rows)
+    if not loose or not spans:
+        return
+    starts = [s for s, _, _ in spans]
+    kept = []
+    for row in loose:
+        k = bisect.bisect_right(starts, row["ts_us"]) - 1
+        if row["source"] == "unattributed" and k >= 0 and row["ts_us"] + row["dur_us"] <= spans[k][1]:
+            key = spans[k][2]
+            row["step"], row["step_inferred"] = list(key), "device_time"
+            by_step[key].append(row)
+        else:
+            kept.append(row)
+    by_step[None] = kept
+    if not kept:
+        del by_step[None]
+
+
 def stream_successors(trace):
     """device activity index -> start of the next activity on its device and stream."""
     streams = collections.defaultdict(list)