Supply both completed block hashes in NPU prefill coverage

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-10-04 13:45:08 +08:00
1 parent 4414701161
commit bd6ab4d50e
1 file changed
+2 -1
+2 -1
View File
@@ -126,7 +126,8 @@ def test_partial_prefill_accumulates_rows_and_no_forward_step_tears_down(runner)
pending = runner._prepare_omni_aux_output()
assert runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]), [0]) == {}
step = _step(runner, hashes={"r": PackedBlockHashes(b"h0", 2)})
# Both completed four-token blocks need their scheduler-provided hashes.
step = _step(runner, hashes={"r": PackedBlockHashes(b"h0h1", 2)})
runner.input_batch.num_computed_tokens_cpu[:] = 4
suffix = _capture(runner, [[4, 5], [5, 6], [6, 7], [7, 8]])
pending = runner._prepare_omni_aux_output()