Skip to content

test(nanbeige): qualify chunked prefill, priority preemption, and prefix caching - #138

Open
cs-ai-li wants to merge 1 commit into
InferMatrixLab:mainfrom
cs-ai-li:feat/nanbeige-adaptation
Open

cs-ai-li wants to merge 1 commit into
InferMatrixLab:mainfrom
cs-ai-li:feat/nanbeige-adaptation

Conversation

@cs-ai-li

Copy link
Copy Markdown
Contributor

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-3B recurrent inference across advanced scheduling and KV cache
features, 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_unsegmented
parameterizes chunk sizes [2, 3, 4], ensuring segmented prefill matches unsegmented
prefill token-for-token.
2. Priority Preemption & Recovery:
test_nanbeige_priority_scheduling_and_preemption tests 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_layout verifies Nanbeige's fixed 2-
loop recurrent semantics under CacheConfig(layout="last_exited").
4. Prefix Caching: test_nanbeige_prefix_caching_hit_and_identity tests KV prefix
reuse between requests sharing prompt prefixes, verifying exact match against clean non-
cached inference.

Validation

- `pytest tests/test_nanbeige.py -k "not parity"`: 17 passed in 2.44s
- Code style: clean on `ruff check` and `ruff format`.
- Zero changes to production code (pure qualification & regression suite).

@cs-ai-li cs-ai-li changed the title test(nanbeige): qualify chunked prefill, priority preemption, and pre… test(nanbeige): qualify chunked prefill, priority preemption, and prefix caching Oct 22, 2026
@cs-ai-li
cs-ai-li marked this pull request as ready for review October 22, 2026 11:43

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread tests/test_nanbeige.py
eng_pc = LLMEngine(
model,
attention_backend="torch",
cache_config=CacheConfig(enable_prefix_caching=True),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Comment thread tests/test_nanbeige.py
"""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")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants