diff --git a/Cargo.lock b/Cargo.lock index 9e73e68a4..6ccbdbdf9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3805,6 +3805,7 @@ dependencies = [ "serde_json", "sha2 0.11.0", "tempfile", + "thiserror 2.0.18", "tokio", "tokio-util", "vllm-text", diff --git a/pegainfer-qwen35/Cargo.toml b/pegainfer-qwen35/Cargo.toml index 562627764..f308934bf 100644 --- a/pegainfer-qwen35/Cargo.toml +++ b/pegainfer-qwen35/Cargo.toml @@ -19,6 +19,7 @@ rand = { workspace = true } safetensors = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +thiserror = { workspace = true } tokio = { workspace = true, features = ["sync"] } [dev-dependencies] diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index a778d3781..8d1c373c6 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -102,9 +102,9 @@ impl Qwen35Model { bufs: &mut BatchDecodeBuffers35, ) -> Result<()> { let eps = self.config.rms_norm_eps; - let tp = self.tensor_parallel; - let num_attention_heads = self.config.local_num_attention_heads(tp); - let num_key_value_heads = self.config.local_num_key_value_heads(tp); + let geom = self.geometry; + let num_attention_heads = geom.local_num_attention_heads(); + let num_key_value_heads = geom.local_num_key_value_heads(); ops::gemm_into(&self.ctx, &attn.q_proj, &bufs.normed, &mut bufs.q_full); ops::gemm_into(&self.ctx, &attn.k_proj, &bufs.normed, &mut bufs.k_attn); diff --git a/pegainfer-qwen35/src/batch_decode_graph.rs b/pegainfer-qwen35/src/batch_decode_graph.rs index 27bcbdcab..eb88ba691 100644 --- a/pegainfer-qwen35/src/batch_decode_graph.rs +++ b/pegainfer-qwen35/src/batch_decode_graph.rs @@ -6,7 +6,7 @@ use pegainfer_core::kv_pool::KvPool; use pegainfer_core::tensor::DeviceContext; use super::config::Config35; -use super::config::TensorParallelConfig; +use super::config::LocalGeometry; use super::decode_buffers::BatchDecodeBuffers35; use super::recurrent_state::LinearStatePointerTables; use super::recurrent_state::RecurrentState; @@ -61,7 +61,7 @@ impl BatchDecodeGraphState { pub(crate) fn with_capacity( ctx: &DeviceContext, config: &Config35, - tensor_parallel: TensorParallelConfig, + geometry: LocalGeometry, kv_pool: &KvPool, max_batch: usize, ) -> Result { @@ -71,7 +71,7 @@ impl BatchDecodeGraphState { let buffers = BatchDecodeBuffers35::new( ctx, config, - tensor_parallel, + geometry, max_batch, max_total_pages, padding_page_id, diff --git a/pegainfer-qwen35/src/config.rs b/pegainfer-qwen35/src/config.rs deleted file mode 100644 index ee40a7d60..000000000 --- a/pegainfer-qwen35/src/config.rs +++ /dev/null @@ -1,683 +0,0 @@ -use std::collections::HashSet; -use std::fs; - -use anyhow::Result; -use anyhow::bail; -use anyhow::ensure; -use log::warn; -use serde::Deserialize; -use serde_json::Value; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) struct TensorParallelConfig { - pub(crate) rank: usize, - pub(crate) world_size: usize, -} - -impl Default for TensorParallelConfig { - fn default() -> Self { - Self { - rank: 0, - world_size: 1, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum LayerType { - FullAttention, - LinearAttention, -} - -#[derive(Debug, Deserialize)] -struct RopeParameters { - rope_theta: f64, - partial_rotary_factor: f64, -} - -#[derive(Debug, Deserialize)] -struct TextConfig { - hidden_size: usize, - intermediate_size: usize, - num_hidden_layers: usize, - num_attention_heads: usize, - num_key_value_heads: usize, - head_dim: usize, - vocab_size: usize, - rms_norm_eps: f64, - layer_types: Vec, - linear_conv_kernel_dim: usize, - linear_key_head_dim: usize, - linear_num_key_heads: usize, - linear_num_value_heads: usize, - linear_value_head_dim: usize, - rope_parameters: RopeParameters, - max_position_embeddings: Option, - tie_word_embeddings: Option, - eos_token_id: u32, -} - -#[derive(Debug, Deserialize)] -struct RawConfig { - text_config: TextConfig, - max_position_embeddings: Option, - tie_word_embeddings: Option, -} - -/// Qwen3.5 model configuration (text-only). -#[derive(Debug)] -pub(crate) struct Config35 { - // Common - pub(crate) hidden_size: usize, - pub(crate) intermediate_size: usize, - pub(crate) num_hidden_layers: usize, - pub(crate) vocab_size: usize, - pub(crate) rms_norm_eps: f32, - pub(crate) eos_token_id: u32, - - // Full attention params - pub(crate) num_attention_heads: usize, - pub(crate) num_key_value_heads: usize, - pub(crate) head_dim: usize, - - // Linear attention params - pub(crate) linear_num_key_heads: usize, - pub(crate) linear_key_head_dim: usize, - pub(crate) linear_num_value_heads: usize, - pub(crate) linear_value_head_dim: usize, - pub(crate) linear_conv_kernel_dim: usize, - - // RoPE - pub(crate) rope_theta: f32, - pub(crate) rotary_dim: usize, - pub(crate) max_position_embeddings: usize, - - // Layer layout - pub(crate) layer_types: Vec, - - /// `false` requires a top-level `lm_head.weight`; `true` reuses `embed_tokens`. - pub(crate) tie_word_embeddings: bool, - - /// Token-selection width: `vocab_size` bounded to the frontend-decodable vocab. - pub(crate) selection_vocab: usize, -} - -/// Head dims baked into the kernels; head counts are runtime parameters. -pub(crate) const GDN_AOT_KEY_HEAD_DIM: usize = 128; -pub(crate) const GDN_AOT_VALUE_HEAD_DIM: usize = 128; -pub(crate) const LINEAR_CONV_MAX_KERNEL_DIM: usize = 4; -const FULL_ATTN_HEAD_DIM: usize = 256; - -impl Config35 { - pub(crate) fn from_file(model_path: &str) -> Result { - let config_path = format!("{}/config.json", model_path); - let content = fs::read_to_string(&config_path)?; - let raw: RawConfig = serde_json::from_str(&content)?; - let root_max_position_embeddings = raw.max_position_embeddings; - let root_tie_word_embeddings = raw.tie_word_embeddings; - let t = raw.text_config; - - let tie_word_embeddings = t - .tie_word_embeddings - .or(root_tie_word_embeddings) - .ok_or_else(|| anyhow::anyhow!("Qwen3.5 config missing tie_word_embeddings"))?; - - let layer_types: Vec = t - .layer_types - .iter() - .map(|s| match s.as_str() { - "full_attention" => Ok(LayerType::FullAttention), - "linear_attention" => Ok(LayerType::LinearAttention), - other => Err(anyhow::anyhow!("Unknown layer type: {}", other)), - }) - .collect::>()?; - - anyhow::ensure!( - layer_types.len() == t.num_hidden_layers, - "layer_types length {} != num_hidden_layers {}", - layer_types.len(), - t.num_hidden_layers - ); - - let rotary_dim = (t.head_dim as f64 * t.rope_parameters.partial_rotary_factor) as usize; - anyhow::ensure!(rotary_dim > 0, "Qwen3.5 rotary_dim must be positive"); - let max_position_embeddings = t - .max_position_embeddings - .or(root_max_position_embeddings) - .ok_or_else(|| anyhow::anyhow!("Qwen3.5 config missing max_position_embeddings"))?; - anyhow::ensure!( - max_position_embeddings > 0, - "Qwen3.5 max_position_embeddings must be positive" - ); - - anyhow::ensure!( - t.linear_key_head_dim == GDN_AOT_KEY_HEAD_DIM - && t.linear_value_head_dim == GDN_AOT_VALUE_HEAD_DIM, - "Qwen3.5 GDN Triton-AOT kernels are baked for key/value head dim {}/{}; \ - config has {}/{} (dims are baked into the AOT signatures in pegainfer-kernels/build.rs).", - GDN_AOT_KEY_HEAD_DIM, - GDN_AOT_VALUE_HEAD_DIM, - t.linear_key_head_dim, - t.linear_value_head_dim, - ); - anyhow::ensure!( - t.head_dim == FULL_ATTN_HEAD_DIM, - "Qwen3.5 full-attention kernels are baked for head_dim {}; config has {}.", - FULL_ATTN_HEAD_DIM, - t.head_dim, - ); - anyhow::ensure!( - (1..=LINEAR_CONV_MAX_KERNEL_DIM).contains(&t.linear_conv_kernel_dim), - "Qwen3.5 linear conv decode kernels support kernel_dim in 1..={}; config has {}.", - LINEAR_CONV_MAX_KERNEL_DIM, - t.linear_conv_kernel_dim, - ); - anyhow::ensure!( - t.linear_num_key_heads > 0 - && t.linear_num_value_heads - .is_multiple_of(t.linear_num_key_heads), - "Qwen3.5 GDN kernels require linear_num_value_heads ({}) divisible by \ - linear_num_key_heads ({})", - t.linear_num_value_heads, - t.linear_num_key_heads, - ); - anyhow::ensure!( - t.num_key_value_heads > 0 - && t.num_attention_heads.is_multiple_of(t.num_key_value_heads), - "Qwen3.5 num_attention_heads ({}) must be a positive multiple of \ - num_key_value_heads ({})", - t.num_attention_heads, - t.num_key_value_heads, - ); - - let config = Self { - hidden_size: t.hidden_size, - intermediate_size: t.intermediate_size, - num_hidden_layers: t.num_hidden_layers, - vocab_size: t.vocab_size, - rms_norm_eps: t.rms_norm_eps as f32, - eos_token_id: t.eos_token_id, - num_attention_heads: t.num_attention_heads, - num_key_value_heads: t.num_key_value_heads, - head_dim: t.head_dim, - linear_num_key_heads: t.linear_num_key_heads, - linear_key_head_dim: t.linear_key_head_dim, - linear_num_value_heads: t.linear_num_value_heads, - linear_value_head_dim: t.linear_value_head_dim, - linear_conv_kernel_dim: t.linear_conv_kernel_dim, - rope_theta: t.rope_parameters.rope_theta as f32, - rotary_dim, - max_position_embeddings, - layer_types, - tie_word_embeddings, - selection_vocab: t.vocab_size, - }; - Ok(config) - } - - /// Number of full attention layers in the model. - pub(crate) fn num_full_attention_layers(&self) -> usize { - self.layer_types - .iter() - .filter(|&&t| t == LayerType::FullAttention) - .count() - } - - /// Q dimension for full attention (without gate). - pub(crate) fn full_attn_q_dim(&self) -> usize { - self.num_attention_heads * self.head_dim - } - - /// KV dimension for full attention. - pub(crate) fn full_attn_kv_dim(&self) -> usize { - self.num_key_value_heads * self.head_dim - } - - pub(crate) fn decode_group_is_compiled(&self) -> bool { - // Uncompiled GQA groups use the batched hybrid eager fallback. - pegainfer_core::ops::SUPPORTED_GQA_GROUP_SIZES - .contains(&(self.num_attention_heads / self.num_key_value_heads)) - } - - /// QKV projection output dimension for linear attention. - pub(crate) fn linear_attn_qkv_dim(&self) -> usize { - let q_dim = self.linear_num_key_heads * self.linear_key_head_dim; - let k_dim = q_dim; - let v_dim = self.linear_num_value_heads * self.linear_value_head_dim; - q_dim + k_dim + v_dim - } - - /// Z projection output dimension for linear attention. - pub(crate) fn linear_attn_z_dim(&self) -> usize { - self.linear_num_value_heads * self.linear_value_head_dim - } - - pub(crate) fn local_num_attention_heads(&self, tp: TensorParallelConfig) -> usize { - self.num_attention_heads / tp.world_size - } - - pub(crate) fn local_num_key_value_heads(&self, tp: TensorParallelConfig) -> usize { - self.num_key_value_heads / tp.world_size - } - - pub(crate) fn local_intermediate_size(&self, tp: TensorParallelConfig) -> usize { - self.intermediate_size / tp.world_size - } - - pub(crate) fn local_full_attn_q_dim(&self, tp: TensorParallelConfig) -> usize { - self.local_num_attention_heads(tp) * self.head_dim - } - - pub(crate) fn local_full_attn_kv_dim(&self, tp: TensorParallelConfig) -> usize { - self.local_num_key_value_heads(tp) * self.head_dim - } - - /// Local gated full-attention q projection output dimension. - pub(crate) fn local_full_attn_gated_q_dim(&self, tp: TensorParallelConfig) -> usize { - self.local_full_attn_q_dim(tp) * 2 - } -} - -impl TensorParallelConfig { - pub(crate) fn validate_for(self, config: &Config35, enable_cuda_graph: bool) -> Result<()> { - if self.world_size == 0 { - return Err(anyhow::anyhow!("tensor_parallel.world_size must be >= 1")); - } - if self.rank >= self.world_size { - return Err(anyhow::anyhow!( - "tensor_parallel.rank {} must be < world_size {}", - self.rank, - self.world_size - )); - } - if self.is_sharded() && enable_cuda_graph { - return Err(anyhow::anyhow!( - "Qwen3.5 tensor parallelism is eager-only in Phase 1; disable CUDA Graph for tp world_size={}", - self.world_size - )); - } - if !config.num_attention_heads.is_multiple_of(self.world_size) { - return Err(anyhow::anyhow!( - "num_attention_heads={} not divisible by tp world_size={}", - config.num_attention_heads, - self.world_size - )); - } - if !config.num_key_value_heads.is_multiple_of(self.world_size) { - return Err(anyhow::anyhow!( - "num_key_value_heads={} not divisible by tp world_size={}", - config.num_key_value_heads, - self.world_size - )); - } - if !config.intermediate_size.is_multiple_of(self.world_size) { - return Err(anyhow::anyhow!( - "intermediate_size={} not divisible by tp world_size={}", - config.intermediate_size, - self.world_size - )); - } - Ok(()) - } - - pub(crate) fn shard_range(self, total: usize) -> (usize, usize) { - let shard_len = total / self.world_size; - (self.rank * shard_len, shard_len) - } - - pub(crate) fn is_sharded(self) -> bool { - self.world_size > 1 - } -} - -#[cfg(test)] -mod tp_tests { - use super::*; - - fn test_config() -> Config35 { - Config35 { - hidden_size: 2560, - intermediate_size: 9216, - num_hidden_layers: 32, - vocab_size: 248_320, - selection_vocab: 248_320, - rms_norm_eps: 1e-6, - eos_token_id: 151_645, - num_attention_heads: 16, - num_key_value_heads: 4, - head_dim: 256, - linear_num_key_heads: 16, - linear_key_head_dim: 128, - linear_num_value_heads: 32, - linear_value_head_dim: 128, - linear_conv_kernel_dim: 4, - rope_theta: 10_000.0, - rotary_dim: 64, - max_position_embeddings: 262_144, - tie_word_embeddings: true, - layer_types: vec![LayerType::LinearAttention; 32], - } - } - - #[test] - fn default_tensor_parallel_is_tp1() { - let config = test_config(); - let tp = TensorParallelConfig::default(); - - tp.validate_for(&config, true).unwrap(); - assert!(!tp.is_sharded()); - assert_eq!(tp.shard_range(config.full_attn_q_dim()), (0, 4096)); - assert_eq!(config.local_num_attention_heads(tp), 16); - assert_eq!(config.local_num_key_value_heads(tp), 4); - assert_eq!(config.local_intermediate_size(tp), 9216); - assert_eq!(config.local_full_attn_q_dim(tp), 4096); - assert_eq!(config.local_full_attn_kv_dim(tp), 1024); - assert_eq!(config.local_full_attn_gated_q_dim(tp), 8192); - } - - #[test] - fn computes_tp2_dense_local_dimensions() { - let config = test_config(); - let tp = TensorParallelConfig { - rank: 1, - world_size: 2, - }; - - tp.validate_for(&config, false).unwrap(); - assert!(tp.is_sharded()); - assert_eq!(tp.shard_range(config.full_attn_q_dim()), (2048, 2048)); - assert_eq!(config.local_num_attention_heads(tp), 8); - assert_eq!(config.local_num_key_value_heads(tp), 2); - assert_eq!(config.local_intermediate_size(tp), 4608); - assert_eq!(config.local_full_attn_q_dim(tp), 2048); - assert_eq!(config.local_full_attn_kv_dim(tp), 512); - assert_eq!(config.local_full_attn_gated_q_dim(tp), 4096); - } - - #[test] - fn rejects_invalid_world_size_and_rank() { - let config = test_config(); - - let err = TensorParallelConfig { - rank: 0, - world_size: 0, - } - .validate_for(&config, false) - .unwrap_err() - .to_string(); - assert!(err.contains("world_size must be >= 1")); - - let err = TensorParallelConfig { - rank: 2, - world_size: 2, - } - .validate_for(&config, false) - .unwrap_err() - .to_string(); - assert!(err.contains("rank 2 must be < world_size 2")); - } - - #[test] - fn rejects_indivisible_dense_dimensions() { - let tp = TensorParallelConfig { - rank: 0, - world_size: 3, - }; - - let mut config = test_config(); - let err = tp.validate_for(&config, false).unwrap_err().to_string(); - assert!(err.contains("num_attention_heads=16 not divisible")); - - config.num_attention_heads = 15; - config.num_key_value_heads = 4; - let err = tp.validate_for(&config, false).unwrap_err().to_string(); - assert!(err.contains("num_key_value_heads=4 not divisible")); - - config.num_key_value_heads = 3; - config.intermediate_size = 9217; - let err = tp.validate_for(&config, false).unwrap_err().to_string(); - assert!(err.contains("intermediate_size=9217 not divisible")); - } - - #[test] - fn rejects_tensor_parallel_cuda_graph_phase1() { - let config = test_config(); - let tp = TensorParallelConfig { - rank: 0, - world_size: 2, - }; - - let err = tp.validate_for(&config, true).unwrap_err().to_string(); - assert!(err.contains("eager-only in Phase 1")); - } - - #[test] - fn phase1_does_not_require_linear_attention_divisibility() { - let mut config = test_config(); - config.linear_num_key_heads = 17; - config.linear_num_value_heads = 31; - let tp = TensorParallelConfig { - rank: 1, - world_size: 2, - }; - - tp.validate_for(&config, false).unwrap(); - } -} - -/// Schema kept identical to the pinned vLLM frontend; unread fields exist for -/// payload type-checking. -#[allow(dead_code)] -// The tokenizer_config schema is bool-heavy by design. -#[allow(clippy::struct_excessive_bools)] -#[derive(Deserialize)] -struct AddedTokenConfig { - #[serde(default)] - id: Option, - content: String, - #[serde(default)] - single_word: bool, - #[serde(default)] - lstrip: bool, - #[serde(default)] - rstrip: bool, - #[serde(default)] - normalized: bool, - #[serde(default)] - special: bool, -} - -#[derive(Deserialize)] -struct TokenizerJsonIds { - model: TokenizerModelIds, - #[serde(default)] - added_tokens: Vec, -} - -#[derive(Deserialize)] -struct TokenizerModelIds { - vocab: std::collections::HashMap, -} - -#[derive(Deserialize)] -struct TokenizerConfigIds { - #[serde(default)] - added_tokens_decoder: std::collections::HashMap, -} - -/// Width of the frontend-decodable id space, mirroring the pinned frontend's -/// merge: `tokenizer.json` vocab and added_tokens (fatal on parse failure - -/// the frontend cannot serve without it) plus `tokenizer_config.json` -/// added_tokens_decoder (whole-file typed parse; failure drops all decoder -/// tokens with a warning, unparsable keys are skipped per entry). The ids -/// must form a dense prefix - a row-range selection bound cannot mask holes - -/// so a sparse id space fails the load instead of silently truncating the -/// output space. -pub(crate) fn tokenizer_effective_vocab(model_path: &str) -> Result { - let path = format!("{}/tokenizer.json", model_path); - let content = - fs::read_to_string(&path).map_err(|e| anyhow::anyhow!("cannot read {path}: {e}"))?; - let tj: TokenizerJsonIds = - serde_json::from_str(&content).map_err(|e| anyhow::anyhow!("cannot parse {path}: {e}"))?; - anyhow::ensure!(!tj.model.vocab.is_empty(), "{path} model.vocab is empty"); - let mut ids: HashSet = tj.model.vocab.into_values().collect(); - ids.extend(tj.added_tokens.iter().filter_map(|t| t.id)); - - let config_path = format!("{}/tokenizer_config.json", model_path); - if let Ok(text) = fs::read_to_string(&config_path) { - match serde_json::from_str::(&text) { - Ok(cfg) => ids.extend( - cfg.added_tokens_decoder - .keys() - .filter_map(|k| k.parse::().ok()), - ), - Err(e) => warn!( - "cannot parse {config_path}: {e}; skipping its added tokens like the frontend does" - ), - } - } - - let width = ids.len(); - let max_id = *ids.iter().max().expect("vocab checked non-empty") as usize; - anyhow::ensure!( - max_id + 1 == width, - "tokenizer id space is not dense (max id {max_id}, {width} distinct ids); \ - a row-range selection bound cannot mask holes" - ); - Ok(width) -} - -/// Identity check that `json` is a Qwen3.5 config; size and shape validation belong to the config loader. -pub(crate) fn probe_config_json(json: &Value) -> Result<()> { - let model_type = json.get("model_type").and_then(Value::as_str).unwrap_or(""); - if model_type != "qwen3_5" { - bail!("not a Qwen3.5 config: model_type={model_type}"); - } - let architectures: Vec<&str> = json - .get("architectures") - .and_then(Value::as_array) - .map(|arr| arr.iter().filter_map(Value::as_str).collect()) - .unwrap_or_default(); - ensure!( - architectures.contains(&"Qwen3_5ForConditionalGeneration"), - "Qwen3.5 architectures must contain Qwen3_5ForConditionalGeneration" - ); - let text_model_type = json - .get("text_config") - .and_then(|tc| tc.get("model_type")) - .and_then(Value::as_str) - .unwrap_or(""); - ensure!( - text_model_type == "qwen3_5_text", - "Qwen3.5 text_config.model_type must be qwen3_5_text, got {text_model_type}" - ); - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::Config35; - - #[test] - fn guard_accepts_48_value_heads() { - let dir = tempfile::tempdir().unwrap(); - let json = r#"{ - "max_position_embeddings": 4096, - "tie_word_embeddings": true, - "text_config": { - "hidden_size": 512, - "intermediate_size": 1024, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 256, - "vocab_size": 1000, - "rms_norm_eps": 1e-6, - "layer_types": ["linear_attention", "full_attention"], - "linear_conv_kernel_dim": 4, - "linear_key_head_dim": 128, - "linear_num_key_heads": 16, - "linear_num_value_heads": 48, - "linear_value_head_dim": 128, - "rope_parameters": { "rope_theta": 10000.0, "partial_rotary_factor": 0.25 }, - "eos_token_id": 0 - } -}"#; - std::fs::write(dir.path().join("config.json"), json).unwrap(); - Config35::from_file(dir.path().to_str().unwrap()).expect("48 value heads must load"); - } - - #[test] - fn guard_rejects_wide_linear_conv_decode_kernel() { - let dir = tempfile::tempdir().unwrap(); - let json = r#"{ - "max_position_embeddings": 4096, - "tie_word_embeddings": true, - "text_config": { - "hidden_size": 512, - "intermediate_size": 1024, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 256, - "vocab_size": 1000, - "rms_norm_eps": 1e-6, - "layer_types": ["linear_attention", "full_attention"], - "linear_conv_kernel_dim": 5, - "linear_key_head_dim": 128, - "linear_num_key_heads": 16, - "linear_num_value_heads": 48, - "linear_value_head_dim": 128, - "rope_parameters": { "rope_theta": 10000.0, "partial_rotary_factor": 0.25 }, - "eos_token_id": 0 - } -}"#; - std::fs::write(dir.path().join("config.json"), json).unwrap(); - - let err = Config35::from_file(dir.path().to_str().unwrap()) - .expect_err("wide conv decode kernels must be rejected"); - assert!( - err.to_string().contains("linear conv decode kernels"), - "unexpected error: {err}" - ); - } - - #[test] - fn effective_vocab_is_the_dense_decodable_width() { - let dir = tempfile::tempdir().unwrap(); - let json = r#"{ - "model": { "vocab": { "a": 0, "b": 1, "c": 2 } }, - "added_tokens": [ { "id": 3, "content": "" } ] -}"#; - std::fs::write(dir.path().join("tokenizer.json"), json).unwrap(); - let cfg = r#"{ "added_tokens_decoder": { "4": { "content": "" }, "5": { "content": "" }, "x": { "content": "" } } }"#; - std::fs::write(dir.path().join("tokenizer_config.json"), cfg).unwrap(); - assert_eq!( - super::tokenizer_effective_vocab(dir.path().to_str().unwrap()).unwrap(), - 6 - ); - } - - #[test] - fn effective_vocab_fails_on_a_sparse_id_space() { - let dir = tempfile::tempdir().unwrap(); - let json = r#"{ "model": { "vocab": { "a": 0, "b": 1 } } }"#; - std::fs::write(dir.path().join("tokenizer.json"), json).unwrap(); - let cfg = r#"{ "added_tokens_decoder": { "5": { "content": "" } } }"#; - std::fs::write(dir.path().join("tokenizer_config.json"), cfg).unwrap(); - assert!(super::tokenizer_effective_vocab(dir.path().to_str().unwrap()).is_err()); - } - - #[test] - fn one_invalid_decoder_entry_drops_all_decoder_tokens() { - let dir = tempfile::tempdir().unwrap(); - let json = r#"{ "model": { "vocab": { "a": 0, "b": 1 } } }"#; - std::fs::write(dir.path().join("tokenizer.json"), json).unwrap(); - let cfg = r#"{ "added_tokens_decoder": { "2": { "content": "" }, "3": { "content": "", "special": "not-a-bool" } } }"#; - std::fs::write(dir.path().join("tokenizer_config.json"), cfg).unwrap(); - assert_eq!( - super::tokenizer_effective_vocab(dir.path().to_str().unwrap()).unwrap(), - 2 - ); - } -} diff --git a/pegainfer-qwen35/src/config/error.rs b/pegainfer-qwen35/src/config/error.rs new file mode 100644 index 000000000..b44a5b996 --- /dev/null +++ b/pegainfer-qwen35/src/config/error.rs @@ -0,0 +1,81 @@ +//! Typed configuration-load errors for the Qwen3.5 model line. +//! +//! Every failure mode here is a *static* invariant violation that is detected +//! before any weights are mapped or any CUDA memory is allocated. Keeping these +//! as variants (instead of free-form `anyhow!` strings) lets callers and tests +//! branch on the exact violation and lets the loader reject an invalid model +//! file fail-closed with a precise message. + +/// A Qwen3.5 config that could not be turned into a validated `Config35`. +/// +/// These cover the model config (`RawConfig -> Config35`), the tensor-parallel +/// assignment and derived local geometry (`TensorParallelConfig`/`LocalGeometry`), +/// and the cross-field + kernel/TP compatibility rules that cut across both. +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub(crate) enum ConfigError { + // ---- model config (cross-field + kernel AOT invariants) ---- + #[error("Qwen3.5 config missing tie_word_embeddings")] + MissingTieWordEmbeddings, + #[error("Qwen3.5 config missing max_position_embeddings")] + MissingMaxPositionEmbeddings, + #[error("unknown layer type: {0}")] + UnknownLayerType(String), + #[error("layer_types length {actual} != num_hidden_layers {expected}")] + LayerTypeCountMismatch { actual: usize, expected: usize }, + #[error("Qwen3.5 rotary_dim must be positive")] + ZeroRotaryDim, + #[error("Qwen3.5 max_position_embeddings must be positive")] + NonPositiveMaxPositionEmbeddings, + #[error( + "Qwen3.5 GDN Triton-AOT kernels are baked for key/value head dim {expected_key}/{expected_value}; \ + config has {key}/{value} (dims are baked into the AOT signatures in pegainfer-kernels/build.rs)" + )] + GdnAotHeadDimMismatch { + key: usize, + value: usize, + expected_key: usize, + expected_value: usize, + }, + #[error( + "Qwen3.5 full-attention kernels are baked for head_dim {expected}; config has {actual}" + )] + FullAttnHeadDimMismatch { expected: usize, actual: usize }, + #[error( + "Qwen3.5 linear conv decode kernels support kernel_dim in 1..={max}; config has {actual}" + )] + LinearConvKernelDim { max: usize, actual: usize }, + #[error( + "Qwen3.5 GDN kernels require linear_num_value_heads ({value_heads}) divisible by \ + linear_num_key_heads ({key_heads})" + )] + LinearHeadDivisibility { + key_heads: usize, + value_heads: usize, + }, + #[error( + "Qwen3.5 num_attention_heads ({attention_heads}) must be a positive multiple of \ + num_key_value_heads ({key_value_heads})" + )] + AttentionHeadDivisibility { + attention_heads: usize, + key_value_heads: usize, + }, + #[error("tokenizer defines ids up to {used} but checkpoint vocab_size is {vocab_size}")] + EffectiveVocabExceedsCheckpoint { used: usize, vocab_size: usize }, + + // ---- tensor-parallel assignment invariants ---- + #[error("tensor_parallel.world_size must be >= 1")] + TpZeroWorldSize, + #[error("tensor_parallel.rank {rank} must be < world_size {world_size}")] + TpRankOutOfRange { rank: usize, world_size: usize }, + #[error( + "Qwen3.5 tensor parallelism is eager-only; disable CUDA Graph for tp world_size={world_size}" + )] + TpRequiresEager { world_size: usize }, + #[error("{field}={value} not divisible by tp world_size={world_size}")] + TpIndivisible { + field: &'static str, + value: usize, + world_size: usize, + }, +} diff --git a/pegainfer-qwen35/src/config/mod.rs b/pegainfer-qwen35/src/config/mod.rs new file mode 100644 index 000000000..fa4bf208a --- /dev/null +++ b/pegainfer-qwen35/src/config/mod.rs @@ -0,0 +1,104 @@ +//! Qwen3.5 configuration: validated model geometry, tensor-parallel local +//! geometry, and the tokenizer schema. +//! +//! Ownership model (directional boundaries): +//! - [`model`] deserializes raw config and validates into a TP-agnostic +//! [`Config35`]; +//! - [`tp`] depends on [`model`] and produces the validated +//! [`LocalGeometry`] that downstream code accepts; +//! - [`tokenizer`] owns the frontend tokenizer schema and decodable-vocab width; +//! - [`error`] owns the typed [`ConfigError`] variants. + +use anyhow::Result; + +mod error; +mod model; +mod tokenizer; +mod tp; + +pub(crate) use model::Config35; +pub(crate) use model::GDN_AOT_KEY_HEAD_DIM; +pub(crate) use model::GDN_AOT_VALUE_HEAD_DIM; +pub(crate) use model::LINEAR_CONV_MAX_KERNEL_DIM; +pub(crate) use model::LayerType; +pub(crate) use tokenizer::tokenizer_effective_vocab; +pub(crate) use tp::LocalGeometry; +pub(crate) use tp::TensorParallelConfig; + +/// Identity check that `json` is a Qwen3.5 config; size and shape validation +/// belong to the config loader. +pub(crate) fn probe_config_json(json: &serde_json::Value) -> Result<()> { + let model_type = json + .get("model_type") + .and_then(serde_json::Value::as_str) + .unwrap_or(""); + if model_type != "qwen3_5" { + anyhow::bail!("not a Qwen3.5 config: model_type={model_type}"); + } + let architectures: Vec<&str> = json + .get("architectures") + .and_then(serde_json::Value::as_array) + .map(|arr| arr.iter().filter_map(serde_json::Value::as_str).collect()) + .unwrap_or_default(); + anyhow::ensure!( + architectures.contains(&"Qwen3_5ForConditionalGeneration"), + "Qwen3.5 architectures must contain Qwen3_5ForConditionalGeneration" + ); + let text_model_type = json + .get("text_config") + .and_then(|tc| tc.get("model_type")) + .and_then(serde_json::Value::as_str) + .unwrap_or(""); + anyhow::ensure!( + text_model_type == "qwen3_5_text", + "Qwen3.5 text_config.model_type must be qwen3_5_text, got {text_model_type}" + ); + Ok(()) +} + +#[cfg(test)] +mod tokenizer_tests { + use std::fs; + + use super::tokenizer_effective_vocab; + + fn dir_with(json: &str, config: &str) -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + fs::write(dir.path().join("tokenizer.json"), json).unwrap(); + fs::write(dir.path().join("tokenizer_config.json"), config).unwrap(); + dir + } + + #[test] + fn effective_vocab_is_the_dense_decodable_width() { + let dir = dir_with( + r#"{ "model": { "vocab": { "a": 0, "b": 1, "c": 2 } }, "added_tokens": [ { "id": 3, "content": "" } ] }"#, + r#"{ "added_tokens_decoder": { "4": { "content": "" }, "5": { "content": "" }, "x": { "content": "" } } }"#, + ); + assert_eq!( + tokenizer_effective_vocab(dir.path().to_str().unwrap()).unwrap(), + 6 + ); + } + + #[test] + fn effective_vocab_fails_on_a_sparse_id_space() { + let dir = dir_with( + r#"{ "model": { "vocab": { "a": 0, "b": 1 } } }"#, + r#"{ "added_tokens_decoder": { "5": { "content": "" } } }"#, + ); + assert!(tokenizer_effective_vocab(dir.path().to_str().unwrap()).is_err()); + } + + #[test] + fn one_invalid_decoder_entry_drops_all_decoder_tokens() { + let dir = dir_with( + r#"{ "model": { "vocab": { "a": 0, "b": 1 } } }"#, + r#"{ "added_tokens_decoder": { "2": { "content": "" }, "3": { "content": "", "special": "not-a-bool" } } }"#, + ); + assert_eq!( + tokenizer_effective_vocab(dir.path().to_str().unwrap()).unwrap(), + 2 + ); + } +} diff --git a/pegainfer-qwen35/src/config/model.rs b/pegainfer-qwen35/src/config/model.rs new file mode 100644 index 000000000..64c7fdecd --- /dev/null +++ b/pegainfer-qwen35/src/config/model.rs @@ -0,0 +1,396 @@ +//! Validated Qwen3.5 model config. +//! +//! This module owns the *model* geometry: deserialized-from-disk [`RawConfig`] +//! is converted once into a validated [`Config35`] via [`TryFrom`]. It has no +//! knowledge of tensor parallelism — sharded/local dimensions, shard ranges and +//! the TP assignment live in [`super::tp`] as a separate directional boundary. + +use std::fs; + +use anyhow::Result; +use serde::Deserialize; + +use super::error::ConfigError; + +/// Which attention variant a layer uses. Model-level (TP-agnostic) layout. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum LayerType { + FullAttention, + LinearAttention, +} + +/// RoPE block inside `text_config`. Only these two keys are read. +#[derive(Debug, Deserialize)] +struct RopeParameters { + rope_theta: f64, + partial_rotary_factor: f64, +} + +/// Raw `text_config` block, exactly as it appears on disk. +#[derive(Debug, Deserialize)] +pub(crate) struct TextConfig { + hidden_size: usize, + intermediate_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + vocab_size: usize, + rms_norm_eps: f64, + layer_types: Vec, + linear_conv_kernel_dim: usize, + linear_key_head_dim: usize, + linear_num_key_heads: usize, + linear_num_value_heads: usize, + linear_value_head_dim: usize, + rope_parameters: RopeParameters, + max_position_embeddings: Option, + tie_word_embeddings: Option, + eos_token_id: u32, +} + +/// Raw Qwen3.5 config as deserialized from `config.json`. Held only long enough +/// to be validated into a [`Config35`]. +#[derive(Debug, Deserialize)] +pub(crate) struct RawConfig { + text_config: TextConfig, + max_position_embeddings: Option, + tie_word_embeddings: Option, +} + +/// Head dims baked into the kernels; head counts are runtime parameters. +pub(crate) const GDN_AOT_KEY_HEAD_DIM: usize = 128; +pub(crate) const GDN_AOT_VALUE_HEAD_DIM: usize = 128; +pub(crate) const LINEAR_CONV_MAX_KERNEL_DIM: usize = 4; +const FULL_ATTN_HEAD_DIM: usize = 256; + +/// Validated Qwen3.5 model configuration (text-only). +/// +/// Every cross-field and kernel-AOT rule is enforced here, at the single +/// [`Self::try_from`]/[`Self::from_file`] boundary, so downstream code can only +/// observe a model that is known to be loadable. Tensor-parallel sharding is +/// never a property of this type. +#[derive(Debug)] +pub(crate) struct Config35 { + // Common + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) vocab_size: usize, + pub(crate) rms_norm_eps: f32, + pub(crate) eos_token_id: u32, + + // Full attention params + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) head_dim: usize, + + // Linear attention params + pub(crate) linear_num_key_heads: usize, + pub(crate) linear_key_head_dim: usize, + pub(crate) linear_num_value_heads: usize, + pub(crate) linear_value_head_dim: usize, + pub(crate) linear_conv_kernel_dim: usize, + + // RoPE + pub(crate) rope_theta: f32, + pub(crate) rotary_dim: usize, + pub(crate) max_position_embeddings: usize, + + // Layer layout + pub(crate) layer_types: Vec, + + /// `false` requires a top-level `lm_head.weight`; `true` reuses `embed_tokens`. + pub(crate) tie_word_embeddings: bool, + + /// Token-selection width: `vocab_size` bounded to the frontend-decodable vocab. + pub(crate) selection_vocab: usize, +} + +impl Config35 { + /// Load and validate `config.json` from a model directory. + pub(crate) fn from_file(model_path: &str) -> Result { + let config_path = format!("{}/config.json", model_path); + let content = fs::read_to_string(&config_path)?; + let raw: RawConfig = serde_json::from_str(&content)?; + let config = Self::try_from(raw)?; + Ok(config) + } + + /// Number of full attention layers in the model. + pub(crate) fn num_full_attention_layers(&self) -> usize { + self.layer_types + .iter() + .filter(|&&t| t == LayerType::FullAttention) + .count() + } + + /// Q dimension for full attention (without gate). + pub(crate) fn full_attn_q_dim(&self) -> usize { + self.num_attention_heads * self.head_dim + } + + /// KV dimension for full attention. + pub(crate) fn full_attn_kv_dim(&self) -> usize { + self.num_key_value_heads * self.head_dim + } + + pub(crate) fn decode_group_is_compiled(&self) -> bool { + // Uncompiled GQA groups use the batched hybrid eager fallback. + pegainfer_core::ops::SUPPORTED_GQA_GROUP_SIZES + .contains(&(self.num_attention_heads / self.num_key_value_heads)) + } + + /// QKV projection output dimension for linear attention. + pub(crate) fn linear_attn_qkv_dim(&self) -> usize { + let q_dim = self.linear_num_key_heads * self.linear_key_head_dim; + let k_dim = q_dim; + let v_dim = self.linear_num_value_heads * self.linear_value_head_dim; + q_dim + k_dim + v_dim + } + + /// Z projection output dimension for linear attention. + pub(crate) fn linear_attn_z_dim(&self) -> usize { + self.linear_num_value_heads * self.linear_value_head_dim + } + + /// Bound the output-selection width to the frontend-decodable vocab. + /// + /// The frontend decodes a dense prefix of the vocab; the checkpoint may pad + /// beyond it. Refusing a tokenizer wider than the checkpoint is the + /// fail-closed rule, and it is checked here at the validation boundary + /// rather than scattered through the loader. + pub(crate) fn bound_selection_vocab( + &mut self, + effective_vocab: usize, + ) -> Result<(), ConfigError> { + if effective_vocab > self.vocab_size { + return Err(ConfigError::EffectiveVocabExceedsCheckpoint { + used: effective_vocab, + vocab_size: self.vocab_size, + }); + } + self.selection_vocab = effective_vocab; + Ok(()) + } +} + +/// Validate a raw Qwen3.5 config into a [`Config35`]. +/// +/// Every check here is a static, load-time invariant: an invalid model file is +/// rejected before any weight is mapped or any CUDA buffer is allocated. The +/// kernel head dims are compile-time AOT constants, so a config that disagrees +/// cannot be served and is rejected here rather than at first kernel launch. +impl TryFrom for Config35 { + type Error = ConfigError; + + fn try_from(raw: RawConfig) -> std::result::Result { + let root_max_position_embeddings = raw.max_position_embeddings; + let root_tie_word_embeddings = raw.tie_word_embeddings; + let t = raw.text_config; + + let tie_word_embeddings = t + .tie_word_embeddings + .or(root_tie_word_embeddings) + .ok_or(ConfigError::MissingTieWordEmbeddings)?; + + let layer_types: Vec = t + .layer_types + .iter() + .map(|s| match s.as_str() { + "full_attention" => Ok(LayerType::FullAttention), + "linear_attention" => Ok(LayerType::LinearAttention), + other => Err(ConfigError::UnknownLayerType(other.to_string())), + }) + .collect::>()?; + + if layer_types.len() != t.num_hidden_layers { + return Err(ConfigError::LayerTypeCountMismatch { + actual: layer_types.len(), + expected: t.num_hidden_layers, + }); + } + + let rotary_dim = (t.head_dim as f64 * t.rope_parameters.partial_rotary_factor) as usize; + if rotary_dim == 0 { + return Err(ConfigError::ZeroRotaryDim); + } + + let max_position_embeddings = t + .max_position_embeddings + .or(root_max_position_embeddings) + .ok_or(ConfigError::MissingMaxPositionEmbeddings)?; + if max_position_embeddings == 0 { + return Err(ConfigError::NonPositiveMaxPositionEmbeddings); + } + + if t.linear_key_head_dim != GDN_AOT_KEY_HEAD_DIM + || t.linear_value_head_dim != GDN_AOT_VALUE_HEAD_DIM + { + return Err(ConfigError::GdnAotHeadDimMismatch { + key: t.linear_key_head_dim, + value: t.linear_value_head_dim, + expected_key: GDN_AOT_KEY_HEAD_DIM, + expected_value: GDN_AOT_VALUE_HEAD_DIM, + }); + } + if t.head_dim != FULL_ATTN_HEAD_DIM { + return Err(ConfigError::FullAttnHeadDimMismatch { + expected: FULL_ATTN_HEAD_DIM, + actual: t.head_dim, + }); + } + if !(1..=LINEAR_CONV_MAX_KERNEL_DIM).contains(&t.linear_conv_kernel_dim) { + return Err(ConfigError::LinearConvKernelDim { + max: LINEAR_CONV_MAX_KERNEL_DIM, + actual: t.linear_conv_kernel_dim, + }); + } + if t.linear_num_key_heads == 0 + || !t + .linear_num_value_heads + .is_multiple_of(t.linear_num_key_heads) + { + return Err(ConfigError::LinearHeadDivisibility { + key_heads: t.linear_num_key_heads, + value_heads: t.linear_num_value_heads, + }); + } + if t.num_key_value_heads == 0 + || !t.num_attention_heads.is_multiple_of(t.num_key_value_heads) + { + return Err(ConfigError::AttentionHeadDivisibility { + attention_heads: t.num_attention_heads, + key_value_heads: t.num_key_value_heads, + }); + } + + Ok(Self { + hidden_size: t.hidden_size, + intermediate_size: t.intermediate_size, + num_hidden_layers: t.num_hidden_layers, + vocab_size: t.vocab_size, + rms_norm_eps: t.rms_norm_eps as f32, + eos_token_id: t.eos_token_id, + num_attention_heads: t.num_attention_heads, + num_key_value_heads: t.num_key_value_heads, + head_dim: t.head_dim, + linear_num_key_heads: t.linear_num_key_heads, + linear_key_head_dim: t.linear_key_head_dim, + linear_num_value_heads: t.linear_num_value_heads, + linear_value_head_dim: t.linear_value_head_dim, + linear_conv_kernel_dim: t.linear_conv_kernel_dim, + rope_theta: t.rope_parameters.rope_theta as f32, + rotary_dim, + max_position_embeddings, + layer_types, + tie_word_embeddings, + selection_vocab: t.vocab_size, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const VALID_RAW: &str = r#"{ + "max_position_embeddings": 4096, + "tie_word_embeddings": true, + "text_config": { + "hidden_size": 512, + "intermediate_size": 1024, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 256, + "vocab_size": 1000, + "rms_norm_eps": 1e-6, + "layer_types": ["linear_attention", "full_attention"], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "rope_parameters": { "rope_theta": 10000.0, "partial_rotary_factor": 0.25 }, + "eos_token_id": 0 + } +}"#; + + fn parse(raw: &str) -> RawConfig { + serde_json::from_str(raw).expect("fixture must deserialize") + } + + fn config(raw: &str) -> Config35 { + Config35::try_from(parse(raw)).expect("fixture must validate") + } + + #[test] + fn valid_fixture_loads() { + let config = config(VALID_RAW); + assert_eq!(config.num_full_attention_layers(), 1); + assert_eq!(config.vocab_size, 1000); + } + + #[test] + fn missing_tie_word_embeddings_is_typed() { + let json = VALID_RAW.replace("\"tie_word_embeddings\": true,", ""); + let err = Config35::try_from(parse(&json)).unwrap_err(); + assert_eq!(err, ConfigError::MissingTieWordEmbeddings); + } + + #[test] + fn unknown_layer_type_is_typed() { + let json = VALID_RAW.replace("linear_attention", "window_attention"); + let err = Config35::try_from(parse(&json)).unwrap_err(); + assert_eq!( + err, + ConfigError::UnknownLayerType("window_attention".to_string()) + ); + } + + #[test] + fn layer_type_count_mismatch_is_typed() { + let json = VALID_RAW.replace( + "[\"linear_attention\", \"full_attention\"]", + "[\"linear_attention\"]", + ); + let err = Config35::try_from(parse(&json)).unwrap_err(); + assert_eq!( + err, + ConfigError::LayerTypeCountMismatch { + actual: 1, + expected: 2, + } + ); + } + + #[test] + fn missing_max_position_embeddings_is_typed() { + let json = VALID_RAW.replace("\"max_position_embeddings\": 4096,", ""); + let err = Config35::try_from(parse(&json)).unwrap_err(); + assert_eq!(err, ConfigError::MissingMaxPositionEmbeddings); + } + + #[test] + fn wide_linear_conv_kernel_is_typed() { + let json = VALID_RAW.replace( + "\"linear_conv_kernel_dim\": 4", + "\"linear_conv_kernel_dim\": 5", + ); + let err = Config35::try_from(parse(&json)).unwrap_err(); + assert_eq!( + err, + ConfigError::LinearConvKernelDim { + max: LINEAR_CONV_MAX_KERNEL_DIM, + actual: 5, + } + ); + } + + #[test] + fn acceptance_of_48_value_heads() { + // Regression guard: 48 value heads with 16 key heads must load. + config(VALID_RAW); + } +} diff --git a/pegainfer-qwen35/src/config/tokenizer.rs b/pegainfer-qwen35/src/config/tokenizer.rs new file mode 100644 index 000000000..f34fe5495 --- /dev/null +++ b/pegainfer-qwen35/src/config/tokenizer.rs @@ -0,0 +1,104 @@ +//! Frontend tokenizer schema and decodable-vocab width for Qwen3.5. +//! +//! The pinned frontend serves from the join of `tokenizer.json` (vocab + +//! `added_tokens`) and `tokenizer_config.json` (`added_tokens_decoder`). We +//! mirror that merge to compute the widest decodable token id, then bound the +//! output selection to a dense prefix of that id space. + +use std::collections::HashSet; +use std::fs; + +use anyhow::Result; +use log::warn; +use serde::Deserialize; + +/// One `added_tokens_decoder` entry in `tokenizer_config.json`. +/// +/// An explicitly-owned frontend compatibility contract. The frontend's decoder +/// schema is bool-heavy, and we parse it with these typed fields purely to make +/// the whole-file typed parse fail-closed on a malformed entry (e.g. a `special` +/// that is not a bool): then every decoder token is dropped, mirroring the +/// frontend. None of the fields are read by [`tokenizer_effective_vocab`] — +/// the contract's entire job is the fail-closed type-check, so this schema is +/// intentionally dead in the code path and its fields exist to shape the parse. +#[allow(clippy::struct_excessive_bools)] +#[allow(dead_code)] +#[derive(Deserialize)] +struct FrontendAddedToken { + #[serde(default)] + id: Option, + content: String, + #[serde(default)] + single_word: bool, + #[serde(default)] + lstrip: bool, + #[serde(default)] + rstrip: bool, + #[serde(default)] + normalized: bool, + #[serde(default)] + special: bool, +} + +#[derive(Deserialize)] +struct TokenizerJsonIds { + model: TokenizerModelIds, + #[serde(default)] + added_tokens: Vec, +} + +#[derive(Deserialize)] +struct TokenizerModelIds { + vocab: std::collections::HashMap, +} + +#[derive(Deserialize)] +struct TokenizerConfigIds { + #[serde(default)] + added_tokens_decoder: std::collections::HashMap, +} + +/// Width of the frontend-decodable id space, mirroring the pinned frontend's +/// merge: `tokenizer.json` vocab and added_tokens (fatal on parse failure - +/// the frontend cannot serve without it) plus `tokenizer_config.json` +/// added_tokens_decoder (whole-file typed parse; failure drops all decoder +/// tokens with a warning, unparsable keys are skipped per entry). The ids +/// must form a dense prefix - a row-range selection bound cannot mask holes - +/// so a sparse id space fails the load instead of silently truncating the +/// output space. +pub(crate) fn tokenizer_effective_vocab(model_path: &str) -> Result { + let path = format!("{}/tokenizer.json", model_path); + let content = + fs::read_to_string(&path).map_err(|e| anyhow::anyhow!("cannot read {path}: {e}"))?; + let tj: TokenizerJsonIds = + serde_json::from_str(&content).map_err(|e| anyhow::anyhow!("cannot parse {path}: {e}"))?; + anyhow::ensure!(!tj.model.vocab.is_empty(), "{path} model.vocab is empty"); + let mut ids: HashSet = tj.model.vocab.into_values().collect(); + ids.extend(tj.added_tokens.iter().filter_map(|t| t.id)); + + let config_path = format!("{}/tokenizer_config.json", model_path); + if let Ok(text) = fs::read_to_string(&config_path) { + match serde_json::from_str::(&text) { + Ok(cfg) => ids.extend( + cfg.added_tokens_decoder + .keys() + .filter_map(|k| k.parse::().ok()), + ), + Err(e) => warn!( + "cannot parse {config_path}: {e}; skipping its added tokens like the frontend does" + ), + } + } + + let width = ids.len(); + let max_id = + *ids.iter() + .max() + .ok_or_else(|| anyhow::anyhow!("tokenizer vocab must be non-empty"))? as usize; + anyhow::ensure!( + max_id + 1 == width, + "tokenizer id space is not dense (max id {max_id}, {width} distinct ids); \ + a row-range selection bound cannot mask holes" + ); + Ok(width) +} diff --git a/pegainfer-qwen35/src/config/tp.rs b/pegainfer-qwen35/src/config/tp.rs new file mode 100644 index 000000000..9f2a97e1a --- /dev/null +++ b/pegainfer-qwen35/src/config/tp.rs @@ -0,0 +1,318 @@ +//! Tensor-parallel assignment and validated local geometry for Qwen3.5. +//! +//! Directional boundary: this module depends on [`super::model`] (it reads a +//! validated [`Config35`] to derive shard dimensions), but [`super::model`] is +//! TP-agnostic and never depends on it. Downstream code accepts the validated +//! [`LocalGeometry`] type produced here instead of re-deriving shards from a +//! raw `(rank, world_size)` pair. +//! +//! Both [`TensorParallelConfig`] and [`LocalGeometry`] validate at construction: +//! an invalid assignment is unrepresentable, and a model that disagrees with the +//! TP world is rejected before any weights are mapped. + +use super::error::ConfigError; +use super::model::Config35; + +/// A validated tensor-parallel assignment. +/// +/// Fields are private so an invalid state (`world_size == 0`, `rank >= +/// world_size`) cannot be constructed. The only way to obtain one is +/// [`Self::try_from`] or [`Default`] (single-rank). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct TensorParallelConfig { + rank: usize, + world_size: usize, +} + +impl Default for TensorParallelConfig { + fn default() -> Self { + Self { + rank: 0, + world_size: 1, + } + } +} + +impl TensorParallelConfig { + pub(crate) fn rank(self) -> usize { + self.rank + } + + pub(crate) fn world_size(self) -> usize { + self.world_size + } + + pub(crate) fn is_sharded(self) -> bool { + self.world_size > 1 + } + + /// Split a global row count into this rank's `(offset, len)`. + pub(crate) fn shard_range(self, total: usize) -> (usize, usize) { + let shard_len = total / self.world_size; + (self.rank * shard_len, shard_len) + } +} + +/// Validate an assignment independently of any model geometry. +impl TryFrom<(usize, usize)> for TensorParallelConfig { + type Error = ConfigError; + + fn try_from((rank, world_size): (usize, usize)) -> Result { + if world_size == 0 { + return Err(ConfigError::TpZeroWorldSize); + } + if rank >= world_size { + return Err(ConfigError::TpRankOutOfRange { rank, world_size }); + } + Ok(Self { rank, world_size }) + } +} + +/// Validated per-rank shard geometry for a Qwen3.5 model under a TP assignment. +/// +/// Built once from a validated model config + TP assignment + cuda-graph mode; +/// every model cross-field and TP/kernel compatibility rule is enforced here. +/// Downstream code sizes buffers and computes shards from this type, never from +/// a raw rank/world-size pair. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct LocalGeometry { + tp: TensorParallelConfig, + local_num_attention_heads: usize, + local_num_key_value_heads: usize, + local_intermediate_size: usize, + local_full_attn_q_dim: usize, + local_full_attn_kv_dim: usize, + local_full_attn_gated_q_dim: usize, +} + +impl LocalGeometry { + /// Validate `config` against `tp` and the runtime execution mode, then derive + /// this rank's local dimensions. + /// + /// Fails on unsupported combinations before expensive loading: + /// - sharded TP demands eager execution (`enable_cuda_graph` off); + /// - every sharded model dimension must divide evenly by `world_size` + /// (linear-attention head counts are intentionally exempt); + /// - `rank < world_size` and `world_size >= 1` are guaranteed by + /// `TensorParallelConfig::try_from`. + pub(crate) fn try_new( + config: &Config35, + tp: TensorParallelConfig, + enable_cuda_graph: bool, + ) -> Result { + if tp.is_sharded() && enable_cuda_graph { + return Err(ConfigError::TpRequiresEager { + world_size: tp.world_size(), + }); + } + if !config.num_attention_heads.is_multiple_of(tp.world_size()) { + return Err(ConfigError::TpIndivisible { + field: "num_attention_heads", + value: config.num_attention_heads, + world_size: tp.world_size(), + }); + } + if !config.num_key_value_heads.is_multiple_of(tp.world_size()) { + return Err(ConfigError::TpIndivisible { + field: "num_key_value_heads", + value: config.num_key_value_heads, + world_size: tp.world_size(), + }); + } + if !config.intermediate_size.is_multiple_of(tp.world_size()) { + return Err(ConfigError::TpIndivisible { + field: "intermediate_size", + value: config.intermediate_size, + world_size: tp.world_size(), + }); + } + + let local_num_attention_heads = config.num_attention_heads / tp.world_size(); + let local_num_key_value_heads = config.num_key_value_heads / tp.world_size(); + let local_intermediate_size = config.intermediate_size / tp.world_size(); + let local_full_attn_q_dim = local_num_attention_heads * config.head_dim; + let local_full_attn_kv_dim = local_num_key_value_heads * config.head_dim; + + Ok(Self { + tp, + local_num_attention_heads, + local_num_key_value_heads, + local_intermediate_size, + local_full_attn_q_dim, + local_full_attn_kv_dim, + local_full_attn_gated_q_dim: local_full_attn_q_dim * 2, + }) + } + + pub(crate) fn rank(&self) -> usize { + self.tp.rank() + } + + pub(crate) fn world_size(&self) -> usize { + self.tp.world_size() + } + + pub(crate) fn is_sharded(&self) -> bool { + self.tp.is_sharded() + } + + pub(crate) fn shard_range(&self, total: usize) -> (usize, usize) { + self.tp.shard_range(total) + } + + pub(crate) fn local_num_attention_heads(&self) -> usize { + self.local_num_attention_heads + } + + pub(crate) fn local_num_key_value_heads(&self) -> usize { + self.local_num_key_value_heads + } + + pub(crate) fn local_intermediate_size(&self) -> usize { + self.local_intermediate_size + } + + pub(crate) fn local_full_attn_q_dim(&self) -> usize { + self.local_full_attn_q_dim + } + + pub(crate) fn local_full_attn_kv_dim(&self) -> usize { + self.local_full_attn_kv_dim + } + + /// Local gated full-attention q projection output dimension. + pub(crate) fn local_full_attn_gated_q_dim(&self) -> usize { + self.local_full_attn_gated_q_dim + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> Config35 { + let raw: crate::config::model::RawConfig = serde_json::from_str( + r#"{ + "max_position_embeddings": 4096, + "tie_word_embeddings": true, + "text_config": { + "hidden_size": 2560, + "intermediate_size": 9216, + "num_hidden_layers": 1, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 256, + "vocab_size": 248320, + "rms_norm_eps": 1e-6, + "layer_types": ["linear_attention"], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "rope_parameters": { "rope_theta": 10000.0, "partial_rotary_factor": 0.25 }, + "eos_token_id": 151645 + } +}"#, + ) + .expect("fixture parses"); + Config35::try_from(raw).expect("fixture validates") + } + + #[test] + fn default_tensor_parallel_is_tp1() { + let tp = TensorParallelConfig::default(); + assert_eq!((tp.rank(), tp.world_size()), (0, 1)); + assert!(!tp.is_sharded()); + assert_eq!(tp.shard_range(4096), (0, 4096)); + } + + #[test] + fn tp2_local_geometry_matches_dense_dims() { + let cfg = config(); + let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); + let geom = LocalGeometry::try_new(&cfg, tp, false).unwrap(); + assert!(geom.is_sharded()); + assert_eq!(geom.shard_range(4096), (2048, 2048)); + assert_eq!(geom.local_num_attention_heads(), 8); + assert_eq!(geom.local_num_key_value_heads(), 2); + assert_eq!(geom.local_intermediate_size(), 4608); + assert_eq!(geom.local_full_attn_q_dim(), 2048); + assert_eq!(geom.local_full_attn_kv_dim(), 512); + assert_eq!(geom.local_full_attn_gated_q_dim(), 4096); + } + + #[test] + fn rejects_zero_world_size_and_rank_out_of_range() { + assert_eq!( + TensorParallelConfig::try_from((0, 0)), + Err(ConfigError::TpZeroWorldSize) + ); + assert_eq!( + TensorParallelConfig::try_from((2, 2)), + Err(ConfigError::TpRankOutOfRange { + rank: 2, + world_size: 2 + }) + ); + } + + #[test] + fn rejects_indivisible_dense_dimensions() { + let tp = TensorParallelConfig::try_from((0, 3)).unwrap(); + let cfg = config(); + let mut err = LocalGeometry::try_new(&cfg, tp, false).unwrap_err(); + assert_eq!( + err, + ConfigError::TpIndivisible { + field: "num_attention_heads", + value: 16, + world_size: 3, + } + ); + + let mut broken = cfg; + broken.num_attention_heads = 15; + broken.num_key_value_heads = 4; + err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + assert_eq!( + err, + ConfigError::TpIndivisible { + field: "num_key_value_heads", + value: 4, + world_size: 3, + } + ); + + broken.num_key_value_heads = 3; + broken.intermediate_size = 9217; + err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + assert_eq!( + err, + ConfigError::TpIndivisible { + field: "intermediate_size", + value: 9217, + world_size: 3, + } + ); + } + + #[test] + fn rejects_tensor_parallel_with_cuda_graph() { + let cfg = config(); + let tp = TensorParallelConfig::try_from((0, 2)).unwrap(); + assert_eq!( + LocalGeometry::try_new(&cfg, tp, true), + Err(ConfigError::TpRequiresEager { world_size: 2 }) + ); + } + + #[test] + fn linear_attention_heads_need_not_divide_world_size() { + let mut cfg = config(); + cfg.linear_num_key_heads = 17; + cfg.linear_num_value_heads = 31; + let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); + LocalGeometry::try_new(&cfg, tp, false).unwrap(); + } +} diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs index 3e5bbfad1..2275d2167 100644 --- a/pegainfer-qwen35/src/decode_buffers.rs +++ b/pegainfer-qwen35/src/decode_buffers.rs @@ -7,7 +7,7 @@ use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::HiddenStates; use super::config::Config35; -use super::config::TensorParallelConfig; +use super::config::LocalGeometry; /// Pre-allocated GPU buffers for Qwen3.5 batch decode (N requests, 1 token each). pub(crate) struct BatchDecodeBuffers35 { @@ -65,21 +65,21 @@ impl BatchDecodeBuffers35 { pub(crate) fn new( ctx: &DeviceContext, config: &Config35, - tensor_parallel: TensorParallelConfig, + geometry: LocalGeometry, max_batch_size: usize, max_total_pages: usize, padding_page_id: i32, ) -> Result { let h = config.hidden_size; let bs = max_batch_size; - let q_proj_dim = config.local_full_attn_gated_q_dim(tensor_parallel); - let q_dim = config.local_full_attn_q_dim(tensor_parallel); - let kv_dim = config.local_full_attn_kv_dim(tensor_parallel); + let q_proj_dim = geometry.local_full_attn_gated_q_dim(); + let q_dim = geometry.local_full_attn_q_dim(); + let kv_dim = geometry.local_full_attn_kv_dim(); let qkv_dim = config.linear_attn_qkv_dim(); let z_dim = config.linear_attn_z_dim(); let b_dim = config.linear_num_value_heads; let a_dim = b_dim; - let intermediate = config.local_intermediate_size(tensor_parallel); + let intermediate = geometry.local_intermediate_size(); Ok(Self { max_batch_size: bs, diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 0638d222c..f46db580c 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -2,6 +2,8 @@ // `pegainfer-kernels/qwen35`, which need Python + Triton at build time. // Without the feature this compiles to an empty crate so plain workspace // builds stay Python-free. +// Submodule `use self::*` globs are the established module-tree convention. +#![allow(clippy::wildcard_imports)] #![cfg(feature = "qwen35")] mod batch_decode; diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index f641983a9..6b65d9459 100644 --- a/pegainfer-qwen35/src/prefill.rs +++ b/pegainfer-qwen35/src/prefill.rs @@ -176,14 +176,14 @@ impl Qwen35Model { kv_state.ensure_capacity(end_pos)?; kv_state.advance(seq_len); let kv_desc = kv_state.desc(); - let tp = self.tensor_parallel; + let geom = self.geometry; let prefill_plan = PrefillPagedPlan::new( &self.ctx, &kv_desc, base_pos, seq_len, - c.local_num_attention_heads(tp), - c.local_num_key_value_heads(tp), + geom.local_num_attention_heads(), + geom.local_num_key_value_heads(), c.head_dim, )?; @@ -237,9 +237,9 @@ impl Qwen35Model { self.batched_rms_norm_offset(hidden_batch, &layer.input_layernorm, eps)?; // 2. Attention / Linear attention — per-token for correctness - let tp = self.tensor_parallel; + let geom = self.geometry; let attn_out_dim = match &layer.attn { - LayerKind::FullAttention(_) => c.local_full_attn_q_dim(tp), + LayerKind::FullAttention(_) => geom.local_full_attn_q_dim(), LayerKind::LinearAttention(_) => c.linear_attn_z_dim(), }; @@ -273,7 +273,7 @@ impl Qwen35Model { // 4. MLP (batched) let gate_up_out = ops::gemm(&self.ctx, &layer.mlp.gate_up_proj, &normed_batch)?; - let mut act_out = HiddenStates::zeros(&self.ctx, c.local_intermediate_size(tp), seq_len)?; + let mut act_out = HiddenStates::zeros(&self.ctx, geom.local_intermediate_size(), seq_len)?; ops::silu_mul_fused_batch_into(&self.ctx, &gate_up_out, &mut act_out)?; let mut mlp_out = ops::gemm(&self.ctx, &layer.mlp.down_proj, &act_out)?; self.all_reduce_hidden(&mut mlp_out)?; @@ -294,10 +294,10 @@ impl Qwen35Model { seq_len: usize, ) -> Result { let c = &self.config; - let tp = self.tensor_parallel; - let num_attention_heads = c.local_num_attention_heads(tp); - let num_key_value_heads = c.local_num_key_value_heads(tp); - let attn_out_dim = c.local_full_attn_q_dim(tp); + let geom = self.geometry; + let num_attention_heads = geom.local_num_attention_heads(); + let num_key_value_heads = geom.local_num_key_value_heads(); + let attn_out_dim = geom.local_full_attn_q_dim(); let eps = c.rms_norm_eps; let q_full_batch = ops::gemm(&self.ctx, &attn.q_proj, normed_batch)?; let k_batch = ops::gemm(&self.ctx, &attn.k_proj, normed_batch)?; diff --git a/pegainfer-qwen35/src/scheduler/backend.rs b/pegainfer-qwen35/src/scheduler/backend.rs new file mode 100644 index 000000000..d5abe8d0f --- /dev/null +++ b/pegainfer-qwen35/src/scheduler/backend.rs @@ -0,0 +1,567 @@ +//! Qwen3.5 scheduler backend abstraction (single-GPU + TP). + +use super::*; + +pub(super) struct SingleGpuBackend { + model: Qwen35Model, + graph_state: BatchDecodeGraphState, + prefill_stream: Option>, +} + +// One instance per scheduler; the size asymmetry costs nothing here. +#[allow(clippy::large_enum_variant)] +pub(super) enum SchedulerBackend { + Single(SingleGpuBackend), + Tp(TpSchedulerBackend), +} + +pub(super) struct AsyncPrefillOutput { + logits: Option, + done: CudaEvent, + stream: Arc, + completed: bool, +} + +impl AsyncPrefillOutput { + pub(super) fn is_ready(&mut self) -> bool { + match unsafe { sys::cuEventQuery(self.done.cu_event()) } { + sys::CUresult::CUDA_SUCCESS => { + self.completed = true; + true + } + sys::CUresult::CUDA_ERROR_NOT_READY => false, + err => fatal_cuda_lifecycle(&format!( + "query Qwen3.5 async prefill event failed: {err:?}" + )), + } + } + + pub(super) fn into_logits(mut self) -> HiddenStates { + if !self.completed { + if let Err(err) = self.done.synchronize() { + fatal_cuda_lifecycle(&format!("wait for Qwen3.5 async prefill failed: {err}")); + } + self.completed = true; + } + self.logits + .take() + .expect("async prefill logits must be consumed exactly once") + } +} + +impl Drop for AsyncPrefillOutput { + fn drop(&mut self) { + if self.completed { + return; + } + if let Err(err) = self.stream.synchronize() { + fatal_cuda_lifecycle(&format!( + "drain Qwen3.5 async prefill during cleanup failed: {err}" + )); + } + } +} + +pub(super) fn fatal_cuda_lifecycle(message: &str) -> ! { + log::error!("FATAL: {message}; aborting before CUDA-referenced state is released"); + std::process::abort(); +} + +pub(super) struct TpSchedulerBackend { + executor: Qwen35TpExecutor, + next_request_id: u64, +} + +impl SingleGpuBackend { + pub(super) fn new( + model: Qwen35Model, + max_batch: usize, + decode_overlap: Qwen35DecodeOverlap, + ) -> Result { + anyhow::ensure!(max_batch > 0, "Qwen3.5 max_batch must be > 0"); + let graph_capacity = crate::batch_decode_graph::bucket_for(max_batch); + let graph_state = model.create_batch_decode_graph_state_with_capacity(graph_capacity)?; + let prefill_stream = match decode_overlap { + Qwen35DecodeOverlap::Off => None, + Qwen35DecodeOverlap::SharedSm => Some( + model + .device_ctx() + .ctx + .new_stream() + .map_err(|err| anyhow::anyhow!("create Qwen3.5 prefill stream: {err}"))?, + ), + }; + Ok(Self { + model, + graph_state, + prefill_stream, + }) + } + + pub(super) fn model(&self) -> &Qwen35Model { + &self.model + } + + pub(super) fn max_batch(&self) -> usize { + // #470: admit the requested `--max-batch`, which may sit below the loaded + // graph bucket (e.g. 5 on bucket 8); never exceed the physical slots. + self.model + .decode_admission_batch + .min(self.graph_state.slot_states.len()) + .max(1) + } + + pub(super) fn page_size(&self) -> usize { + self.model.kv_pool().layout().page_size + } + + pub(super) fn available_pages(&self) -> usize { + self.model.kv_pool().available_pages() + } + + pub(super) fn capacity_pages_for_requests(&self) -> usize { + self.model.kv_pool().capacity_pages().saturating_sub(1) + } + + pub(super) fn max_position_embeddings(&self) -> usize { + self.model.config().max_position_embeddings + } + + pub(super) fn alloc_kv(&self) -> KvState { + self.model.alloc_kv() + } + + pub(super) fn alloc_recurrent(&self) -> Result { + RecurrentState::new(self.model.device_ctx(), self.model.config()) + } + + pub(super) fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU prefill received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + self.model + .batch_prefill_logits(&window_refs, kvs, &mut rec_refs) + } + + pub(super) fn overlap_enabled(&self) -> bool { + self.prefill_stream.is_some() + } + + pub(super) fn launch_async_prefill( + &mut self, + chunk: &mut ScheduledChunk, + ) -> Result { + let prefill_stream = self + .prefill_stream + .clone() + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 decode overlap is disabled"))?; + + // Request KV/recurrent state was allocated on the model stream. Order + // those producers before the prefill stream without blocking the host. + prefill_stream + .join(&self.model.device_ctx().stream) + .map_err(|err| anyhow::anyhow!("join Qwen3.5 prefill stream: {err}"))?; + + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU async prefill received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + let logits = match self.model.batch_prefill_logits_on_stream( + Arc::clone(&prefill_stream), + &window_refs, + kvs, + &mut rec_refs, + ) { + Ok(logits) => logits, + Err(err) => { + if let Err(sync_err) = prefill_stream.synchronize() { + fatal_cuda_lifecycle(&format!( + "Qwen3.5 async prefill failed ({err}); stream drain failed: {sync_err}" + )); + } + return Err(err); + } + }; + let done = match prefill_stream.record_event(None) { + Ok(done) => done, + Err(err) => { + if let Err(sync_err) = prefill_stream.synchronize() { + fatal_cuda_lifecycle(&format!( + "record Qwen3.5 async prefill event failed ({err}); stream drain failed: {sync_err}" + )); + } + return Err(anyhow::anyhow!("record Qwen3.5 async prefill event: {err}")); + } + }; + Ok(AsyncPrefillOutput { + logits: Some(logits), + done, + stream: prefill_stream, + completed: false, + }) + } + + pub(super) fn unified_step( + &mut self, + chunk: &mut ScheduledChunk, + active: &mut [ActiveRequest35], + ) -> Result { + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU unified step received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); + let mut decode_kv_refs: Vec<&mut KvState> = active + .iter_mut() + .map(|r| match &mut r.backend_state { + ActiveBackendState::Single { kv, .. } => kv, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU unified step received TP active state") + } + }) + .collect(); + self.model.unified_step( + &window_refs, + kvs, + &mut rec_refs, + &decode_tokens, + &mut decode_kv_refs, + &mut self.graph_state, + ) + } + + pub(super) fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { + let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); + let mut kv_refs: Vec<&mut KvState> = active + .iter_mut() + .map(|r| match &mut r.backend_state { + ActiveBackendState::Single { kv, .. } => kv, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU decode received TP active state") + } + }) + .collect(); + self.model + .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) + } + + pub(super) fn sample_prefill_logits( + &mut self, + pending: &[SchedulerRequest], + logits: &HiddenStates, + sample_seed: u64, + ) -> Result<(Vec, Vec>)> { + debug_assert_eq!( + logits.seq_len, + pending.len(), + "Qwen3.5 prefill logits rows must preserve pending request order" + ); + let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect(); + let cpu_logits = + snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?; + let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); + let tokens = self.model.select_tokens_from_logits_varied( + logits, + &mut self.graph_state.buffers, + ¶ms_refs, + sample_seed, + )?; + + let logprobs = cpu_logits + .into_iter() + .enumerate() + .map(|(i, logits_opt)| { + logits_opt.and_then(|logits_f32| { + pegainfer_sample::token_logprob_from_row( + &logits_f32, + tokens[i], + pending[i].logprobs, + ) + }) + }) + .collect(); + Ok((tokens, logprobs)) + } + + pub(super) fn sample_decode_logits( + &mut self, + active: &[ActiveRequest35], + sample_seed: u64, + ) -> Result<(Vec, Vec>)> { + let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); + let cpu_logits = snapshot_requested_logprobs( + self.model.device_ctx(), + &self.graph_state.buffers.logits, + &requested_logprobs, + )?; + let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); + let tokens = self.model.select_tokens_batch_varied( + &mut self.graph_state.buffers, + ¶ms_refs, + sample_seed, + )?; + + let logprobs = cpu_logits + .into_iter() + .enumerate() + .map(|(i, logits_opt)| { + logits_opt.and_then(|logits_f32| { + pegainfer_sample::token_logprob_from_row( + &logits_f32, + tokens[i], + active[i].logprobs, + ) + }) + }) + .collect(); + Ok((tokens, logprobs)) + } + + pub(super) fn is_stop_token(&self, token: u32) -> bool { + self.model.is_stop_token(token) + } + + pub(super) fn copy_recurrent_to_slot( + &mut self, + recurrent: &RecurrentState, + slot_idx: usize, + ) -> Result<()> { + self.graph_state + .copy_state_to_slot(self.model.device_ctx(), recurrent, slot_idx) + } + + pub(super) fn compact_slot( + &mut self, + active: &mut [ActiveRequest35], + compaction: plan::SlotCompaction, + ) { + let src_slot = match active[compaction.moved_to].backend_state { + ActiveBackendState::Single { graph_slot_idx, .. } => graph_slot_idx, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU slot compaction received TP active state") + } + }; + debug_assert_eq!(src_slot, compaction.moved_from); + + let ctx = self.model.device_ctx(); + let src = &self.graph_state.slot_states[compaction.moved_from]; + for layer_idx in 0..src.layers.len() { + let (src_part, dst_part) = if compaction.moved_to < compaction.moved_from { + let (left, right) = self + .graph_state + .slot_states + .split_at_mut(compaction.moved_from); + ( + &right[0].layers[layer_idx], + &mut left[compaction.moved_to].layers[layer_idx], + ) + } else { + unreachable!("idx < active.len() <= last"); + }; + + ctx.stream + .memcpy_dtod(&src_part.state, &mut dst_part.state) + .expect("compact slot state copy failed"); + ctx.stream + .memcpy_dtod(&src_part.conv_state.data, &mut dst_part.conv_state.data) + .expect("compact slot conv_state copy failed"); + } + self.graph_state.slot_states[compaction.moved_to].seq_len = + self.graph_state.slot_states[compaction.moved_from].seq_len; + + match &mut active[compaction.moved_to].backend_state { + ActiveBackendState::Single { graph_slot_idx, .. } => { + *graph_slot_idx = compaction.moved_to; + } + ActiveBackendState::Tp { .. } => { + panic!("single-GPU slot compaction received TP active state") + } + } + } +} + +impl TpSchedulerBackend { + pub(super) fn new( + model_path: &str, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, + ) -> Result { + let executor = Qwen35TpExecutor::from_runtime_with_limits( + model_path, + false, + device_ordinals, + max_batch, + max_prefill_tokens, + )?; + Ok(Self { + executor, + next_request_id: 1, + }) + } + + pub(super) fn alloc_request_id(&mut self) -> RequestId { + let id = RequestId::new(self.next_request_id); + self.next_request_id = self.next_request_id.wrapping_add(1).max(1); + id + } + + pub(super) fn max_batch(&self) -> usize { + self.executor.max_batch() + } + + pub(super) fn page_size(&self) -> usize { + self.executor.page_size() + } + + pub(super) fn capacity_pages_for_requests(&self) -> usize { + self.executor.capacity_pages_for_requests() + } + + pub(super) fn max_position_embeddings(&self) -> usize { + self.executor.max_position_embeddings() + } + + pub(super) fn is_stop_token(&self, token: u32) -> bool { + self.executor.is_stop_token(token) + } + + pub(super) fn available_pages( + &self, + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + ) -> usize { + let page_size = self.page_size(); + let active_pages: usize = active + .iter() + .map(|req| pages_needed(current_active_tokens(req), page_size)) + .sum(); + let prefilling_pages: usize = prefilling + .iter() + .map(|req| pages_needed(req.cursor, page_size)) + .sum(); + self.capacity_pages_for_requests() + .saturating_sub(active_pages.saturating_add(prefilling_pages)) + } + + pub(super) fn execute_prefill_chunk( + &self, + chunk: &ScheduledChunk, + sample_seed: u64, + ) -> Result>> { + let items = tp_prefill_items(chunk)?; + let result = self + .executor + .execute_prefill_chunks_with_seed(&items, sample_seed)?; + align_prefill_results(chunk, &result) + .map_err(|err| self.executor.poison_artifact_contract("prefill", &err)) + } + + pub(super) fn execute_decode( + &self, + active: &[ActiveRequest35], + sample_seed: u64, + ) -> Result> { + let items = tp_decode_items(active)?; + let result = self.executor.execute_decode_items(&items, sample_seed)?; + align_decode_results(active, &result) + .map_err(|err| self.executor.poison_artifact_contract("decode", &err)) + } + + pub(super) fn execute_unified( + &self, + chunk: &ScheduledChunk, + active: &[ActiveRequest35], + decode_sample_seed: u64, + prefill_sample_seed: u64, + ) -> Result { + let plan = TpUnifiedPlan { + prefill: tp_prefill_items(chunk)?, + decode: tp_decode_items(active)?, + prefill_sample_seed, + decode_sample_seed, + }; + let result = self.executor.execute_unified(&plan)?; + let prefill = align_prefill_results(chunk, &result.prefill).map_err(|err| { + self.executor + .poison_artifact_contract("unified prefill", &err) + })?; + let decode = align_decode_results(active, &result.decode).map_err(|err| { + self.executor + .poison_artifact_contract("unified decode", &err) + })?; + Ok(AlignedUnifiedArtifacts { prefill, decode }) + } + + pub(super) fn drop_request( + &self, + request_id: RequestId, + expectation: DropExpectation, + ) -> Result<()> { + self.executor.drop_request(request_id, expectation) + } +} + +impl SchedulerBackend { + pub(super) fn max_batch(&self) -> usize { + match self { + Self::Single(backend) => backend.max_batch(), + Self::Tp(backend) => backend.max_batch(), + } + } + + pub(super) fn page_size(&self) -> usize { + match self { + Self::Single(backend) => backend.page_size(), + Self::Tp(backend) => backend.page_size(), + } + } + + pub(super) fn available_pages( + &self, + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + ) -> usize { + match self { + Self::Single(backend) => backend.available_pages(), + Self::Tp(backend) => backend.available_pages(active, prefilling), + } + } + + pub(super) fn capacity_pages_for_requests(&self) -> usize { + match self { + Self::Single(backend) => backend.capacity_pages_for_requests(), + Self::Tp(backend) => backend.capacity_pages_for_requests(), + } + } + + pub(super) fn max_position_embeddings(&self) -> usize { + match self { + Self::Single(backend) => backend.max_position_embeddings(), + Self::Tp(backend) => backend.max_position_embeddings(), + } + } + + pub(super) fn alloc_prefill_state(&mut self) -> Result { + match self { + Self::Single(backend) => Ok(PrefillBackendState::Single { + kv: backend.alloc_kv(), + rec: backend.alloc_recurrent()?, + }), + Self::Tp(backend) => Ok(PrefillBackendState::Tp { + request_id: backend.alloc_request_id(), + }), + } + } + + pub(super) fn is_stop_token(&self, token: u32) -> bool { + match self { + Self::Single(backend) => backend.is_stop_token(token), + Self::Tp(backend) => backend.is_stop_token(token), + } + } +} diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler/mod.rs similarity index 79% rename from pegainfer-qwen35/src/scheduler.rs rename to pegainfer-qwen35/src/scheduler/mod.rs index c2552c23c..212b67257 100644 --- a/pegainfer-qwen35/src/scheduler.rs +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -4,8 +4,8 @@ //! - `RecurrentState` alongside `KvState` (linear attention layers) //! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) +mod backend; mod plan; - use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; @@ -42,6 +42,7 @@ use rand::rngs::StdRng; use tokio::sync::mpsc; use tokio::sync::watch; +use self::backend::*; use self::plan::ActiveDecodeState; use self::plan::ActiveKvBudget; use self::plan::ExecutionPlan; @@ -462,559 +463,6 @@ pub(crate) fn start_tp_with_capacity( ) } -struct SingleGpuBackend { - model: Qwen35Model, - graph_state: BatchDecodeGraphState, - prefill_stream: Option>, -} - -// One instance per scheduler; the size asymmetry costs nothing here. -#[allow(clippy::large_enum_variant)] -enum SchedulerBackend { - Single(SingleGpuBackend), - Tp(TpSchedulerBackend), -} - -struct AsyncPrefillOutput { - logits: Option, - done: CudaEvent, - stream: Arc, - completed: bool, -} - -impl AsyncPrefillOutput { - fn is_ready(&mut self) -> bool { - match unsafe { sys::cuEventQuery(self.done.cu_event()) } { - sys::CUresult::CUDA_SUCCESS => { - self.completed = true; - true - } - sys::CUresult::CUDA_ERROR_NOT_READY => false, - err => fatal_cuda_lifecycle(&format!( - "query Qwen3.5 async prefill event failed: {err:?}" - )), - } - } - - fn into_logits(mut self) -> HiddenStates { - if !self.completed { - if let Err(err) = self.done.synchronize() { - fatal_cuda_lifecycle(&format!("wait for Qwen3.5 async prefill failed: {err}")); - } - self.completed = true; - } - self.logits - .take() - .expect("async prefill logits must be consumed exactly once") - } -} - -impl Drop for AsyncPrefillOutput { - fn drop(&mut self) { - if self.completed { - return; - } - if let Err(err) = self.stream.synchronize() { - fatal_cuda_lifecycle(&format!( - "drain Qwen3.5 async prefill during cleanup failed: {err}" - )); - } - } -} - -fn fatal_cuda_lifecycle(message: &str) -> ! { - log::error!("FATAL: {message}; aborting before CUDA-referenced state is released"); - std::process::abort(); -} - -struct TpSchedulerBackend { - executor: Qwen35TpExecutor, - next_request_id: u64, -} - -impl SingleGpuBackend { - fn new( - model: Qwen35Model, - max_batch: usize, - decode_overlap: Qwen35DecodeOverlap, - ) -> Result { - anyhow::ensure!(max_batch > 0, "Qwen3.5 max_batch must be > 0"); - let graph_capacity = crate::batch_decode_graph::bucket_for(max_batch); - let graph_state = model.create_batch_decode_graph_state_with_capacity(graph_capacity)?; - let prefill_stream = match decode_overlap { - Qwen35DecodeOverlap::Off => None, - Qwen35DecodeOverlap::SharedSm => Some( - model - .device_ctx() - .ctx - .new_stream() - .map_err(|err| anyhow::anyhow!("create Qwen3.5 prefill stream: {err}"))?, - ), - }; - Ok(Self { - model, - graph_state, - prefill_stream, - }) - } - - fn model(&self) -> &Qwen35Model { - &self.model - } - - fn max_batch(&self) -> usize { - // #470: admit the requested `--max-batch`, which may sit below the loaded - // graph bucket (e.g. 5 on bucket 8); never exceed the physical slots. - self.model - .decode_admission_batch - .min(self.graph_state.slot_states.len()) - .max(1) - } - - fn page_size(&self) -> usize { - self.model.kv_pool().layout().page_size - } - - fn available_pages(&self) -> usize { - self.model.kv_pool().available_pages() - } - - fn capacity_pages_for_requests(&self) -> usize { - self.model.kv_pool().capacity_pages().saturating_sub(1) - } - - fn max_position_embeddings(&self) -> usize { - self.model.config().max_position_embeddings - } - - fn alloc_kv(&self) -> KvState { - self.model.alloc_kv() - } - - fn alloc_recurrent(&self) -> Result { - RecurrentState::new(self.model.device_ctx(), self.model.config()) - } - - fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { - anyhow::bail!("single-GPU prefill received TP chunk state"); - }; - let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - self.model - .batch_prefill_logits(&window_refs, kvs, &mut rec_refs) - } - - fn overlap_enabled(&self) -> bool { - self.prefill_stream.is_some() - } - - fn launch_async_prefill(&mut self, chunk: &mut ScheduledChunk) -> Result { - let prefill_stream = self - .prefill_stream - .clone() - .ok_or_else(|| anyhow::anyhow!("Qwen3.5 decode overlap is disabled"))?; - - // Request KV/recurrent state was allocated on the model stream. Order - // those producers before the prefill stream without blocking the host. - prefill_stream - .join(&self.model.device_ctx().stream) - .map_err(|err| anyhow::anyhow!("join Qwen3.5 prefill stream: {err}"))?; - - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { - anyhow::bail!("single-GPU async prefill received TP chunk state"); - }; - let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - let logits = match self.model.batch_prefill_logits_on_stream( - Arc::clone(&prefill_stream), - &window_refs, - kvs, - &mut rec_refs, - ) { - Ok(logits) => logits, - Err(err) => { - if let Err(sync_err) = prefill_stream.synchronize() { - fatal_cuda_lifecycle(&format!( - "Qwen3.5 async prefill failed ({err}); stream drain failed: {sync_err}" - )); - } - return Err(err); - } - }; - let done = match prefill_stream.record_event(None) { - Ok(done) => done, - Err(err) => { - if let Err(sync_err) = prefill_stream.synchronize() { - fatal_cuda_lifecycle(&format!( - "record Qwen3.5 async prefill event failed ({err}); stream drain failed: {sync_err}" - )); - } - return Err(anyhow::anyhow!("record Qwen3.5 async prefill event: {err}")); - } - }; - Ok(AsyncPrefillOutput { - logits: Some(logits), - done, - stream: prefill_stream, - completed: false, - }) - } - - fn unified_step( - &mut self, - chunk: &mut ScheduledChunk, - active: &mut [ActiveRequest35], - ) -> Result { - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { - anyhow::bail!("single-GPU unified step received TP chunk state"); - }; - let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); - let mut decode_kv_refs: Vec<&mut KvState> = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU unified step received TP active state") - } - }) - .collect(); - self.model.unified_step( - &window_refs, - kvs, - &mut rec_refs, - &decode_tokens, - &mut decode_kv_refs, - &mut self.graph_state, - ) - } - - fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { - let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); - let mut kv_refs: Vec<&mut KvState> = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU decode received TP active state") - } - }) - .collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) - } - - fn sample_prefill_logits( - &mut self, - pending: &[SchedulerRequest], - logits: &HiddenStates, - sample_seed: u64, - ) -> Result<(Vec, Vec>)> { - debug_assert_eq!( - logits.seq_len, - pending.len(), - "Qwen3.5 prefill logits rows must preserve pending request order" - ); - let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect(); - let cpu_logits = - snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?; - let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); - let tokens = self.model.select_tokens_from_logits_varied( - logits, - &mut self.graph_state.buffers, - ¶ms_refs, - sample_seed, - )?; - - let logprobs = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - pegainfer_sample::token_logprob_from_row( - &logits_f32, - tokens[i], - pending[i].logprobs, - ) - }) - }) - .collect(); - Ok((tokens, logprobs)) - } - - fn sample_decode_logits( - &mut self, - active: &[ActiveRequest35], - sample_seed: u64, - ) -> Result<(Vec, Vec>)> { - let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); - let cpu_logits = snapshot_requested_logprobs( - self.model.device_ctx(), - &self.graph_state.buffers.logits, - &requested_logprobs, - )?; - let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); - let tokens = self.model.select_tokens_batch_varied( - &mut self.graph_state.buffers, - ¶ms_refs, - sample_seed, - )?; - - let logprobs = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - pegainfer_sample::token_logprob_from_row( - &logits_f32, - tokens[i], - active[i].logprobs, - ) - }) - }) - .collect(); - Ok((tokens, logprobs)) - } - - fn is_stop_token(&self, token: u32) -> bool { - self.model.is_stop_token(token) - } - - fn copy_recurrent_to_slot( - &mut self, - recurrent: &RecurrentState, - slot_idx: usize, - ) -> Result<()> { - self.graph_state - .copy_state_to_slot(self.model.device_ctx(), recurrent, slot_idx) - } - - fn compact_slot(&mut self, active: &mut [ActiveRequest35], compaction: plan::SlotCompaction) { - let src_slot = match active[compaction.moved_to].backend_state { - ActiveBackendState::Single { graph_slot_idx, .. } => graph_slot_idx, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU slot compaction received TP active state") - } - }; - debug_assert_eq!(src_slot, compaction.moved_from); - - let ctx = self.model.device_ctx(); - let src = &self.graph_state.slot_states[compaction.moved_from]; - for layer_idx in 0..src.layers.len() { - let (src_part, dst_part) = if compaction.moved_to < compaction.moved_from { - let (left, right) = self - .graph_state - .slot_states - .split_at_mut(compaction.moved_from); - ( - &right[0].layers[layer_idx], - &mut left[compaction.moved_to].layers[layer_idx], - ) - } else { - unreachable!("idx < active.len() <= last"); - }; - - ctx.stream - .memcpy_dtod(&src_part.state, &mut dst_part.state) - .expect("compact slot state copy failed"); - ctx.stream - .memcpy_dtod(&src_part.conv_state.data, &mut dst_part.conv_state.data) - .expect("compact slot conv_state copy failed"); - } - self.graph_state.slot_states[compaction.moved_to].seq_len = - self.graph_state.slot_states[compaction.moved_from].seq_len; - - match &mut active[compaction.moved_to].backend_state { - ActiveBackendState::Single { graph_slot_idx, .. } => { - *graph_slot_idx = compaction.moved_to; - } - ActiveBackendState::Tp { .. } => { - panic!("single-GPU slot compaction received TP active state") - } - } - } -} - -impl TpSchedulerBackend { - fn new( - model_path: &str, - device_ordinals: &[usize], - max_batch: usize, - max_prefill_tokens: usize, - ) -> Result { - let executor = Qwen35TpExecutor::from_runtime_with_limits( - model_path, - false, - device_ordinals, - max_batch, - max_prefill_tokens, - )?; - Ok(Self { - executor, - next_request_id: 1, - }) - } - - fn alloc_request_id(&mut self) -> RequestId { - let id = RequestId::new(self.next_request_id); - self.next_request_id = self.next_request_id.wrapping_add(1).max(1); - id - } - - fn max_batch(&self) -> usize { - self.executor.max_batch() - } - - fn page_size(&self) -> usize { - self.executor.page_size() - } - - fn capacity_pages_for_requests(&self) -> usize { - self.executor.capacity_pages_for_requests() - } - - fn max_position_embeddings(&self) -> usize { - self.executor.max_position_embeddings() - } - - fn is_stop_token(&self, token: u32) -> bool { - self.executor.is_stop_token(token) - } - - fn available_pages( - &self, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - ) -> usize { - let page_size = self.page_size(); - let active_pages: usize = active - .iter() - .map(|req| pages_needed(current_active_tokens(req), page_size)) - .sum(); - let prefilling_pages: usize = prefilling - .iter() - .map(|req| pages_needed(req.cursor, page_size)) - .sum(); - self.capacity_pages_for_requests() - .saturating_sub(active_pages.saturating_add(prefilling_pages)) - } - - fn execute_prefill_chunk( - &self, - chunk: &ScheduledChunk, - sample_seed: u64, - ) -> Result>> { - let items = tp_prefill_items(chunk)?; - let result = self - .executor - .execute_prefill_chunks_with_seed(&items, sample_seed)?; - align_prefill_results(chunk, &result) - .map_err(|err| self.executor.poison_artifact_contract("prefill", &err)) - } - - fn execute_decode( - &self, - active: &[ActiveRequest35], - sample_seed: u64, - ) -> Result> { - let items = tp_decode_items(active)?; - let result = self.executor.execute_decode_items(&items, sample_seed)?; - align_decode_results(active, &result) - .map_err(|err| self.executor.poison_artifact_contract("decode", &err)) - } - - fn execute_unified( - &self, - chunk: &ScheduledChunk, - active: &[ActiveRequest35], - decode_sample_seed: u64, - prefill_sample_seed: u64, - ) -> Result { - let plan = TpUnifiedPlan { - prefill: tp_prefill_items(chunk)?, - decode: tp_decode_items(active)?, - prefill_sample_seed, - decode_sample_seed, - }; - let result = self.executor.execute_unified(&plan)?; - let prefill = align_prefill_results(chunk, &result.prefill).map_err(|err| { - self.executor - .poison_artifact_contract("unified prefill", &err) - })?; - let decode = align_decode_results(active, &result.decode).map_err(|err| { - self.executor - .poison_artifact_contract("unified decode", &err) - })?; - Ok(AlignedUnifiedArtifacts { prefill, decode }) - } - - fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { - self.executor.drop_request(request_id, expectation) - } -} - -impl SchedulerBackend { - fn max_batch(&self) -> usize { - match self { - Self::Single(backend) => backend.max_batch(), - Self::Tp(backend) => backend.max_batch(), - } - } - - fn page_size(&self) -> usize { - match self { - Self::Single(backend) => backend.page_size(), - Self::Tp(backend) => backend.page_size(), - } - } - - fn available_pages( - &self, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - ) -> usize { - match self { - Self::Single(backend) => backend.available_pages(), - Self::Tp(backend) => backend.available_pages(active, prefilling), - } - } - - fn capacity_pages_for_requests(&self) -> usize { - match self { - Self::Single(backend) => backend.capacity_pages_for_requests(), - Self::Tp(backend) => backend.capacity_pages_for_requests(), - } - } - - fn max_position_embeddings(&self) -> usize { - match self { - Self::Single(backend) => backend.max_position_embeddings(), - Self::Tp(backend) => backend.max_position_embeddings(), - } - } - - fn alloc_prefill_state(&mut self) -> Result { - match self { - Self::Single(backend) => Ok(PrefillBackendState::Single { - kv: backend.alloc_kv(), - rec: backend.alloc_recurrent()?, - }), - Self::Tp(backend) => Ok(PrefillBackendState::Tp { - request_id: backend.alloc_request_id(), - }), - } - } - - fn is_stop_token(&self, token: u32) -> bool { - match self { - Self::Single(backend) => backend.is_stop_token(token), - Self::Tp(backend) => backend.is_stop_token(token), - } - } -} - fn current_active_tokens(req: &ActiveRequest35) -> usize { req.prompt_len .saturating_add(req.generated_count.saturating_sub(1)) diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 80eafc675..e35a05d36 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -321,7 +321,7 @@ impl Qwen35TpExecutor { model_path, ModelRuntimeConfig { enable_cuda_graph: false, - tensor_parallel: Some(TensorParallelConfig { rank, world_size }), + tensor_parallel: Some(TensorParallelConfig::try_from((rank, world_size))?), device_ordinal, }, )?); diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index ddb55c34d..70cb6bb80 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -26,6 +26,7 @@ use safetensors::SafeTensors; use super::config::Config35; use super::config::LayerType; +use super::config::LocalGeometry; use super::config::TensorParallelConfig; /// Full attention layer weights (8 layers in Qwen3.5-4B). @@ -107,7 +108,7 @@ impl Default for ModelRuntimeConfig { pub struct Qwen35Model { pub(super) ctx: DeviceContext, pub(super) config: Config35, - pub(super) tensor_parallel: TensorParallelConfig, + pub(super) geometry: LocalGeometry, pub(super) embed_tokens: DeviceMatrix, lm_head: Option, pub(super) layers: Vec, @@ -208,7 +209,8 @@ impl Qwen35Model { let mut config = Config35::from_file(model_path)?; let tensor_parallel = runtime.tensor_parallel.unwrap_or_default(); - tensor_parallel.validate_for(&config, runtime.enable_cuda_graph)?; + let geometry = LocalGeometry::try_new(&config, tensor_parallel, runtime.enable_cuda_graph) + .map_err(anyhow::Error::from)?; debug!( "Config: hidden_size={}, num_layers={}, full_attn={}, linear_attn={}, max_position_embeddings={}, tp_rank={}, tp_world_size={}", config.hidden_size, @@ -216,28 +218,23 @@ impl Qwen35Model { config.num_full_attention_layers(), config.num_hidden_layers - config.num_full_attention_layers(), config.max_position_embeddings, - tensor_parallel.rank, - tensor_parallel.world_size, + geometry.rank(), + geometry.world_size(), ); let effective_vocab = super::config::tokenizer_effective_vocab(model_path)?; - anyhow::ensure!( - effective_vocab <= config.vocab_size, - "tokenizer defines ids up to {} but checkpoint vocab_size is {}", - effective_vocab - 1, - config.vocab_size, - ); - if effective_vocab < config.vocab_size { - config.selection_vocab = effective_vocab; + config + .bound_selection_vocab(effective_vocab) + .map_err(anyhow::Error::from)?; + if config.selection_vocab < config.vocab_size { info!( "output projection: selection bounded to decodable vocab {} (checkpoint pads to {})", - effective_vocab, config.vocab_size + config.selection_vocab, config.vocab_size ); } let (shard_paths, weight_map) = load_shard_info_fixed(model_path)?; debug!("Loading {} safetensor shard(s)", shard_paths.len()); - let prefetch = - (tensor_parallel.world_size == 1).then(|| WeightPrefetch::spawn(&shard_paths)); + let prefetch = (geometry.world_size() == 1).then(|| WeightPrefetch::spawn(&shard_paths)); let mmaps = mmap_shards(&shard_paths)?; let shards = deserialize_shards(&mmaps)?; @@ -279,9 +276,9 @@ impl Qwen35Model { config.num_hidden_layers ); let mut layers = Vec::with_capacity(config.num_hidden_layers); - let (_, q_rows) = tensor_parallel.shard_range(config.full_attn_q_dim()); - let (kv_row_offset, kv_rows) = tensor_parallel.shard_range(config.full_attn_kv_dim()); - let (inter_row_offset, inter_rows) = tensor_parallel.shard_range(config.intermediate_size); + let (_, q_rows) = geometry.shard_range(config.full_attn_q_dim()); + let (kv_row_offset, kv_rows) = geometry.shard_range(config.full_attn_kv_dim()); + let (inter_row_offset, inter_rows) = geometry.shard_range(config.intermediate_size); for i in 0..config.num_hidden_layers { let prefix = format!("{}.layers.{}", wp, i); let layer_type = config.layer_types[i]; @@ -296,14 +293,14 @@ impl Qwen35Model { &weight_map, &format!("{}.q_proj.weight", attn_prefix), &config, - tensor_parallel, + geometry, )?, k_proj: load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.k_proj.weight", attn_prefix), - tensor_parallel, + geometry, kv_row_offset, kv_rows, )?, @@ -312,7 +309,7 @@ impl Qwen35Model { &shards, &weight_map, &format!("{}.v_proj.weight", attn_prefix), - tensor_parallel, + geometry, kv_row_offset, kv_rows, )?, @@ -321,8 +318,8 @@ impl Qwen35Model { &shards, &weight_map, &format!("{}.o_proj.weight", attn_prefix), - tensor_parallel, - tensor_parallel.shard_range(config.full_attn_q_dim()).0, + geometry, + geometry.shard_range(config.full_attn_q_dim()).0, q_rows, )?, q_norm: load_tensor_1d( @@ -405,7 +402,7 @@ impl Qwen35Model { &shards, &weight_map, &format!("{}.mlp.gate_proj.weight", prefix), - tensor_parallel, + geometry, inter_row_offset, inter_rows, )?; @@ -414,7 +411,7 @@ impl Qwen35Model { &shards, &weight_map, &format!("{}.mlp.up_proj.weight", prefix), - tensor_parallel, + geometry, inter_row_offset, inter_rows, )?; @@ -443,7 +440,7 @@ impl Qwen35Model { &shards, &weight_map, &format!("{}.mlp.down_proj.weight", prefix), - tensor_parallel, + geometry, inter_row_offset, inter_rows, )?, @@ -491,7 +488,7 @@ impl Qwen35Model { let num_full_layers = config.num_full_attention_layers(); let layout = pegainfer_core::kv_pool::KvLayout::new( num_full_layers, - config.local_num_key_value_heads(tensor_parallel), + geometry.local_num_key_value_heads(), config.head_dim, page_size, ) @@ -532,7 +529,7 @@ impl Qwen35Model { let kv_pool = pegainfer_core::kv_pool::KvPool::new( &ctx, num_full_layers, - config.local_num_key_value_heads(tensor_parallel), + geometry.local_num_key_value_heads(), config.head_dim, page_size, num_pages, @@ -541,7 +538,7 @@ impl Qwen35Model { Ok(Self { ctx, config, - tensor_parallel, + geometry, embed_tokens, lm_head, layers, @@ -610,13 +607,13 @@ impl Qwen35Model { let ctx = &self.ctx; let hidden = self.config.hidden_size; let vocab = self.config.selection_vocab; - let tp = self.tensor_parallel; - let full_q = self.config.local_full_attn_gated_q_dim(tp); - let full_kv = self.config.local_full_attn_kv_dim(tp); + let geom = self.geometry; + let full_q = geom.local_full_attn_gated_q_dim(); + let full_kv = geom.local_full_attn_kv_dim(); let linear_qkv = self.config.linear_attn_qkv_dim(); let linear_z = self.config.linear_attn_z_dim(); let linear_ba = self.config.linear_num_value_heads; - let intermediate = self.config.local_intermediate_size(tp); + let intermediate = geom.local_intermediate_size(); let full_q_samples: Vec<_> = self .layers @@ -732,7 +729,7 @@ impl Qwen35Model { super::batch_decode_graph::BatchDecodeGraphState::with_capacity( &self.ctx, &self.config, - self.tensor_parallel, + self.geometry, &self.kv_pool, max_batch, ) @@ -745,7 +742,7 @@ impl Qwen35Model { super::decode_buffers::BatchDecodeBuffers35::new( &self.ctx, &self.config, - self.tensor_parallel, + self.geometry, max_batch, self.kv_pool.capacity_pages(), self.kv_pool.padding_page_id(), @@ -777,12 +774,12 @@ struct GatedQShardRange { fn full_attention_gated_q_shard_range( config: &Config35, - tensor_parallel: TensorParallelConfig, + geometry: LocalGeometry, ) -> GatedQShardRange { // HF/PegaInfer kernels interpret q_proj rows as per-head [q, gate] chunks. // Keep each local head's q rows adjacent to its gate rows. - let local_heads = config.local_num_attention_heads(tensor_parallel); - let head_start = tensor_parallel.rank * local_heads; + let local_heads = geometry.local_num_attention_heads(); + let head_start = geometry.rank() * local_heads; GatedQShardRange { row_offset: head_start * config.head_dim * 2, rows: local_heads * config.head_dim * 2, @@ -795,13 +792,13 @@ fn load_full_attention_gated_q_proj( weight_map: &HashMap, name: &str, config: &Config35, - tensor_parallel: TensorParallelConfig, + geometry: LocalGeometry, ) -> Result { - if !tensor_parallel.is_sharded() { + if !geometry.is_sharded() { return load_tensor_2d(ctx, shards, weight_map, name); } - let range = full_attention_gated_q_shard_range(config, tensor_parallel); + let range = full_attention_gated_q_shard_range(config, geometry); load_tensor_2d_row_shard(ctx, shards, weight_map, name, range.row_offset, range.rows) } @@ -810,11 +807,11 @@ fn load_tensor_2d_row_shard_if_needed( shards: &[SafeTensors], weight_map: &HashMap, name: &str, - tensor_parallel: TensorParallelConfig, + geometry: LocalGeometry, row_offset: usize, rows: usize, ) -> Result { - if tensor_parallel.is_sharded() { + if geometry.is_sharded() { load_tensor_2d_row_shard(ctx, shards, weight_map, name, row_offset, rows) } else { load_tensor_2d(ctx, shards, weight_map, name) @@ -826,11 +823,11 @@ fn load_tensor_2d_col_shard_if_needed( shards: &[SafeTensors], weight_map: &HashMap, name: &str, - tensor_parallel: TensorParallelConfig, + geometry: LocalGeometry, col_offset: usize, cols: usize, ) -> Result { - if tensor_parallel.is_sharded() { + if geometry.is_sharded() { load_tensor_2d_col_shard(ctx, shards, weight_map, name, col_offset, cols) } else { load_tensor_2d(ctx, shards, weight_map, name) @@ -842,41 +839,47 @@ mod tests { use super::*; fn test_config() -> Config35 { - Config35 { - hidden_size: 2560, - intermediate_size: 9216, - num_hidden_layers: 32, - vocab_size: 248_320, - selection_vocab: 248_320, - rms_norm_eps: 1e-6, - eos_token_id: 151_645, - num_attention_heads: 16, - num_key_value_heads: 4, - head_dim: 256, - linear_num_key_heads: 16, - linear_key_head_dim: 128, - linear_num_value_heads: 32, - linear_value_head_dim: 128, - linear_conv_kernel_dim: 4, - rope_theta: 10_000.0, - rotary_dim: 64, - max_position_embeddings: 262_144, - tie_word_embeddings: true, - layer_types: vec![LayerType::LinearAttention; 32], - } + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{ + "max_position_embeddings": 262144, + "tie_word_embeddings": true, + "text_config": { + "hidden_size": 2560, + "intermediate_size": 9216, + "num_hidden_layers": 1, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 256, + "vocab_size": 248320, + "rms_norm_eps": 1e-6, + "layer_types": ["linear_attention"], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "rope_parameters": { "rope_theta": 10000.0, "partial_rotary_factor": 0.25 }, + "eos_token_id": 151645 + } +}"#, + ) + .unwrap(); + Config35::from_file(dir.path().to_str().unwrap()).expect("fixture validates") + } + + fn test_geometry(rank: usize, world_size: usize) -> LocalGeometry { + let config = test_config(); + let tp = TensorParallelConfig::try_from((rank, world_size)).unwrap(); + LocalGeometry::try_new(&config, tp, false).unwrap() } #[test] fn gated_q_shard_range_keeps_matching_q_and_gate_rows() { let config = test_config(); - let rank0 = full_attention_gated_q_shard_range( - &config, - TensorParallelConfig { - rank: 0, - world_size: 2, - }, - ); + let rank0 = full_attention_gated_q_shard_range(&config, test_geometry(0, 2)); assert_eq!( rank0, GatedQShardRange { @@ -885,13 +888,7 @@ mod tests { } ); - let rank1 = full_attention_gated_q_shard_range( - &config, - TensorParallelConfig { - rank: 1, - world_size: 2, - }, - ); + let rank1 = full_attention_gated_q_shard_range(&config, test_geometry(1, 2)); assert_eq!( rank1, GatedQShardRange { @@ -904,14 +901,11 @@ mod tests { #[test] fn mlp_tp2_uses_matching_gate_up_rows_and_down_cols() { let config = test_config(); - let tp = TensorParallelConfig { - rank: 1, - world_size: 2, - }; + let geom = test_geometry(1, 2); - let (inter_offset, inter_rows) = tp.shard_range(config.intermediate_size); + let (inter_offset, inter_rows) = geom.shard_range(config.intermediate_size); assert_eq!((inter_offset, inter_rows), (4608, 4608)); - assert_eq!(config.local_intermediate_size(tp), inter_rows); + assert_eq!(geom.local_intermediate_size(), inter_rows); let local_gate_up_rows = 2 * inter_rows; let local_down_cols = inter_rows;