spark_model/
mistral_loader.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Weight loader for Mistral Small 4 (MLA + MoE architecture).
4//!
5//! GQA fallback: MLA LoRA projections are expanded to dense at load time
6//! via GPU matmul. Q = wq_b @ wq_a, K/V split from `wkv_b @ wkv_a[:kv_lora]`.
7//! Loses MLA's 12.8x KV cache compression but produces coherent output.
8
9use anyhow::Result;
10use spark_runtime::gpu::{DevicePtr, GpuBackend};
11
12use crate::layers::ops;
13use crate::weight_map::DenseWeight;
14
15pub struct MistralWeightLoader;
16
17/// GPU matmul: `C[M,N] = A[M,K] × B[K,N]` using dense_gemm_bf16 kernel.
18/// Allocate GPU memory, falling back to managed (UVM) if device alloc fails.
19/// Uses a static flag to avoid retrying device alloc after the first failure
20/// (which wastes time and fragments memory).
21pub(crate) fn gpu_alloc_or_managed(gpu: &dyn GpuBackend, bytes: usize) -> Result<DevicePtr> {
22    // Latched on the BACKEND, not in a static: after a model is unloaded the
23    // memory pressure that caused the fallback is gone, and the next load
24    // should try device memory again instead of inheriting a UVM sentence
25    // from a model that is no longer resident.
26    if gpu.op_cache().alloc_fell_back() {
27        return gpu.alloc_managed(bytes);
28    }
29    match gpu.alloc(bytes) {
30        Ok(p) => Ok(p),
31        Err(_) => {
32            tracing::warn!(
33                "GPU alloc failed ({bytes} bytes) — switching to managed for remaining allocations"
34            );
35            gpu.op_cache().note_alloc_fallback();
36            gpu.alloc_managed(bytes)
37        }
38    }
39}
40
41#[allow(dead_code)]
42fn gpu_matmul(
43    a: DevicePtr,
44    b: DevicePtr,
45    m: usize,
46    n: usize,
47    k: usize,
48    gpu: &dyn GpuBackend,
49) -> Result<DevicePtr> {
50    let bf16 = 2usize;
51    let c = gpu_alloc_or_managed(gpu, m * n * bf16)?;
52    let stream = gpu.default_stream();
53    let gemm_k = gpu.kernel("gemm", "dense_gemm_bf16")?;
54    let b_dense = DenseWeight { weight: b };
55    ops::dense_gemm(
56        gpu, gemm_k, a, &b_dense, c, m as u32, n as u32, k as u32, stream,
57    )?;
58    gpu.synchronize(stream)?;
59    Ok(c)
60}
61
62pub(crate) mod loader_impl;
spark_model/mistral_loader.rs

spark_model/
mistral_loader.rs
