From 0da9c57154ce8e7b85b6841c36653d7f3f3851e5 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 9 Nov 2026 15:03:55 -0800
Subject: [PATCH] Skip balanced split sizing when one chunk fits

With a dynamic item count, the balanced-chunk arithmetic adds guards. When the count is statically known to fit, return it directly.
---
 sdm/_memory.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/sdm/_memory.py b/sdm/_memory.py
index ae55f6da4..240409866 100644
--- a/sdm/_memory.py
+++ b/sdm/_memory.py
@@ -69,5 +69,8 @@ def split_size(num_items: int, item_bytes: int, device: torch.device) -> int:
         return max(num_items, 1)
     limit = max(chunk_memory_limit(device), 1)
     capacity = max(limit // item_bytes, 1)
+    # Avoid guards on the balanced-chunk arithmetic when one chunk fits.
+    if statically_known_true(num_items <= capacity):
+        return max(num_items, 1)
     num_chunks = max(-(-num_items // capacity), 1)
     return max(-(-num_items // num_chunks), 1)