[gspark-etri]

Background

GraphNPUBackend.run pads the token count up to token_buckets, but _bucket returns T itself above the largest bucket (256 on RNGD, 64 on ATOM). Prefill rows reached run() unchunked, so every new prompt length compiled another graph (one per expert in baked residency) and _graphs grew without bound. With a stub backend, prefill lengths 300-303 gave 4 compiles in inputs mode and 32 in baked mode.

Changes

• opt_prime/opt_prime/accel/platforms/npu_graph.py:249-252: run() splits inputs longer than the largest token bucket into chunks of that bucket and concatenates the chunk outputs, so long prefills reuse the bucket graphs.
• opt_prime/tests/test_moe_review.py: test_npu_graph_chunks_rows_above_the_largest_bucket[inputs|baked] uses buckets (1, 4), feeds T = 5, 6, 7, 9, checks the result against an unchunked backend (bucket 64) and that only bucket-sized graphs are requested.

Verification

CPU only (CUDA_VISIBLE_DEVICES=1, no NPU), Python 3.10, torch 2.10.0+cu128, base optimus-rm @253be20.

New test on the unfixed code (from opt_prime/):

TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests/test_moe_review.py -k chunks
E  assert {5, 6, 7, 9} <= {1, 4}        (inputs)
E  assert {1, 5, 6, 7, 9} <= {1, 4}     (baked)
2 failed

(outputs already matched; only the compiled shapes were wrong.) After the fix: 2 passed.

Stub-backend repro from the issue (Fake _compile counter, prefill lengths 300-303, default buckets): inputs mode 4 -> 2 compiles (buckets 256 and 64, once each); baked mode 32 -> 16 (8 experts x the same two buckets), and no further compiles for other lengths.

Whole suite:

TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests
before (253be20): 400 passed, 23 skipped, 135 subtests passed
after:            402 passed, 23 skipped, 135 subtests passed

Compatibility and limits

• Outputs are unchanged (rows are independent); prefills above the largest bucket now run several device calls instead of one large one. Not measured on RNGD/ATOM hardware.
• RngdOptBackend.run (rngd/experts_opt.py) is a separate code path and is not changed.

Related Issue

Closes #171

🤖 Generated with Claude Code