perf(closed_loop): torch.compile the evaluation forward pass

scenario_generation/closed_loop_evaluation.py

@@ -30,6 +30,7 @@
     segment_row_for_json,
     tdigest_sidecar_row,
 )
+from scenario_generation.inference_compile import compiled_for_inference
 from scenario_generation.perf_timer import Timers
 from scenario_generation.reproducer_rollout import render_segment
 from scenario_generation.route_timeline import RouteTimeline
@@ -167,7 +168,8 @@ def run(self) -> dict:
                 f"{len(jobs)} job(s) -> {[j.job_id for j in jobs]}"
             )
 
-        result = self.execute_jobs(jobs)
+        with compiled_for_inference(self.model):
+            result = self.execute_jobs(jobs)
         elapsed_sec = time.perf_counter() - t0
 
         if self.ddp_world_size > 1:

scenario_generation/inference_compile.py

@@ -0,0 +1,81 @@
+"""``torch.compile`` for the closed-loop evaluation forward pass.
+
+``nn.MultiheadAttention`` takes a fused fastpath in eval mode that dynamo cannot trace, so a
+compiled model always runs the decomposed path instead. Both are correct; they round differently
+in the last float32 bit, and a closed loop feeds that difference back into its own next input.
+Closed-loop metrics therefore shift once. No backend choice avoids this -- it is structural.
+
+``cudagraphs`` replays the same kernels in the same order and was the faster of the two
+candidates on a real rollout. Inductor generates its own reductions and was no faster.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class _CloneEncoderOutput(torch.nn.Module):
+    """Clone the encoder's output so the DiT's next cudagraph replay cannot overwrite it.
+
+    The encoding is computed once per inference and read by every DPM-solver step, but a
+    cudagraph-managed output buffer belongs to the graph that produced it and is reused on the
+    next replay.
+    """
+
+    def __init__(self, inner: torch.nn.Module) -> None:
+        super().__init__()
+        self.inner = inner
+
+    def forward(self, *args, **kwargs):
+        return self.inner(*args, **kwargs).clone()
+
+
+@contextlib.contextmanager
+def compiled_for_inference(model):
+    """Compile the encoder and the DiT for the duration, then put the model back as it was.
+
+    Scoped rather than permanent because the training loop hands over its LIVE model: compiling
+    in place would leave ``model.encoder`` wrapped and rename the decoder's parameter keys
+    (``_orig_mod.`` prefixes), which would then land in the next checkpoint.
+
+    The fastpath switch is entered here rather than left to the caller because compilation is
+    lazy and the flag is not in dynamo's guard set: setting it afterwards would leave the graph
+    traced under the old value while reading back as changed.
+
+    Callers must run inference under ``no_grad`` -- an output carrying an autograd graph drops
+    compile off the CUDA-graph fast path silently -- and call :func:`mark_inference_step` once
+    per inference.
+    """
+    encoder = getattr(model, "encoder", None)
+    decoder = getattr(model, "decoder", None)
+    dit = getattr(decoder, "dit", None)
+    if not isinstance(encoder, torch.nn.Module) or not isinstance(dit, torch.nn.Module):
+        # e.g. the ONNX adapter from simulate.load_onnx_model, which runs its own runtime.
+        logger.info("%s has nothing to compile; running it as-is.", type(model).__name__)
+        yield model
+        return
+
+    fastpath = torch.backends.mha.get_fastpath_enabled()
+    torch.backends.mha.set_fastpath_enabled(False)
+    model.encoder = _CloneEncoderOutput(torch.compile(encoder, backend="cudagraphs"))
+    decoder.dit = torch.compile(dit, backend="cudagraphs")
+    try:
+        yield model
+    finally:
+        model.encoder = encoder
+        decoder.dit = dit
+        torch.backends.mha.set_fastpath_enabled(fastpath)
+
+
+def mark_inference_step() -> None:
+    """Open a new cudagraph step.
+
+    Unconditional on purpose: the rollout should not have to know whether the model it was handed
+    is compiled, and skipping it on a compiled model corrupts the replayed buffers.
+    """
+    torch.compiler.cudagraph_mark_step_begin()

scenario_generation/reproducer_rollout.py

@@ -31,6 +31,7 @@
 
 from planner_metrics.scene_format import future_to_4col
 from scenario_generation.danger_event_selection import OnlineEventSelector
+from scenario_generation.inference_compile import mark_inference_step
 from scenario_generation.metrics import (
     score_object_step,
     score_object_step_batched,
@@ -1660,6 +1661,8 @@ def render_segment(
         override = None
         if plan_world is None or offset == 0:
             data = _to_torch_batch([np_dict], model_args, device)
+            # No-op unless the model was compiled with cudagraphs; one inference is one step.
+            mark_inference_step()
             _, outputs = model(data)
             pred = outputs["prediction"][0, 0].cpu().numpy()
             plan_world = _ego_pred_to_world(
@@ -1968,6 +1971,9 @@ def run_segments_batched(
                                 else:
                                     nxt = s.cursor.max_idx_reached + 1
                                 pool.submit(s.tl.prefetch, range(nxt, nxt + prefetch_ahead))
+                        # No-op unless the model was compiled with cudagraphs; one inference is
+                        # one step.
+                        mark_inference_step()
                         _, outputs = model(data)
                         preds = outputs["prediction"][:, 0].cpu().numpy()  # (B,80,4)
                         # Model's predicted turn indicator per segment, decoded with the SAME

scenario_generation/tests/test_inference_compile.py

@@ -0,0 +1,52 @@
+"""Smoke tests for scenario_generation.inference_compile.
+
+The compiled forward itself needs a GPU, so what is worth pinning here is the scoping: the
+training loop hands over its live model, and a wrapper left behind would rename the decoder's
+parameter keys in the next checkpoint.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from scenario_generation.inference_compile import compiled_for_inference, mark_inference_step
+
+
+class _Decoder(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.dit = nn.Linear(4, 4)
+
+
+class _Model(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.encoder = nn.Linear(4, 4)
+        self.decoder = _Decoder()
+
+
+def test_model_and_fastpath_are_restored_on_exit():
+    model = _Model()
+    encoder, dit = model.encoder, model.decoder.dit
+    fastpath = torch.backends.mha.get_fastpath_enabled()
+
+    with compiled_for_inference(model):
+        assert model.encoder is not encoder
+        assert not torch.backends.mha.get_fastpath_enabled()
+
+    assert model.encoder is encoder
+    assert model.decoder.dit is dit
+    assert torch.backends.mha.get_fastpath_enabled() == fastpath
+    assert set(model.state_dict()) == set(_Model().state_dict())
+
+
+def test_a_model_with_nothing_to_compile_passes_through():
+    """The ONNX adapter runs its own runtime and has no encoder/dit to hand to dynamo."""
+    model = object()
+    with compiled_for_inference(model) as handed_back:
+        assert handed_back is model
+
+
+def test_mark_inference_step_is_safe_without_a_compiled_model():
+    mark_inference_step()