From 536aa1a6ed72cbeaf0b54adaf40dc61b99b805dd Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Sun, 5 Jul 2026 12:46:52 -0400 Subject: [PATCH] feat(inference): parse Qwen3.5 vision config fields (ADR-069 S1) Parse and store the vision-language checkpoint fields that were previously dropped by qwen35_config.rs, without wiring them into the forward pass: - New VisionConfig struct (ViT depth/hidden_size/heads/patch geometry etc.), parsed from the top-level vision_config object; None for text-only checkpoints. - Four vision token ids (image/video/vision_start/vision_end), threaded from the top-level config.json the same way tie_word_embeddings already is. - RopeParams gains mrope_section and mrope_interleaved, parsed from text_config.rope_parameters alongside the existing rope_theta and partial_rotary_factor handling (unchanged). All five new Qwen35Config fields default to None and are additive; every preset constructor and existing config.json parse path is unaffected. Adding them to the struct required a mechanical None-field addition at existing exhaustive struct-literal test helpers across the crate (no logic changes) so the crate keeps compiling under default, metal-gpu, and all-features builds. Round-trips the real 0.8B checkpoint fixture (vision_config dims, all four token ids, mrope_section == [11, 11, 10], mrope_interleaved == true) and adds a text-only-config negative test proving the plain decoder path is unaffected. Consumption by the forward pass (ViT encoder, merger, decoder M-RoPE) lands in later ADR-069 stages. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/inference/src/forward/batch_prefill.rs | 5 + crates/inference/src/forward/cpu_f16.rs | 15 ++ crates/inference/src/forward/cpu_q8.rs | 20 ++ crates/inference/src/forward/metal_qwen35.rs | 20 ++ crates/inference/src/forward/neon_forward.rs | 30 +++ crates/inference/src/model/qwen35/eval.rs | 5 + .../inference/src/model/qwen35/generation.rs | 10 + crates/inference/src/model/qwen35/tests.rs | 5 + crates/inference/src/model/qwen35_config.rs | 185 +++++++++++++++++- crates/inference/src/stop_token_contract.rs | 5 + 10 files changed, 296 insertions(+), 4 deletions(-) diff --git a/crates/inference/src/forward/batch_prefill.rs b/crates/inference/src/forward/batch_prefill.rs index 36d9ea78ed..8f115c0a46 100644 --- a/crates/inference/src/forward/batch_prefill.rs +++ b/crates/inference/src/forward/batch_prefill.rs @@ -1663,6 +1663,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } diff --git a/crates/inference/src/forward/cpu_f16.rs b/crates/inference/src/forward/cpu_f16.rs index 3cb58dc81c..f8c696a5bf 100644 --- a/crates/inference/src/forward/cpu_f16.rs +++ b/crates/inference/src/forward/cpu_f16.rs @@ -1120,6 +1120,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope_dim = cfg.rope_dim(); // = 16 @@ -1505,6 +1510,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; // embed_tokens is [vocab * hidden] packed u16 (f16 zeros = 0u16). @@ -1652,6 +1662,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; // The negative gamma at dim-0 flips the sign of that component after diff --git a/crates/inference/src/forward/cpu_q8.rs b/crates/inference/src/forward/cpu_q8.rs index cf6e8713ad..e17ee3c594 100644 --- a/crates/inference/src/forward/cpu_q8.rs +++ b/crates/inference/src/forward/cpu_q8.rs @@ -934,6 +934,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope_dim = cfg.rope_dim(); // = 16 @@ -1204,6 +1209,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope = RopeTable::new(cfg.rope_dim(), 512, cfg.rope_theta); @@ -1377,6 +1387,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; // embed_tokens is [vocab, hidden] and also serves as the tied LM head. @@ -1620,6 +1635,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope_dim = (head_dim as f32 * cfg.partial_rotary_factor) as usize; // 16 diff --git a/crates/inference/src/forward/metal_qwen35.rs b/crates/inference/src/forward/metal_qwen35.rs index 0a2aec20a9..012708099a 100644 --- a/crates/inference/src/forward/metal_qwen35.rs +++ b/crates/inference/src/forward/metal_qwen35.rs @@ -289,6 +289,11 @@ mod mtp_resolve_tests { eos_token_id: 63, max_position_embeddings: 128, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } @@ -19617,6 +19622,11 @@ kernel void decode_attention_reference( eos_token_id: (vocab - 1) as u32, max_position_embeddings: 128, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let mut embed_tokens = vec![0.0f32; vocab * hidden]; @@ -22375,6 +22385,11 @@ kernel void decode_attention_reference( eos_token_id: (vocab - 1) as u32, max_position_embeddings: 128, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let make_common = || CommonLayerWeights { @@ -22721,6 +22736,11 @@ kernel void decode_attention_reference( eos_token_id: (vocab - 1) as u32, max_position_embeddings: 128, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let embed_tokens = vec![0.0f32; vocab * hidden]; diff --git a/crates/inference/src/forward/neon_forward.rs b/crates/inference/src/forward/neon_forward.rs index 96381eadd2..11ea390ad9 100644 --- a/crates/inference/src/forward/neon_forward.rs +++ b/crates/inference/src/forward/neon_forward.rs @@ -1153,6 +1153,11 @@ pub mod bench_support { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope_dim = (head_dim as f32 * cfg.partial_rotary_factor) as usize; // 32 @@ -1832,6 +1837,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope_dim = cfg.rope_dim(); // = 16 @@ -2025,6 +2035,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope_dim = (head_dim as f32 * cfg.partial_rotary_factor) as usize; // 16 @@ -2292,6 +2307,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let rope = RopeTable::new(cfg.rope_dim(), 512, cfg.rope_theta); @@ -2529,6 +2549,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; // All-zero packed weights: scale=0, all i8=0 → logits all 0 → greedy picks 0. @@ -2632,6 +2657,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; // The negative gamma at dim-0 creates a "bounce" each decode step. diff --git a/crates/inference/src/model/qwen35/eval.rs b/crates/inference/src/model/qwen35/eval.rs index 3fc2451a3c..a040f7442a 100644 --- a/crates/inference/src/model/qwen35/eval.rs +++ b/crates/inference/src/model/qwen35/eval.rs @@ -545,6 +545,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } diff --git a/crates/inference/src/model/qwen35/generation.rs b/crates/inference/src/model/qwen35/generation.rs index 435da89106..da705d8adb 100644 --- a/crates/inference/src/model/qwen35/generation.rs +++ b/crates/inference/src/model/qwen35/generation.rs @@ -1302,6 +1302,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; // Deterministic xorshift for reproducible synthetic weights. @@ -1524,6 +1529,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, }; let z = |len: usize| vec![0.0_f32; len]; diff --git a/crates/inference/src/model/qwen35/tests.rs b/crates/inference/src/model/qwen35/tests.rs index 7f5438ee24..fcae56dacb 100644 --- a/crates/inference/src/model/qwen35/tests.rs +++ b/crates/inference/src/model/qwen35/tests.rs @@ -482,6 +482,11 @@ mod lora_serving { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } diff --git a/crates/inference/src/model/qwen35_config.rs b/crates/inference/src/model/qwen35_config.rs index 299246c6a3..7c054bb42d 100644 --- a/crates/inference/src/model/qwen35_config.rs +++ b/crates/inference/src/model/qwen35_config.rs @@ -42,6 +42,37 @@ pub enum LayerType { FullAttention, } +/// **Unstable**: ViT geometry for the vision-language checkpoint variants, parsed from the +/// top-level `vision_config` object in `config.json`. `None` on [`Qwen35Config`] for text-only +/// checkpoints (e.g. the bare 2B/35B-A3B text decoders). +/// +/// Parsed but not yet consumed (ADR-069 stage S1) — the Metal ViT forward pass, patch merger, +/// and image-token injection land in later stages. +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct VisionConfig { + /// Number of ViT transformer blocks. + pub depth: usize, + /// ViT hidden dimension. + pub hidden_size: usize, + /// ViT attention head count. + pub num_heads: usize, + /// Patch size in pixels (square patches). + pub patch_size: usize, + /// Spatial merge factor applied by the patch merger before projecting into decoder space. + pub spatial_merge_size: usize, + /// Merger output dimension (equals the text decoder `hidden_size` for this checkpoint). + pub out_hidden_size: usize, + /// Temporal patch size (frames per patch) for video input. + pub temporal_patch_size: usize, + /// Learned position-embedding table size. + pub num_position_embeddings: usize, + /// Input image channel count. + pub in_channels: usize, + /// DeepStack visual layer indexes; empty for checkpoints that don't use DeepStack fusion. + #[serde(default)] + pub deepstack_visual_indexes: Vec, +} + /// **Unstable**: Qwen hybrid attention model configuration; fields evolving with model variants. #[derive(Debug, Clone, serde::Deserialize)] #[serde(default)] @@ -122,6 +153,24 @@ pub struct Qwen35Config { /// preset's context length. #[serde(default = "default_max_position_embeddings")] pub max_position_embeddings: usize, + + // --- Vision-language extension (ADR-069 S1) --- + /// ViT geometry from the top-level `vision_config` object; `None` for text-only checkpoints. + /// Parsed but not yet consumed by the forward pass (S3/S5 wire it in). + #[serde(default)] + pub vision_config: Option, + /// Placeholder token id marking an image-patch span in the input sequence. + #[serde(default)] + pub image_token_id: Option, + /// Placeholder token id marking a video-frame span in the input sequence. + #[serde(default)] + pub video_token_id: Option, + /// Token id opening a vision content span (wraps image/video token runs). + #[serde(default)] + pub vision_start_token_id: Option, + /// Token id closing a vision content span. + #[serde(default)] + pub vision_end_token_id: Option, } fn default_tie_word_embeddings() -> bool { @@ -140,15 +189,38 @@ pub struct RopeParams { pub rope_theta: f64, #[serde(default)] pub partial_rotary_factor: Option, + /// Per-axis (temporal, height, width) M-RoPE section sizes, e.g. `[11, 11, 10]`. + /// `None` when the checkpoint has no vision M-RoPE config (text-only decoders use plain + /// 1-D partial RoPE). Parsed but not yet consumed by the forward pass (ADR-069 S1; wired + /// in S5). + #[serde(default)] + pub mrope_section: Option>, + /// Whether M-RoPE frequencies are interleaved `[T, H, W, T, H, W, ...]` (true for + /// Qwen3.5) rather than block-concatenated. `None` when absent (text-only decoders). + #[serde(default)] + pub mrope_interleaved: Option, } -// Private helper for HF config.json structure (outer wrapper with text_config). +// Private helper for HF config.json structure (outer wrapper with text_config). The +// vision-language fields are siblings of text_config at the top level of config.json, not +// nested inside it, so they are threaded onto the parsed Qwen35Config in from_config_json_str +// the same way tie_word_embeddings already is. #[derive(Debug, serde::Deserialize)] struct HfQwenConfigFile { #[serde(default)] text_config: Option, #[serde(default)] tie_word_embeddings: Option, + #[serde(default)] + vision_config: Option, + #[serde(default)] + image_token_id: Option, + #[serde(default)] + video_token_id: Option, + #[serde(default)] + vision_start_token_id: Option, + #[serde(default)] + vision_end_token_id: Option, } impl Default for Qwen35Config { @@ -203,12 +275,20 @@ impl Qwen35Config { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + // Vision-language extension (this preset is the bare text decoder, no vision tower) + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } /// **Unstable**: Qwen3.5-0.8B configuration. Same hybrid architecture as the /// 2B (24 layers, `[linear, linear, linear, full] x 6`), scaled down. The - /// released checkpoint is a vision-language model; this is its text decoder. + /// released checkpoint is a vision-language model; this preset is its bare text decoder + /// (no vision_config) — parse the real config.json via `from_config_json` to get the + /// vision-language fields (ADR-069 S1). pub fn qwen35_0_8b() -> Self { let num_hidden_layers = 24; let full_attention_interval = 4; @@ -253,6 +333,12 @@ impl Qwen35Config { mtp_num_hidden_layers: 1, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + // Vision-language extension (this preset is the bare text decoder, no vision tower) + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } @@ -295,6 +381,12 @@ impl Qwen35Config { mtp_num_hidden_layers: 1, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + // Vision-language extension (this preset is text-only, no vision tower) + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } @@ -345,6 +437,12 @@ impl Qwen35Config { mtp_num_hidden_layers: 1, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + // Vision-language extension (this preset is text-only, no vision tower) + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } } @@ -365,6 +463,14 @@ impl Qwen35Config { if let Some(tie) = parsed.tie_word_embeddings { cfg.tie_word_embeddings = tie; } + // Vision-language fields (ADR-069 S1) are siblings of text_config at the top level of + // config.json, not nested inside it — thread them onto cfg the same way as + // tie_word_embeddings above. Parsed but not yet consumed by the forward pass. + cfg.vision_config = parsed.vision_config; + cfg.image_token_id = parsed.image_token_id; + cfg.video_token_id = parsed.video_token_id; + cfg.vision_start_token_id = parsed.vision_start_token_id; + cfg.vision_end_token_id = parsed.vision_end_token_id; // Many models nest rope_theta and partial_rotary_factor under rope_parameters // instead of at the text_config level — extract when the flat fields are unset. if let Some(rp) = &cfg.rope_parameters { @@ -1530,8 +1636,9 @@ mod tests { assert_eq!(cfg.max_position_embeddings, 262_144); assert_eq!(cfg.mtp_num_hidden_layers, 1); - // Released checkpoint is a dense vision-language model — neither the MoE - // fields nor the vision wrapper may leak into the text config. + // Released checkpoint is a dense vision-language model: the MoE fields must still + // resolve to None (0.8B has no MoE), while the vision_config / token-id fields below + // are now parsed onto the config (ADR-069 S1) instead of being dropped. assert!(!cfg.is_moe(), "Qwen3.5-0.8B is dense, not MoE"); // rope_theta and partial_rotary_factor are nested under rope_parameters @@ -1555,6 +1662,76 @@ mod tests { // tie_word_embeddings is taken from the outer wrapper. assert!(cfg.tie_word_embeddings); + + // ADR-069 S1: vision_config, the four vision token ids, and the M-RoPE section + // fields are top-level / text_config.rope_parameters siblings, respectively, and + // must now round-trip onto the config (parsed but not yet consumed by forward). + let vision = cfg + .vision_config + .as_ref() + .expect("released checkpoint has a vision_config"); + assert_eq!(vision.depth, 12); + assert_eq!(vision.hidden_size, 768); + assert_eq!(vision.num_heads, 12); + assert_eq!(vision.patch_size, 16); + assert_eq!(vision.spatial_merge_size, 2); + assert_eq!(vision.out_hidden_size, 1024); + assert_eq!(vision.temporal_patch_size, 2); + assert_eq!(vision.num_position_embeddings, 2304); + assert_eq!(vision.in_channels, 3); + assert!(vision.deepstack_visual_indexes.is_empty()); + + assert_eq!(cfg.image_token_id, Some(248_056)); + assert_eq!(cfg.video_token_id, Some(248_057)); + assert_eq!(cfg.vision_start_token_id, Some(248_053)); + assert_eq!(cfg.vision_end_token_id, Some(248_054)); + + let rope_params = cfg + .rope_parameters + .as_ref() + .expect("released checkpoint nests rope_parameters under text_config"); + assert_eq!(rope_params.mrope_section, Some(vec![11, 11, 10])); + assert_eq!(rope_params.mrope_interleaved, Some(true)); + } + + #[test] + fn test_text_only_config_has_no_vision_fields() { + // A text-only config.json (no vision_config, no token ids, no mrope fields) must + // still parse cleanly, with every vision-language field resolving to None — proving + // the ADR-069 S1 additions don't regress the plain-text decoder path. + let json = r#"{ + "text_config": { + "hidden_size": 1024, + "num_hidden_layers": 4, + "vocab_size": 1000, + "intermediate_size": 2048, + "rms_norm_eps": 1e-6, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "head_dim": 256, + "rope_theta": 10000000.0, + "partial_rotary_factor": 0.25, + "linear_num_key_heads": 16, + "linear_num_value_heads": 16, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "full_attention_interval": 4, + "eos_token_id": 999, + "max_position_embeddings": 4096 + } + }"#; + let cfg = + Qwen35Config::from_config_json_str(json).expect("text-only config.json still parses"); + + assert!(cfg.vision_config.is_none()); + assert!(cfg.image_token_id.is_none()); + assert!(cfg.video_token_id.is_none()); + assert!(cfg.vision_start_token_id.is_none()); + assert!(cfg.vision_end_token_id.is_none()); + // No rope_parameters object at all in this fixture (rope_theta/partial_rotary_factor + // are flat), so mrope_section has nothing to nest under either. + assert!(cfg.rope_parameters.is_none()); } // ────────────────────────────────────────────────────────────────────── diff --git a/crates/inference/src/stop_token_contract.rs b/crates/inference/src/stop_token_contract.rs index 8e2992f048..4950671360 100644 --- a/crates/inference/src/stop_token_contract.rs +++ b/crates/inference/src/stop_token_contract.rs @@ -150,6 +150,11 @@ mod tests { mtp_num_hidden_layers: 0, mtp_use_dedicated_embeddings: false, quarot_rotation_seed: None, + vision_config: None, + image_token_id: None, + video_token_id: None, + vision_start_token_id: None, + vision_end_token_id: None, } }