diff --git a/openinfer-dynamo-backend/src/engine.rs b/openinfer-dynamo-backend/src/engine.rs index 89449510..4a748ce9 100644 --- a/openinfer-dynamo-backend/src/engine.rs +++ b/openinfer-dynamo-backend/src/engine.rs @@ -30,7 +30,7 @@ use openinfer_engine::engine::{ EngineHandle, GenerateRequest, KvBlockEvent, LoadSnapshot, TokenSink, TokenStreamReceiver, }; use openinfer_qwen3::{ - DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, + DEFAULT_GPU_MEMORY_UTILIZATION, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; @@ -71,6 +71,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 +115,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-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..94e0a2c2 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, @@ -1066,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 ); @@ -1208,4 +1227,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..969a9c64 100644 --- a/openinfer-qwen3/tests/dflash_speculative_gate.rs +++ b/openinfer-qwen3/tests/dflash_speculative_gate.rs @@ -45,8 +45,8 @@ 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, + 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 +109,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..e4a9d9e6 100644 --- a/openinfer-qwen3/tests/dflash_speculative_perf.rs +++ b/openinfer-qwen3/tests/dflash_speculative_perf.rs @@ -21,8 +21,8 @@ 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, + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_KV_PAGE_SIZE, DEFAULT_MAX_PREFILL_TOKENS, + DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; mod common; @@ -59,9 +59,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-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 5a220f3d..82aa63c8 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")] + #[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 /// 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