Fix AttnRes silently dropping blocks past 32

src/liger_kernel/ops/attn_res.py

@@ -211,11 +211,14 @@ def _next_pow2(n):
 
 
 def _get_max_blocks(n_blocks):
-    """Round up to constexpr-friendly value."""
-    for mb in [4, 8, 16, 32]:
-        if n_blocks <= mb:
-            return mb
-    return 32
+    """Round n_blocks up to a power-of-two constexpr.
+
+    MAX_BLOCKS is the static trip count of the `tl.static_range(0, MAX_BLOCKS)`
+    loops and the width of the register-held `scores`/`alpha` vectors, so it must
+    be >= n_blocks or the trailing blocks are silently skipped. It also has to be
+    a power of two for `tl.arange(0, MAX_BLOCKS)`.
+    """
+    return max(4, _next_pow2(n_blocks))
 
 
 def attn_res_forward(blocks, w_query, w_norm, eps=1e-6):

test/transformers/test_attn_res.py

@@ -43,6 +43,8 @@ def pytorch_attn_res(V, w_query, w_norm, eps=1e-6):
         (4, 2, 64, 4096),
         (8, 2, 64, 4096),
         (4, 2, 64, 8192),
+        # more blocks than the old MAX_BLOCKS ceiling of 32
+        (40, 1, 32, 128),
         # weird shapes
         (3, 5, 37, 123),
     ],
@@ -79,6 +81,8 @@ def test_correctness(N, B, T, D, dtype, atol, rtol):
     [
         (4, 2, 64, 4096),
         (8, 2, 64, 4096),
+        # more blocks than the old MAX_BLOCKS ceiling of 32
+        (40, 1, 32, 128),
         # weird shapes
         (3, 5, 37, 123),
     ],