fix: chunk NPU graph rows to the largest token bucket (#171)

opt_prime/opt_prime/accel/platforms/npu_graph.py

@@ -246,7 +246,10 @@ def _slot_graph(self, spec, T_b, K_b, with_bias, host_sample=None):
 
     def run(self, h: Handle, hidden, local_index, weights):
         spec = h.spec
-        T = hidden.shape[0]
+        T, cap = hidden.shape[0], self.token_buckets[-1]
+        if T > cap:                                # longer prefills run in chunks of the largest bucket (no graph per T)
+            return torch.cat([self.run(h, hidden[i:i + cap], local_index[i:i + cap], weights[i:i + cap])
+                              for i in range(0, T, cap)])
         T_b = _bucket(T, self.token_buckets)
         x = torch.zeros(T_b, spec.hidden, dtype=spec.dtype)
         x[:T] = hidden.to(spec.dtype)

opt_prime/tests/test_moe_review.py

@@ -67,6 +67,29 @@ def weights(bias):
     torch.testing.assert_close(b, a + 3)
 
 
+@pytest.mark.parametrize('residency', ['inputs', 'baked'])
+def test_npu_graph_chunks_rows_above_the_largest_bucket(residency):
+    class Counting(GraphNPUBackend):
+        def __init__(self):
+            super().__init__(emulate=True, token_buckets=(1, 4))
+            self.residency, self.keys = residency, set()
+        def _cached_compile(self, key, module, samples, place=True):
+            self.keys.add(key)
+            return module
+    backend, whole = Counting(), GraphNPUBackend(emulate=True, token_buckets=(64,))   # whole: one unchunked graph
+    whole.residency = residency
+    spec = ExpertSpec(2, 2, 2, dtype=torch.float32)
+    w = ExpertWeights(spec, [0, 1], torch.randn(2, 4, 2), torch.randn(2, 2, 2))
+    h, ref = backend.load(w), whole.load(w)
+    loaded = len(backend.keys)
+    torch.manual_seed(0)
+    for T in (5, 6, 7, 9):                    # prefill lengths above the largest bucket (4)
+        args = (torch.randn(T, 2), torch.randint(0, 2, (T, 2)), torch.rand(T, 2))
+        torch.testing.assert_close(backend.run(h, *args), whole.run(ref, *args))
+    assert {k[3 if residency == 'baked' else 1] for k in backend.keys} <= {1, 4}
+    assert len(backend.keys) - loaded <= (2 if residency == 'baked' else 4)   # bucket graphs only, not one per T
+
+
 def test_adaptation_waits_for_a_forward_using_old_mapping():
     m = Experts()
     adapter = GenericAdapter()