Fix/kernel op identity

.github/workflows/ci-l0-checks.yml

@@ -89,6 +89,8 @@ jobs:
             e2e_workflow/scripts/tests/test_parse_profile.py \
             e2e_workflow/scripts/tests/test_capture_shapes.py \
             e2e_workflow/scripts/tests/test_overlay_setup.py \
+            e2e_workflow/scripts/tests/test_seam_trace.py \
+            e2e_workflow/scripts/tests/test_kernel_selection.py \
             e2e_workflow/scripts/tests/test_attribute_weights_edges.py \
             e2e_workflow/scripts/tests/test_op_bench.py \
             e2e_workflow/scripts/tests/test_harness_lib.py \

e2e_workflow/roles/kernel_extractor.md

@@ -122,26 +122,35 @@ If the live regime genuinely cannot be reproduced offline (op only exists fused
 routing-dependent MoE token counts), say so in `notes` and report `editable:false`/drop rather than
 freeze an out-of-regime oracle nobody should trust.
 
-1. **Locate the source.** **If `KERNEL.source_hint`/`KERNEL.launcher_hint` is provided (TraceLens
+1. **Locate the source and select the live launcher.** `KERNEL.device_kernel` is the profiled GPU
+   symbol this extraction must reach. `KERNEL.target_callable` is only a hint, and
+   `KERNEL.live_call_seam` is prose that must never become a machine target. Start from
+   `KERNEL.seam_candidates[]`. If source/runtime inspection finds a missing inner launcher, append it
+   to the returned `seam_candidates[]` with its exact role, depth, matching device kernels, and evidence.
+   **If `KERNEL.source_hint`/`KERNEL.launcher_hint` is provided (TraceLens
    pre-resolved the file/seam), look there FIRST** — but always CONFIRM by importing the package +
    grepping the `short_name`/`module:attr` target; never trust the hint blindly (it may point at a
    launcher/wrapper rather than the true defining file). If no hint, resolve as usual
    (`python3 -c "import sglang,os;print(os.path.dirname(sglang.__file__))"`, then grep the
    `short_name` / the `module:attr` target).
    **OP-IDENTITY IS THE RULE: extract the op the LIVE kernel actually is, at the seam it is actually called
-   from — never a different op.** Two cases:
+   from — never a different op.** Three cases:
    - **Standalone LIBRARY op** (a discrete hipBLASLt/rocBLAS `gemm(...)` / library attention whose only
      call site is that library call, no editable body) → STOP, report `editable=false`, `target_callable=""`;
      it belongs to the config/tune-hook track (per-shape DB tune / backend env), not a source rewrite. Do
      NOT synthesize a standalone-GEMM proxy just to make it look extractable.
    - **FUSED / monolithic op** (fused-MoE, grouped-expert GEMM, asm/CK fused kernel — `KERNEL` arrives with
      `op_kind=moe` and `GEMM_SYNTH=false`): **extract the FUSED op** (capture its live I/O oracle), NOT its
-     constituent standalone GEMMs. Set `target_callable` to the **dispatcher** actually called at runtime —
-     use `KERNEL.target_callable`/`KERNEL.live_call_seam` if provided (e.g. the vLLM `fused_moe`/
-     `fused_experts` dispatcher), which is editable Python EVEN WHEN the underlying kernel is a non-editable
+     constituent standalone GEMMs. Select the bindable whole-operation **`op_seam`** from
+     `KERNEL.seam_candidates` (or an exact `KERNEL.target_callable` hint), which is editable Python EVEN
+     WHEN the underlying kernel is a non-editable
      library/asm `.so`. That dispatcher seam is what lets a fused op be BACKEND-SWAPPED (aiter/flydsl/triton
      fused) or AUTHOR-fused-replaced regardless of the underlying kernel's editability. Report
      `editable=true` (the seam is rebindable). NEVER decompose it into a dense A·Bᵀ GEMM — no live call site.
+   - **OUTER WRAPPER / dispatcher** → do not stop after rejecting it. Descend to the deepest safe Python
+     `inner_launcher` or `op_seam` that launches `KERNEL.device_kernel`. A native or Triton
+     `kernel_entry` remains source evidence, not a monkeypatch target. If no safe callable can be found,
+     report `editable=false`; never claim selection success from rejection alone.
 2. **Capture shapes + oracle** from a live server using `scripts/capture_shapes.py` via a temporary
    capture overlay, driven by the SAME workload as the profile so shapes match the regime:
    ```bash
@@ -152,17 +161,37 @@ freeze an out-of-regime oracle nobody should trust.
    # PRISTINE install instead of the server the accepted kernels actually built.
    python3 "$SKILL_DIR/scripts/overlay_setup.py" add-capture \
      --overlay "$TASK/_capture_overlay" --from "$CURRENT_OVERLAY" \
-     --target "<module:attr>" --out "$TASK" --max 5 \
+     --target "<selected module:attr>" --out "$TASK" --max 5 \
      --capture-file "$SKILL_DIR/scripts/capture_shapes.py"
+   # Repeat add-marker for every relevant safe Python candidate, deepest first. The shim always
+   # installs captures BEFORE markers, so a marker can never freeze an alias of a function that the
+   # capture hook had not wrapped yet.
+   python3 "$SKILL_DIR/scripts/overlay_setup.py" add-marker \
+     --overlay "$TASK/_capture_overlay" --target "<candidate module:attr>" \
+     --marker-file "$SKILL_DIR/scripts/seam_trace.py"
    cp -r "$CURRENT_OVERLAY"/. "$TASK/baseline_overlay"/ 2>/dev/null || \
      python3 -c "import sys;sys.path.insert(0,'$SKILL_DIR/scripts');import overlay_setup as o;o._ensure_overlay('$TASK/baseline_overlay')"
 
    BACKEND="<backend>" OUT_DIR="$TASK/_capture" GPU="$GPU_ID" MODEL="$MODEL_PATH" \
    ISL=<WORKLOAD.isl> OSL=<WORKLOAD.osl> CONC=<WORKLOAD.conc> REPEATS=0 PROFILE=0 \
    OVERLAY_PYTHONPATH="$TASK/_capture_overlay" \
-   EXTRA_ENV="CAPTURE_TARGET=<module:attr> CAPTURE_OUT=$TASK CAPTURE_MAX=5" \
+   EXTRA_ENV="CAPTURE_TARGET=<selected module:attr> CAPTURE_OUT=$TASK CAPTURE_MAX=5 GEAK_SELECTION_TRACE=$TASK/selection_trace.json" \
      bash "$EVAL_DIR/bench_e2e.sh" 2>&1 | tee "$EVAL_DIR/logs/capture_<short_name>.log"
+   python3 "$SKILL_DIR/scripts/kernel_selection.py" \
+     --target "<selected module:attr>" --device-kernel "<KERNEL.device_kernel>" \
+     --capture-meta "$TASK"/capture.pid-*.rank-*/meta.json \
+     --torch-trace "$TASK"/selection_trace.pid-*.rank-*.call-*.json \
+     --candidate-target "<candidate module:attr>" \
+     --out "$TASK/selection_validation.json"
    ```
+   Repeat `--candidate-target` for every relevant candidate. `seam_trace` writes an atomic trace for
+   each PID/rank/root call, and `capture_shapes` writes atomic process-local artifacts. The verifier
+   merges all calls for each PID, isolates External ids between trace files, and requires every capture
+   PID to pass. Installation proof is distinct from execution: a marked mutually exclusive branch may
+   stay inactive, while a selected outer target fails if a deeper marked candidate launches the same
+   device kernel in any call/rank. Require `deepest_verified:true`, then copy the selected process's
+   `meta.json` and `reference_io.pt` into the task root.
+
    🔴 **Capture on the CURRENT stack, not the install.** The oracle you freeze is the truth source the
    candidate is judged against, and the baseline you time against is `baseline_overlay/`. Both must be
    the server as it runs RIGHT NOW (config + every accepted kernel). Capturing on the pristine install
@@ -451,7 +480,14 @@ Return JSON:
   "editable": true,
   "task_dir": "<EVAL_DIR>/kernels/<short_name>_task",
   "source_path_in_sglang": "<abs path under site-packages>",
+  "device_kernel": "<KERNEL.device_kernel verbatim>",
   "target_callable": "<module:attr>",
+  "seam_candidates": [
+    {"target_callable": "<module:attr>",
+     "role": "outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry",
+     "device_kernels": ["<KERNEL.device_kernel>"], "depth": 0, "evidence": "source/runtime evidence"}
+  ],
+  "selection_validation": {"contract": "kernel_selection", "ok": true, "deepest_verified": true},
   "candidate_bind": {"kind": "module|rebind", "module": "<dotted>", "file": "kernel_src/<f>.py"},
   "baseline_overlay": "<task_dir>/baseline_overlay",
   "baseline_frozen": true,
@@ -807,7 +843,14 @@ Return JSON:
   "regimes_captured": ["prefill"],
   "candidate_backends": ["aiter","hipblaslt","triton","ck"],
   "reference_io_sha256": "<or '' if synthesized>",
+  "device_kernel": "<KERNEL.device_kernel verbatim>",
   "target_callable": "<module:attr rebind seam if one exists, else ''>",
+  "seam_candidates": [
+    {"target_callable": "<module:attr>",
+     "role": "outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry",
+     "device_kernels": ["<KERNEL.device_kernel>"], "depth": 0, "evidence": "source/runtime evidence"}
+  ],
+  "selection_validation": {"contract": "kernel_selection", "ok": true, "deepest_verified": true},
   "baseline_callable": "<module:attr of the live default backend, resolved OUTSIDE the task dir>",
   "smoke": "pass|fail",
   "notes": "transpose/bias inference, regime, whether oracle was synthesized vs captured"

e2e_workflow/roles/profiler.md

@@ -100,8 +100,21 @@ An upstream orchestrator may already have profiled the SAME baseline workload wi
   `<br>` args, `classification`←map from `kernel_category`/`bound_type` (MoE/grouped-GEMM→library_gemm
   or triton per `kernel_kind`; attention→library_attn; etc.), `editable`←`op_to_source_patchable`. Carry
   `source_file`/`kernel_path` into each entry's `notes` (the Architect/Extractor reuse them). Write
-  `profile_topN.json` + `.md` via your own Write (you may shell out to `parse_profile.py` only if you
-  also have a trace; otherwise assemble the JSON yourself) and set `source:"tracelens"`.
+  `profile_topN.json` + `.md` via your own Write and set `source:"tracelens"`.
+  **Then annotate the assembled rows from profiler evidence; never hand-write `entity_kind`:**
+  ```bash
+  TLT=$(ls -1 "$TRACELENS_TRACE_FILE"/*rank0*.pt.trace.json.gz 2>/dev/null | head -1)
+  [ -z "$TLT" ] && TLT=$(ls -1 "$TRACELENS_TRACE_FILE"/*.pt.trace.json.gz \
+    "$TRACELENS_TRACE_FILE"/*.json.gz "$TRACELENS_TRACE_FILE"/*.json 2>/dev/null | head -1)
+  python3 "$EVAL_DIR/parse_profile.py" \
+    --annotate "$EVAL_DIR/profile/round_${ROUND}/profile_topN.json" \
+    --torch-trace "$TLT" \
+    --annotate-out "$EVAL_DIR/profile/round_${ROUND}/profile_topN.json"
+  ```
+  If no trace is available, fall back to the normal collection below instead of guessing a row's
+  entity kind. Annotation expands an outer dispatcher/custom-op row through torch-profiler External-id
+  edges into its concrete device children. Preserve the resulting `device_kernel`, `profile_parent`,
+  and split GPU percentages: rejecting a dispatcher without discovering its children is not success.
 - **If `TRACELENS_TRACE_FILE` is also a non-empty path that EXISTS → run an ADDITIONAL trace-analysis
   pass on top of analysis.md to sharpen the picture** (this is required by contract when the trace is
   present). `TRACELENS_TRACE_FILE` is a `torch_trace` **directory** that holds one steady-state serving

e2e_workflow/roles/system_architect.md

@@ -244,6 +244,20 @@ OPTIONAL upstream TraceLens prior (may be empty strings — treat empty/missing
    e2e by MORE than the noise band. Otherwise drop it — say so.
 5. Write `EVAL_DIR/strategy.md` (human-readable plan) and return the routing.
 
+> **Every head candidate must identify a device kernel and a structured callable chain.**
+> - Copy `entity_kind` and `device_kernel` from the profiled row. Route only `gpu_kernel` rows; if a
+>   dispatcher was expanded, route its device children rather than recreating the outer aggregate.
+> - `live_call_seam` is prose context only. Never copy arrows, signatures, paths, or prose into
+>   `target_callable`.
+> - Build `seam_candidates[]` from source and the baseline server log. Each entry has an exact importable
+>   `target_callable` (`module:attr`), `role` (`outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry`),
+>   matching `device_kernels`, `depth`, and evidence. Include every plausible callable on the live path.
+> - Keep native/JIT `kernel_entry` objects as source evidence only; replacing them can break their
+>   `.run`, `.warmup`, or cache protocols. For non-fused heads prefer the deepest safe
+>   `inner_launcher`/`op_seam`; fused heads must select the whole-operation `op_seam`.
+> - `target_callable` is only an initial hint. The Extractor may add a missing inner launcher and must
+>   prove the final choice with runtime markers; merely rejecting an outer wrapper is not discovery.
+
 Return JSON:
 ```json
 {
@@ -256,6 +270,15 @@ Return JSON:
   "head_candidates": [
     {"id": "h0", "short_name": "...", "op_kind": "gemm|attn", "pct_gpu_time": 0.0,
      "shapes": "[[1024,5120],[5120,34816]]", "dtype": "bf16", "regime": "prefill|decode|both",
+     "entity_kind": "gpu_kernel",
+     "device_kernel": "<exact GPU symbol copied from profile>",
+     "target_callable": "<exact module:attr hint selected from seam_candidates, or ''>",
+     "seam_candidates": [
+       {"target_callable": "<exact module:attr>",
+        "role": "outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry",
+        "device_kernels": ["<exact profiled GPU symbol>"], "depth": 0,
+        "runtime_verified": false, "evidence": "source/log evidence"}
+     ],
      "transpose_b": true, "bias": false,
      "candidate_backends": ["aiter","hipblaslt","triton","ck"],
      "is_fused_kernel": false,

e2e_workflow/scripts/capture_shapes.py

@@ -54,6 +54,22 @@
 }
 
 
+def _rank():
+    return next((str(os.environ[key]) for key in
+                 ("RANK", "LOCAL_RANK", "TP_RANK", "SLURM_PROCID")
+                 if key in os.environ), "unknown")
+
+
+def _process_out_dir(out_dir):
+    """Isolate selection-capture artifacts so TP workers cannot corrupt each other."""
+    unique = (os.environ.get("CAPTURE_PROCESS_UNIQUE") == "1"
+              or bool(os.environ.get("GEAK_SELECTION_TRACE")))
+    if not unique:
+        return out_dir
+    return os.path.join(
+        out_dir, f"capture.pid-{os.getpid()}.rank-{_rank()}")
+
+
 def _shapes_dtypes(args, kwargs):
     """Light shape/dtype walk (no clone) so we can catalog EVERY distinct shape cheaply, independent of
     the memory-bounded oracle capture."""
@@ -153,7 +169,20 @@ def _sig(args, kwargs):
 
 def _wrapper(*args, **kwargs):
     s = _STATE
-    out = s["orig"](*args, **kwargs)
+    # This marker is consumed by kernel_selection.py from the capture run's torch trace.  It turns
+    # "the hook saw calls" into stronger evidence: the GPU kernel selected from the baseline profile
+    # must actually execute while THIS callable is active.  record_function is effectively a no-op
+    # when no profiler is collecting, so existing non-profiled capture users keep the same behavior.
+    marker = f"GEAK_TARGET::{s['target']}"
+    try:
+        record_function = _torch().profiler.record_function
+    except Exception:
+        record_function = None
+    if record_function:
+        with record_function(marker):
+            out = s["orig"](*args, **kwargs)
+    else:
+        out = s["orig"](*args, **kwargs)
     s["calls"] += 1
     in_graph = _capturing()
     try:
@@ -239,7 +268,9 @@ def _flush(write_oracle=True):
     # workload capture (< max_cases distinct shapes) and a late regime-coverage case (appended past
     # max_cases) land on disk; records is bounded, so this rewrites only a handful of times.
     if write_oracle and records and len(records) > s["oracle_records"]:
-        torch.save({"target": s["target"], "records": records}, io_path)
+        tmp_io = f"{io_path}.tmp-{os.getpid()}-{threading.get_ident()}"
+        torch.save({"target": s["target"], "records": records}, tmp_io)
+        os.replace(tmp_io, io_path)
         import hashlib
         h = hashlib.sha256()
         with open(io_path, "rb") as fh:
@@ -274,6 +305,8 @@ def walk(o):
     # not just single-shape h.check_correct_multi.
     meta = {
         "target": s["target"],
+        "process_id": os.getpid(),
+        "rank": _rank(),
         "module": s["mod"].__name__ if s["mod"] else None,
         "attr": s["attr"],
         "num_cases": len(records),
@@ -291,8 +324,11 @@ def walk(o):
         "build": False,  # default: pure-python/triton; Extractor flips to True for HIP/CK/asm tasks
         "note": "Oracle captured from baseline. Do NOT edit unittest.py or reference_io.pt during opt.",
     }
-    with open(os.path.join(out_dir, "meta.json"), "w") as fh:
+    meta_path = os.path.join(out_dir, "meta.json")
+    tmp_meta = f"{meta_path}.tmp-{os.getpid()}-{threading.get_ident()}"
+    with open(tmp_meta, "w") as fh:
         json.dump(meta, fh, indent=2)
+    os.replace(tmp_meta, meta_path)
     sys.stderr.write(f"[capture_shapes] flushed {len(records)} case(s) "
                      f"(regimes={sorted(s['regime_seen'])}), "
                      f"oracle_complete={s['oracle_written']} -> {out_dir}\n")
@@ -365,8 +401,15 @@ def install(target, out_dir, max_cases=5):
             f"({type(orig).__module__}.{type(orig).__name__}): a plain-function stand-in for a native/"
             f"triton-JIT callable SIGSEGVs the server (e.g. mxfp4 matmul_ogs). Hook a Python-level seam "
             f"(its caller) instead, or set CAPTURE_WRAP_UNSAFE=1 to force.")
+    out_dir = _process_out_dir(out_dir)
     s.update(target=target, out_dir=out_dir, max_cases=int(max_cases),
              orig=orig, mod=mod, attr=attr, installed=True)
+    if os.environ.get("GEAK_SELECTION_TRACE"):
+        # Selection runs are intentionally short and server teardown may use SIGTERM, which skips
+        # Python atexit. Persist the first eager capture immediately.
+        s["flush_every"] = 1
+    elif os.environ.get("CAPTURE_FLUSH_EVERY"):
+        s["flush_every"] = max(1, int(os.environ["CAPTURE_FLUSH_EVERY"]))
     setattr(mod, attr, _make_wrapper(orig))
     atexit.register(_flush)
     sys.stderr.write(f"[capture_shapes] hooked {target}; recording up to {max_cases} cases -> {out_dir}\n")