From 731c92f7adc0f8e727daec3eab3d30732d85dfca Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 12 Nov 2026 17:55:41 -0800
Subject: [PATCH] Skip automatic chunking in compiled code

A compiled graph records the number of chunks of the inputs it was traced
with, so it is only valid for inputs that need the same number. Compiled
code now processes its inputs in one chunk wherever the chunk size is
derived from the memory budget: `TransformerBlock` without an explicit
`batch_size_limit`, `CellEmbedding` with `batch_size_limit="auto"` and
`split_size`. Explicit limits are honoured as before, and eager code chunks
as before.

The check is `torch.compiler.is_dynamo_compiling()`: PyTorch 2.14 sets
`torch.compiler.is_compiling()` for the whole process while any thread
compiles, which would switch chunking off for eager calls in other threads.
---
 sdm/_memory.py                     | 12 +++++++++---
 sdm/models/tabfm/cell_embedding.py |  3 ++-
 sdm/nn/attention.py                |  7 ++++++-
 test/nn/test_compile.py            | 16 ++++++++++++++++
 4 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/sdm/_memory.py b/sdm/_memory.py
index ae55f6da4..b78329f22 100644
--- a/sdm/_memory.py
+++ b/sdm/_memory.py
@@ -62,10 +62,16 @@ def split_size(num_items: int, item_bytes: int, device: torch.device) -> int:
     r"""Split size of balanced chunks of items of ``item_bytes`` bytes each.
 
     On CUDA devices, chunks fit :func:`chunk_memory_limit`, unless a single
-    item exceeds the limit. Autograd keeps the memory of every chunk, so one
-    chunk holds all items elsewhere or while gradients are enabled.
+    item exceeds the limit. Autograd keeps the memory of every chunk, and a
+    compiled graph records the number of chunks, so one chunk holds all items
+    elsewhere, while gradients are enabled or in compiled code.
     """
-    if device.type != "cuda" or torch.is_grad_enabled() or item_bytes == 0:
+    if (
+        device.type != "cuda"
+        or torch.is_grad_enabled()
+        or torch.compiler.is_dynamo_compiling()
+        or item_bytes == 0
+    ):
         return max(num_items, 1)
     limit = max(chunk_memory_limit(device), 1)
     capacity = max(limit // item_bytes, 1)
diff --git a/sdm/models/tabfm/cell_embedding.py b/sdm/models/tabfm/cell_embedding.py
index babe20c5d..e38bf47d7 100644
--- a/sdm/models/tabfm/cell_embedding.py
+++ b/sdm/models/tabfm/cell_embedding.py
@@ -135,7 +135,8 @@ def forward(
 
         if batch_size_limit == "auto":
             batch_size_limit = None
-            if x.is_cuda:
+            # A compiled graph records the number of chunks.
+            if x.is_cuda and not torch.compiler.is_dynamo_compiling():
                 if torch.is_autocast_enabled(x.device.type):
                     element_size = torch.empty(
                         size=(),
diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index a7e53d84a..ec542a975 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -726,7 +726,12 @@ def forward(
                 "already cached"
             )
 
-        if torch.is_grad_enabled():
+        # Autograd keeps the memory of every chunk. A compiled graph records
+        # the number of chunks, so it only chunks by an explicit limit.
+        if torch.is_grad_enabled() or (
+            not isinstance(batch_size_limit, int)
+            and torch.compiler.is_dynamo_compiling()
+        ):
             return self._forward(
                 query=query,
                 key_value=key_value,
diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py
index 8f2f2b2ee..0d70acf70 100644
--- a/test/nn/test_compile.py
+++ b/test/nn/test_compile.py
@@ -219,6 +219,22 @@ def test_sdpa_compile_dynamic_batch_autocast(dtype: torch.dtype) -> None:
     torch.testing.assert_close(out, expected)
 
 
+def test_transformer_block_compile_dynamic_batch() -> None:
+    module = TransformerBlock(
+        channels=8, num_query_heads=2, mlp=torch.nn.Identity()
+    ).eval()
+    counter = CompileCounter()
+    compiled = torch.compile(module, fullgraph=True, backend=counter)
+    with torch.inference_mode():
+        for batch_size in [3, 65_536]:
+            query = torch.randn(batch_size, 4, 8)
+            torch._dynamo.mark_dynamic(query, 0)
+            expected = module(query=query, batch_size_limit="auto")
+            out = compiled(query=query, batch_size_limit="auto")
+            torch.testing.assert_close(out, expected)
+    assert counter.frame_count == 1
+
+
 @withCUDA
 @pytest.mark.parametrize(
     ("num_key_value_heads", "qassmax", "self_attn"),