perf(pt): pass is_inference to compile_fx_inner for compiled eval graphs

deepmd/pt/model/model/sezm_model.py

@@ -2322,11 +2322,13 @@ def compute_fn(  # type: ignore[misc]
             def _inductor_inference_compiler(
                 fx_gm: torch.fx.GraphModule, fx_inputs: list[Any]
             ) -> Any:
+                # This callback lowers the forward-only AOTAutograd inference
+                # graph (the eval branch below owns it; training compiles
+                # through the separate torch.compile branch above).
                 # ``compile_options`` already includes ``triton.max_tiles=1``
-                # from ``build_inductor_compile_options``, so the 1D launch-grid
-                # constraint applies to both the training and evaluation graphs.
+                # from ``build_inductor_compile_options``.
                 with _ind_cfg.patch(compile_options):
-                    return compile_fx_inner(fx_gm, fx_inputs)
+                    return compile_fx_inner(fx_gm, fx_inputs, is_inference=True)
 
             # select_decomp_table keeps the decomposition set aligned with
             # Inductor's fallback set (a mismatch raises an aten._to_copy

source/tests/pt/model/test_sezm_model.py

@@ -492,6 +492,53 @@ def test_eval_compile_first_frame_nloc_matches_trace_edge_count(self) -> None:
             msg="eval force mismatch when first-frame nloc == trace edge count",
         )
 
+    @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
+    def test_eval_compile_pins_is_inference_on_the_inference_slot(self) -> None:
+        """The eval slot must lower through ``compile_fx_inner(is_inference=True)``.
+
+        The existing suite (cache-slot presence, eager parity) holds with the
+        flag set, cleared, or absent; only this call-shape pin notices a
+        refactor that drops the argument and quietly reinstates the
+        peak-memory cost.  ``sezm_model`` imports ``compile_fx_inner`` from
+        its home module at call time, so the spy patches the source module.
+        """
+        coord, atype, box = self._make_frame_with_natoms(7)
+
+        model_dyn = get_sezm_model(self._build_model_params(use_compile=False))
+        self._randomize_params(model_dyn)
+        model_dyn.eval()
+        with mock.patch.dict(
+            os.environ,
+            {
+                "DP_COMPILE_INFER": "1",
+                # A cache hit would skip the lowering this test observes.
+                "TORCHINDUCTOR_FX_GRAPH_CACHE": "0",
+                "TORCHINDUCTOR_AUTOGRAD_CACHE": "0",
+            },
+            clear=False,
+        ):
+            model_cmp = get_sezm_model(self._build_model_params(use_compile=True))
+            model_cmp.load_state_dict(model_dyn.state_dict())
+            model_cmp.eval()
+
+            import torch._inductor.compile_fx as _compile_fx_mod
+
+            calls: list[dict] = []
+            real = _compile_fx_mod.compile_fx_inner
+
+            def _spy(*args, **kwargs):
+                calls.append(kwargs)
+                return real(*args, **kwargs)
+
+            with mock.patch.object(_compile_fx_mod, "compile_fx_inner", _spy):
+                model_cmp(coord, atype, box=box)
+
+        self.assertTrue(
+            any(kwargs.get("is_inference") is True for kwargs in calls),
+            "the eval slot must lower through compile_fx_inner(is_inference=True); "
+            f"observed call kwargs: {calls}",
+        )
+
     @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
     def test_compile_cache_slots_and_eval_shape_change(self) -> None:
         """Compile cache slots should coexist while eval handles batch-size growth."""