chore: import upstream snapshot with attribution

utils/torch_compile_utils.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+import torch.distributed as dist
+
+
+def _is_main_process() -> bool:
+    return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
+
+
+def _log_once(message: str) -> None:
+    if _is_main_process():
+        print(message)
+
+
+class SafeCompiledCallable:
+    """Lazy torch.compile wrapper that falls back to eager on compile/runtime errors."""
+
+    def __init__(
+        self,
+        fn,
+        *,
+        name: str,
+        backend: str = "inductor",
+        mode: str | None = "max-autotune-no-cudagraphs",
+        fullgraph: bool = False,
+        dynamic: bool | None = False,
+        options: dict | None = None,
+        suppress_errors: bool = True,
+    ) -> None:
+        self.fn = fn
+        self.name = name
+        self.enabled = True
+        self.failed = False
+        self.failure_reason = None
+
+        if suppress_errors:
+            try:
+                import torch._dynamo as torch_dynamo
+
+                torch_dynamo.config.suppress_errors = True
+            except Exception as exc:
+                _log_once(f"[torch.compile] Could not enable suppress_errors: {exc}")
+
+        compile_kwargs = {
+            "backend": backend,
+            "fullgraph": fullgraph,
+            "dynamic": dynamic,
+        }
+        if mode:
+            compile_kwargs["mode"] = mode
+        if options:
+            compile_kwargs["options"] = options
+
+        _log_once(
+            "[torch.compile] Preparing "
+            f"{name}: backend={backend}, mode={mode}, "
+            f"fullgraph={fullgraph}, dynamic={dynamic}"
+        )
+        self.compiled_fn = torch.compile(fn, **compile_kwargs)
+
+    def __call__(self, *args, **kwargs):
+        if not self.enabled:
+            return self.fn(*args, **kwargs)
+
+        try:
+            return self.compiled_fn(*args, **kwargs)
+        except Exception as exc:
+            self.enabled = False
+            self.failed = True
+            self.failure_reason = repr(exc)
+            _log_once(
+                f"[torch.compile][warn] {self.name} failed; "
+                f"falling back to eager. reason={exc}"
+            )
+            return self.fn(*args, **kwargs)
+
+
+def configure_module_call_torch_compile(
+    module,
+    *,
+    name: str,
+    backend: str = "inductor",
+    mode: str | None = "max-autotune-no-cudagraphs",
+    fullgraph: bool = False,
+    dynamic: bool | None = False,
+    options: dict | None = None,
+    suppress_errors: bool = True,
+):
+    if not torch.cuda.is_available():
+        _log_once(f"[torch.compile] Skipping {name}: CUDA is not available")
+        return None
+
+    try:
+        return SafeCompiledCallable(
+            module,
+            name=name,
+            backend=backend,
+            mode=mode,
+            fullgraph=fullgraph,
+            dynamic=dynamic,
+            options=options,
+            suppress_errors=suppress_errors,
+        )
+    except Exception as exc:
+        _log_once(
+            f"[torch.compile][warn] Could not prepare {name}; "
+            f"continuing in eager mode. reason={exc}"
+        )
+        return None