[Dynamic shapes 5/n] Skip automatic chunking in compiled code

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)

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=(),

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,

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"),