metrale_model_arch/
mistral_loader.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Weight loader for Mistral Small 4 (`model_type = "mistral"`): MLA attention with MoE FFNs.
4//!
5//! Owner: model-arch (Mistral loader).
6//! Invariants: none beyond the types.
7
8use anyhow::Result;
9use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
10
11use metrale_model_layers::layers::ops;
12use metrale_model_layers::weight_map::DenseWeight;
13
14pub struct MistralWeightLoader;
15
16/// 2026-09-25: Allocate device memory, or managed (UVM) memory once a device
17/// allocation has failed. The first failure latches
18/// `OpCache::note_alloc_fallback` on this backend; every later call on the
19/// same backend then goes straight to managed memory.
20pub(crate) fn gpu_alloc_or_managed(gpu: &dyn GpuBackend, bytes: usize) -> Result<DevicePtr> {
21    if gpu.op_cache().alloc_fell_back() {
22        return gpu.alloc_managed(bytes);
23    }
24    match gpu.alloc(bytes) {
25        Ok(p) => Ok(p),
26        Err(_) => {
27            tracing::warn!(
28                "GPU alloc failed ({bytes} bytes) — switching to managed for remaining allocations"
29            );
30            gpu.op_cache().note_alloc_fallback();
31            gpu.alloc_managed(bytes)
32        }
33    }
34}
35
36#[allow(dead_code)]
37fn gpu_matmul(
38    a: DevicePtr,
39    b: DevicePtr,
40    m: usize,
41    n: usize,
42    k: usize,
43    gpu: &dyn GpuBackend,
44) -> Result<DevicePtr> {
45    let bf16 = 2usize;
46    let c = gpu_alloc_or_managed(gpu, m * n * bf16)?;
47    let stream = gpu.default_stream();
48    let gemm_k = gpu.kernel("gemm", "dense_gemm_bf16")?;
49    let b_dense = DenseWeight { weight: b };
50    ops::dense_gemm(
51        gpu, gemm_k, a, &b_dense, c, m as u32, n as u32, k as u32, stream,
52    )?;
53    gpu.synchronize(stream)?;
54    Ok(c)
55}
56
57pub(crate) mod loader_impl;
metrale_model_arch/mistral_loader.rs

metrale_model_arch/
mistral_loader.rs
