From 8512b8115193c4a54172b2f0517f8347412bb625 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 10:52:50 -0700
Subject: [PATCH 1/4] Compile compact ICL cache copies as nested regions

---
 sdm/models/kumo/tabular/icl.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py
index 067d5f6191..f472fa5838 100644
--- a/sdm/models/kumo/tabular/icl.py
+++ b/sdm/models/kumo/tabular/icl.py
@@ -98,9 +98,12 @@ def forward(
 
                 if cache is not None and cache.is_recording:
                     x, (key, value) = result
-                    if self.kv_heads is not None:
-                        key = key[..., : self.kv_heads, :].contiguous()
-                        value = value[..., : self.kv_heads, :].contiguous()
+                    kv_heads = self.kv_heads
+                    if kv_heads is not None and kv_heads < key.size(-2):
+                        key, value = _copy_key_value(
+                            key[..., :kv_heads, :],
+                            value[..., :kv_heads, :],
+                        )
                     cache[cache_key] = KVCacheEntry(key, value)
                     del key, value
                 else:
@@ -134,3 +137,15 @@ def forward(
                 x = torch.cat([x_context, x_query], -2)
 
         return self.head(self.norm(x))
+
+
+def _copy_key_value(key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]:
+    return (
+        key.clone(memory_format=torch.contiguous_format),
+        value.clone(memory_format=torch.contiguous_format),
+    )
+
+
+if hasattr(torch.compiler, "nested_compile_region"):
+    # Compile cache copies separately to free each layer's full K/V tensors.
+    _copy_key_value = torch.compiler.nested_compile_region(_copy_key_value)

From 1940095ac4158984999a08a25ee6cea8cff618e8 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 15:20:52 -0700
Subject: [PATCH 2/4] Pass cache-copy tensors by keyword

---
 sdm/models/kumo/tabular/icl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py
index f472fa5838..0f8567dfdd 100644
--- a/sdm/models/kumo/tabular/icl.py
+++ b/sdm/models/kumo/tabular/icl.py
@@ -101,8 +101,8 @@ def forward(
                     kv_heads = self.kv_heads
                     if kv_heads is not None and kv_heads < key.size(-2):
                         key, value = _copy_key_value(
-                            key[..., :kv_heads, :],
-                            value[..., :kv_heads, :],
+                            key=key[..., :kv_heads, :],
+                            value=value[..., :kv_heads, :],
                         )
                     cache[cache_key] = KVCacheEntry(key, value)
                     del key, value

From 3f307cd8fd7a39aa861fdc8ad1f2239a8b25188c Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Fri, 23 Oct 2026 07:49:00 -0700
Subject: [PATCH 3/4] Explain nested compilation compatibility check

---
 sdm/models/kumo/tabular/icl.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py
index 0f8567dfdd..2602901bb3 100644
--- a/sdm/models/kumo/tabular/icl.py
+++ b/sdm/models/kumo/tabular/icl.py
@@ -146,6 +146,7 @@ def _copy_key_value(key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]:
     )
 
 
+# Available since PyTorch 2.8; keep the plain function for PyTorch 2.7.
 if hasattr(torch.compiler, "nested_compile_region"):
     # Compile cache copies separately to free each layer's full K/V tensors.
     _copy_key_value = torch.compiler.nested_compile_region(_copy_key_value)

From d4bd43e3e1950bbed129cfcd9542011575bfc5e3 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Fri, 23 Oct 2026 07:57:03 -0700
Subject: [PATCH 4/4] Clarify why nested compilation needs a compatibility
 guard

---
 sdm/models/kumo/tabular/icl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py
index 2602901bb3..efb75ff3f0 100644
--- a/sdm/models/kumo/tabular/icl.py
+++ b/sdm/models/kumo/tabular/icl.py
@@ -146,7 +146,7 @@ def _copy_key_value(key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]:
     )
 
 
-# Available since PyTorch 2.8; keep the plain function for PyTorch 2.7.
+# Keep the guard for PyTorch 2.7; 2.8+ can use the decorator directly.
 if hasattr(torch.compiler, "nested_compile_region"):
     # Compile cache copies separately to free each layer's full K/V tensors.
     _copy_key_value = torch.compiler.nested_compile_region(_copy_key_value)