diff --git a/tile_kernels/engram/engram_grad_w_reduce_kernel.py b/tile_kernels/engram/engram_grad_w_reduce_kernel.py
index 5d13a29..85b7a9b 100644
--- a/tile_kernels/engram/engram_grad_w_reduce_kernel.py
+++ b/tile_kernels/engram/engram_grad_w_reduce_kernel.py
@@ -21,8 +21,9 @@ def get_engram_grad_w_reduce_kernel(
     blk_d = 512
     assert hidden_size % blk_d == 0
     num_tiles = hidden_size // blk_d
-    num_batches = 4
-    assert num_persistent_blocks % num_batches == 0
+    # SM90/SM100 devices can have SM counts such as 114, so avoid assuming
+    # that the persistent-block count is always divisible by four.
+    num_batches = max(batch for batch in range(4, 0, -1) if num_persistent_blocks % batch == 0)
     num_rows = num_persistent_blocks // num_batches
 
     @T.prim_func