feat(qwen3): make KV page size configurable

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:#}")))?;

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)]

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<KvBudget> {
         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}"
+            );
+        }
+    }
 }

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<PathBuf>) -> 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,

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<PathBuf>) -> 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,

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,

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

openinfer-server/src/main.rs

@@ -217,6 +217,7 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result<EngineHandl
                     memory: openinfer_qwen3::Qwen3MemoryOptions::new(
                         args.gpu_memory_utilization,
                         kv_cache_memory_margin_bytes,
+                        args.kv_page_size,
                     )
                     .validate()?,
                     lora,