From 6ecf7567e947f3f1469178f72a3bb3d7cd657bfc Mon Sep 17 00:00:00 2001 From: KingManiya Date: Sat, 18 Apr 2026 17:40:19 +0400 Subject: [PATCH] feat: add memory mode --- README.md | 6 ++++-- llama_cli.py | 39 +++++++++++++++++++++++++++++++++++++-- nodes.py | 27 ++++++++++++++++++++++----- 3 files changed, 63 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index c284eaa..2efe929 100644 --- a/README.md +++ b/README.md @@ -125,8 +125,10 @@ Each top-level `.txt` file appears in the `system_prompt` dropdown. Choose | `top_k` | Top-K sampling cutoff. | | `repeat_penalty` | Penalty for repeated tokens. | | `ctx_size` | llama.cpp context size. | -| `n_gpu_layers` | Layers offloaded to GPU. Use `0` for CPU. | -| `seed` | Random seed. | +| `memory_mode` | Advanced memory placement mode. `auto` passes no layer flags. | +| `n_gpu_layers` | Used in GPU layer modes. Passes `--gpu-layers` / `-ngl`. | +| `n_cpu_moe_layers` | Used in CPU MoE modes. Passes `--n-cpu-moe`. | +| `seed` | Random seed. Use `-1` for a random seed. | | `enable_thinking` | Uses `/think` when enabled and `/no_think` when disabled. | | `timeout_seconds` | Maximum runtime before the subprocess is stopped. | | `image` | Optional ComfyUI image input. Uses the first image in a batch. | diff --git a/llama_cli.py b/llama_cli.py index 66c144f..3398fae 100644 --- a/llama_cli.py +++ b/llama_cli.py @@ -9,6 +9,13 @@ from pathlib import Path THINK_BLOCK_RE = re.compile(r"]*>.*?", flags=re.IGNORECASE | re.DOTALL) +MAX_LLAMA_SEED = 2**32 - 1 +MEMORY_MODES = ( + "auto", + "gpu_layers", + "cpu_moe_layers", + "gpu_and_cpu_moe_layers", +) def tensor_to_temp_png(image) -> Path: @@ -51,6 +58,17 @@ def _split_extra_args(extra_args: str) -> list[str]: return [part.strip("\"'") for part in parts] +def normalize_seed(seed: int) -> int: + # llama.cpp uses -1 as "random seed"; otherwise Windows builds parse the + # value as an unsigned 32-bit integer. + seed = int(seed) + if seed == -1: + return -1 + if not 0 <= seed <= MAX_LLAMA_SEED: + raise ValueError(f"seed must be -1 or between 0 and {MAX_LLAMA_SEED}") + return seed + + def build_command( cli_path: Path, model_path: Path, @@ -63,10 +81,17 @@ def build_command( top_k: int, repeat_penalty: float, ctx_size: int, + memory_mode: str, n_gpu_layers: int, + n_cpu_moe_layers: int, seed: int, extra_args: str = "", ) -> list[str]: + if memory_mode not in MEMORY_MODES: + raise ValueError(f"Unsupported memory_mode: {memory_mode}") + if memory_mode in {"cpu_moe_layers", "gpu_and_cpu_moe_layers"} and n_cpu_moe_layers < 1: + raise ValueError("n_cpu_moe_layers must be at least 1 when CPU MoE layers are enabled.") + command = [ str(cli_path), "-m", str(model_path), @@ -77,10 +102,16 @@ def build_command( "--top-k", str(top_k), "--repeat-penalty", str(repeat_penalty), "-c", str(ctx_size), - "-ngl", str(n_gpu_layers), - "--seed", str(seed), + "--seed", str(normalize_seed(seed)), ] + # In auto mode llama.cpp receives neither flag and uses its own placement + # defaults. Other modes pass only the user-selected memory controls. + if memory_mode in {"gpu_layers", "gpu_and_cpu_moe_layers"}: + command.extend(["-ngl", str(n_gpu_layers)]) + if memory_mode in {"cpu_moe_layers", "gpu_and_cpu_moe_layers"}: + command.extend(["--n-cpu-moe", str(n_cpu_moe_layers)]) + if mmproj_path is not None: command.extend(["--mmproj", str(mmproj_path)]) if image_path is not None: @@ -103,7 +134,9 @@ def run_llama_cli( top_k: int, repeat_penalty: float, ctx_size: int, + memory_mode: str, n_gpu_layers: int, + n_cpu_moe_layers: int, seed: int, enable_thinking: bool, timeout_seconds: int, @@ -130,7 +163,9 @@ def run_llama_cli( top_k=top_k, repeat_penalty=repeat_penalty, ctx_size=ctx_size, + memory_mode=memory_mode, n_gpu_layers=n_gpu_layers, + n_cpu_moe_layers=n_cpu_moe_layers, seed=seed, extra_args=extra_args, ) diff --git a/nodes.py b/nodes.py index 557bf10..7c1029b 100644 --- a/nodes.py +++ b/nodes.py @@ -10,7 +10,7 @@ from .folder_registry import ( system_prompt_options, ) from .llama_binary import ensure_llama_cli -from .llama_cli import extract_thinking, run_llama_cli +from .llama_cli import MAX_LLAMA_SEED, MEMORY_MODES, extract_thinking, run_llama_cli class QwenGGUF: @@ -74,17 +74,30 @@ class QwenGGUF: "step": 512, "tooltip": "llama.cpp context window size in tokens.", }), + "memory_mode": (list(MEMORY_MODES), { + "default": "auto", + "tooltip": "auto passes no layer flags; other modes pass GPU layers, CPU MoE layers, or both.", + "advanced": True, + }), "n_gpu_layers": ("INT", { "default": 99, "min": -1, "max": 999, - "tooltip": "Number of layers to offload to GPU. 0 runs on CPU.", + "tooltip": "Used only in gpu_layers and gpu_and_cpu_moe_layers modes. Number of layers to offload to GPU.", + "advanced": True, + }), + "n_cpu_moe_layers": ("INT", { + "default": 1, + "min": 1, + "max": 999, + "tooltip": "Used only in cpu_moe_layers and gpu_and_cpu_moe_layers modes. Keeps MoE weights of the first N layers on CPU.", + "advanced": True, }), "seed": ("INT", { "default": 1, - "min": 0, - "max": 0xffffffffffffffff, - "tooltip": "Random seed used by llama.cpp.", + "min": -1, + "max": MAX_LLAMA_SEED, + "tooltip": "Random seed used by llama.cpp. Use -1 for a random seed.", }), "enable_thinking": ("BOOLEAN", { "default": False, @@ -131,7 +144,9 @@ class QwenGGUF: top_k: int, repeat_penalty: float, ctx_size: int, + memory_mode: str, n_gpu_layers: int, + n_cpu_moe_layers: int, seed: int, enable_thinking: bool, timeout_seconds: int, @@ -159,7 +174,9 @@ class QwenGGUF: top_k=top_k, repeat_penalty=repeat_penalty, ctx_size=ctx_size, + memory_mode=memory_mode, n_gpu_layers=n_gpu_layers, + n_cpu_moe_layers=n_cpu_moe_layers, seed=seed, enable_thinking=enable_thinking, timeout_seconds=timeout_seconds,