moe: MXFP4-resident experts on GPU (single-card gpt-oss)

Experts now stay MXFP4-packed on GPU (~10GB whole model, fits one 32GB
card) instead of dequantized to ~38GB BF16. loader::load_model_dir_split
returns BF16 tensors + raw U8 (_blocks/_scales) in one pass; GptOss slices
each expert's MXFP4 bytes to a GpuBuffer at load, and expert_forward
dequantizes the selected expert to a BF16 scratch (dequant_mxfp4) right
before its GEMM — no per-token CPU->GPU upload, no 38GB BF16 dir.

Verified: gptoss-logits on the original MXFP4 dir
(/opt/wjh/models/gpt-oss-20b) gives logits byte-identical to the BF16 path
— top-1 token 12650 = " Paris" @ 15.3125, full top-10 unchanged — running
on a single GPU. Build green on dash5 (release).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xserv-model/src/loader.rs

@@ -63,6 +63,50 @@ pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
     all_tensors
 }
 
+/// Load a model dir splitting tensors by dtype: float tensors (F32/F16/BF16)
+/// become `Tensor`s on `device`; U8 tensors (gpt-oss MXFP4 `_blocks`/`_scales`,
+/// which are not an xserv Tensor dtype) are returned as raw `(bytes, shape)`.
+/// One pass over the shards (the 13GB MXFP4 file is read once).
+pub fn load_model_dir_split(
+    dir: &Path, device: Device,
+) -> (HashMap<String, Tensor>, HashMap<String, (Vec<u8>, Vec<usize>)>) {
+    let mut files: Vec<std::path::PathBuf> = Vec::new();
+    let single = dir.join("model.safetensors");
+    if single.exists() {
+        files.push(single);
+    } else {
+        let mut entries: Vec<_> = std::fs::read_dir(dir).unwrap()
+            .filter_map(|e| e.ok())
+            .filter(|e| e.path().file_name()
+                .map(|f| f.to_string_lossy().ends_with(".safetensors")).unwrap_or(false))
+            .collect();
+        entries.sort_by_key(|e| e.file_name());
+        files.extend(entries.into_iter().map(|e| e.path()));
+    }
+    assert!(!files.is_empty(), "no safetensors files in {}", dir.display());
+
+    let mut floats = HashMap::new();
+    let mut u8s = HashMap::new();
+    for path in &files {
+        let data = std::fs::read(path)
+            .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
+        let st = SafeTensors::deserialize(&data)
+            .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display()));
+        for (name, view) in st.tensors() {
+            let shape: Vec<usize> = view.shape().to_vec();
+            let raw = view.data();
+            match view.dtype() {
+                safetensors::Dtype::F32 => { floats.insert(name, make_tensor(raw, &shape, DType::F32).to_device(device)); }
+                safetensors::Dtype::F16 => { floats.insert(name, make_tensor(raw, &shape, DType::F16).to_device(device)); }
+                safetensors::Dtype::BF16 => { floats.insert(name, make_tensor(raw, &shape, DType::BF16).to_device(device)); }
+                safetensors::Dtype::U8 => { u8s.insert(name, (raw.to_vec(), shape)); }
+                other => eprintln!("load_model_dir_split: skipping {name}: dtype {other:?}"),
+            }
+        }
+    }
+    (floats, u8s)
+}
+
 fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
     match dtype {
         DType::F32 => {