From f47c264769724a3355efaa43c189b10ddd1f9cb7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 7 Jun 2026 16:01:25 +0900 Subject: [PATCH] test: extend paged KV cache parity to llama3 (#121 sub-step c) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generalize the scheduler and prefix-share real-model parity tests to cover both pool-backed families — qwen3 and llama3 (Fp16 dense-natural backend) — via shared per-family helpers, completing sub-step (c) of #121's a/b/c decomposition. The scheduler test runs single-sequence and batched (B=2) paged-vs-dense parity per family; the prefix-share DoD test runs the radix-store adopt/share-vs-cold parity per family. All six cases produce byte-identical greedy tokens. Model-owned-state families (gemma3/llama4/qwen3.5) and quantized-KV (Int8/Turbo) paths keep the dense backend and are out of scope for pool parity; the pool-backing gate in cache.rs already documents this, and both test module docs now reinforce it. --- tests/paged_prefix_share_parity.rs | 85 ++++++++++------- tests/paged_scheduler_parity.rs | 144 +++++++++++++++++++---------- 2 files changed, 148 insertions(+), 81 deletions(-) diff --git a/tests/paged_prefix_share_parity.rs b/tests/paged_prefix_share_parity.rs index 1cdcddfd0..dff305826 100644 --- a/tests/paged_prefix_share_parity.rs +++ b/tests/paged_prefix_share_parity.rs @@ -13,12 +13,15 @@ // limitations under the License. //! Real-model parity for the unified radix prompt-prefix cache over the paged -//! block pool (#121 sub-step b). +//! block pool (#121). //! -//! Proves the headline behaviour of the sub-step: two requests sharing a prompt -//! prefix store that prefix's KV **once** in the paged pool and the second -//! request reuses it without re-prefilling — and still decodes bit-identically -//! to a cold (no-cache) paged run. +//! Proves the headline behaviour: two requests sharing a prompt prefix store +//! that prefix's KV **once** in the paged pool and the second request reuses it +//! without re-prefilling — and still decodes bit-identically to a cold +//! (no-cache) paged run. The case runs for each pool-backed family with a small +//! checkpoint — **qwen3** and **llama3** (Fp16 dense-natural backend). The +//! adopt/donate machinery operates on the pool + block table, not model +//! internals, so it is otherwise model-agnostic. //! //! ## What it drives //! @@ -39,17 +42,18 @@ //! //! ## Running //! -//! `#[ignore]` (loads qwen3-0.6b-4bit and runs real GPU forwards). Run with: +//! `#[ignore]` (loads a real checkpoint and runs real GPU forwards). Run with: //! //! ```text //! cargo test --test paged_prefix_share_parity --release \ -//! --features metal,accelerate -- --ignored --nocapture +//! --features metal,accelerate -- --ignored --nocapture --test-threads=1 //! ``` //! -//! Soft-skips when `models/qwen3-0.6b-4bit` is absent. Fetch with: +//! Each case soft-skips when its model directory is absent. Fetch with: //! //! ```text //! ./target/release/mlxcel download mlx-community/Qwen3-0.6B-4bit +//! ./target/release/mlxcel download mlx-community/Llama-3.2-1B-Instruct-4bit //! ``` mod common; @@ -64,8 +68,10 @@ use mlxcel::server::prompt_cache::{ use mlxcel::{LanguageModel, initialize_runtime, load_model}; use mlxcel_core::cache::{CachePool, PagedKvLayout, SequenceStateLayout}; -/// Model directory name (present at `models/qwen3-0.6b-4bit`). -const MODEL_DIR_NAME: &str = "qwen3-0.6b-4bit"; +/// qwen3 checkpoint directory name (pool-backed family). +const QWEN3_DIR: &str = "qwen3-0.6b-4bit"; +/// llama3 checkpoint directory name (pool-backed family). +const LLAMA3_DIR: &str = "llama-3.2-1b-4bit"; /// Paged block size (tokens per physical block) — matches the scheduler's /// `DEFAULT_PAGED_BLOCK_SIZE`. @@ -131,30 +137,27 @@ fn prefill_and_decode( decoded } -fn load_or_skip() -> Option { - let model_dir = repo_model_dir(MODEL_DIR_NAME); +fn load_or_skip(model_dir_name: &str, fetch_repo: &str) -> Option { + let model_dir = repo_model_dir(model_dir_name); if !model_dir.exists() { eprintln!( - "Skipping {MODEL_DIR_NAME}: model directory not found at {}.\n\ - Fetch with: ./target/release/mlxcel download mlx-community/Qwen3-0.6B-4bit", + "Skipping {model_dir_name}: model directory not found at {}.\n\ + Fetch with: ./target/release/mlxcel download {fetch_repo}", model_dir.display() ); return None; } - let (model, _tokenizer) = load_model(&model_dir).expect("load qwen3-0.6b-4bit"); + let (model, _tokenizer) = + load_model(&model_dir).unwrap_or_else(|e| panic!("load {model_dir_name}: {e:?}")); Some(model) } /// Sequence A stores its prefix in the radix store; sequence B adopts it via the /// paged pool, shares A's prefix blocks (no re-prefill), and decodes identically -/// to a cold no-cache run of the same `prefix + suffix` prompt. -#[test] -#[ignore = "loads qwen3-0.6b-4bit and runs real GPU forwards; run with --ignored"] -fn paged_prefix_share_matches_cold_run() { - let _runtime = initialize_runtime(); - let Some(model) = load_or_skip() else { - return; - }; +/// to a cold no-cache run of the same `prefix + suffix` prompt. `cache_model_key` +/// is the `PromptCacheKey` model field (arbitrary, but must match between A's +/// insert and B's lookup). +fn assert_prefix_share_parity(model: &mlxcel::LoadedModel, label: &str, cache_model_key: &str) { let num_layers = model.num_layers(); let prefix_len = PREFIX_TOKENS.len() as i32; @@ -162,7 +165,7 @@ fn paged_prefix_share_matches_cold_run() { full_prompt.extend_from_slice(SUFFIX_TOKENS); eprintln!( - "\n=== paged prefix-share parity: {MODEL_DIR_NAME} ({num_layers} layers, \ + "\n=== paged prefix-share parity: {label} ({num_layers} layers, \ prefix={prefix_len}, suffix={}) ===", SUFFIX_TOKENS.len() ); @@ -171,10 +174,10 @@ fn paged_prefix_share_matches_cold_run() { let cold_tokens = { let mut pool = CachePool::new(2); let id = pool - .allocate_with_layout(&model, Some(scheduler_paged_layout(num_layers))) + .allocate_with_layout(model, Some(scheduler_paged_layout(num_layers))) .expect("cold paged allocate"); let caches = pool.get_caches_mut(id).unwrap(); - prefill_and_decode(&model, caches, &full_prompt, 0) + prefill_and_decode(model, caches, &full_prompt, 0) }; eprintln!("cold decoded: {cold_tokens:?}"); @@ -190,7 +193,7 @@ fn paged_prefix_share_matches_cold_run() { // A prefills ONLY the shared prefix (positions [0, prefix_len)). let seq_a = pool - .allocate_with_layout(&model, Some(scheduler_paged_layout(num_layers))) + .allocate_with_layout(model, Some(scheduler_paged_layout(num_layers))) .expect("A paged allocate"); { let caches = pool.get_caches_mut(seq_a).unwrap(); @@ -223,7 +226,7 @@ fn paged_prefix_share_matches_cold_run() { ); let entry = CacheEntry::new(PREFIX_TOKENS.to_vec(), DetachedKvSet::Paged(set_a)); let insert_key = PromptCacheKey::new_full( - "qwen3", + cache_model_key, None, "tpl-v1", Some("sess"), @@ -234,7 +237,7 @@ fn paged_prefix_share_matches_cold_run() { // B arrives with [prefix + suffix]; the store returns A's prefix entry. let lookup_key = PromptCacheKey::new_full( - "qwen3", + cache_model_key, None, "tpl-v1", Some("sess"), @@ -256,7 +259,7 @@ fn paged_prefix_share_matches_cold_run() { DetachedKvSet::Paged(p) => p, DetachedKvSet::Dense(_) => panic!("paged backend must store a paged set"), }; - let seq_b = pool.adopt_paged(&model, set_b).expect("B adopt_paged"); + let seq_b = pool.adopt_paged(model, set_b).expect("B adopt_paged"); // B SHARES A's prefix blocks (the prefix is stored once, not re-prefilled). let b_blocks = pool @@ -284,7 +287,7 @@ fn paged_prefix_share_matches_cold_run() { caches.iter().all(|c| c.is_paged_backed()), "adopted B must have pool-backed caches" ); - prefill_and_decode(&model, caches, SUFFIX_TOKENS, prefix_len) + prefill_and_decode(model, caches, SUFFIX_TOKENS, prefix_len) }; eprintln!("cached decoded: {cached_tokens:?}"); @@ -298,3 +301,23 @@ fn paged_prefix_share_matches_cold_run() { a_blocks.len() ); } + +#[test] +#[ignore = "loads qwen3-0.6b-4bit and runs real GPU forwards; run with --ignored"] +fn paged_prefix_share_matches_cold_run_qwen3() { + let _runtime = initialize_runtime(); + let Some(model) = load_or_skip(QWEN3_DIR, "mlx-community/Qwen3-0.6B-4bit") else { + return; + }; + assert_prefix_share_parity(&model, QWEN3_DIR, "qwen3"); +} + +#[test] +#[ignore = "loads llama-3.2-1b-4bit and runs real GPU forwards; run with --ignored"] +fn paged_prefix_share_matches_cold_run_llama3() { + let _runtime = initialize_runtime(); + let Some(model) = load_or_skip(LLAMA3_DIR, "mlx-community/Llama-3.2-1B-Instruct-4bit") else { + return; + }; + assert_prefix_share_parity(&model, LLAMA3_DIR, "llama3"); +} diff --git a/tests/paged_scheduler_parity.rs b/tests/paged_scheduler_parity.rs index 07deac26f..1e1eec9b3 100644 --- a/tests/paged_scheduler_parity.rs +++ b/tests/paged_scheduler_parity.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Real-model parity for the scheduler-driven paged KV cache (#121 sub-step a). +//! Real-model parity for the scheduler-driven paged KV cache (#121). //! //! #152 proved a hand-built pool-backed `KVCache` matches a dense cache through //! `forward`. This test goes one layer up: it drives the **scheduler's** paged @@ -22,6 +22,16 @@ //! the scheduler's dense path (`CachePool::allocate` with the model's natural //! dense layout). //! +//! # Scope: the pool-backed model families +//! +//! Pool-backing is gated to **Fp16 dense-natural-backend** sequences. The two +//! such families with small checkpoints are exercised here: **qwen3** and +//! **llama3**. Model-owned-state families (gemma3, llama4, qwen3.5) and +//! quantized-KV (Int8 / Turbo) paths keep the dense backend and are therefore +//! out of scope for this pool-parity test — their decode path is unchanged by +//! #121. Each pool-backed family runs both the single-sequence and the batched +//! parity case below. +//! //! # Why `CachePool` and not `BatchScheduler` //! //! `BatchScheduler::new` is `pub(crate)`, so an integration test (a separate @@ -33,34 +43,34 @@ //! `CachePool` hands back therefore faithfully reproduces the scheduler's //! per-sequence paged path without standing up the async worker loop. //! -//! # What each test covers +//! # What each parity case covers //! -//! * [`paged_scheduler_single_sequence_matches_dense`] — the scheduler's -//! **single-sequence** path. A lone request decodes via `decode_single_step` -//! (`scheduler.rs`: `if seq_ids.len() <= 1 { decode_single_step }`), i.e. -//! single-sequence `model.forward`, whose pool intercept (`update_and_fetch` -//! / `update`) writes to and gathers from the pool. This is the path the -//! task's single-sequence parity requirement targets. -//! * [`paged_scheduler_batched_decode_matches_dense`] — the **batched** decode -//! wiring (`B == 2`) via `forward_batched_with_context_and_ids` + a paged -//! [`DecodeBatchContext`], exactly as `execute_batched_decode` dispatches it. -//! This exercises the per-model `is_paged_backed()` decode guard that routes -//! pool-backed caches through the per-sequence `update_and_fetch` loop instead -//! of the dense-pointer native kernel. +//! * [`assert_single_sequence_parity`] — the scheduler's **single-sequence** +//! path. A lone request decodes via `decode_single_step` (`scheduler.rs`: `if +//! seq_ids.len() <= 1 { decode_single_step }`), i.e. single-sequence +//! `model.forward`, whose pool intercept (`update_and_fetch` / `update`) +//! writes to and gathers from the pool. +//! * [`assert_batched_decode_parity`] — the **batched** decode wiring (`B == 2`) +//! via `forward_batched_with_context_and_ids` + a paged [`DecodeBatchContext`], +//! exactly as `execute_batched_decode` dispatches it. This exercises the +//! per-model `is_paged_backed()` decode guard that routes pool-backed caches +//! through the per-sequence `update_and_fetch` loop instead of the +//! dense-pointer native kernel. //! //! # Running //! -//! `#[ignore]` (loads qwen3-0.6b-4bit and runs a real GPU forward). Run with: +//! `#[ignore]` (loads a real checkpoint and runs a real GPU forward). Run with: //! //! ```text //! cargo test --test paged_scheduler_parity --release \ -//! --features metal,accelerate -- --ignored --nocapture +//! --features metal,accelerate -- --ignored --nocapture --test-threads=1 //! ``` //! -//! Soft-skips when `models/qwen3-0.6b-4bit` is absent. Fetch with: +//! Each case soft-skips when its model directory is absent. Fetch with: //! //! ```text //! ./target/release/mlxcel download mlx-community/Qwen3-0.6B-4bit +//! ./target/release/mlxcel download mlx-community/Llama-3.2-1B-Instruct-4bit //! ``` mod common; @@ -69,8 +79,10 @@ use common::repo_model_dir; use mlxcel::{DecodeBatchContext, LanguageModel, initialize_runtime, load_model}; use mlxcel_core::cache::{CachePool, PagedKvLayout, SequenceStateLayout}; -/// Model directory name (present at `models/qwen3-0.6b-4bit`). -const MODEL_DIR_NAME: &str = "qwen3-0.6b-4bit"; +/// qwen3 checkpoint directory name (pool-backed family). +const QWEN3_DIR: &str = "qwen3-0.6b-4bit"; +/// llama3 checkpoint directory name (pool-backed family). +const LLAMA3_DIR: &str = "llama-3.2-1b-4bit"; /// Paged block size (tokens per physical block) — matches the scheduler's /// `DEFAULT_PAGED_BLOCK_SIZE`. @@ -80,7 +92,8 @@ const BLOCK_SIZE: usize = 32; const DECODE_STEPS: usize = 16; /// Fixed prompt token ids (deterministic; no tokenizer needed). Identical bytes -/// feed every run so the comparison is purely about the cache backend. +/// feed every run so the comparison is purely about the cache backend. All ids +/// are < 128k, valid for both the qwen3 (~151k) and llama3 (128k) vocabularies. const PROMPT_TOKENS: &[i32] = &[ 9707, 11, 358, 1079, 264, 4128, 1614, 13, 5209, 3291, 752, 911, 697, 7990, 13, 358, ]; @@ -132,17 +145,20 @@ fn run_single_sequence( decoded } -fn load_or_skip() -> Option { - let model_dir = repo_model_dir(MODEL_DIR_NAME); +/// Load `model_dir_name` from the repo's `models/` directory, or return `None` +/// (and print a fetch hint) when it is absent so the case soft-skips. +fn load_or_skip(model_dir_name: &str, fetch_repo: &str) -> Option { + let model_dir = repo_model_dir(model_dir_name); if !model_dir.exists() { eprintln!( - "Skipping {MODEL_DIR_NAME}: model directory not found at {}.\n\ - Fetch with: ./target/release/mlxcel download mlx-community/Qwen3-0.6B-4bit", + "Skipping {model_dir_name}: model directory not found at {}.\n\ + Fetch with: ./target/release/mlxcel download {fetch_repo}", model_dir.display() ); return None; } - let (model, _tokenizer) = load_model(&model_dir).expect("load qwen3-0.6b-4bit"); + let (model, _tokenizer) = + load_model(&model_dir).unwrap_or_else(|e| panic!("load {model_dir_name}: {e:?}")); Some(model) } @@ -150,27 +166,21 @@ fn load_or_skip() -> Option { /// request allocates pool-backed caches (`decode_storage_backend == Paged`), /// prefills + decodes through them, and produces the same greedy tokens as the /// dense backend. -#[test] -#[ignore = "loads qwen3-0.6b-4bit and runs a real GPU forward; run with --ignored"] -fn paged_scheduler_single_sequence_matches_dense() { - let _runtime = initialize_runtime(); - let Some(model) = load_or_skip() else { - return; - }; +fn assert_single_sequence_parity(model: &mlxcel::LoadedModel, label: &str) { let num_layers = model.num_layers(); eprintln!( - "\n=== scheduler paged vs dense (single sequence): {MODEL_DIR_NAME} ({num_layers} layers) ===" + "\n=== scheduler paged vs dense (single sequence): {label} ({num_layers} layers) ===" ); // Dense backend: scheduler `decode_storage_backend == Dense` → no layout // override → `CachePool::allocate` with the model's natural dense layout. let mut dense_pool = CachePool::new(2); - let dense_id = dense_pool.allocate(&model).expect("dense allocate"); + let dense_id = dense_pool.allocate(model).expect("dense allocate"); assert!( !dense_pool.get_caches_mut(dense_id).unwrap()[0].is_paged_backed(), "dense backend must not pool-back caches" ); - let dense_tokens = run_single_sequence(&model, dense_pool.get_caches_mut(dense_id).unwrap()); + let dense_tokens = run_single_sequence(model, dense_pool.get_caches_mut(dense_id).unwrap()); eprintln!("dense decoded: {dense_tokens:?}"); // Paged backend: scheduler `decode_storage_backend == Paged` + @@ -178,7 +188,7 @@ fn paged_scheduler_single_sequence_matches_dense() { // which pool-backs the per-layer caches for this dense-natural Fp16 model. let mut paged_pool = CachePool::new(2); let paged_id = paged_pool - .allocate_with_layout(&model, Some(scheduler_paged_layout(num_layers))) + .allocate_with_layout(model, Some(scheduler_paged_layout(num_layers))) .expect("paged allocate"); assert!( paged_pool @@ -188,7 +198,7 @@ fn paged_scheduler_single_sequence_matches_dense() { .all(|c| c.is_paged_backed()), "paged backend must pool-back every layer cache for a dense-natural Fp16 model" ); - let paged_tokens = run_single_sequence(&model, paged_pool.get_caches_mut(paged_id).unwrap()); + let paged_tokens = run_single_sequence(model, paged_pool.get_caches_mut(paged_id).unwrap()); eprintln!("paged decoded: {paged_tokens:?}"); assert_eq!( @@ -206,36 +216,30 @@ fn paged_scheduler_single_sequence_matches_dense() { /// `forward_batched_with_context_and_ids` + a paged `DecodeBatchContext` — /// exactly the dispatch `execute_batched_decode` performs — and both rows must /// reproduce the dense single-sequence reference. -#[test] -#[ignore = "loads qwen3-0.6b-4bit and runs a real GPU forward; run with --ignored"] -fn paged_scheduler_batched_decode_matches_dense() { - let _runtime = initialize_runtime(); - let Some(model) = load_or_skip() else { - return; - }; +fn assert_batched_decode_parity(model: &mlxcel::LoadedModel, label: &str) { if !model.supports_paged_decode_backend() { - eprintln!("Skipping: {MODEL_DIR_NAME} does not support the paged decode backend"); + eprintln!("Skipping: {label} does not support the paged decode backend"); return; } let num_layers = model.num_layers(); let prompt_len = PROMPT_TOKENS.len() as i32; eprintln!( - "\n=== scheduler paged batched decode (B=2) vs dense: {MODEL_DIR_NAME} ({num_layers} layers) ===" + "\n=== scheduler paged batched decode (B=2) vs dense: {label} ({num_layers} layers) ===" ); // Dense single-sequence reference. let mut dense_pool = CachePool::new(4); - let dense_id = dense_pool.allocate(&model).expect("dense allocate"); - let dense_tokens = run_single_sequence(&model, dense_pool.get_caches_mut(dense_id).unwrap()); + let dense_id = dense_pool.allocate(model).expect("dense allocate"); + let dense_tokens = run_single_sequence(model, dense_pool.get_caches_mut(dense_id).unwrap()); eprintln!("dense decoded: {dense_tokens:?}"); // Two pool-backed paged sequences fed the identical prompt. let mut paged_pool = CachePool::new(4); let id0 = paged_pool - .allocate_with_layout(&model, Some(scheduler_paged_layout(num_layers))) + .allocate_with_layout(model, Some(scheduler_paged_layout(num_layers))) .expect("paged allocate 0"); let id1 = paged_pool - .allocate_with_layout(&model, Some(scheduler_paged_layout(num_layers))) + .allocate_with_layout(model, Some(scheduler_paged_layout(num_layers))) .expect("paged allocate 1"); // Prefill each sequence with the single-sequence forward (as @@ -288,3 +292,43 @@ fn paged_scheduler_batched_decode_matches_dense() { "OK: {DECODE_STEPS} batched (B=2) paged decode steps identical to the dense reference." ); } + +#[test] +#[ignore = "loads qwen3-0.6b-4bit and runs a real GPU forward; run with --ignored"] +fn paged_scheduler_single_sequence_matches_dense_qwen3() { + let _runtime = initialize_runtime(); + let Some(model) = load_or_skip(QWEN3_DIR, "mlx-community/Qwen3-0.6B-4bit") else { + return; + }; + assert_single_sequence_parity(&model, QWEN3_DIR); +} + +#[test] +#[ignore = "loads qwen3-0.6b-4bit and runs a real GPU forward; run with --ignored"] +fn paged_scheduler_batched_decode_matches_dense_qwen3() { + let _runtime = initialize_runtime(); + let Some(model) = load_or_skip(QWEN3_DIR, "mlx-community/Qwen3-0.6B-4bit") else { + return; + }; + assert_batched_decode_parity(&model, QWEN3_DIR); +} + +#[test] +#[ignore = "loads llama-3.2-1b-4bit and runs a real GPU forward; run with --ignored"] +fn paged_scheduler_single_sequence_matches_dense_llama3() { + let _runtime = initialize_runtime(); + let Some(model) = load_or_skip(LLAMA3_DIR, "mlx-community/Llama-3.2-1B-Instruct-4bit") else { + return; + }; + assert_single_sequence_parity(&model, LLAMA3_DIR); +} + +#[test] +#[ignore = "loads llama-3.2-1b-4bit and runs a real GPU forward; run with --ignored"] +fn paged_scheduler_batched_decode_matches_dense_llama3() { + let _runtime = initialize_runtime(); + let Some(model) = load_or_skip(LLAMA3_DIR, "mlx-community/Llama-3.2-1B-Instruct-4bit") else { + return; + }; + assert_batched_decode_parity(&model, LLAMA3_DIR); +}