feat(v8-r06): compile.trace RPC + CompileTracePanel

apstenku123 · apstenku123 · commit 3dc3dc1f05a7 · 2026-05-23T19:42:39.000+02:00
Closes cppmega-mlx-yz6n.6. New RPC compile.trace(spec, backend) reads
the FusionRegionPlan output of plan_fusion_regions and emits a
per-op compile trace with fused/dlpack/materialised chips +
aggregate counters (fused_groups list, dlpack_crossings count,
materialised_ops list).

The backend selector ('tilelang' | 'torch_inductor' | 'mlx') is
echoed through but the underlying planner is unified — each op's
per-row backend label is the actual backend the planner chose
(path_c / metal_inline / dlpack_handoff).

UI: CompileTracePanel renders the trace as a chip table + an
aggregate header (fused / dlpack / materialised counters), one chip
per op.

Tests: 5 new pytest (per-brick op count, fused flag, backend label,
dispatch envelope, aggregate consistency), 2 new vitest (panel mount,
error banner). Regression: 112 pytest fusion+compile+mxfp4+catalog +
468 vitest green.
diff --git a/cppmega_v4/jsonrpc/compile_trace_method.py b/cppmega_v4/jsonrpc/compile_trace_method.py
@@ -0,0 +1,132 @@
+"""V8-R06: ``compile.trace`` RPC — surface the fusion plan as a wire-
+form compile trace that the UI can render with per-op chips.
+
+Each op in the trace carries:
+  - name: the brick name
+  - fused: True iff its region has 2+ bricks
+  - group: the region label (e.g. ``"gemm_softmax_0"``)
+  - materialised: True iff the region is the ``dlpack_handoff`` backend
+    (a materialised boundary tensor is unavoidable across DLPack)
+  - dlpack_boundary: True iff this op crosses a backend boundary
+  - backend: tilelang | mlx | torch_inductor | metal_inline | dlpack_handoff
+
+For now there is one shared planner — ``plan_fusion_regions`` — and
+the same trace is returned for all three ``backend`` selectors. The
+backend label in each op is taken from the FusionRegionPlan rather
+than the request, so the UI can render the actual backend chosen by
+the planner. The ``backend`` request parameter is reserved for the
+upcoming inductor-specific path.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from cppmega_v4.fusion.auto_planner import plan_fusion_regions
+from cppmega_v4.fusion.brick_graph import from_block_specs
+from cppmega_v4.jsonrpc.cache import LRUCache
+from cppmega_v4.jsonrpc.methods import (
+    _cache_lookup, _cache_store, _graph_to_specs,
+)
+from cppmega_v4.jsonrpc.schema import VerifyParams
+
+
+__all__ = [
+    "CompileTraceParams",
+    "CompileTraceOp",
+    "CompileTraceResult",
+    "compile_trace",
+]
+
+
+CompileBackend = Literal["tilelang", "mlx", "torch_inductor"]
+
+
+class CompileTraceParams(BaseModel):
+    """Input — VerifyParams payload + the backend perspective."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    spec: VerifyParams
+    backend: CompileBackend = "mlx"
+
+
+class CompileTraceOp(BaseModel):
+    """One row in the compile trace — chip-renderable in the UI."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: str
+    fused: bool
+    group: str
+    materialised: bool
+    dlpack_boundary: bool
+    backend: str
+
+
+class CompileTraceResult(BaseModel):
+    """Full trace + aggregate counters that surface in extras.train."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    ops: list[CompileTraceOp] = Field(default_factory=list)
+    fused_groups: list[str] = Field(default_factory=list)
+    dlpack_crossings: int = 0
+    materialised_ops: list[str] = Field(default_factory=list)
+    compile_artifact_path: str | None = None
+    backend: str
+
+
+def compile_trace(
+    params: CompileTraceParams, *, cache: LRUCache | None = None,
+) -> CompileTraceResult:
+    """Plan fusion regions for ``params.spec`` and render the trace."""
+    key, hit = _cache_lookup(cache, "compile.trace", params)
+    if hit is not None:
+        return hit
+
+    specs = _graph_to_specs(params.spec.graph)
+    hidden = params.spec.dim_env.get("H", 64)
+    graph = from_block_specs(specs, hidden_size=hidden, instantiate=False)
+    plans = list(plan_fusion_regions(graph))
+
+    ops: list[CompileTraceOp] = []
+    fused_groups: list[str] = []
+    materialised: list[str] = []
+    dlpack_crossings = 0
+    last_backend: str | None = None
+
+    for idx, plan in enumerate(plans):
+        group_label = f"region_{idx:02d}_{plan.backend}"
+        is_materialised = plan.backend == "dlpack_handoff"
+        if is_materialised and last_backend is not None and \
+                last_backend != plan.backend:
+            dlpack_crossings += 1
+        if plan.is_fused:
+            fused_groups.append(group_label)
+        for i, name in enumerate(plan.brick_names):
+            ops.append(CompileTraceOp(
+                name=name,
+                fused=plan.is_fused,
+                group=group_label,
+                materialised=is_materialised,
+                dlpack_boundary=(is_materialised and i == 0
+                                  and idx > 0),
+                backend=plan.backend,
+            ))
+            if is_materialised and name not in materialised:
+                materialised.append(name)
+        last_backend = plan.backend
+
+    out = CompileTraceResult(
+        ops=ops,
+        fused_groups=fused_groups,
+        dlpack_crossings=dlpack_crossings,
+        materialised_ops=materialised,
+        compile_artifact_path=None,   # set by tilelang codegen path
+        backend=params.backend,
+    )
+    _cache_store(cache, key, out)
+    return out
diff --git a/cppmega_v4/jsonrpc/dispatcher.py b/cppmega_v4/jsonrpc/dispatcher.py
@@ -87,6 +87,10 @@
     HfQuickstartParams,
     hf_quickstart_method,
 )
+from cppmega_v4.jsonrpc.compile_trace_method import (
+    CompileTraceParams,
+    compile_trace,
+)
 from cppmega_v4.jsonrpc.tokenizer_roundtrip_text_method import (
     TokenizerRoundtripTextParams,
     roundtrip_text,
@@ -153,6 +157,10 @@
         HfQuickstartParams,
         lambda p, c: hf_quickstart_method(p, cache=c),
     ),
+    "compile.trace": (
+        CompileTraceParams,
+        lambda p, c: compile_trace(p, cache=c),
+    ),
     "probe.run": (
         ProbeRunParams,
         lambda p, c: probe_run(p, cache=c),
diff --git a/cppmega_v4/jsonrpc/schema.py b/cppmega_v4/jsonrpc/schema.py
@@ -811,6 +811,7 @@ class PipelineRunResult(BaseModel):
     "memory.matrix",
     "architectures.auto_fit",
     "data.hf_quickstart",
+    "compile.trace",
 })
 
 
diff --git a/tests/v4/test_compile_trace.py b/tests/v4/test_compile_trace.py
@@ -0,0 +1,81 @@
+"""V8-R06 pytest: compile.trace RPC."""
+
+from __future__ import annotations
+
+import pytest
+
+from cppmega_v4.jsonrpc.compile_trace_method import (
+    CompileTraceParams, compile_trace,
+)
+from cppmega_v4.jsonrpc.dispatcher import dispatch
+from cppmega_v4.jsonrpc.schema import JsonRpcRequest, VerifyParams
+
+
+def _spec() -> dict:
+    return {
+        "graph": {
+            "nodes": [
+                {"id": "a", "kind": "attention", "params": {}},
+                {"id": "b", "kind": "mlp", "params": {}},
+                {"id": "c", "kind": "attention", "params": {}},
+                {"id": "d", "kind": "mlp", "params": {}},
+            ],
+            "edges": [
+                {"src": "a", "dst": "b"},
+                {"src": "b", "dst": "c"},
+                {"src": "c", "dst": "d"},
+            ],
+        },
+        "dim_env": {"H": 128},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["d"]},
+        "optim": {"kind": "adamw", "groups": [
+            {"matcher": "all", "lr": 1e-3, "weight_decay": 0.01,
+             "betas": [0.9, 0.95]}]},
+        "sharding": None,
+        "training": True,
+    }
+
+
+def test_returns_one_op_per_brick():
+    r = compile_trace(CompileTraceParams(spec=VerifyParams(**_spec())))
+    assert len(r.ops) == 4
+    assert {op.name for op in r.ops} == {"a", "b", "c", "d"}
+
+
+def test_fused_flag_matches_region_size():
+    r = compile_trace(CompileTraceParams(spec=VerifyParams(**_spec())))
+    # All 4 land in one big path_c region for this contrived spec, so
+    # fused is True everywhere. The aggregate fused_groups list has
+    # one entry.
+    assert all(op.fused for op in r.ops)
+    assert len(r.fused_groups) >= 1
+
+
+def test_backend_label_round_trips():
+    r = compile_trace(CompileTraceParams(
+        spec=VerifyParams(**_spec()), backend="tilelang"))
+    assert r.backend == "tilelang"
+
+
+def test_dispatch_end_to_end():
+    req = JsonRpcRequest(
+        jsonrpc="2.0", id="t-r06",
+        method="compile.trace",
+        params={"spec": _spec(), "backend": "mlx"},
+    )
+    resp = dispatch(req)
+    assert resp.error is None, resp.error
+    r = resp.result
+    assert "ops" in r and len(r["ops"]) == 4
+    for op in r["ops"]:
+        assert {"name", "fused", "group", "materialised",
+                "dlpack_boundary", "backend"} <= set(op)
+
+
+def test_aggregate_counters_consistent():
+    r = compile_trace(CompileTraceParams(spec=VerifyParams(**_spec())))
+    # Every materialised op-name appears in materialised_ops exactly once.
+    mat_set = {op.name for op in r.ops if op.materialised}
+    assert mat_set == set(r.materialised_ops)
+    # dlpack_crossings counts backend transitions, not per-op flags.
+    assert r.dlpack_crossings >= 0
diff --git a/vbgui/src/components/CompileTracePanel.tsx b/vbgui/src/components/CompileTracePanel.tsx
@@ -0,0 +1,151 @@
+/**
+ * V8-R06 CompileTracePanel — renders the compile.trace RPC output as
+ * a per-op chip table with fused / dlpack-crossing / materialised
+ * markers and an aggregate header line.
+ */
+
+import { useEffect, useState } from "react";
+import type { RpcClient } from "@/lib/rpc";
+
+interface CompileTraceOp {
+  name: string;
+  fused: boolean;
+  group: string;
+  materialised: boolean;
+  dlpack_boundary: boolean;
+  backend: string;
+}
+
+interface CompileTraceResult {
+  ops: CompileTraceOp[];
+  fused_groups: string[];
+  dlpack_crossings: number;
+  materialised_ops: string[];
+  compile_artifact_path: string | null;
+  backend: string;
+}
+
+export interface CompileTracePanelProps {
+  rpc: RpcClient;
+  specPayload: unknown;
+  backend?: "tilelang" | "torch_inductor" | "mlx";
+}
+
+export function CompileTracePanel({
+  rpc, specPayload, backend = "mlx",
+}: CompileTracePanelProps): JSX.Element {
+  const [trace, setTrace] = useState<CompileTraceResult | null>(null);
+  const [err, setErr] = useState<string | null>(null);
+
+  useEffect(() => {
+    let cancelled = false;
+    (async () => {
+      setErr(null);
+      try {
+        const r = await rpc.call<CompileTraceResult>(
+          "compile.trace",
+          { spec: specPayload, backend },
+        );
+        if (!cancelled) setTrace(r);
+      } catch (e) {
+        if (!cancelled) {
+          setErr(e instanceof Error ? e.message : String(e));
+        }
+      }
+    })();
+    return () => { cancelled = true; };
+  }, [rpc, specPayload, backend]);
+
+  if (err) {
+    return <div data-testid="compile-trace-error"
+                style={{ color: "#b91c1c", padding: 12 }}>{err}</div>;
+  }
+  if (!trace) {
+    return <div data-testid="compile-trace-loading"
+                style={{ padding: 12, color: "#6b7280" }}>loading…</div>;
+  }
+
+  return (
+    <div data-testid="compile-trace" style={{ padding: 12,
+         fontFamily: "system-ui, sans-serif", fontSize: 12 }}>
+      <header style={{ display: "flex", gap: 12, marginBottom: 8 }}>
+        <span>backend: <strong>{trace.backend}</strong></span>
+        <span data-testid="compile-trace-fused-count">
+          fused groups: {trace.fused_groups.length}
+        </span>
+        <span data-testid="compile-trace-dlpack-crossings">
+          dlpack crossings: {trace.dlpack_crossings}
+        </span>
+        <span data-testid="compile-trace-materialised-count">
+          materialised: {trace.materialised_ops.length}
+        </span>
+      </header>
+      {trace.fused_groups.map((g) => (
+        <span key={g}
+              data-testid={`compile-trace-fused-group-${g}`}
+              style={{ display: "inline-block", marginRight: 6,
+                       background: "#dbeafe", color: "#1e3a8a",
+                       borderRadius: 4, padding: "2px 6px" }}>
+          {g}
+        </span>
+      ))}
+      <table style={{ marginTop: 8, borderCollapse: "collapse",
+                      width: "100%" }}>
+        <thead>
+          <tr style={{ borderBottom: "1px solid #e5e7eb" }}>
+            <th style={th}>op</th>
+            <th style={th}>backend</th>
+            <th style={th}>group</th>
+            <th style={th}>chips</th>
+          </tr>
+        </thead>
+        <tbody>
+          {trace.ops.map((op, i) => (
+            <tr key={`${op.name}-${i}`}
+                data-testid={`compile-trace-op-${i}`}>
+              <td style={td}>{op.name}</td>
+              <td style={td}>{op.backend}</td>
+              <td style={td}>{op.group}</td>
+              <td style={td}>
+                {op.fused && (
+                  <span style={chip("#dbeafe", "#1e3a8a")}>fused</span>
+                )}
+                {op.materialised && (
+                  <span style={chip("#fef3c7", "#92400e")}
+                        data-testid={
+                          `compile-trace-op-${i}-materialised`}>
+                    materialised
+                  </span>
+                )}
+                {op.dlpack_boundary && (
+                  <span style={chip("#fee2e2", "#7f1d1d")}
+                        data-testid={
+                          `compile-trace-op-${i}-dlpack`}>
+                    dlpack
+                  </span>
+                )}
+              </td>
+            </tr>
+          ))}
+        </tbody>
+      </table>
+    </div>
+  );
+}
+
+const th: React.CSSProperties = {
+  textAlign: "left", padding: "4px 6px", color: "#374151",
+  fontWeight: 600,
+};
+
+const td: React.CSSProperties = {
+  padding: "4px 6px", color: "#1f2937",
+  borderBottom: "1px solid #f3f4f6",
+};
+
+function chip(bg: string, fg: string): React.CSSProperties {
+  return {
+    display: "inline-block", marginRight: 4, padding: "1px 6px",
+    borderRadius: 4, background: bg, color: fg, fontSize: 10,
+  };
+}
diff --git a/vbgui/tests/CompileTracePanel.test.tsx b/vbgui/tests/CompileTracePanel.test.tsx
