[Feature Request] Add support for `n_cpu_moe` parameter in Llama class

[KLL535]

Description

This parameter allows users to offload the expert layers of MOE models directly to the CPU/RAM while keeping the Attention layers in the GPU.

This is incredibly useful for running large MoE models on systems with limited VRAM. Currently, this parameter cannot be passed directly when initializing the Llama class.

Thank you!

[JamePeng]

It seems that ctypes needs to implement quite a few GGML API adaptations to achieve this functionality. I'll look into it later when I have time.

[KLL535]

I managed to get this mode to work:

1. Patch llama.h and llama.cpp

LLAMA_API ggml_backend_buffer_type_t llama_get_cpu_buffer_type(void);

ggml_backend_buffer_type_t llama_get_cpu_buffer_type(void) {
    return ggml_backend_cpu_buffer_type();
}

2. Rebuild and Check:

H:\ComfyUI128\python_embeded\Lib\site-packages\llama_cpp\lib>dumpbin /exports llama.dll | findstr "llama_get_cpu_buffer_type"
      17628 44DB 00033240 llama_get_cpu_buffer_type

3. Patch llama_cpp.py

# llama.h bindings

_lib.llama_get_cpu_buffer_type.argtypes = []
_lib.llama_get_cpu_buffer_type.restype = ctypes.c_void_p
def llama_get_cpu_buffer_type():
    return _lib.llama_get_cpu_buffer_type()

4. Patch llama.py

class Llama:
...
    def __init__(
...
        n_cpu_moe: int = 0,  
...
    ):
...
        if not os.path.exists(model_path):
            raise ValueError(f"Model path does not exist: {model_path}")

        # === EMULATION n_cpu_moe via tensor_buft_overrides ===
        if n_cpu_moe > 0:

            OverrideStruct = llama_cpp.llama_model_tensor_buft_override

            TOTAL_LAYERS = 40   # So far only for this one model
            start_layer = max(0, TOTAL_LAYERS - n_cpu_moe)

            cpu_buft_ptr = llama_cpp.llama_get_cpu_buffer_type()   # ctypes.c_void_p

            overrides = []
            for layer_idx in range(start_layer, TOTAL_LAYERS):
                pattern = f"blk\\.{layer_idx}\\.ffn_(down|gate|up)_exps\\.weight" # So far only for this one model
                pattern_bytes = pattern.encode('utf-8')
                ov = OverrideStruct()
                ov.pattern = pattern_bytes
                ov.buft = cpu_buft_ptr   # real pointer
                overrides.append(ov)

            null_ov = OverrideStruct()
            null_ov.pattern = None
            null_ov.buft = ctypes.c_void_p(0)
            overrides.append(null_ov)

            overrides_array = (OverrideStruct * len(overrides))(*overrides)
            self.model_params.tensor_buft_overrides = ctypes.cast(
                overrides_array,
                ctypes.POINTER(OverrideStruct)
            )

        # === END OF EMULATION ===

        self._model = self._stack.enter_context(
            contextlib.closing(
                internals.LlamaModel(

5. Test in Qwen3.6-35B-A3B-Q4_K_M + image 1024x1024 + thinking on RTX5080 16 Gb
   n_cpu_moe = 20
   n_gpu_layers = -1
   n_threads = 10

LM Studio: 60 tok/sec 1237 tok
llama_cpp_python: 50.76 tok/sec 1328 tok

It's still unclear where the 10 tok/sec disappear to.
use_mmap = false - Increases the speed to 54.
There's also a peak and drop in performance at the beginning. LM Studio doesn't have this and has a smooth trapezoid.

[JamePeng]

You can try turning off the hybrid cache saving mechanism and then try again.

[KLL535]

ctx_checkpoints = 0 not changed speed.
But I found settings that eliminates the "gap" in the trapezoid. split_mode = 0 changes something (single GPU) and gives 57-60 tokens/sec.

            "max_tokens": 2048,
            "image_min_tokens": 1024,
            "image_max_tokens": 2048,
            "n_ctx": 8192,
            "n_batch": 2048,
            "n_ubatch": 512,
            "n_gpu_layers": -1,
            "n_threads": 8,
            "n_cpu_moe": 20,
            "use_mmap": false,
            "split_mode": 0,
            //"type_k": 8,
            //"type_v": 8,
            "temperature": 0.8,
            "top_p": 0.95,
            "min_p": 0.05,
            "repeat_penalty": 1.1,
            "presence_penalty": 0.0,
            "top_k": 40,
            "ctx_checkpoints": 0,

[DEBUG] config read: 0.000s | mode direct_clean
[DEBUG] clear memory: unload_all_models: 0.000s
[DEBUG] clear memory: cuda.empty_cache: 0.000s
[DEBUG] process_images: 0.017s
[DEBUG] import llama_cpp: 0.000s
[DEBUG] create_chat_handler: 0.014s
[DEBUG] load_model: 3.783s
[DEBUG] create message (with 1/0/0 image/audio/video): 0.006s
[DEBUG] inference 59.14 tok/sec, 1555 tokens: 26.292s
[DEBUG] unload_llama_model: 0.834s
[DEBUG] total time: 30.951s

For example last n_gpu_layers offload mode:

            "max_tokens": 2048,
            "image_min_tokens": 1024,
            "image_max_tokens": 2048,
            "n_ctx": 8192,
            "n_batch": 2048,
            "n_ubatch": 512,
            "n_gpu_layers": 22,
            "n_threads": 8,
            //"n_cpu_moe": 20,
            "use_mmap": false,
            "split_mode": 0
            //"type_k": 8,
            //"type_v": 8,
            "temperature": 0.8,
            "top_p": 0.95,
            "min_p": 0.05,
            "repeat_penalty": 1.1,
            "presence_penalty": 0.0,
            "top_k": 40,

[DEBUG] config read: 0.000s | mode direct_clean
[DEBUG] clear memory: unload_all_models: 0.000s
[DEBUG] clear memory: cuda.empty_cache: 0.000s
[DEBUG] process_images: 0.023s
[DEBUG] import llama_cpp: 0.000s
[DEBUG] create_chat_handler: 0.021s
[DEBUG] load_model: 3.741s
[DEBUG] create message (with 1/0/0 image/audio/video): 0.006s
[DEBUG] inference 29.00 tok/sec, 1320 tokens: 45.516s
[DEBUG] unload_llama_model: 0.839s
[DEBUG] total time: 50.150s

[JamePeng]

@KLL535 I push the commit of fine-grained MoE CPU offloading controls, there is no need to patch the llama.h and llama.cpp

• Introduce cpu_moe (bool) and n_cpu_moe (int) parameters to Llama.__init__ for precise Mixture of Experts (MoE) weight offloading.
• cpu_moe=True forces all MoE expert weights to the CPU memory, regardless of n_gpu_layers.
• n_cpu_moe=N offloads the expert weights of the first N layers to the CPU, while keeping attention and router weights on the GPU.

[KLL535]

Everything seems to be working.
cpu_moe = true - lower speed, but minimal VRAM consumption.
n_cpu_moe = 20 (choose the best number) - all available VRAM is full, but higher speed.

A message started appearing in the log that spams thousands of lines and doesn't let me see what happened next. How can I disable it?

CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused
CUDA Graph id 60 reused