use super::{ bf16_word_to_f32, fnv1a64_u32_words, gemma4_qproj_case_input_bf16_words, MlxDType, MlxFamilySnapshot, MlxIndexedSafetensors, MlxModelFamily, MlxModelManifest, MlxModelSnapshot, MlxQwen35MoeIndexedSafetensors, MlxQwen35MoeLayerKind, MlxQwen35MoeSnapshot, MlxQwen35MoeTensors, MlxSafetensorsHeader, MlxTokenizer, GEMMA4_QPROJ_CASE_INNER_DIM, GEMMA4_QPROJ_CASE_OUTPUT_DIM, GEMMA4_QPROJ_CASE_OUTPUT_FNV1A64, }; use crate::chat::{MlxChatDecodeMode, MlxChatRole, MlxChatSession}; use crate::{ extract_qwen35moe_assistant_response_text, format_qwen35moe_chat_prompt_with_options, MlxModelLayerRole, MlxQwen35MoeRuntimeSession, QwenChatMessage, QwenChatPromptOptions, QwenChatRole, QwenThinkingMode, }; use std::fs; use std::path::PathBuf; use std::time::{SystemTime, UNIX_EPOCH}; fn local_model_dir() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../local/models/gemma-4-26b-mlx") } fn local_model_shard_1() -> PathBuf { local_model_dir().join("model-00001-of-00003.safetensors") } fn temp_model_dir(prefix: &str) -> PathBuf { let stamp = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_nanos(); let path = std::env::temp_dir().join(format!("makepad_mlx_{prefix}_{stamp}")); fs::create_dir_all(&path).unwrap(); path } fn write_text(path: &PathBuf, text: &str) { fs::write(path, text).unwrap(); } fn write_fake_safetensors(path: &PathBuf, names: &[String]) { let mut header = String::from("{\"__metadata__\":{\"format\":\"mlx\"}"); let mut payload = Vec::with_capacity(names.len() * 2); let mut offset = 0u64; for name in names { header.push_str(&format!( ",\"{}\":{{\"dtype\":\"BF16\",\"shape\":[1],\"data_offsets\":[{},{}]}}", name, offset, offset + 2 )); payload.extend_from_slice(&[0u8, 0u8]); offset += 2; } header.push('}'); let mut bytes = Vec::with_capacity(8 + header.len() + payload.len()); bytes.extend_from_slice(&(header.len() as u64).to_le_bytes()); bytes.extend_from_slice(header.as_bytes()); bytes.extend_from_slice(&payload); fs::write(path, bytes).unwrap(); } fn write_qwen35moe_test_model(root: &PathBuf) { write_text( &root.join("config.json"), r#"{ "architectures":["Qwen3_5MoeForConditionalGeneration"], "image_token_id":9, "model_type":"qwen3_5_moe", "quantization":{ "bits":4, "group_size":64, "mode":"affine", "model.language_model.layers.0.linear_attn.in_proj_qkv":{ "bits":8, "group_size":64, "mode":"affine" } }, "text_config":{ "attention_bias":false, "attention_dropout":0.0, "attn_output_gate":true, "bos_token_id":1, "dtype":"bfloat16", "eos_token_id":2, "full_attention_interval":4, "head_dim":256, "hidden_act":"silu", "hidden_size":2048, "initializer_range":0.02, "layer_types":["linear_attention","linear_attention","linear_attention","full_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, "mamba_ssm_dtype":"float32", "max_position_embeddings":262144, "model_type":"qwen3_5_moe_text", "moe_intermediate_size":512, "mtp_num_hidden_layers":1, "mtp_use_dedicated_embeddings":false, "num_attention_heads":16, "num_experts":256, "num_experts_per_tok":8, "num_hidden_layers":4, "num_key_value_heads":2, "output_router_logits":false, "pad_token_id":null, "partial_rotary_factor":0.25, "rms_norm_eps":0.000001, "rope_parameters":{ "mrope_interleaved":true, "mrope_section":[11,11,10], "partial_rotary_factor":0.25, "rope_theta":10000000.0, "rope_type":"default" }, "router_aux_loss_coef":0.001, "shared_expert_intermediate_size":512, "tie_word_embeddings":false, "use_cache":true, "vocab_size":248320 }, "tie_word_embeddings":false, "transformers_version":"4.52.3", "video_token_id":10, "vision_config":{ "deepstack_visual_indexes":[], "depth":1, "hidden_act":"gelu_pytorch_tanh", "hidden_size":1152, "in_channels":3, "initializer_range":0.02, "intermediate_size":4304, "model_type":"qwen3_5_moe", "num_heads":16, "num_position_embeddings":2304, "out_hidden_size":2048, "patch_size":16, "spatial_merge_size":2, "temporal_patch_size":2 }, "vision_end_token_id":11, "vision_start_token_id":12 }"#, ); write_text( &root.join("generation_config.json"), r#"{ "bos_token_id":1, "do_sample":true, "eos_token_id":[2], "pad_token_id":1, "temperature":1.0, "top_k":20, "top_p":0.95, "transformers_version":"4.52.3" }"#, ); write_text( &root.join("preprocessor_config.json"), r#"{ "size":{"longest_edge":16777216,"shortest_edge":65536}, "patch_size":16, "temporal_patch_size":2, "merge_size":2, "image_mean":[0.5,0.5,0.5], "image_std":[0.5,0.5,0.5], "processor_class":"Qwen3VLProcessor", "image_processor_type":"Qwen2VLImageProcessorFast" }"#, ); write_text( &root.join("tokenizer_config.json"), r#"{ "tokenizer_class":"Qwen2Tokenizer", "model_max_length":262144, "extra_special_tokens":{ "image_token":"<|image_pad|>", "vision_bos_token":"<|vision_start|>", "vision_eos_token":"<|vision_end|>" } }"#, ); write_text( &root.join("tokenizer.json"), r#"{ "normalizer":{"type":"NFC"}, "pre_tokenizer":{"type":"Sequence"}, "decoder":{"type":"ByteLevel"}, "model":{ "type":"BPE", "byte_fallback":false, "vocab":{"Ġ":0,"h":1,"i":2,"Ġh":3,"Ġhi":4}, "merges":["Ġ h","Ġh i"] }, "added_tokens":[{"id":5,"content":"<|image_pad|>","special":true}] }"#, ); let mut names = vec![ "lm_head.weight".to_string(), "model.language_model.embed_tokens.weight".to_string(), "model.language_model.norm.weight".to_string(), "model.visual.patch_embed.proj.weight".to_string(), "model.visual.patch_embed.proj.bias".to_string(), "model.visual.pos_embed.weight".to_string(), "model.visual.blocks.0.norm1.weight".to_string(), "model.visual.blocks.0.norm1.bias".to_string(), "model.visual.blocks.0.attn.qkv.weight".to_string(), "model.visual.blocks.0.attn.qkv.bias".to_string(), "model.visual.blocks.0.attn.proj.weight".to_string(), "model.visual.blocks.0.attn.proj.bias".to_string(), "model.visual.blocks.0.norm2.weight".to_string(), "model.visual.blocks.0.norm2.bias".to_string(), "model.visual.blocks.0.mlp.linear_fc1.weight".to_string(), "model.visual.blocks.0.mlp.linear_fc1.bias".to_string(), "model.visual.blocks.0.mlp.linear_fc2.weight".to_string(), "model.visual.blocks.0.mlp.linear_fc2.bias".to_string(), "model.visual.merger.norm.weight".to_string(), "model.visual.merger.norm.bias".to_string(), "model.visual.merger.linear_fc1.weight".to_string(), "model.visual.merger.linear_fc1.bias".to_string(), "model.visual.merger.linear_fc2.weight".to_string(), "model.visual.merger.linear_fc2.bias".to_string(), ]; for layer in 0..4u32 { names.push(format!( "model.language_model.layers.{layer}.input_layernorm.weight" )); names.push(format!( "model.language_model.layers.{layer}.post_attention_layernorm.weight" )); names.push(format!( "model.language_model.layers.{layer}.mlp.gate.weight" )); names.push(format!( "model.language_model.layers.{layer}.mlp.experts.gate_up_proj" )); names.push(format!( "model.language_model.layers.{layer}.mlp.experts.down_proj" )); names.push(format!( "model.language_model.layers.{layer}.mlp.shared_expert_gate.weight" )); names.push(format!( "model.language_model.layers.{layer}.mlp.shared_expert.gate_proj.weight" )); names.push(format!( "model.language_model.layers.{layer}.mlp.shared_expert.up_proj.weight" )); names.push(format!( "model.language_model.layers.{layer}.mlp.shared_expert.down_proj.weight" )); if layer == 3 { names.push(format!( "model.language_model.layers.{layer}.self_attn.q_proj.weight" )); names.push(format!( "model.language_model.layers.{layer}.self_attn.k_proj.weight" )); names.push(format!( "model.language_model.layers.{layer}.self_attn.v_proj.weight" )); names.push(format!( "model.language_model.layers.{layer}.self_attn.o_proj.weight" )); names.push(format!( "model.language_model.layers.{layer}.self_attn.q_norm.weight" )); names.push(format!( "model.language_model.layers.{layer}.self_attn.k_norm.weight" )); } else { names.push(format!( "model.language_model.layers.{layer}.linear_attn.in_proj_qkv.weight" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.in_proj_z.weight" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.conv1d.weight" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.dt_bias" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.A_log" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.in_proj_a.weight" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.in_proj_b.weight" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.norm.weight" )); names.push(format!( "model.language_model.layers.{layer}.linear_attn.out_proj.weight" )); } } names.sort(); let mut weight_index = String::from("{\"metadata\":{\"total_size\":1},\"weight_map\":{"); for (index, name) in names.iter().enumerate() { if index != 0 { weight_index.push(','); } weight_index.push_str(&format!( "\"{}\":\"model-00001-of-00001.safetensors\"", name )); } weight_index.push_str("}}"); write_text(&root.join("model.safetensors.index.json"), &weight_index); write_fake_safetensors(&root.join("model-00001-of-00001.safetensors"), &names); } #[test] fn loads_local_gemma4_mlx_snapshot() { let snapshot = MlxModelSnapshot::load(local_model_dir()).unwrap(); assert_eq!( snapshot.config.architectures, vec!["Gemma4ForConditionalGeneration".to_string()] ); assert_eq!(snapshot.config.model_type, "gemma4"); assert_eq!(snapshot.config.text_config.num_hidden_layers, 30); assert_eq!(snapshot.config.vision_config.num_hidden_layers, 27); assert_eq!(snapshot.config.quantization.bits, 4); assert_eq!(snapshot.config.quantization.group_size, 64); assert_eq!(snapshot.config.quantization.mode, "affine"); assert_eq!(snapshot.processor_config.image_processor.size.height, 224); assert_eq!(snapshot.processor_config.image_processor.size.width, 224); assert_eq!( snapshot.tokenizer_config.model_max_length, 1000000000000000019884624838656u128 ); assert_eq!(snapshot.weight_index.metadata.total_size, 15_335_574_684); assert_eq!(snapshot.unique_weight_shards().len(), 3); assert_eq!( snapshot .weight_index .weight_map .get("language_model.model.layers.0.self_attn.q_proj.weight") .map(String::as_str), Some("model-00001-of-00003.safetensors") ); assert_eq!( snapshot .weight_index .weight_map .get("embed_vision.embedding_projection.weight") .map(String::as_str), Some("model-00003-of-00003.safetensors") ); } #[test] fn reads_local_safetensors_header_without_touching_payload() { let header = MlxSafetensorsHeader::load(local_model_shard_1()).unwrap(); assert_eq!(header.header_len, 64_065); assert_eq!( header.metadata.get("format").map(String::as_str), Some("mlx") ); assert_eq!(header.tensors.len(), 488); let embed_weight = header .tensor("language_model.model.embed_tokens.weight") .unwrap(); assert_eq!(embed_weight.dtype, MlxDType::U32); assert_eq!(embed_weight.shape, vec![262_144, 352]); assert_eq!(embed_weight.data_offsets, [3_116_339_724, 3_485_438_476]); assert_eq!( embed_weight.data_len_bytes(), embed_weight.expected_len_bytes() ); let q_proj = header .tensor("language_model.model.layers.0.self_attn.q_proj.weight") .unwrap(); assert_eq!(q_proj.dtype, MlxDType::U32); assert_eq!(q_proj.shape, vec![4_096, 352]); assert_eq!(q_proj.data_offsets, [3_612_518_924, 3_618_286_092]); assert_eq!(q_proj.data_len_bytes(), q_proj.expected_len_bytes()); let q_proj_scales = header .tensor("language_model.model.layers.0.self_attn.q_proj.scales") .unwrap(); assert_eq!(q_proj_scales.dtype, MlxDType::BF16); assert_eq!(q_proj_scales.shape, vec![4_096, 44]); assert_eq!(q_proj_scales.data_offsets, [2_536_973_576, 2_537_334_024]); assert_eq!( q_proj_scales.file_offsets(header.payload_base_offset())[0], 2_537_037_649 ); } #[test] fn indexed_safetensors_resolves_late_layer_to_correct_shard() { let indexed = MlxIndexedSafetensors::load(local_model_dir()).unwrap(); assert_eq!( indexed .shard_name_for_tensor("language_model.model.layers.29.self_attn.q_proj.weight") .unwrap(), "model-00003-of-00003.safetensors" ); let header = indexed .header_for_tensor("language_model.model.layers.29.self_attn.q_proj.weight") .unwrap(); assert!(header.path.ends_with("model-00003-of-00003.safetensors")); let entry = indexed .tensor("language_model.model.layers.29.self_attn.q_proj.weight") .unwrap(); assert_eq!(entry.dtype, MlxDType::U32); assert_eq!(entry.shape, vec![8_192, 352]); } #[test] fn loads_local_tokenizer_metadata() { let tokenizer = MlxTokenizer::load(local_model_dir()).unwrap(); assert!(tokenizer.vocab_size() > 260_000); assert!(tokenizer.merge_count() > 500_000); assert_eq!(tokenizer.token_to_id(""), Some(2)); assert_eq!(tokenizer.token_to_id(""), Some(1)); assert_eq!(tokenizer.token_to_id("<|video|>"), Some(258_884)); assert_eq!(tokenizer.token_to_id("say"), Some(30_468)); assert_eq!(tokenizer.token_to_id("▁hi"), Some(5_631)); assert_eq!(tokenizer.token_to_id("▁("), Some(568)); assert_eq!(tokenizer.id_to_token(2), Some("")); assert_eq!(tokenizer.id_to_token(568), Some("▁(")); } #[test] fn loads_qwen_manifest_with_preprocessor_fallback() { let root = temp_model_dir("qwen_manifest"); write_text(&root.join("config.json"), r#"{"model_type":"qwen3_5_moe"}"#); write_text(&root.join("generation_config.json"), "{}"); write_text(&root.join("preprocessor_config.json"), "{}"); write_text( &root.join("tokenizer_config.json"), r#"{ "tokenizer_class":"Qwen2Tokenizer", "model_max_length":262144, "extra_special_tokens":{ "image_token":"<|image_pad|>", "vision_bos_token":"<|vision_start|>", "vision_eos_token":"<|vision_end|>" } }"#, ); write_text( &root.join("tokenizer.json"), r#"{ "normalizer":{"type":"NFC"}, "pre_tokenizer":{"type":"Sequence"}, "decoder":{"type":"ByteLevel"}, "model":{ "type":"BPE", "byte_fallback":false, "vocab":{"Ġ":0,"h":1,"i":2,"Ġh":3,"Ġhi":4}, "merges":["Ġ h","Ġh i"] }, "added_tokens":[{"id":5,"content":"<|image_pad|>","special":true}] }"#, ); write_text( &root.join("model.safetensors.index.json"), r#"{"metadata":{"total_size":71903645408.0},"weight_map":{}}"#, ); let manifest = MlxModelManifest::load(&root).unwrap(); assert_eq!(manifest.family, MlxModelFamily::Qwen35Moe); assert!(manifest .paths .processor_config_json .as_ref() .unwrap() .ends_with("preprocessor_config.json")); assert_eq!(manifest.tokenizer_config.image_token, "<|image_pad|>"); assert_eq!(manifest.tokenizer_config.boi_token, "<|vision_start|>"); assert_eq!(manifest.tokenizer_config.eoi_token, "<|vision_end|>"); fs::remove_dir_all(root).unwrap(); } #[test] fn qwen_bytelevel_tokenizer_roundtrips_minimal_prompt() { let root = temp_model_dir("qwen_tokenizer"); write_text(&root.join("config.json"), r#"{"model_type":"qwen3_5_moe"}"#); write_text(&root.join("generation_config.json"), "{}"); write_text(&root.join("preprocessor_config.json"), "{}"); write_text( &root.join("tokenizer_config.json"), r#"{ "tokenizer_class":"Qwen2Tokenizer", "model_max_length":262144, "extra_special_tokens":{"image_token":"<|image_pad|>"} }"#, ); write_text( &root.join("tokenizer.json"), r#"{ "normalizer":{"type":"NFC"}, "pre_tokenizer":{"type":"Sequence"}, "decoder":{"type":"ByteLevel"}, "model":{ "type":"BPE", "byte_fallback":false, "vocab":{"Ġ":0,"h":1,"i":2,"Ġh":3,"Ġhi":4}, "merges":["Ġ h","Ġh i"] }, "added_tokens":[{"id":5,"content":"<|image_pad|>","special":true}] }"#, ); write_text( &root.join("model.safetensors.index.json"), r#"{"metadata":{"total_size":1},"weight_map":{}}"#, ); let tokenizer = MlxTokenizer::load(&root).unwrap(); assert_eq!(tokenizer.encode(" hi").unwrap(), vec![4]); assert_eq!(tokenizer.decode(&[4]).unwrap(), " hi"); assert_eq!(tokenizer.encode("<|image_pad|>").unwrap(), vec![5]); fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_snapshot_builds_canonical_tensor_map() { let root = temp_model_dir("qwen_snapshot"); write_qwen35moe_test_model(&root); let snapshot = MlxQwen35MoeSnapshot::load(&root).unwrap(); assert_eq!( snapshot.layer_kind(0).unwrap(), MlxQwen35MoeLayerKind::Recurrent ); assert_eq!( snapshot.layer_kind(3).unwrap(), MlxQwen35MoeLayerKind::Attention ); let aliases = snapshot.canonical_weight_map().unwrap(); assert_eq!( aliases.get("token_embd.weight").map(String::as_str), Some("model.language_model.embed_tokens.weight") ); assert_eq!( aliases.get("blk.0.attn_qkv.weight").map(String::as_str), Some("model.language_model.layers.0.linear_attn.in_proj_qkv.weight") ); assert_eq!( aliases.get("blk.3.attn_q.weight").map(String::as_str), Some("model.language_model.layers.3.self_attn.q_proj.weight") ); assert_eq!( aliases .get("blk.0.ffn_gate_up_exps.weight") .map(String::as_str), Some("model.language_model.layers.0.mlp.experts.gate_up_proj") ); assert_eq!( aliases .get("visual.blk.0.attn_qkv.weight") .map(String::as_str), Some("model.visual.blocks.0.attn.qkv.weight") ); fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_indexed_safetensors_resolves_canonical_names() { let root = temp_model_dir("qwen_indexed"); write_qwen35moe_test_model(&root); let indexed = MlxQwen35MoeIndexedSafetensors::load(&root).unwrap(); assert_eq!( indexed.actual_tensor_name("blk.3.attn_k.weight").unwrap(), "model.language_model.layers.3.self_attn.k_proj.weight" ); assert_eq!( indexed .shard_name_for_tensor("visual.merger.fc1.weight") .unwrap(), "model-00001-of-00001.safetensors" ); let tensor = indexed.tensor("blk.0.ssm_out.weight").unwrap(); assert_eq!(tensor.dtype, MlxDType::BF16); assert_eq!(tensor.shape, vec![1]); assert_eq!( indexed .quantization_for_tensor("blk.0.attn_qkv.weight") .unwrap() .unwrap() .bits, 8 ); assert_eq!( indexed .quantization_for_tensor("blk.0.ssm_out.weight") .unwrap() .unwrap() .bits, 4 ); fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_indexed_builds_typed_text_inventory() { let root = temp_model_dir("qwen_inventory"); write_qwen35moe_test_model(&root); let indexed = MlxQwen35MoeIndexedSafetensors::load(&root).unwrap(); let tensors = MlxQwen35MoeTensors::from_indexed(&indexed).unwrap(); assert_eq!(tensors.globals.token_embd, "token_embd.weight"); assert_eq!(tensors.layers.len(), 4); assert_eq!(tensors.layers[0].kind, MlxQwen35MoeLayerKind::Recurrent); assert!(tensors.layers[0].attention.is_none()); assert!(tensors.layers[0].recurrent.is_some()); assert!(tensors.layers[0].moe.uses_merged_gate_up()); assert_eq!(tensors.layers[3].kind, MlxQwen35MoeLayerKind::Attention); assert!(tensors.layers[3].attention.is_some()); assert!(tensors.layers[3].recurrent.is_none()); fs::remove_dir_all(root).unwrap(); } #[test] fn family_snapshot_front_door_loads_qwen_model() { let root = temp_model_dir("qwen_front_door"); write_qwen35moe_test_model(&root); let family = MlxFamilySnapshot::load(&root).unwrap(); assert_eq!(family.family(), MlxModelFamily::Qwen35Moe); match family { MlxFamilySnapshot::Gemma4(_) => panic!("expected qwen family snapshot"), MlxFamilySnapshot::Qwen35Moe(snapshot) => { assert_eq!(snapshot.config.model_type, "qwen3_5_moe"); } } fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_runtime_session_loads_dims_and_cache_template() { let root = temp_model_dir("qwen_runtime"); write_qwen35moe_test_model(&root); let session = MlxQwen35MoeRuntimeSession::load(&root).unwrap(); assert_eq!(session.dims.vocab_size, 1); assert_eq!(session.dims.block_count, 4); assert_eq!(session.dims.embedding_length, 2048); assert_eq!(session.dims.attention_head_count, 16); assert_eq!(session.dims.attention_head_count_kv, 2); assert_eq!(session.dims.ssm_state_size, 128); assert_eq!(session.dims.ssm_group_count, 16); assert_eq!(session.dims.ssm_time_step_rank, 32); assert_eq!(session.dims.recurrent_value_head_dim().unwrap(), 128); assert_eq!(session.cache_template.attention_layers, vec![3]); assert_eq!(session.cache_template.recurrent_layers, vec![0, 1, 2]); assert_eq!(session.cache_template.attention_k_width, 512); assert_eq!(session.cache_template.attention_v_width, 512); assert_eq!(session.cache_template.recurrent_r_width, 24576); assert_eq!(session.cache_template.recurrent_s_width, 524288); assert_eq!( session.stop_tokens.iter().copied().collect::>(), vec![2] ); fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_runtime_session_builds_decode_state_layout() { let root = temp_model_dir("qwen_decode_state"); write_qwen35moe_test_model(&root); let session = MlxQwen35MoeRuntimeSession::load(&root).unwrap(); let decode_state = session.new_decode_state().unwrap(); assert_eq!(decode_state.token_count, 0); assert_eq!(decode_state.layers.len(), 4); match &decode_state.layers[0] { crate::MlxQwen35MoeLayerDecodeState::Recurrent(layer) => { assert_eq!(layer.layer_index, 0); assert_eq!(layer.conv_state.len(), 24_576); assert_eq!(layer.ssm_state.len(), 524_288); } other => panic!("expected recurrent layer state, got {other:?}"), } match &decode_state.layers[3] { crate::MlxQwen35MoeLayerDecodeState::Attention(layer) => { assert_eq!(layer.layer_index, 3); assert!(layer.key_cache.is_empty()); assert!(layer.value_cache.is_empty()); } other => panic!("expected attention layer state, got {other:?}"), } fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_chat_prompt_and_extractors_match_expected_control_tokens() { let root = temp_model_dir("qwen_prompt"); write_qwen35moe_test_model(&root); let session = MlxQwen35MoeRuntimeSession::load(&root).unwrap(); let prompt = session .format_chat_prompt( &[ QwenChatMessage::new(QwenChatRole::System, "You are terse."), QwenChatMessage::new(QwenChatRole::User, "Describe the duck."), ], true, ) .unwrap(); assert!(prompt.contains("<|im_start|>system\nYou are terse.<|im_end|>\n")); assert!(prompt.contains("<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the duck.<|im_end|>\n")); assert!(prompt.ends_with("<|im_start|>assistant\n\n\n\n\n")); let think_prompt = format_qwen35moe_chat_prompt_with_options( session.tokenizer_config(), &[ QwenChatMessage::new(QwenChatRole::System, "You are terse."), QwenChatMessage::new(QwenChatRole::User, "Describe the duck."), ], QwenChatPromptOptions { thinking_mode: QwenThinkingMode::Enabled, ..QwenChatPromptOptions::default() }, ) .unwrap(); assert!(think_prompt.ends_with("<|im_start|>assistant\n\n")); let text = extract_qwen35moe_assistant_response_text( session.tokenizer_config(), "<|im_start|>assistant\n\ninternal\n\nFinal answer<|im_end|>", ); assert_eq!(text, "Final answer"); fs::remove_dir_all(root).unwrap(); } #[test] fn mlx_chat_session_routes_qwen_without_gemma_specific_cli_logic() { let root = temp_model_dir("qwen_generic_chat"); write_qwen35moe_test_model(&root); let mut session = MlxChatSession::load_with_mode(&root, Some(32), MlxChatDecodeMode::Greedy).unwrap(); assert_eq!(session.family(), MlxModelFamily::Qwen35Moe); assert_eq!(session.decode_mode(), MlxChatDecodeMode::Greedy); assert!(session.backend_label().starts_with("qwen-")); assert!(session.messages().is_empty()); match session.send_user_message_streaming("Describe the duck.", |_| Ok(())) { Ok(output) => { assert!(!output.generated_token_ids.is_empty()); assert_eq!(session.messages().len(), 2); assert_eq!(session.messages()[0].role, MlxChatRole::User); assert_eq!(session.messages()[1].role, MlxChatRole::Assistant); } Err(err) => { let message = err.to_string(); assert!(message.contains("missing tokenizer piece"), "{message}"); assert!(session.messages().is_empty()); } } fs::remove_dir_all(root).unwrap(); } #[test] fn qwen35moe_execution_plan_builds_tensor_inventory() { let root = temp_model_dir("qwen_execution_plan"); write_qwen35moe_test_model(&root); let session = MlxQwen35MoeRuntimeSession::load(&root).unwrap(); let plan = session.execution_plan().unwrap(); assert_eq!(plan.layer_count(), 4); assert_eq!( plan.inventory .count_layers_with_role(MlxModelLayerRole::Attention), 1 ); assert_eq!( plan.inventory .count_layers_with_role(MlxModelLayerRole::Recurrent), 3 ); assert_eq!(plan.inventory.unique_tensor_count(), 72); assert_eq!(plan.inventory.total_tensor_bytes(), 144); assert_eq!(plan.tail_probe.output.actual_name, "lm_head.weight"); assert_eq!( plan.tail_probe.output_norm.actual_name, "model.language_model.norm.weight" ); assert_eq!( plan.inventory.globals["token_embd"].actual_name, "model.language_model.embed_tokens.weight" ); assert_eq!(plan.inventory.layers[0].role, MlxModelLayerRole::Recurrent); assert_eq!(plan.inventory.layers[3].role, MlxModelLayerRole::Attention); assert_eq!( plan.inventory.layers[0].tensors["ssm_conv1d"].actual_name, "model.language_model.layers.0.linear_attn.conv1d.weight" ); assert_eq!( plan.inventory.layers[3].tensors["attn_q"].actual_name, "model.language_model.layers.3.self_attn.q_proj.weight" ); fs::remove_dir_all(root).unwrap(); } #[test] fn tokenizer_encodes_and_decodes_simple_phrase() { let tokenizer = MlxTokenizer::load(local_model_dir()).unwrap(); assert_eq!(tokenizer.encode("say hi").unwrap(), vec![30_468, 5_631]); assert_eq!(tokenizer.encode(" hi").unwrap(), vec![5_631]); assert_eq!(tokenizer.decode(&[30_468, 5_631]).unwrap(), "say hi"); assert_eq!(tokenizer.decode(&[1_879, 5_631]).unwrap(), " say hi"); } #[test] fn embeds_and_norms_local_text_token_rows() { let weights = MlxIndexedSafetensors::load(local_model_dir()).unwrap(); let embed = weights.embed_token_bf16_words(30_468).unwrap(); assert_eq!(embed.len(), 2_816); let final_norm = weights.final_text_norm_bf16_words(&embed).unwrap(); assert_eq!(final_norm.len(), 2_816); } #[test] #[ignore] fn embed_rows_report_hashes_for_two_token_prompt() { let weights = MlxIndexedSafetensors::load(local_model_dir()).unwrap(); for token_id in [30_468u32, 5_631u32] { let bits = weights .embed_token_bf16_words(token_id) .unwrap() .into_iter() .map(|word| (word as u32) << 16) .collect::>(); println!( "token_id={} embed_fnv1a64=0x{:016X}", token_id, fnv1a64_u32_words(&bits) ); println!( "token_id={} embed_first16_f32_bits={}", token_id, bits.iter() .take(16) .map(|bits| format!("0x{bits:08X}")) .collect::>() .join(",") ); } } #[test] fn reads_local_tensor_payload_words() { let header = MlxSafetensorsHeader::load(local_model_shard_1()).unwrap(); let q_proj_weight = header .read_u32_tensor_words("language_model.model.layers.0.self_attn.q_proj.weight") .unwrap(); assert_eq!( &q_proj_weight[..8], &[ 2_259_126_473, 1_283_001_501, 2_291_701_430, 1_970_953_151, 1_283_929_482, 2_027_333_543, 934_918_473, 3_033_893_010, ] ); let q_proj_scales = header .read_bf16_tensor_words("language_model.model.layers.0.self_attn.q_proj.scales") .unwrap(); assert_eq!( &q_proj_scales[..16], &[ 15_321, 48_110, 48_135, 48_112, 15_290, 48_057, 15_308, 15_307, 15_254, 15_260, 15_397, 15_275, 15_300, 15_297, 48_099, 15_299, ] ); } #[test] fn dequantizes_one_local_q_proj_row_matches_mlx_oracle() { let header = MlxSafetensorsHeader::load(local_model_shard_1()).unwrap(); let row = header .affine_dequantize_row_f32( "language_model.model.layers.0.self_attn.q_proj.weight", "language_model.model.layers.0.self_attn.q_proj.scales", "language_model.model.layers.0.self_attn.q_proj.biases", 0, 64, 4, ) .unwrap(); assert_eq!(row.len(), 2_816); assert_eq!( fnv1a64_u32_words(&row.iter().map(|value| value.to_bits()).collect::>()), 0x2D44_4223_7EE7_C10F ); assert_eq!( &row[..16] .iter() .map(|value| value.to_bits()) .collect::>(), &[ 0x3BD9_0000, 0x3CD9_0000, 0x0000_0000, 0x0000_0000, 0xBBD9_0000, 0x3C59_0000, 0xBC59_0000, 0x0000_0000, 0x3D08_0000, 0x3BD9_0000, 0x3CD9_0000, 0xBD59_0000, 0x3BD9_0000, 0xBBD9_0000, 0x3CD9_0000, 0xBCD9_0000, ] ); } #[test] fn quantized_matmul_one_local_q_proj_case_matches_mlx_oracle() { let header = MlxSafetensorsHeader::load(local_model_shard_1()).unwrap(); let x = gemma4_qproj_case_input_bf16_words(GEMMA4_QPROJ_CASE_INNER_DIM); assert_eq!( x[..16] .iter() .copied() .map(bf16_word_to_f32) .map(f32::to_bits) .collect::>(), vec![ 0xBF80_0000, 0xBF40_0000, 0xBF00_0000, 0xBE80_0000, 0x0000_0000, 0x3E80_0000, 0x3F00_0000, 0x3F40_0000, 0x3F80_0000, 0x3F00_0000, 0x0000_0000, 0xBF00_0000, 0xBF80_0000, 0x3E00_0000, 0x3EC0_0000, 0x3F20_0000, ] ); let out = header .affine_quantized_matmul_t_f32( &x, "language_model.model.layers.0.self_attn.q_proj.weight", "language_model.model.layers.0.self_attn.q_proj.scales", "language_model.model.layers.0.self_attn.q_proj.biases", 64, 4, ) .unwrap(); assert_eq!(out.len(), GEMMA4_QPROJ_CASE_OUTPUT_DIM); let out_bits = out.iter().map(|value| value.to_bits()).collect::>(); assert_eq!( &out_bits[..16], &[ 0xBF59_0000, 0x4029_0000, 0x3F3B_0000, 0xBF63_0000, 0x3DAF_0000, 0xBF51_0000, 0xBF49_0000, 0x3FCE_0000, 0x3D8B_0000, 0xBEB9_0000, 0x3F0F_0000, 0x3E8E_0000, 0x3DF2_0000, 0x3E80_0000, 0x3E89_0000, 0xC022_0000, ] ); assert_eq!( fnv1a64_u32_words(&out_bits), GEMMA4_QPROJ_CASE_OUTPUT_FNV1A64 ); } #[test] fn rms_norm_one_local_layer0_input_case_matches_mlx_gpu_oracle() { let header = MlxSafetensorsHeader::load(local_model_shard_1()).unwrap(); let x = gemma4_qproj_case_input_bf16_words(2_816); let out = header .rms_norm_weighted_f32( &x, "language_model.model.layers.0.input_layernorm.weight", 1e-6, ) .unwrap(); assert_eq!(out.len(), 2_816); let out_bits = out.iter().map(|value| value.to_bits()).collect::>(); assert_eq!( &out_bits[..16], &[ 0xC0A2_0000, 0xC080_0000, 0xC033_0000, 0xBFBC_0000, 0x0000_0000, 0x3FB6_0000, 0x402F_0000, 0x40E3_0000, 0x4126_0000, 0x4041_0000, 0x0000_0000, 0xC048_0000, 0xC11C_0000, 0x3F6C_0000, 0x4081_0000, 0x40A4_0000, ] ); assert_eq!(fnv1a64_u32_words(&out_bits), 0xBF5E_A05B_53DF_E923); }