From 6d75a8aad4c31502638c39db9f82a05f6aae98b5 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:59:31 +0900 Subject: [PATCH] accel.platforms: chunk NPU graph rows to the largest token bucket (#171) _bucket returns T itself above the largest bucket, so each new prefill length compiled and cached another graph (one per expert in baked mode). run() now splits longer inputs into chunks of the largest bucket, which reuse the bucket graphs. Co-Authored-By: Claude Opus 5.5 --- .../opt_prime/accel/platforms/npu_graph.py | 5 +++- opt_prime/tests/test_moe_review.py | 23 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/opt_prime/opt_prime/accel/platforms/npu_graph.py b/opt_prime/opt_prime/accel/platforms/npu_graph.py index 69653bda..8492e3a9 100644 --- a/opt_prime/opt_prime/accel/platforms/npu_graph.py +++ b/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) diff --git a/opt_prime/tests/test_moe_review.py b/opt_prime/tests/test_moe_review.py index da86b590..a107a9fa 100644 --- a/opt_prime/tests/test_moe_review.py +++ b/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()