fix offline batch chunked prefill

[vegetabledoww]

Summary

• apply an effective token budget to offline LLMEngine.generate_batch prefill
• preserve one normal batched prefill dispatch when the complete input fits the budget
• process over-budget long prompts one request at a time in round-robin order with a 256-token safety threshold
• preserve absolute positions, cumulative sequence lengths, KV allocations, and final per-request logits across chunks
• enable device greedy sampling only on final chunks
• keep Qwen warmup within the same effective budget, including when the budget is smaller than the configured batch
• support both host embedding tensors and the device-embedding path where input_embeddings is None

Root cause

The offline batch path built one PrefillBatch for all prompts and submitted it through a single run_prefill call. max_num_batched_tokens constrained warmup but did not constrain the real offline prefill dispatch.

With 16 prompts of 3,338 tokens, the old path submitted 53,408 packed tokens at once and failed with AICPU error 507018.

Validation

• pre-commit run --all-files
  • check-headers passed
  • check-english-only passed
  • ruff check passed
• focused batching regression: 23 passed
• broader CPU regression covering batching, DeepSeek, package, parallel, and serving worker paths: 64 passed
• compileall and git diff --check passed
• NPU A/B on the original issue commit matrix:
  • fixed matrix-aligned worktree completed 224 chunked prefill calls and all decode steps, exit 0
  • baseline dispatched 53,408 tokens in one prefill call and reproduced 507018
• Exact workload rerun on the current rebased PR (2bca273, pypto-lib@306519dc) still reproduces 507018 before decode:
  • NPU 3: failed on prefill chunk 6
  • NPU 2: completed 41 chunks and failed on chunk 42 at cumulative sequence length 768

Refs #91

[coderabbitai]

📝 Walkthrough

Walkthrough

Prefill execution now derives an effective token budget from runtime settings, uses chunked dispatch when prompts exceed that budget, and adds a configurable long-prefill threshold. Qwen warmup and batching tests apply and validate the same budgeting behavior. The pypto-lib reference was updated.

Changes

Prefill Budgeting

Layer / File(s)	Summary
Runtime budget contract and warmup alignment python/core/types.py, examples/model/qwen3_14b/runner/npu_runner.py, pypto-lib	RuntimeConfig adds long_prefill_token_threshold; Qwen warmup incorporates it into effective token budgeting and startup logging; the pypto-lib reference is updated.
Budgeted engine prefill flow python/core/engine.py	Batch generation uses the effective prefill budget for fast-path selection or chunked executor prefill, collecting final logits and sampled token IDs across chunks.
Batching behavior validation tests/test_batching.py	New tokenizer, model parameters, recording executor, and batching tests verify constrained prefill positions and request rotation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• hw-native-sys/pypto-serving#84: Both changes modify prefill input handling in python/core/engine.py.

Poem

I’m a bunny who chunks through the queue,
With token-sized carrots in view.
Budgets now guide every hop,
Warmups know when they should stop.
Logs and tests keep paths bright—
Prefill bounds are set just right!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly matches the main change: offline batch chunked prefill.
Description check	✅ Passed	The description directly describes the batching prefill changes, root cause, and validation.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces a chunked prefill mechanism to handle offline batches within an effective token budget, adding a long_prefill_token_threshold configuration parameter and updating the NPU runner warmup logic and unit tests accordingly. Feedback on the changes suggests optimizing the chunked prefill implementation to co-schedule chunks from multiple requests for better NPU utilization, avoiding redundant greedy sampling on intermediate chunks, and simplifying the tensor stacking logic.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

tests/test_batching.py (1)

59-61: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert the prompt slices too. The batching checks only track positions, so a chunk from the wrong request can still pass. Add token_ids assertions for the recorded PrefillBatch rows in both batching tests.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/test_batching.py` around lines 59 - 61, Update both batching tests to
assert the token_ids values on each recorded PrefillBatch row, not just
positional fields. Use the expected encoded prompt slices from
_CharacterTokenizer so chunks from an incorrect request cannot pass, while
preserving the existing position assertions.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/model/qwen3_14b/runner/npu_runner.py`:
- Around line 397-405: The token-budget calculation around effective_budget,
step_tokens, and per_req must not dispatch more than the configured budget when
0 < effective_budget < batch. Reject this configuration or reduce the active
rows during warmup so total_tokens never exceeds effective_budget, and add a
regression test covering batch=4 with a threshold of 2.

In `@pypto-lib`:
- Line 1: Update the pypto-lib submodule pointer from the unreachable revision
4ee8cf44e4af815ba8cc15c0486ad8edfe73b1af to a commit that is published and
fetchable from the submodule’s remote repository.

In `@python/core/engine.py`:
- Around line 231-245: Update the prefill flow around _run_prefill_in_chunks so
an in-budget batch continues through batched prefill when try_generate_batch
returns None, rather than being reduced to one row. Use chunking only when the
total prompt token count exceeds prefill_token_budget, and apply the same
behavior at the corresponding later flow.

---

Nitpick comments:
In `@tests/test_batching.py`:
- Around line 59-61: Update both batching tests to assert the token_ids values
on each recorded PrefillBatch row, not just positional fields. Use the expected
encoded prompt slices from _CharacterTokenizer so chunks from an incorrect
request cannot pass, while preserving the existing position assertions.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: bb45bfb6-ca4f-4e79-bc88-de3ad6e2c904

📥 Commits

Reviewing files that changed from the base of the PR and between a77f42c and 2ffea2a.

📒 Files selected for processing (5)

• examples/model/qwen3_14b/runner/npu_runner.py
• pypto-lib
• python/core/engine.py
• python/core/types.py
• tests/test_batching.py

[vegetabledoww]

Review follow-up completed locally:

• preserved one batched prefill call when the full input fits the effective budget
• limited device greedy sampling to final chunks
• kept warmup dispatch within the effective budget when the threshold is smaller than the configured batch
• strengthened both chunking tests to reconstruct and assert token ID slices per request, in addition to positions and sequence lengths

Validation:

• pre-commit run --all-files
• relevant CPU regression suite: 63 passed, 1 deselected
• compileall and git diff --check passed

The one deselected test reads the unavailable local checkout of the pypto-lib commit referenced by upstream main.

[vegetabledoww]

Correction to my previous validation note: both pypto-lib@4ee8cf4 and the current main pointer 306519dc are fetchable from hw-native-sys/pypto-lib. The earlier failure came from this worktree's submodule origin being overridden to a local issue-matrix mirror that did not contain 306519dc, not from the public remote. After checking out the main pointer, the previously excluded test passes: batching is now 23 passed and the broader CPU regression is 64 passed.