Add Qwen3.6 MoE inference support

crates/xserv-model/src/loader.rs

@@ -1,10 +1,20 @@
 use half::{bf16, f16};
 use safetensors::SafeTensors;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::path::Path;
 use xserv_tensor::{DType, Device, Tensor};
 
 pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
+    load_safetensors_filtered(path, device, |_| true)
+}
+
+/// Load only tensors accepted by `keep`. The safetensors file is still read as
+/// one byte buffer, but unneeded tensor payloads are not copied into Tensors.
+pub fn load_safetensors_filtered(
+    path: &Path,
+    device: Device,
+    keep: impl Fn(&str) -> bool,
+) -> HashMap<String, Tensor> {
     let data =
         std::fs::read(path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
     let st = SafeTensors::deserialize(&data)
@@ -13,6 +23,9 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
     let mut tensors = HashMap::new();
 
     for (name, view) in st.tensors() {
+        if !keep(&name) {
+            continue;
+        }
         let shape: Vec<usize> = view.shape().to_vec();
         let raw_bytes = view.data();
         let dtype = match view.dtype() {
@@ -36,27 +49,64 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
 
 /// Load from a directory containing model.safetensors (or sharded files) + config.json.
 pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
+    load_model_dir_filtered(dir, device, |_| true)
+}
+
+/// Load a filtered subset of a model directory. For indexed sharded models,
+/// consult `model.safetensors.index.json` first so shards containing no wanted
+/// tensors are never read. This is critical for pipeline parallelism on large
+/// models: each stage should read only its layer range, not the full checkpoint.
+pub fn load_model_dir_filtered(
+    dir: &Path,
+    device: Device,
+    keep: impl Fn(&str) -> bool,
+) -> HashMap<String, Tensor> {
     let single = dir.join("model.safetensors");
     if single.exists() {
-        return load_safetensors(&single, device);
+        return load_safetensors_filtered(&single, device, keep);
     }
 
+    let index_path = dir.join("model.safetensors.index.json");
+    let wanted_shards: Option<HashSet<String>> = if index_path.exists() {
+        let text = std::fs::read_to_string(&index_path)
+            .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
+        let index: serde_json::Value = serde_json::from_str(&text)
+            .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
+        let map = index
+            .get("weight_map")
+            .and_then(|v| v.as_object())
+            .unwrap_or_else(|| panic!("{} has no weight_map", index_path.display()));
+        Some(
+            map.iter()
+                .filter(|(name, _)| keep(name))
+                .filter_map(|(_, shard)| shard.as_str().map(str::to_string))
+                .collect(),
+        )
+    } else {
+        None
+    };
+
     // Try sharded: model-00001-of-NNNNN.safetensors
     let mut all_tensors = HashMap::new();
     let mut entries: Vec<_> = std::fs::read_dir(dir)
         .unwrap()
         .filter_map(|e| e.ok())
         .filter(|e| {
-            e.path()
+            let is_safetensors = e
+                .path()
                 .file_name()
                 .map(|f| f.to_string_lossy().ends_with(".safetensors"))
-                .unwrap_or(false)
+                .unwrap_or(false);
+            let is_wanted = wanted_shards.as_ref().is_none_or(|wanted| {
+                wanted.contains(&e.file_name().to_string_lossy().to_string())
+            });
+            is_safetensors && is_wanted
         })
         .collect();
     entries.sort_by_key(|e| e.file_name());
 
     for entry in entries {
-        let tensors = load_safetensors(&entry.path(), device);
+        let tensors = load_safetensors_filtered(&entry.path(), device, &keep);
         all_tensors.extend(tensors);
     }