From c41c229e81369ca269e361da63e84e200b926323 Mon Sep 17 00:00:00 2001 From: sparkzky Date: Sat, 4 Jul 2026 15:13:57 +0800 Subject: [PATCH 1/2] feat(qwen3): make KV page size configurable Replace the hardcoded page_size = 16 in kv_budget_geometry() with a configurable knob threaded through Qwen3MemoryOptions, exposed via the --kv-page-size CLI flag (default 16). FlashInfer's paged attention only accepts 16 and 64, so validate() rejects anything else. Closes #545 --- openinfer-dynamo-backend/src/engine.rs | 11 +++- openinfer-kernels/third_party/flashinfer | 2 +- openinfer-qwen3/src/lib.rs | 3 +- openinfer-qwen3/src/weights.rs | 57 +++++++++++++++++-- .../tests/dflash_speculative_gate.rs | 17 ++++-- .../tests/dflash_speculative_perf.rs | 17 ++++-- openinfer-server/src/config.rs | 6 ++ openinfer-server/src/main.rs | 1 + 8 files changed, 93 insertions(+), 21 deletions(-) diff --git a/openinfer-dynamo-backend/src/engine.rs b/openinfer-dynamo-backend/src/engine.rs index 89449510..cc92a70f 100644 --- a/openinfer-dynamo-backend/src/engine.rs +++ b/openinfer-dynamo-backend/src/engine.rs @@ -29,8 +29,11 @@ use futures::stream::BoxStream; use openinfer_engine::engine::{ EngineHandle, GenerateRequest, KvBlockEvent, LoadSnapshot, TokenSink, TokenStreamReceiver, }; -use openinfer_qwen3::{ - DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, +use openinfer_qwen3_4b::{ + DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, + DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, + Qwen3OffloadOptions, +}; DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; @@ -71,6 +74,9 @@ struct Args { /// Fraction of GPU memory the engine may use (weights + KV cache). #[arg(long, default_value_t = DEFAULT_GPU_MEMORY_UTILIZATION)] gpu_memory_utilization: f64, + /// KV cache page (block) size in tokens. FlashInfer only accepts 16 or 64. + #[arg(long, default_value_t = DEFAULT_KV_PAGE_SIZE)] + kv_page_size: usize, } pub struct OpeninferBackend { @@ -112,6 +118,7 @@ impl OpeninferBackend { let memory = Qwen3MemoryOptions::new( args.gpu_memory_utilization, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, + args.kv_page_size, ) .validate() .map_err(|e| convert::invalid_argument(format!("invalid memory options: {e:#}")))?; diff --git a/openinfer-kernels/third_party/flashinfer b/openinfer-kernels/third_party/flashinfer index 57ba7eeb..d768c14e 160000 --- a/openinfer-kernels/third_party/flashinfer +++ b/openinfer-kernels/third_party/flashinfer @@ -1 +1 @@ -Subproject commit 57ba7eeb7ea3003a2d6ad5d9a057c4f952709bac +Subproject commit d768c14e7cf5dd5df45a8a1de78ae815879f108a diff --git a/openinfer-qwen3/src/lib.rs b/openinfer-qwen3/src/lib.rs index cec8010e..5ee78d37 100644 --- a/openinfer-qwen3/src/lib.rs +++ b/openinfer-qwen3/src/lib.rs @@ -28,7 +28,8 @@ use openinfer_core::engine::{EngineHandle, EngineLoadOptions, EpBackend, ModelIn pub use kernel_plan::kernel_plan; pub use scheduler::DEFAULT_MAX_PREFILL_TOKENS; pub use weights::{ - DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, Qwen3MemoryOptions, + DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, + Qwen3MemoryOptions, }; #[derive(Clone, Copy, Debug, Eq, PartialEq)] diff --git a/openinfer-qwen3/src/weights.rs b/openinfer-qwen3/src/weights.rs index 8ca85bb9..23477c89 100644 --- a/openinfer-qwen3/src/weights.rs +++ b/openinfer-qwen3/src/weights.rs @@ -24,6 +24,10 @@ use crate::batch_decode_buffers::BatchDecodeBuffers; pub const DEFAULT_GPU_MEMORY_UTILIZATION: f64 = 0.90; pub const DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES: usize = 150 * 1024 * 1024; +/// Default KV cache page (block) size in tokens. +pub const DEFAULT_KV_PAGE_SIZE: usize = 16; +/// Page sizes FlashInfer's paged attention kernels accept (see #545). +pub(crate) const VALID_KV_PAGE_SIZES: &[usize] = &[16, 64]; #[derive(Clone, Copy, Debug, PartialEq)] pub struct Qwen3MemoryOptions { @@ -34,13 +38,21 @@ pub struct Qwen3MemoryOptions { /// Extra bytes held back after the profile result to cover allocator /// fragmentation and small unprofiled runtime drift. pub kv_cache_memory_margin_bytes: usize, + /// KV cache page (block) size in tokens (`--kv-page-size`). FlashInfer + /// constrains this to [`VALID_KV_PAGE_SIZES`]; 16 by default. + pub page_size: usize, } impl Qwen3MemoryOptions { - pub const fn new(gpu_memory_utilization: f64, kv_cache_memory_margin_bytes: usize) -> Self { + pub const fn new( + gpu_memory_utilization: f64, + kv_cache_memory_margin_bytes: usize, + page_size: usize, + ) -> Self { Self { gpu_memory_utilization, kv_cache_memory_margin_bytes, + page_size, } } @@ -50,6 +62,12 @@ impl Qwen3MemoryOptions { "gpu_memory_utilization must be in (0, 1], got {}", self.gpu_memory_utilization ); + anyhow::ensure!( + VALID_KV_PAGE_SIZES.contains(&self.page_size), + "page_size must be one of {:?}, got {}", + VALID_KV_PAGE_SIZES, + self.page_size + ); Ok(self) } } @@ -59,6 +77,7 @@ impl Default for Qwen3MemoryOptions { Self { gpu_memory_utilization: DEFAULT_GPU_MEMORY_UTILIZATION, kv_cache_memory_margin_bytes: DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, + page_size: DEFAULT_KV_PAGE_SIZE, } } } @@ -879,7 +898,7 @@ impl Qwen3Model { /// KV cache geometry and budget for kernel-call tracing. #[cfg(feature = "kernel-call-trace")] pub(crate) fn kv_budget(&self) -> KvBudget { - let geometry = self.kv_budget_geometry(); + let geometry = self.kv_budget_geometry(DEFAULT_KV_PAGE_SIZE); let bytes_per_block = self.kv_bytes_per_block(&geometry); let (free_bytes, _) = cudarc::driver::result::mem_get_info().expect("cuMemGetInfo failed"); let kv_budget_bytes = (free_bytes as f64 * 0.85) as usize; @@ -901,7 +920,7 @@ impl Qwen3Model { memory_options: Qwen3MemoryOptions, ) -> Result { let memory_options = memory_options.validate()?; - let geometry = self.kv_budget_geometry(); + let geometry = self.kv_budget_geometry(memory_options.page_size); let bytes_per_block = self.kv_bytes_per_block(&geometry); let (initial_free_bytes, total_bytes) = mem_info_bytes()?; let requested_bytes = @@ -1027,8 +1046,7 @@ impl Qwen3Model { )) } - fn kv_budget_geometry(&self) -> KvBudget { - let page_size = 16; + fn kv_budget_geometry(&self, page_size: usize) -> KvBudget { let num_kv_heads = self.local_num_key_value_heads(); KvBudget { num_layers: self.config.num_hidden_layers, @@ -1208,4 +1226,33 @@ mod tests { .expect("released slot should be reused"); assert_eq!(slot_c, 0); } + + #[test] + fn memory_options_default_page_size_is_16() { + // #545: default behavior is unchanged when the flag is omitted. + assert_eq!(Qwen3MemoryOptions::default().page_size, 16); + } + + #[test] + fn memory_options_accepts_valid_page_sizes() { + for &valid in VALID_KV_PAGE_SIZES { + Qwen3MemoryOptions::new(0.9, 0, valid) + .validate() + .unwrap_or_else(|_| panic!("page_size {valid} should be valid")); + } + } + + #[test] + fn memory_options_rejects_invalid_page_sizes() { + // FlashInfer only accepts 16 and 64; anything else must fail validation. + for &invalid in &[0usize, 1, 15, 17, 32, 63, 65, 128] { + let err = Qwen3MemoryOptions::new(0.9, 0, invalid) + .validate() + .expect_err(&format!("page_size {invalid} should be rejected")); + assert!( + err.to_string().contains("page_size"), + "error should mention page_size, got: {err}" + ); + } + } } diff --git a/openinfer-qwen3/tests/dflash_speculative_gate.rs b/openinfer-qwen3/tests/dflash_speculative_gate.rs index b1d850ed..0f874263 100644 --- a/openinfer-qwen3/tests/dflash_speculative_gate.rs +++ b/openinfer-qwen3/tests/dflash_speculative_gate.rs @@ -44,9 +44,10 @@ use std::time::Duration; use openinfer_core::engine::{EngineHandle, GenerateRequest, TokenEvent, TokenSink}; use openinfer_core::sampler::SamplingParams; -use openinfer_qwen3::{ - DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, - Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, +use openinfer_qwen3_4b::{ + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, + DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, +}; }; use vllm_text::tokenizer::DynTokenizer; @@ -109,9 +110,13 @@ fn launch_options(draft: Option) -> Qwen3LaunchOptions { // baseline so both take the same cold prefill path. no_prefix_cache: true, max_prefill_tokens: DEFAULT_MAX_PREFILL_TOKENS, - memory: Qwen3MemoryOptions::new(0.85, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES) - .validate() - .expect("valid memory options"), + memory: Qwen3MemoryOptions::new( + 0.85, + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, + DEFAULT_KV_PAGE_SIZE, + ) + .validate() + .expect("valid memory options"), lora: None, decode_overlap: DecodeOverlap::Off, batch_invariant: false, diff --git a/openinfer-qwen3/tests/dflash_speculative_perf.rs b/openinfer-qwen3/tests/dflash_speculative_perf.rs index 4eaf3c16..3f544ac0 100644 --- a/openinfer-qwen3/tests/dflash_speculative_perf.rs +++ b/openinfer-qwen3/tests/dflash_speculative_perf.rs @@ -20,9 +20,10 @@ use std::time::{Duration, Instant}; use openinfer_core::engine::{EngineHandle, GenerateRequest, TokenEvent, TokenSink}; use openinfer_core::sampler::SamplingParams; -use openinfer_qwen3::{ - DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, - Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, +use openinfer_qwen3_4b::{ + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, + DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, +}; }; mod common; @@ -59,9 +60,13 @@ fn launch_options(draft: Option) -> Qwen3LaunchOptions { offload: Qwen3OffloadOptions::disabled(), no_prefix_cache: true, max_prefill_tokens: DEFAULT_MAX_PREFILL_TOKENS, - memory: Qwen3MemoryOptions::new(0.85, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES) - .validate() - .expect("valid memory options"), + memory: Qwen3MemoryOptions::new( + 0.85, + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, + DEFAULT_KV_PAGE_SIZE, + ) + .validate() + .expect("valid memory options"), lora: None, decode_overlap: DecodeOverlap::Off, batch_invariant: false, diff --git a/openinfer-server/src/config.rs b/openinfer-server/src/config.rs index 5a220f3d..489d605f 100644 --- a/openinfer-server/src/config.rs +++ b/openinfer-server/src/config.rs @@ -139,6 +139,12 @@ pub(crate) struct Args { #[cfg(feature = "qwen3")] #[arg(long, default_value_t = (openinfer_qwen3::DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES >> 20) as usize)] pub kv_cache_memory_margin_mib: usize, + /// KV cache page (block) size in tokens. FlashInfer's paged attention only + /// accepts a restricted set; 16 (default) or 64. Larger pages cut block + /// bookkeeping overhead at the cost of coarser-grained allocation. + #[cfg(feature = "qwen3-4b")] + #[arg(long, default_value_t = openinfer_qwen3_4b::DEFAULT_KV_PAGE_SIZE)] + pub kv_page_size: usize, /// How prefill and decode share the GPU (single-GPU Qwen3 only). /// `off` serializes them on one stream (lowest TTFT); `stream` overlaps on /// two streams sharing all SMs; `green-ctx` pins each to a disjoint Green diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index 7019a76b..11de1593 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -217,6 +217,7 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result Date: Tue, 7 Jul 2026 00:40:31 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(qwen3):=20address=20#554=20review=20?= =?UTF-8?q?=E2=80=94=20make=20branch=20build=20+=20fmt=20clean?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - rename openinfer_qwen3_4b:: -> openinfer_qwen3:: (engine.rs, both dflash tests, config.rs) - gate kv_page_size on qwen3 (not qwen3-4b) so it exists under the default build - fix malformed import blocks (dangling };) + cargo fmt --all - pass page_size as 3rd arg to Qwen3MemoryOptions::new at tp_concurrent_decode.rs - restore flashinfer submodule pointer to origin/main (57ba7eeb) - print page size in the 'KV cache (profiled)' load line --- openinfer-dynamo-backend/src/engine.rs | 5 +---- openinfer-kernels/third_party/flashinfer | 2 +- openinfer-qwen3/src/weights.rs | 3 ++- openinfer-qwen3/tests/dflash_speculative_gate.rs | 3 +-- openinfer-qwen3/tests/dflash_speculative_perf.rs | 3 +-- openinfer-qwen3/tests/tp_concurrent_decode.rs | 14 +++++++++----- openinfer-server/src/config.rs | 4 ++-- 7 files changed, 17 insertions(+), 17 deletions(-) diff --git a/openinfer-dynamo-backend/src/engine.rs b/openinfer-dynamo-backend/src/engine.rs index cc92a70f..4a748ce9 100644 --- a/openinfer-dynamo-backend/src/engine.rs +++ b/openinfer-dynamo-backend/src/engine.rs @@ -29,13 +29,10 @@ use futures::stream::BoxStream; use openinfer_engine::engine::{ EngineHandle, GenerateRequest, KvBlockEvent, LoadSnapshot, TokenSink, TokenStreamReceiver, }; -use openinfer_qwen3_4b::{ +use openinfer_qwen3::{ DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, -}; - DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, - Qwen3OffloadOptions, }; use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; diff --git a/openinfer-kernels/third_party/flashinfer b/openinfer-kernels/third_party/flashinfer index d768c14e..57ba7eeb 160000 --- a/openinfer-kernels/third_party/flashinfer +++ b/openinfer-kernels/third_party/flashinfer @@ -1 +1 @@ -Subproject commit d768c14e7cf5dd5df45a8a1de78ae815879f108a +Subproject commit 57ba7eeb7ea3003a2d6ad5d9a057c4f952709bac diff --git a/openinfer-qwen3/src/weights.rs b/openinfer-qwen3/src/weights.rs index 23477c89..94e0a2c2 100644 --- a/openinfer-qwen3/src/weights.rs +++ b/openinfer-qwen3/src/weights.rs @@ -1084,8 +1084,9 @@ impl Qwen3Model { bytes_per_block + dflash_kv_bytes_per_token * geometry.block_size; let num_blocks = (kv_budget_bytes / effective_bytes_per_block).max(64); let kv_mb = num_blocks * bytes_per_block / (1024 * 1024); + let page_size = geometry.block_size; log::info!( - "KV cache ({source}): {num_blocks} blocks ({kv_mb} MB, {:.0}% of {:.0} MB free)", + "KV cache ({source}): {num_blocks} blocks ({kv_mb} MB, page size {page_size}, {:.0}% of {:.0} MB free)", kv_budget_bytes as f64 / free_bytes as f64 * 100.0, free_bytes as f64 / 1024.0 / 1024.0 ); diff --git a/openinfer-qwen3/tests/dflash_speculative_gate.rs b/openinfer-qwen3/tests/dflash_speculative_gate.rs index 0f874263..969a9c64 100644 --- a/openinfer-qwen3/tests/dflash_speculative_gate.rs +++ b/openinfer-qwen3/tests/dflash_speculative_gate.rs @@ -44,11 +44,10 @@ use std::time::Duration; use openinfer_core::engine::{EngineHandle, GenerateRequest, TokenEvent, TokenSink}; use openinfer_core::sampler::SamplingParams; -use openinfer_qwen3_4b::{ +use openinfer_qwen3::{ DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; -}; use vllm_text::tokenizer::DynTokenizer; mod common; diff --git a/openinfer-qwen3/tests/dflash_speculative_perf.rs b/openinfer-qwen3/tests/dflash_speculative_perf.rs index 3f544ac0..e4a9d9e6 100644 --- a/openinfer-qwen3/tests/dflash_speculative_perf.rs +++ b/openinfer-qwen3/tests/dflash_speculative_perf.rs @@ -20,11 +20,10 @@ use std::time::{Duration, Instant}; use openinfer_core::engine::{EngineHandle, GenerateRequest, TokenEvent, TokenSink}; use openinfer_core::sampler::SamplingParams; -use openinfer_qwen3_4b::{ +use openinfer_qwen3::{ DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; -}; mod common; diff --git a/openinfer-qwen3/tests/tp_concurrent_decode.rs b/openinfer-qwen3/tests/tp_concurrent_decode.rs index b2471da1..21ffbe61 100644 --- a/openinfer-qwen3/tests/tp_concurrent_decode.rs +++ b/openinfer-qwen3/tests/tp_concurrent_decode.rs @@ -8,8 +8,8 @@ use std::time::{Duration, Instant}; use openinfer_core::engine::{GenerateRequest, TokenEvent, TokenSink}; use openinfer_core::sampler::SamplingParams; use openinfer_qwen3::{ - DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, - Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, + DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; use tokio::sync::mpsc::error::TryRecvError; @@ -56,9 +56,13 @@ fn tp2_concurrent_decode_completes() { offload: Qwen3OffloadOptions::disabled(), no_prefix_cache: false, max_prefill_tokens: DEFAULT_MAX_PREFILL_TOKENS, - memory: Qwen3MemoryOptions::new(0.85, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES) - .validate() - .expect("valid memory options"), + memory: Qwen3MemoryOptions::new( + 0.85, + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, + DEFAULT_KV_PAGE_SIZE, + ) + .validate() + .expect("valid memory options"), lora: None, decode_overlap: DecodeOverlap::Off, batch_invariant: false, diff --git a/openinfer-server/src/config.rs b/openinfer-server/src/config.rs index 489d605f..82aa63c8 100644 --- a/openinfer-server/src/config.rs +++ b/openinfer-server/src/config.rs @@ -142,8 +142,8 @@ pub(crate) struct Args { /// KV cache page (block) size in tokens. FlashInfer's paged attention only /// accepts a restricted set; 16 (default) or 64. Larger pages cut block /// bookkeeping overhead at the cost of coarser-grained allocation. - #[cfg(feature = "qwen3-4b")] - #[arg(long, default_value_t = openinfer_qwen3_4b::DEFAULT_KV_PAGE_SIZE)] + #[cfg(feature = "qwen3")] + #[arg(long, default_value_t = openinfer_qwen3::DEFAULT_KV_PAGE_SIZE)] pub kv_page_size: usize, /// How prefill and decode share the GPU (single-GPU Qwen3 only). /// `off` serializes them on one stream (lowest TTFT); `stream` overlaps on