From 54d5babed312330a1e2dc2a7cbd8dd6209928551 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 2 Nov 2026 14:46:04 -0800
Subject: [PATCH 1/3] Preserve asynchronous cache prefetching in compiled
 prediction

---
 sdm/models/base/icl.py   |  4 +++
 sdm/models/base/model.py | 53 ++++++++++++++++++++++++++++++++--------
 2 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py
index 2285346699..0f7a0f5b08 100644
--- a/sdm/models/base/icl.py
+++ b/sdm/models/base/icl.py
@@ -354,6 +354,10 @@ def fit(
 
         self._cache = cache.freeze()
 
+        # Prepare the stream before compiled predict to avoid recompilation.
+        if x.is_cuda and x.device not in self._transfer_streams:
+            self._transfer_streams[x.device] = torch.cuda.Stream(x.device)
+
     def predict(
         self,
         x: Tensor | TableTensor,  # [..., R, D]
diff --git a/sdm/models/base/model.py b/sdm/models/base/model.py
index 1b48d1ecfd..fb78e5e6b7 100644
--- a/sdm/models/base/model.py
+++ b/sdm/models/base/model.py
@@ -3,6 +3,7 @@
 
 import abc
 from collections.abc import Iterator, Sequence
+from functools import partial
 from typing import ClassVar, Generic, Self, TypeVar
 
 import torch
@@ -91,6 +92,38 @@ def _validate_feature_stype(self, x: TableTensor) -> None:
             )
 
 
+def _copy_to_device(
+    tensor: Tensor,
+    device: torch.device,
+    stream: torch.cuda.Stream,
+) -> Tensor:
+    """Keep copy allocations separate from the parent graph's buffer reuse."""
+    with torch.cuda.stream(stream):
+        return tensor.to(device=device, non_blocking=True)
+
+
+if hasattr(torch.compiler, "nested_compile_region"):
+    _copy_to_device = torch.compiler.nested_compile_region(_copy_to_device)
+
+
+def _transfer(
+    item: T,
+    device: torch.device,
+    stream: torch.cuda.Stream,
+) -> T:
+    """Move a tensor or tensor container onto the prefetch stream."""
+    if not torch.compiler.is_compiling():
+        with torch.cuda.stream(stream):
+            return item.to(device=device, non_blocking=True)  # type: ignore
+    if isinstance(item, DeviceMixin):
+        return item._apply_tensor(
+            partial(_transfer, device=device, stream=stream)
+        )
+    if item.device == device:
+        return item
+    return _copy_to_device(item, device, stream)  # type: ignore
+
+
 class _Prefetcher(Generic[T]):
     def __init__(
         self,
@@ -112,11 +145,11 @@ def __enter__(self) -> Self:
         if self._transfer_stream is not None:
             self._compute_stream = torch.cuda.current_stream(self._device)
             try:
-                with torch.cuda.stream(self._transfer_stream):
-                    self._next = self._items[0].to(  # type: ignore
-                        device=self._device,
-                        non_blocking=True,
-                    )
+                self._next = _transfer(
+                    item=self._items[0],
+                    device=self._device,
+                    stream=self._transfer_stream,
+                )
             except BaseException:
                 self._transfer_stream.synchronize()
                 raise
@@ -153,11 +186,11 @@ def __iter__(self) -> Iterator[T]:
 
             if self._next is not None:
                 if self._transfer_stream is not None:
-                    with torch.cuda.stream(self._transfer_stream):
-                        self._next = self._next.to(  # type: ignore
-                            device=self._device,
-                            non_blocking=True,
-                        )
+                    self._next = _transfer(
+                        item=self._next,
+                        device=self._device,
+                        stream=self._transfer_stream,
+                    )
                 else:
                     self._next = self._next.to(self._device)  # type: ignore
 

From f09b39955a56d1790b776d91cf7ee95989e1b92f Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 2 Nov 2026 19:42:11 -0800
Subject: [PATCH 2/3] Explain nested compile region compatibility guard

---
 sdm/models/base/model.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/sdm/models/base/model.py b/sdm/models/base/model.py
index fb78e5e6b7..8182d35df3 100644
--- a/sdm/models/base/model.py
+++ b/sdm/models/base/model.py
@@ -102,6 +102,7 @@ def _copy_to_device(
         return tensor.to(device=device, non_blocking=True)
 
 
+# Keep eager imports working on older supported PyTorch versions.
 if hasattr(torch.compiler, "nested_compile_region"):
     _copy_to_device = torch.compiler.nested_compile_region(_copy_to_device)
 

From d4969b4cf037aae795121db8814386f1113168d0 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 2 Nov 2026 19:43:12 -0800
Subject: [PATCH 3/3] Specify PyTorch versions for nested region compatibility

---
 sdm/models/base/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdm/models/base/model.py b/sdm/models/base/model.py
index 8182d35df3..8950f7aba0 100644
--- a/sdm/models/base/model.py
+++ b/sdm/models/base/model.py
@@ -102,7 +102,7 @@ def _copy_to_device(
         return tensor.to(device=device, non_blocking=True)
 
 
-# Keep eager imports working on older supported PyTorch versions.
+# Keep eager imports working on PyTorch 2.7; this API was added in 2.8.
 if hasattr(torch.compiler, "nested_compile_region"):
     _copy_to_device = torch.compiler.nested_compile_region(_copy_to_device)