From f6b6bdac3648af85a345d53554c700dad431f08c Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 16:01:58 -0700
Subject: [PATCH 1/4] Avoid symbolic batch iteration in single-chunk attention

---
 sdm/nn/attention.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index 3d0275119..e5410d3ca 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -171,7 +171,13 @@ def forward(
         key = key.transpose(-3, -2)  # [B, Hkv, KV, C]
         value = value.transpose(-3, -2)  # [B, Hkv, KV, C]
 
-        # CUDA attention kernels accept at most 65,535 batch entries per call:
+        # CUDA attention kernels accept at most 65,535 batch entries per call.
+        # A fixed loop length also lets torch.compile trace bounded batches.
+        chunk_starts = (
+            (0,)
+            if query.size(0) <= 65_535
+            else range(0, query.size(0), 65_535)
+        )
         outs = [
             F.scaled_dot_product_attention(
                 query=query[i : i + 65_535],
@@ -185,7 +191,7 @@ def forward(
                 is_causal=is_causal,
                 scale=self.scale,
             ).transpose(-3, -2)  # [B, Q, Hq, C]
-            for i in range(0, query.size(0), 65_535)
+            for i in chunk_starts
         ]
         out = outs[0] if len(outs) == 1 else torch.cat(outs)
 

From cfef9e024f04f8ce1ccd81816e24c858de4db85c Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 20:21:04 -0700
Subject: [PATCH 2/4] Inline attention chunk iteration and clarify compilation
 comment

---
 sdm/nn/attention.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index 948c3e736..69444c27e 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -172,12 +172,8 @@ def forward(
         value = value.transpose(-3, -2)  # [B, Hkv, KV, C]
 
         # CUDA attention kernels accept at most 65,535 batch entries per call.
-        # A fixed loop length also lets torch.compile trace bounded batches.
-        chunk_starts = (
-            (0,)
-            if query.size(0) <= 65_535
-            else range(0, query.size(0), 65_535)
-        )
+        # For one attention call, (0,) avoids constructing range() from a
+        # batch size that torch.compile may not know.
         outs = [
             F.scaled_dot_product_attention(
                 query=query[i : i + 65_535],
@@ -191,7 +187,11 @@ def forward(
                 is_causal=is_causal,
                 scale=self.scale,
             ).transpose(-3, -2)  # [B, Q, Hq, C]
-            for i in chunk_starts
+            for i in (
+                (0,)
+                if query.size(0) <= 65_535
+                else range(0, query.size(0), 65_535)
+            )
         ]
         out = outs[0] if len(outs) == 1 else torch.cat(outs)
 

From 85f220b7a7f6bb3b13713df0fca461104a8c47bf Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 20:23:38 -0700
Subject: [PATCH 3/4] Shorten compilation comment and place it beside attention
 loop

---
 sdm/nn/attention.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index 69444c27e..6dafcecba 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -172,8 +172,6 @@ def forward(
         value = value.transpose(-3, -2)  # [B, Hkv, KV, C]
 
         # CUDA attention kernels accept at most 65,535 batch entries per call.
-        # For one attention call, (0,) avoids constructing range() from a
-        # batch size that torch.compile may not know.
         outs = [
             F.scaled_dot_product_attention(
                 query=query[i : i + 65_535],
@@ -188,6 +186,7 @@ def forward(
                 scale=self.scale,
             ).transpose(-3, -2)  # [B, Q, Hq, C]
             for i in (
+                # Avoid range() needing the batch size during compilation.
                 (0,)
                 if query.size(0) <= 65_535
                 else range(0, query.size(0), 65_535)

From 765c509f368902baf243c6a14815ecf28beff484 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 22 Oct 2026 20:24:37 -0700
Subject: [PATCH 4/4] Apply suggestion from @aw471

---
 sdm/nn/attention.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index 6dafcecba..e8a380392 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -186,7 +186,6 @@ def forward(
                 scale=self.scale,
             ).transpose(-3, -2)  # [B, Q, Hq, C]
             for i in (
-                # Avoid range() needing the batch size during compilation.
                 (0,)
                 if query.size(0) <= 65_535
                 else range(0, query.size(0), 65_535)