Repository navigation
Conversation
hsliuustc0106
left a comment
There was a problem hiding this comment.
Model-specific qualification is useful, but two of the new tests do not establish their claimed contracts: the prefix-cache case cannot reuse any complete block, and the LAST_EXITED case compares identical execution paths. The existing prefix-reuse test pattern and dense Nanbeige oracle provide direct ways to strengthen these checks.
Source tracing confirms that chunked prefill reaches production chunk scheduling and that priority preemption snapshots populated prompt KV and hidden state before restoring them.
Reviewed head: 2903409; base: 4a35034; merge base: a32c89a. The remote head was rechecked before submission. Local Ruff lint, format, and diff checks passed. Existing CPU CI reports 806 passed and 356 skipped. PR runtime code was not executed locally because suitable isolation is unavailable; GPU behavior remains unvalidated.
| eng_pc = LLMEngine( | ||
| model, | ||
| attention_backend="torch", | ||
| cache_config=CacheConfig(enable_prefix_caching=True), |
There was a problem hiding this comment.
[P2] Make the prefix-cache test exercise reuse
This configuration retains the default 16-token block size, but each prompt contains only eight tokens and shares six. publish_prefix indexes complete prompt blocks while leaving the final prompt token for recomputation, so no prefix block is published or reused here. The output comparison therefore passes even if prefix lookup/reuse is completely broken. Use a smaller block size (for example, 2) or a longer shared prefix, and assert that the second request actually reuses cached tokens, through an increase in prefix_hits or its scheduled prefill start, alongside output identity.
| """Verify Nanbeige fixed 2-loop inference under LAST_EXITED KV layout.""" | ||
| prompt = [1, 5, 8, 12, 19, 24, 30] | ||
| torch.manual_seed(42) | ||
| eng_ref = LLMEngine(model, attention_backend="torch") |
There was a problem hiding this comment.
[P2] Compare LAST_EXITED against an independent oracle
LLMEngine supplies CacheConfig() when omitted, and its default layout is already last_exited. Both engines therefore run the same implementation with identical settings; a deterministic regression in depth-specific KV addressing can change both token streams equally and leave this assertion green. Derive the expected generation from the existing dense_nanbeige_reference oracle so the test validates LAST_EXITED correctness rather than repeatability.
1. Commit Message
test(nanbeige): qualify chunked prefill, priority preemption, and prefix caching
Nanbeige4.2-3B was integrated in #63 with basic fixed 2-loop inference,but its feature qualification status for advanced scheduling and KV features remained unmarked in the model progress roadmap.
This change adds comprehensive qualification and regression tests in tests/test_nanbeige.py:
- test_nanbeige_chunked_prefill_matches_unsegmented: verifies chunked prefill across
chunk sizes (2, 3, 4) produces bitwise identical tokens to unsegmented prefill.
- test_nanbeige_priority_scheduling_and_preemption: verifies priority-based preemption
and lossless recovery on Nanbeige.
- test_nanbeige_last_exited_kv_layout: confirms exact output under LAST_EXITED KV layout.
- test_nanbeige_prefix_caching_hit_and_identity: confirms that shared prompt prefixes
reuse cached KV blocks with token-identical output compared to fresh execution.
──────
2. PR Description
Summary
Qualifies
Nanbeige4.2-3Brecurrent inference across advanced scheduling and KV cachefeatures, resolving the open qualification cells in the model progress tracker:
- Scheduling: Chunked prefill (⬜ ➔ ✅) and Priority scheduling / preemption (⬜ ➔ ✅)
- KV cache: KV layouts (
last_exited) (⬜ ➔ ✅) and Prefix caching (⬜ ➔ ✅)Background
Nanbeige's basic architecture was merged in #63 and basic PD in #77. However, per the
RFC guidelines ("An existing shared utility does not by itself complete a new model's
integration. Unchanged empty cells require model-specific integration or qualification
rather than inheriting Ouro results"), Nanbeige lacked explicit test suites qualifying
these execution paths.
Changes
Added dedicated end-to-end qualification tests in
tests/test_nanbeige.py:1. Chunked Prefill:
test_nanbeige_chunked_prefill_matches_unsegmentedparameterizes chunk sizes
[2, 3, 4], ensuring segmented prefill matches unsegmentedprefill token-for-token.
2. Priority Preemption & Recovery:
test_nanbeige_priority_scheduling_and_preemptiontests preempting an in-flight low-priority sequence with a high-priority sequence, validating execution order and
uncorrupted output recovery.
3. KV Layouts:
test_nanbeige_last_exited_kv_layoutverifies Nanbeige's fixed 2-loop recurrent semantics under
CacheConfig(layout="last_exited").4. Prefix Caching:
test_nanbeige_prefix_caching_hit_and_identitytests KV prefixreuse between requests sharing prompt prefixes, verifying exact match against clean non-
cached inference.
Validation