diff --git a/README.md b/README.md index 3067a7a4e71e..dbf547ce5c5b 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,7 @@ Everything else is upstream `llama.cpp`. The additions currently carried here: | Speculative checkpoints on device | | `llama-server` keeps speculative-decoding checkpoints in device memory instead of copying them to the host | | ROCmFPx quant types | `llama-quantize` types `Q4_0_ROCMFP4`, `Q4_0_ROCMFP4_FAST`, `Q2/Q3/Q6/Q8_0_ROCMFPX` and the `_LEAN`/`_COHERENT`/`_STRIX` recipes | Loads the ROCmFP4 GGUFs published for Strix Halo. CPU codecs plus Vulkan dequant, mat-vec, matmul and integer-dot kernels. Weight formats only: not accepted as KV-cache types | | Repeatable output at depth | | Freed KV cells are zeroed so masked-out rows never carry stale K/V, and the Vulkan radix top-k assigns output slots deterministically | +| Host-memory watchdog | [`scripts/strix_memory_watchdog.py`](docs/strix-memory-watchdog.md) | Runs a command in a process group, requires zero active swap, and stops before host-wide memory reaches the 120 GiB validation ceiling | Every ROCm/HIP change above is guarded on architecture, shape and layout, so other devices see upstream behaviour. Run `--help`, or see [tools/server/README.md](tools/server/README.md), for the full options. diff --git a/cmake/build-info.cmake b/cmake/build-info.cmake index c7005950c561..edf89ebb1458 100644 --- a/cmake/build-info.cmake +++ b/cmake/build-info.cmake @@ -18,7 +18,7 @@ endif() # Get the commit count and hash if(Git_FOUND) execute_process( - COMMAND ${GIT_EXECUTABLE} rev-parse --short HEAD + COMMAND ${GIT_EXECUTABLE} rev-parse HEAD WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} OUTPUT_VARIABLE HEAD OUTPUT_STRIP_TRAILING_WHITESPACE diff --git a/common/arg.cpp b/common/arg.cpp index 3ac2a09c2990..841d0af22d36 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2808,6 +2808,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.ple_direct_io = value; } ).set_env("LLAMA_ARG_NGRAM_DIRECT_IO")); + add_opt(common_arg( + {"--expert-cache-slots"}, "N", + "DeepSeek V4.1 routed experts resident per layer; requires --expert-cache-mib", + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.expert_cache_slots = value; + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_SLOTS")); + add_opt(common_arg( + {"--expert-cache-mib"}, "MiB", + "aggregate DeepSeek V4.1 fixed expert slot-tensor capacity; requires --expert-cache-slots", + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.expert_cache_mib = value; + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_MIB")); add_opt(common_arg( {"-cmoe", "--cpu-moe"}, "keep all Mixture of Experts (MoE) weights in the CPU", diff --git a/common/common.cpp b/common/common.cpp index a6a5364163bc..b80b7fa5f809 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1697,6 +1697,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.ple_direct_io = params.ple_direct_io; mparams.ple_io_threads = params.ple_io_threads; mparams.ple_cache_mb = params.ple_cache_mb; + mparams.expert_cache_slots = params.expert_cache_slots; + mparams.expert_cache_bytes = params.expert_cache_mib > 0 ? (size_t) params.expert_cache_mib << 20 : 0; if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; diff --git a/common/common.h b/common/common.h index 8c4427e0c719..a838f90b52a8 100644 --- a/common/common.h +++ b/common/common.h @@ -626,6 +626,8 @@ struct common_params { bool ple_direct_io = true; // ... read with O_DIRECT int32_t ple_io_threads = 64; // ... parallel readers (random 4 KiB reads: this NVMe gives 62k IOPS at 16, 130k at 64, ~160k at 128+) int32_t ple_cache_mb = 256; // ... row cache, 0 disables + int32_t expert_cache_slots = 0; // DeepSeek V4.1 routed experts resident per layer + int32_t expert_cache_mib = 0; // aggregate fixed slot-tensor capacity bool single_turn = false; // single turn chat conversation diff --git a/docs/strix-memory-watchdog.md b/docs/strix-memory-watchdog.md new file mode 100644 index 000000000000..f0099b9a55ad --- /dev/null +++ b/docs/strix-memory-watchdog.md @@ -0,0 +1,94 @@ +# Strix host-memory watchdog + +`scripts/strix_memory_watchdog.py` is an external Linux command wrapper for headless Strix Halo validation. It does not change model loading or cache sizing. It measures host-wide memory from procfs and controls the launched command's process group. + +```sh +./scripts/strix_memory_watchdog.py -- ./build/bin/llama-server +``` + +The wrapper performs these checks and actions: + +- It refuses to launch if `/proc/swaps` contains any active entry. +- It calculates used memory as `MemTotal - MemAvailable`. Linux reports these fields in KiB, so the wrapper multiplies each value by 1024 and keeps all accounting as integer bytes. +- It sends `SIGTERM` to the process group at 116 GiB used. +- It sends `SIGKILL` at 118 GiB used or 30 seconds after `SIGTERM`. +- It reports `grace_timeout` if any descendant requires `SIGKILL` after the soft-threshold grace period, even when the direct child exited earlier. +- It sends `SIGKILL` and fails if swap appears or required procfs data becomes unavailable during execution. +- It forwards wrapper `SIGHUP`, `SIGINT`, or `SIGTERM` to the process group, waits the configured grace period, then sends `SIGKILL` if any group member remains. +- It checks the process group after the direct child exits and cleans up remaining descendants before returning the child's classification. +- It applies the same bounded process-group cleanup if an unexpected post-launch error occurs. +- It propagates an unmonitored child exit code. A signal exit uses the shell convention `128 + signal`. + +The 118 GiB emergency threshold leaves a 2 GiB sampling margin below the strict 120 GiB ceiling. The default sample interval is one second. This margin cannot guarantee the ceiling for a workload that can allocate more than 2 GiB between samples. Lower `--emergency-gib` or shorten `--sample-interval-seconds` for such a workload. + +Use `--procfs-root` to select a different procfs mount or a test fixture. `--soft-gib`, `--emergency-gib`, `--grace-seconds`, and `--sample-interval-seconds` override the other defaults. The emergency threshold must remain below 120 GiB. The fail-closed timing bounds are a maximum 30-second grace, maximum one-second sample interval, and maximum five-second heartbeat age. + +The wrapper writes timestamped JSON Lines records to standard error. Preflight, sample, signal, and final records include total, available, used, and peak-used bytes, swap entry count, child status, process-group status, threshold reason, and final classification where applicable. Signal records are written immediately after each process-group signal. Child standard input, standard output, and standard error are inherited unchanged. + +## Watchdog-owned validation lease + +Use all three artifact options together when another process must prove that it is inside the active watchdog process group: + +```sh +./scripts/strix_memory_watchdog.py \ + --lease-path /run/deepseek-v41/watchdog-lease.json \ + --heartbeat-path /run/deepseek-v41/watchdog-heartbeat.json \ + --audit-path /run/deepseek-v41/watchdog-audit.jsonl \ + -- \ + python3 tools/deepseek-v41-trace/run_matrix.py +``` + +The watchdog creates and exclusively locks the persistent audit before launch. It then starts an internal guardian as the new session and process-group leader; the guardian starts the supplied command in that same group without inheriting the private control pipe. After the guardian reports the payload PID, the watchdog atomically creates the lease and heartbeat. Existing artifact paths are rejected rather than overwritten. The payload receives the resolved paths through `STRIX_MEMORY_WATCHDOG_LEASE_PATH`, `STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH`, and `STRIX_MEMORY_WATCHDOG_AUDIT_PATH`. It also receives `STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS`. + +The child can run before the first atomic lease rename. A matching preflight must retry the inherited lease path for a bounded interval and fail closed if a complete valid lease does not appear. It must not accept a lease path supplied separately by the operator. Consumers must require version 2; version 1 does not describe the guardian topology or timing policy and is rejected. + +Lease format `strix-memory-watchdog-lease`, version 2, contains: + +- `lease_id` and active/final `state` +- `watchdog_pid`, `watchdog_start_time_utc`, Linux `watchdog_start_time_ticks`, `watchdog_executable_path`, `watchdog_command_sha256`, `watchdog_script_path`, and `watchdog_script_sha256` +- exact `soft_bytes`, `emergency_bytes`, `strict_ceiling_bytes`, `grace_seconds`, and `sample_interval_seconds` +- `procfs_root` +- `guardian_pid`, payload `child_pid`, `child_process_group_id`, `command`, and `child_command_sha256` +- `heartbeat_path`, `max_heartbeat_age_seconds`, and `audit_path` +- device, inode, owner, and mode identity for atomic JSON artifacts, plus the watchdog-held audit descriptor identity +- the authoritative `final` audit record after termination + +Heartbeat format `strix-memory-watchdog-heartbeat`, version 2, binds `lease_id`, watchdog PID/start ticks, child PID/process group, sequence, state, and update timestamps. Every memory sample first checks swap and memory thresholds, pulses the guardian through the private nonblocking pipe, then atomically replaces the heartbeat with the complete sample audit record and its persistent-audit record hash. It pulses again after persistence succeeds. A blocked audit or heartbeat write cannot delay the emergency signal; if persistence stalls past the guardian deadline, the guardian fails closed. A final heartbeat and final lease update remain on disk with the persistent JSONL audit; the watchdog does not delete this evidence. + +The guardian uses Linux `PR_SET_PDEATHSIG` with a parent-race check. It kills its process group on watchdog death, control-pipe EOF/error, or a missed pulse deadline, including a stopped or wedged watchdog. When the watchdog sends a graceful signal, it also puts the guardian into a bounded grace mode and continues private pulses while it waits. This lets the watchdog own the configured grace deadline and record any `SIGKILL` escalation instead of letting the shorter heartbeat deadline preempt cleanup. If the grace control message or a cleanup pulse fails, the watchdog independently sends `SIGKILL` to the process group and reaps the child before it reports `signal_error`. The payload must call `start_process_group_lease_guard()` before it starts exporter descendants. This validates the lease with bounded startup retries, arms a second parent-death link to the guardian, and starts a thread that kills the process group if any validation or artifact operation fails or the watchdog evidence becomes stale. + +A matching Linux preflight must verify all of the following: + +- The inherited lease, heartbeat, and audit paths match the paths inside the lease. +- `/proc//exe` is the exact expected Python executable and argv position 1 is the exact repository watchdog script. `-c`, `-m`, helper-script, inert-argument, and interpreter-option substitutions are rejected. +- The watchdog command line itself supplies the exact 116/118 GiB thresholds, `/proc`, inherited artifact paths, timing policy, and command after `--`; the lease cannot override those expectations. +- `/proc//stat` start ticks and `/proc//cmdline` SHA-256 match the lease and remain stable across validation. A pidfd is held during validation when Linux provides `pidfd_open`. +- The topology is watchdog parent -> guardian process-group leader -> payload child. The current process must be inside `child_process_group_id`. +- The command identity is expected, the procfs root is `/proc`, and thresholds are exactly 116 GiB soft, 118 GiB emergency, and 120 GiB strict ceiling for the final run. +- Lease and heartbeat files are regular, mode 0600, owned by the current UID, opened with `O_NOFOLLOW`, and match their recorded device/inode identity. +- The heartbeat identity matches the lease, its monotonic timestamp is not older than `max_heartbeat_age_seconds`, and its audit-record hash exists in the persistent audit. +- The persistent audit matches the watchdog-held descriptor device/inode and remains exclusively locked by the live watchdog. + +These checks reject accidental or helper-process substitution and make regular-file heartbeat forgery unable to keep the process group alive after private pulses stop. They are not a security boundary against intentionally hostile code running as the same UID; use a separately owned systemd user service or cgroup if that threat is in scope. + +The guardian controls only the process group. A payload that deliberately calls `setsid()` can escape it. The correctness harness must not do that. If arbitrary payload code is in scope, launch the watchdog in a service/cgroup configured to kill every member when the unit stops. + +Exit classifications are authoritative in the last final JSON record. If final artifact persistence fails after a primary safety failure, the primary classification and exit code remain unchanged and the artifact failure is listed in `secondary_errors`. Operational failures use these exit codes: + +| Exit code | Classification | +| ---: | --- | +| 2 | procfs or configuration error | +| 3 | swap active at startup or detected during execution | +| 4 | soft threshold reached | +| 5 | emergency threshold reached | +| 6 | soft-threshold grace period expired | +| 7 | process-group signaling or termination failure | +| 8 | lease, heartbeat, or persistent audit failure | +| 70 | unexpected post-launch error | +| 127 | command launch failure | + +No model, backend, or ROCm package is required to run the unit tests: + +```sh +python3 tests/test_strix_memory_watchdog.py +``` diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index ba9bc83b9b08..c022dc0d5a09 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -13,7 +13,7 @@ list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/") find_program(GIT_EXE NAMES git git.exe NO_CMAKE_FIND_ROOT_PATH) if(GIT_EXE) # Get current git commit hash - execute_process(COMMAND ${GIT_EXE} rev-parse --short HEAD + execute_process(COMMAND ${GIT_EXE} rev-parse HEAD WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} OUTPUT_VARIABLE GGML_BUILD_COMMIT OUTPUT_STRIP_TRAILING_WHITESPACE diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a6c0d9f43cf5..c619de20bffb 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -558,6 +558,7 @@ class MODEL_ARCH(IntEnum): DEEPSEEK2OCR = auto() DEEPSEEK32 = auto() DEEPSEEK4 = auto() + DEEPSEEK41 = auto() CHATGLM = auto() GLM4 = auto() GLM4_MOE = auto() @@ -822,6 +823,10 @@ class MODEL_TENSOR(IntEnum): PLE_NORM_QUERY = auto() # qwen4exp PLE_NORM_CONV = auto() # qwen4exp PLE_CONV1D = auto() # qwen4exp + ENGRAM_EMBD = auto() + ENGRAM_Q_NORM = auto() + ENGRAM_K_NORM = auto() + ENGRAM_KV = auto() ATTN_COMPRESSOR_WKV = auto() ATTN_COMPRESSOR_WGATE = auto() ATTN_COMPRESSOR_APE = auto() @@ -1315,6 +1320,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr", MODEL_ARCH.DEEPSEEK32: "deepseek32", MODEL_ARCH.DEEPSEEK4: "deepseek4", + MODEL_ARCH.DEEPSEEK41: "deepseek41", MODEL_ARCH.CHATGLM: "chatglm", MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", @@ -1578,6 +1584,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PLE_NORM_QUERY: "blk.{bid}.ple_norm_query", # qwen4exp MODEL_TENSOR.PLE_NORM_CONV: "blk.{bid}.ple_norm_conv", # qwen4exp MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", # qwen4exp + MODEL_TENSOR.ENGRAM_EMBD: "blk.{bid}.engram_embd", + MODEL_TENSOR.ENGRAM_Q_NORM: "blk.{bid}.engram_q_norm", + MODEL_TENSOR.ENGRAM_K_NORM: "blk.{bid}.engram_k_norm", + MODEL_TENSOR.ENGRAM_KV: "blk.{bid}.engram_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", @@ -3921,6 +3931,47 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.DEEPSEEK41: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT_A, + MODEL_TENSOR.ATTN_OUT_B, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.ATTN_COMPRESSOR_WKV, + MODEL_TENSOR.ATTN_COMPRESSOR_WGATE, + MODEL_TENSOR.ATTN_COMPRESSOR_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_EXP_PROBS_B_VL, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.ENGRAM_EMBD, + MODEL_TENSOR.ENGRAM_Q_NORM, + MODEL_TENSOR.ENGRAM_K_NORM, + MODEL_TENSOR.ENGRAM_KV, + ], MODEL_ARCH.ERNIE4_5_MOE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/tests/test_deepseek41_schema.py b/gguf-py/tests/test_deepseek41_schema.py new file mode 100644 index 000000000000..b360e7d02162 --- /dev/null +++ b/gguf-py/tests/test_deepseek41_schema.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 + +import os +import sys +import unittest +from pathlib import Path + +if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / "gguf-py").exists(): + sys.path.insert(0, str(Path(__file__).parent.parent)) + +from gguf.constants import MODEL_ARCH, MODEL_ARCH_NAMES, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES + + +class TestDeepSeek41Schema(unittest.TestCase): + + def test_architecture_name(self): + self.assertEqual(MODEL_ARCH_NAMES[MODEL_ARCH.DEEPSEEK41], "deepseek41") + + def test_engram_tensor_names(self): + expected = { + MODEL_TENSOR.ENGRAM_EMBD: "blk.14.engram_embd", + MODEL_TENSOR.ENGRAM_Q_NORM: "blk.14.engram_q_norm", + MODEL_TENSOR.ENGRAM_K_NORM: "blk.14.engram_k_norm", + MODEL_TENSOR.ENGRAM_KV: "blk.14.engram_kv", + } + + for tensor, name in expected.items(): + self.assertIn(tensor, MODEL_TENSORS[MODEL_ARCH.DEEPSEEK41]) + self.assertEqual(TENSOR_NAMES[tensor].format(bid=14), name) + + def test_output_head_does_not_require_deepseek4_hc_tensors(self): + tensors = MODEL_TENSORS[MODEL_ARCH.DEEPSEEK41] + + self.assertIn(MODEL_TENSOR.OUTPUT_NORM, tensors) + self.assertIn(MODEL_TENSOR.OUTPUT, tensors) + self.assertNotIn(MODEL_TENSOR.HC_HEAD_FN, tensors) + self.assertNotIn(MODEL_TENSOR.HC_HEAD_BASE, tensors) + self.assertNotIn(MODEL_TENSOR.HC_HEAD_SCALE, tensors) + + def test_kv_a_norm_does_not_use_deepseek4_tensor_kind(self): + tensors = MODEL_TENSORS[MODEL_ARCH.DEEPSEEK41] + + self.assertIn(MODEL_TENSOR.ATTN_KV_A_NORM, tensors) + self.assertNotIn(MODEL_TENSOR.ATTN_KV_NORM, tensors) + self.assertEqual(TENSOR_NAMES[MODEL_TENSOR.ATTN_KV_A_NORM].format(bid=14), "blk.14.attn_kv_a_norm") + + +if __name__ == "__main__": + unittest.main() diff --git a/include/llama.h b/include/llama.h index 74c26e0e9dcd..cd03424c61eb 100644 --- a/include/llama.h +++ b/include/llama.h @@ -349,6 +349,10 @@ extern "C" { int32_t ple_io_threads; // parallel pread workers int32_t ple_cache_mb; // in-memory cache of recently read rows, 0 disables + // DeepSeek V4.1 routed-expert cache. Both values must be non-zero. + size_t expert_cache_bytes; + int32_t expert_cache_slots; + // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; diff --git a/scripts/strix_memory_watchdog.py b/scripts/strix_memory_watchdog.py new file mode 100755 index 000000000000..a06c85de3f96 --- /dev/null +++ b/scripts/strix_memory_watchdog.py @@ -0,0 +1,2575 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import ctypes +import fcntl +import hashlib +import json +import math +import os +import re +import secrets +import select +import signal +import stat +import subprocess +import sys +import threading +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import IO, Any, Protocol + + +GIB = 1024**3 +STRICT_CEILING_BYTES = 120 * GIB +DEFAULT_SOFT_BYTES = 116 * GIB +DEFAULT_EMERGENCY_BYTES = 118 * GIB +DEFAULT_GRACE_SECONDS = 30.0 +DEFAULT_SAMPLE_INTERVAL_SECONDS = 1.0 +DEFAULT_HEARTBEAT_MAX_AGE_SECONDS = 5.0 +MAX_GRACE_SECONDS = 30.0 +MAX_SAMPLE_INTERVAL_SECONDS = 1.0 +MAX_HEARTBEAT_MAX_AGE_SECONDS = 5.0 + +LEASE_FORMAT = "strix-memory-watchdog-lease" +LEASE_VERSION = 2 +HEARTBEAT_FORMAT = "strix-memory-watchdog-heartbeat" +HEARTBEAT_VERSION = 2 +PR_SET_PDEATHSIG = 1 +LEASE_GUARD_SIGNAL = signal.SIGUSR1 + +EXIT_PROCFS_ERROR = 2 +EXIT_SWAP_ACTIVE = 3 +EXIT_SOFT_LIMIT = 4 +EXIT_EMERGENCY_LIMIT = 5 +EXIT_GRACE_TIMEOUT = 6 +EXIT_SIGNAL_ERROR = 7 +EXIT_LEASE_ERROR = 8 +EXIT_INTERNAL_ERROR = 70 +EXIT_LAUNCH_ERROR = 127 + +MEMINFO_VALUE_RE = re.compile(r"([0-9]+) kB") +SWAPS_HEADER = ["Filename", "Type", "Size", "Used", "Priority"] +PARENT_SIGNALS = (signal.SIGHUP, signal.SIGINT, signal.SIGTERM) + + +class ProcfsError(RuntimeError): + pass + + +class ProcessGroupError(RuntimeError): + pass + + +class ArtifactError(RuntimeError): + def __init__(self, component: str, detail: str): + self.component = component + super().__init__(detail) + + +class LeaseValidationError(RuntimeError): + pass + + +class ParentSignal(RuntimeError): + def __init__(self, signal_number: int): + self.signal_number = signal_number + super().__init__(signal.Signals(signal_number).name) + + +class ProcessHandle(Protocol): + pid: int + + def poll(self) -> int | None: + ... + + def wait(self, timeout: float | None = None) -> int: + ... + + +@dataclass +class GuardianProcess: + process: subprocess.Popen[bytes] + payload_pid: int + pulse_fd: int + + @property + def pid(self) -> int: + return self.process.pid + + def poll(self) -> int | None: + return self.process.poll() + + def wait(self, timeout: float | None = None) -> int: + return self.process.wait(timeout=timeout) + + def pulse(self) -> None: + self._write_control(b"P") + + def begin_grace(self) -> None: + self._write_control(b"G") + + def _write_control(self, value: bytes) -> None: + try: + os.write(self.pulse_fd, value) + except BlockingIOError as exc: + raise ProcessGroupError( + "guardian pulse pipe is blocked" + ) from exc + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot pulse guardian: {detail}" + ) from exc + + def close(self) -> None: + try: + os.close(self.pulse_fd) + except OSError: + pass + + +@dataclass(frozen=True) +class HostSnapshot: + total_bytes: int + available_bytes: int + active_swaps: tuple[str, ...] + + @property + def used_bytes(self) -> int: + return self.total_bytes - self.available_bytes + + +@dataclass +class RuntimeState: + snapshot: HostSnapshot + peak_used_bytes: int + + +@dataclass(frozen=True) +class ArtifactPaths: + lease: Path + heartbeat: Path + audit: Path + + +@dataclass(frozen=True) +class WatchdogConfig: + command: tuple[str, ...] + procfs_root: Path = Path("/proc") + soft_bytes: int = DEFAULT_SOFT_BYTES + emergency_bytes: int = DEFAULT_EMERGENCY_BYTES + grace_seconds: float = DEFAULT_GRACE_SECONDS + sample_interval_seconds: float = DEFAULT_SAMPLE_INTERVAL_SECONDS + lease_path: Path | None = None + heartbeat_path: Path | None = None + audit_path: Path | None = None + heartbeat_max_age_seconds: float = DEFAULT_HEARTBEAT_MAX_AGE_SECONDS + + @property + def lease_enabled(self) -> bool: + return self.lease_path is not None + + def validate(self) -> ArtifactPaths | None: + if not self.command: + raise ValueError("a command is required after --") + if self.soft_bytes <= 0: + raise ValueError("soft threshold must be greater than zero") + if self.emergency_bytes <= self.soft_bytes: + raise ValueError("emergency threshold must be greater than soft threshold") + if self.emergency_bytes >= STRICT_CEILING_BYTES: + raise ValueError("emergency threshold must be below 120 GiB") + if ( + not math.isfinite(self.grace_seconds) + or self.grace_seconds <= 0 + or self.grace_seconds > MAX_GRACE_SECONDS + ): + raise ValueError( + "grace period must be greater than zero and at most 30 seconds" + ) + if ( + not math.isfinite(self.sample_interval_seconds) + or self.sample_interval_seconds <= 0 + or self.sample_interval_seconds > MAX_SAMPLE_INTERVAL_SECONDS + ): + raise ValueError( + "sample interval must be greater than zero and at most 1 second" + ) + if ( + not math.isfinite(self.heartbeat_max_age_seconds) + or self.heartbeat_max_age_seconds + <= self.sample_interval_seconds + or self.heartbeat_max_age_seconds + > MAX_HEARTBEAT_MAX_AGE_SECONDS + ): + raise ValueError( + "heartbeat max age must be greater than sample interval " + "and at most 5 seconds" + ) + lease_paths = ( + self.lease_path, + self.heartbeat_path, + self.audit_path, + ) + if any(path is not None for path in lease_paths) and not all( + path is not None for path in lease_paths + ): + raise ValueError( + "lease, heartbeat, and audit paths must be specified together" + ) + if self.lease_enabled: + assert self.lease_path is not None + assert self.heartbeat_path is not None + assert self.audit_path is not None + try: + paths = ArtifactPaths( + self.lease_path.expanduser().resolve(), + self.heartbeat_path.expanduser().resolve(), + self.audit_path.expanduser().resolve(), + ) + except (OSError, RuntimeError) as exc: + raise ValueError( + f"cannot resolve watchdog artifact path: {exc}" + ) from exc + if len({paths.lease, paths.heartbeat, paths.audit}) != 3: + raise ValueError( + "lease, heartbeat, and audit paths must be distinct" + ) + return paths + return None + + +class ProcfsReader: + def __init__(self, root: Path): + self.root = root + + def _read_text(self, name: str) -> str: + path = self.root / name + try: + return path.read_text(encoding="utf-8") + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcfsError(f"cannot read {path}: {detail}") from exc + + def read_snapshot(self) -> HostSnapshot: + active_swaps = self._parse_swaps(self._read_text("swaps")) + total_bytes, available_bytes = self._parse_meminfo( + self._read_text("meminfo") + ) + return HostSnapshot(total_bytes, available_bytes, active_swaps) + + @staticmethod + def _parse_meminfo(content: str) -> tuple[int, int]: + values: dict[str, int] = {} + required = {"MemTotal", "MemAvailable"} + for line in content.splitlines(): + key, separator, raw_value = line.partition(":") + if not separator or key not in required: + continue + if key in values: + raise ProcfsError(f"duplicate {key} in meminfo") + match = MEMINFO_VALUE_RE.fullmatch(raw_value.strip()) + if match is None: + raise ProcfsError(f"malformed {key} in meminfo") + values[key] = int(match.group(1)) * 1024 + + missing = sorted(required - values.keys()) + if missing: + raise ProcfsError(f"missing {', '.join(missing)} in meminfo") + if values["MemAvailable"] > values["MemTotal"]: + raise ProcfsError("MemAvailable exceeds MemTotal") + return values["MemTotal"], values["MemAvailable"] + + @staticmethod + def _parse_swaps(content: str) -> tuple[str, ...]: + lines = content.splitlines() + if not lines or lines[0].split() != SWAPS_HEADER: + raise ProcfsError("malformed swaps header") + + entries: list[str] = [] + for line in lines[1:]: + if not line.strip(): + continue + fields = line.split() + if len(fields) != len(SWAPS_HEADER): + raise ProcfsError("malformed swaps entry") + try: + int(fields[2]) + int(fields[3]) + int(fields[4]) + except ValueError as exc: + raise ProcfsError("malformed swaps entry") from exc + entries.append(fields[0]) + return tuple(entries) + + +def _timestamp_utc( + wall_clock: Callable[[], datetime] | None = None, +) -> str: + timestamp = (wall_clock or ( + lambda: datetime.now(timezone.utc) + ))().astimezone(timezone.utc) + return timestamp.isoformat(timespec="milliseconds").replace( + "+00:00", "Z" + ) + + +def _sha256_bytes(content: bytes) -> str: + return hashlib.sha256(content).hexdigest() + + +def _sha256_file(path: Path) -> str: + try: + return _sha256_bytes(path.read_bytes()) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "lease", f"cannot hash {path}: {detail}" + ) from exc + + +def _command_sha256(command: Sequence[str]) -> str: + encoded = json.dumps( + list(command), + ensure_ascii=True, + separators=(",", ":"), + ).encode("utf-8") + return _sha256_bytes(encoded) + + +def _set_parent_death_signal( + signal_number: int, expected_parent_pid: int +) -> None: + if not sys.platform.startswith("linux"): + return + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(PR_SET_PDEATHSIG, signal_number, 0, 0, 0) != 0: + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + if os.getppid() != expected_parent_pid: + os.kill(os.getpid(), signal.SIGKILL) + + +def _kill_own_process_group( + _signal_number: int | None = None, + _frame: object | None = None, +) -> None: + try: + os.killpg(os.getpgrp(), signal.SIGKILL) + except OSError: + os._exit(EXIT_SIGNAL_ERROR) + + +def _guardian_main( + control_fd: int, + status_fd: int, + pulse_timeout_seconds: float, + grace_timeout_seconds: float, + command: tuple[str, ...], +) -> int: + if not sys.platform.startswith("linux"): + return EXIT_LAUNCH_ERROR + os.set_inheritable(control_fd, False) + os.set_inheritable(status_fd, False) + signal.signal(LEASE_GUARD_SIGNAL, _kill_own_process_group) + for signal_number in PARENT_SIGNALS: + signal.signal(signal_number, signal.SIG_IGN) + _set_parent_death_signal(LEASE_GUARD_SIGNAL, os.getppid()) + + def prepare_payload() -> None: + for signal_number in PARENT_SIGNALS: + signal.signal(signal_number, signal.SIG_DFL) + + try: + payload = subprocess.Popen(command, preexec_fn=prepare_payload) + except (OSError, ValueError) as exc: + os.write( + status_fd, + json.dumps( + {"error": getattr(exc, "strerror", None) or str(exc)} + ).encode("utf-8") + + b"\n", + ) + os.close(status_fd) + return EXIT_LAUNCH_ERROR + + os.write( + status_fd, + json.dumps({"payload_pid": payload.pid}).encode("utf-8") + b"\n", + ) + os.close(status_fd) + poller = select.poll() + poller.register( + control_fd, + select.POLLIN | select.POLLHUP | select.POLLERR, + ) + current_timeout_seconds = pulse_timeout_seconds + deadline = time.monotonic() + pulse_timeout_seconds + while True: + remaining = max(0.0, deadline - time.monotonic()) + events = poller.poll(max(1, min(50, int(remaining * 1000)))) + for _, event_mask in events: + if event_mask & (select.POLLHUP | select.POLLERR): + _kill_own_process_group() + try: + pulse = os.read(control_fd, 65536) + except BlockingIOError: + pulse = b"" + if not pulse: + _kill_own_process_group() + if b"G" in pulse: + current_timeout_seconds = grace_timeout_seconds + deadline = time.monotonic() + current_timeout_seconds + if time.monotonic() >= deadline: + _kill_own_process_group() + returncode = payload.poll() + if returncode is not None: + if returncode >= 0: + return returncode + signal_number = -returncode + if signal_number not in (signal.SIGKILL, signal.SIGSTOP): + signal.signal(signal_number, signal.SIG_DFL) + os.kill(os.getpid(), signal_number) + return 128 + signal_number + + +def _read_guardian_status( + descriptor: int, timeout_seconds: float +) -> int: + poller = select.poll() + poller.register(descriptor, select.POLLIN | select.POLLHUP) + deadline = time.monotonic() + timeout_seconds + content = b"" + while time.monotonic() < deadline: + events = poller.poll( + max(1, int((deadline - time.monotonic()) * 1000)) + ) + if not events: + continue + chunk = os.read(descriptor, 4096) + if not chunk: + break + content += chunk + if b"\n" in content: + break + if not content: + raise OSError("guardian did not report payload startup") + try: + status = json.loads(content.splitlines()[0]) + except (UnicodeError, json.JSONDecodeError) as exc: + raise OSError("guardian returned malformed startup status") from exc + if not isinstance(status, dict): + raise OSError("guardian returned malformed startup status") + if "error" in status: + raise OSError(str(status["error"])) + payload_pid = status.get("payload_pid") + if not isinstance(payload_pid, int): + raise OSError("guardian did not report a payload PID") + return payload_pid + + +def _launch_guardian( + command: tuple[str, ...], + environment: dict[str, str], + pulse_timeout_seconds: float, + grace_timeout_seconds: float, + launch_mask: set[signal.Signals], +) -> GuardianProcess: + control_read, control_write = os.pipe() + os.set_blocking(control_read, False) + os.set_blocking(control_write, False) + status_read, status_write = os.pipe() + parent_pid = os.getpid() + + def prepare_guardian() -> None: + signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) + _set_parent_death_signal(signal.SIGKILL, parent_pid) + + guardian_command = ( + sys.executable, + str(Path(__file__).resolve()), + "--internal-guardian", + str(control_read), + str(status_write), + str(pulse_timeout_seconds), + str(grace_timeout_seconds), + "--", + *command, + ) + try: + process = subprocess.Popen( + guardian_command, + start_new_session=True, + pass_fds=(control_read, status_write), + preexec_fn=prepare_guardian, + env=environment, + ) + finally: + os.close(control_read) + os.close(status_write) + try: + payload_pid = _read_guardian_status(status_read, 5.0) + except OSError: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5.0) + os.close(control_write) + raise + finally: + os.close(status_read) + guardian = GuardianProcess(process, payload_pid, control_write) + guardian.pulse() + return guardian + + +def _read_proc_bytes(root: Path, process_id: int, name: str) -> bytes: + path = root / str(process_id) / name + try: + return path.read_bytes() + except OSError as exc: + detail = exc.strerror or str(exc) + raise LeaseValidationError( + f"cannot read {path}: {detail}" + ) from exc + + +def _parse_proc_stat(content: str) -> tuple[int, int, int]: + close_paren = content.rfind(")") + if close_paren < 0: + raise LeaseValidationError("malformed process stat") + fields = content[close_paren + 1:].split() + if len(fields) < 20: + raise LeaseValidationError("malformed process stat") + try: + return int(fields[1]), int(fields[2]), int(fields[19]) + except ValueError as exc: + raise LeaseValidationError("malformed process stat") from exc + + +def _read_proc_stat( + root: Path, process_id: int +) -> tuple[int, int, int]: + content = _read_proc_bytes( + root, process_id, "stat" + ).decode("utf-8") + return _parse_proc_stat(content) + + +def _write_json_atomic( + path: Path, + value: dict[str, object], + *, + create: bool = False, +) -> None: + parent = path.parent + temp_path = parent / ( + f".{path.name}.{os.getpid()}.{secrets.token_hex(8)}.tmp" + ) + try: + descriptor = os.open( + temp_path, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 0o600, + ) + with os.fdopen(descriptor, "wb") as stream: + file_status = os.fstat(stream.fileno()) + record = { + **value, + "file_device": file_status.st_dev, + "file_inode": file_status.st_ino, + "file_uid": file_status.st_uid, + "file_mode": stat.S_IMODE(file_status.st_mode), + } + payload = ( + json.dumps( + record, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ) + + "\n" + ).encode("utf-8") + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + if create: + os.link(temp_path, path) + temp_path.unlink() + else: + os.replace(temp_path, path) + directory_descriptor = os.open(parent, os.O_RDONLY) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + except OSError as exc: + try: + temp_path.unlink() + except FileNotFoundError: + pass + detail = exc.strerror or str(exc) + action = "create" if create else "write" + raise ArtifactError( + "lease", f"cannot atomically {action} {path}: {detail}" + ) from exc + + +class LeaseManager: + def __init__( + self, + config: WatchdogConfig, + paths: ArtifactPaths, + *, + process_procfs_root: Path = Path("/proc"), + wall_clock: Callable[[], datetime] | None = None, + monotonic_ns: Callable[[], int] | None = None, + ): + self.config = config + self.lease_path = paths.lease + self.heartbeat_path = paths.heartbeat + self.audit_path = paths.audit + self.process_procfs_root = process_procfs_root + self.wall_clock = wall_clock + self.monotonic_ns = monotonic_ns or time.monotonic_ns + self.lease_id = secrets.token_hex(16) + self.sequence = 0 + self.lease: dict[str, object] | None = None + + def _watchdog_identity(self) -> dict[str, object]: + script_path = Path(__file__).resolve() + cmdline_path = ( + self.process_procfs_root / str(os.getpid()) / "cmdline" + ) + proc_start_time_ticks: int | None = None + try: + cmdline = cmdline_path.read_bytes() + _, _, proc_start_time_ticks = _read_proc_stat( + self.process_procfs_root, os.getpid() + ) + executable_path = ( + self.process_procfs_root + / str(os.getpid()) + / "exe" + ).resolve() + except (OSError, LeaseValidationError): + if sys.platform.startswith("linux"): + raise ArtifactError( + "lease", + "cannot read watchdog process identity from procfs", + ) + cmdline = b"\0".join( + os.fsencode(argument) for argument in sys.argv + ) + executable_path = Path(sys.executable).resolve() + return { + "pid": os.getpid(), + "start_time_utc": _timestamp_utc(self.wall_clock), + "proc_start_time_ticks": proc_start_time_ticks, + "cmdline_sha256": _sha256_bytes(cmdline), + "executable_path": str(executable_path), + "script_path": str(script_path), + "script_sha256": _sha256_file(script_path), + } + + def _heartbeat_record( + self, + state: str, + sample: dict[str, object] | None = None, + ) -> dict[str, object]: + assert self.lease is not None + self.sequence += 1 + record: dict[str, object] = { + "format": HEARTBEAT_FORMAT, + "version": HEARTBEAT_VERSION, + "lease_id": self.lease_id, + "sequence": self.sequence, + "state": state, + "updated_at": _timestamp_utc(self.wall_clock), + "updated_monotonic_ns": self.monotonic_ns(), + "watchdog_pid": self.lease["watchdog_pid"], + "watchdog_start_time_ticks": ( + self.lease["watchdog_start_time_ticks"] + ), + "child_pid": self.lease["child_pid"], + "child_process_group_id": ( + self.lease["child_process_group_id"] + ), + } + if sample is not None: + record["sample"] = sample + return record + + def start( + self, child: ProcessHandle, audit: AuditLogger + ) -> None: + watchdog_identity = self._watchdog_identity() + audit_identity = audit.persistent_identity() + payload_pid = ( + child.payload_pid + if isinstance(child, GuardianProcess) + else child.pid + ) + self.lease = { + "format": LEASE_FORMAT, + "version": LEASE_VERSION, + "lease_id": self.lease_id, + "state": "active", + "watchdog_pid": watchdog_identity["pid"], + "watchdog_start_time_utc": ( + watchdog_identity["start_time_utc"] + ), + "watchdog_start_time_ticks": ( + watchdog_identity["proc_start_time_ticks"] + ), + "watchdog_command_sha256": ( + watchdog_identity["cmdline_sha256"] + ), + "watchdog_executable_path": ( + watchdog_identity["executable_path"] + ), + "watchdog_script_path": watchdog_identity["script_path"], + "watchdog_script_sha256": ( + watchdog_identity["script_sha256"] + ), + "soft_bytes": self.config.soft_bytes, + "emergency_bytes": self.config.emergency_bytes, + "strict_ceiling_bytes": STRICT_CEILING_BYTES, + "grace_seconds": self.config.grace_seconds, + "sample_interval_seconds": self.config.sample_interval_seconds, + "guardian_pid": child.pid, + "child_pid": payload_pid, + "child_process_group_id": child.pid, + "command": list(self.config.command), + "child_command_sha256": _command_sha256( + self.config.command + ), + "heartbeat_path": str(self.heartbeat_path), + "max_heartbeat_age_seconds": ( + self.config.heartbeat_max_age_seconds + ), + "audit_path": str(self.audit_path), + "audit_device": audit_identity["device"], + "audit_inode": audit_identity["inode"], + "audit_uid": audit_identity["uid"], + "audit_mode": audit_identity["mode"], + "audit_fd": audit_identity["fd"], + "procfs_root": str( + self.config.procfs_root.expanduser().resolve() + ), + } + heartbeat = self._heartbeat_record( + "active", + {"audit_record_sha256": audit.last_record_sha256}, + ) + _write_json_atomic(self.heartbeat_path, heartbeat, create=True) + _write_json_atomic(self.lease_path, self.lease, create=True) + + def update_heartbeat(self, sample: dict[str, object]) -> None: + heartbeat = self._heartbeat_record("active", sample) + _write_json_atomic(self.heartbeat_path, heartbeat) + + def finalize(self, final_record: dict[str, object]) -> None: + if self.lease is None: + return + self.lease["state"] = "final" + self.lease["final"] = final_record + heartbeat = self._heartbeat_record("final") + _write_json_atomic(self.heartbeat_path, heartbeat) + _write_json_atomic(self.lease_path, self.lease) + + +def _read_json_object(path: Path) -> dict[str, object]: + try: + descriptor = os.open( + path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + with os.fdopen(descriptor, "r", encoding="utf-8") as stream: + file_status = os.fstat(stream.fileno()) + if ( + not stat.S_ISREG(file_status.st_mode) + or file_status.st_uid != os.getuid() + or stat.S_IMODE(file_status.st_mode) != 0o600 + ): + raise LeaseValidationError( + f"{path} has unsafe type, owner, or mode" + ) + value = json.load(stream) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise LeaseValidationError( + f"cannot read valid JSON from {path}: {exc}" + ) from exc + if not isinstance(value, dict): + raise LeaseValidationError(f"{path} must contain a JSON object") + if ( + value.get("file_device") != file_status.st_dev + or value.get("file_inode") != file_status.st_ino + or value.get("file_uid") != file_status.st_uid + or value.get("file_mode") != stat.S_IMODE(file_status.st_mode) + ): + raise LeaseValidationError(f"{path} identity does not match") + return value + + +def _require_int(value: object, field: str) -> int: + if not isinstance(value, int) or isinstance(value, bool): + raise LeaseValidationError(f"lease field {field} is invalid") + return value + + +def _require_string(value: object, field: str) -> str: + if not isinstance(value, str) or not value: + raise LeaseValidationError(f"lease field {field} is invalid") + return value + + +def validate_active_lease( + lease_path: Path, + *, + expected_script_path: Path, + expected_executable_path: Path | None = None, + expected_soft_bytes: int = DEFAULT_SOFT_BYTES, + expected_emergency_bytes: int = DEFAULT_EMERGENCY_BYTES, + expected_procfs_root: Path = Path("/proc"), + expected_command: Sequence[str] | None = None, + expected_heartbeat_path: Path | None = None, + expected_audit_path: Path | None = None, + expected_max_heartbeat_age_seconds: float | None = None, + current_process_id: int | None = None, + process_procfs_root: Path = Path("/proc"), + monotonic_ns: Callable[[], int] | None = None, + pidfd_open: Callable[[int], int] | None = getattr( + os, "pidfd_open", None + ), +) -> dict[str, object]: + lease_path = lease_path.expanduser().resolve() + lease = _read_json_object(lease_path) + if ( + lease.get("format") != LEASE_FORMAT + or lease.get("version") != LEASE_VERSION + or lease.get("state") != "active" + ): + raise LeaseValidationError("lease format, version, or state is invalid") + + script_path = Path( + _require_string( + lease.get("watchdog_script_path"), + "watchdog_script_path", + ) + ).resolve() + expected_script_path = expected_script_path.expanduser().resolve() + if script_path != expected_script_path: + raise LeaseValidationError("watchdog script path does not match") + script_sha256 = _require_string( + lease.get("watchdog_script_sha256"), + "watchdog_script_sha256", + ) + if script_sha256 != _sha256_file(expected_script_path): + raise LeaseValidationError("watchdog script SHA does not match") + + if ( + _require_int( + lease.get("soft_bytes"), "soft_bytes" + ) + != expected_soft_bytes + or _require_int( + lease.get("emergency_bytes"), + "emergency_bytes", + ) + != expected_emergency_bytes + or _require_int( + lease.get("strict_ceiling_bytes"), + "strict_ceiling_bytes", + ) + != STRICT_CEILING_BYTES + ): + raise LeaseValidationError("watchdog thresholds do not match") + lease_procfs_root = Path( + _require_string(lease.get("procfs_root"), "procfs_root") + ).resolve() + if lease_procfs_root != expected_procfs_root.expanduser().resolve(): + raise LeaseValidationError("watchdog procfs root does not match") + + watchdog_pid = _require_int( + lease.get("watchdog_pid"), "watchdog_pid" + ) + pidfd: int | None = None + if pidfd_open is not None: + try: + pidfd = pidfd_open(watchdog_pid) + except OSError as exc: + raise LeaseValidationError( + "cannot open watchdog pidfd" + ) from exc + watchdog_start_ticks = _require_int( + lease.get("watchdog_start_time_ticks"), + "watchdog_start_time_ticks", + ) + _, _, live_watchdog_start_ticks = _read_proc_stat( + process_procfs_root, watchdog_pid + ) + if live_watchdog_start_ticks != watchdog_start_ticks: + raise LeaseValidationError("watchdog process start time does not match") + expected_executable = ( + expected_executable_path or Path(sys.executable) + ).expanduser().resolve() + try: + live_executable = ( + process_procfs_root / str(watchdog_pid) / "exe" + ).resolve() + except OSError as exc: + raise LeaseValidationError( + "cannot resolve watchdog executable" + ) from exc + if ( + live_executable != expected_executable + or Path( + _require_string( + lease.get("watchdog_executable_path"), + "watchdog_executable_path", + ) + ).resolve() + != expected_executable + ): + raise LeaseValidationError("watchdog executable does not match") + live_cmdline = _read_proc_bytes( + process_procfs_root, watchdog_pid, "cmdline" + ) + if _sha256_bytes(live_cmdline) != _require_string( + lease.get("watchdog_command_sha256"), + "watchdog_command_sha256", + ): + raise LeaseValidationError("watchdog command line does not match") + argv = [ + os.fsdecode(argument) + for argument in live_cmdline.split(b"\0") + if argument + ] + if len(argv) < 2 or argv[1] in ("-c", "-m"): + raise LeaseValidationError( + "watchdog script is not in executable argv position" + ) + try: + watchdog_cwd = ( + process_procfs_root / str(watchdog_pid) / "cwd" + ).resolve() + except OSError as exc: + raise LeaseValidationError( + "cannot resolve watchdog working directory" + ) from exc + argv_script = Path(argv[1]).expanduser() + if not argv_script.is_absolute(): + argv_script = watchdog_cwd / argv_script + if argv_script.resolve() != expected_script_path: + raise LeaseValidationError( + "watchdog script is not in executable argv position" + ) + try: + live_config = parse_args(argv[2:]) + live_paths = live_config.validate() + except (SystemExit, ValueError) as exc: + raise LeaseValidationError( + "watchdog command line is invalid" + ) from exc + if ( + live_config.soft_bytes != expected_soft_bytes + or live_config.emergency_bytes != expected_emergency_bytes + or live_config.procfs_root.expanduser().resolve() + != expected_procfs_root.expanduser().resolve() + ): + raise LeaseValidationError( + "watchdog command-line policy does not match" + ) + if ( + lease.get("grace_seconds") != live_config.grace_seconds + or lease.get("sample_interval_seconds") + != live_config.sample_interval_seconds + or lease.get("max_heartbeat_age_seconds") + != live_config.heartbeat_max_age_seconds + ): + raise LeaseValidationError( + "watchdog lease timing policy does not match" + ) + if ( + live_paths is None + or live_paths.lease != lease_path + or ( + expected_heartbeat_path is not None + and live_paths.heartbeat + != expected_heartbeat_path.expanduser().resolve() + ) + or ( + expected_audit_path is not None + and live_paths.audit + != expected_audit_path.expanduser().resolve() + ) + ): + raise LeaseValidationError( + "watchdog command-line artifact paths do not match" + ) + if expected_command is not None and tuple( + expected_command + ) != live_config.command: + raise LeaseValidationError("monitored command does not match") + + guardian_pid = _require_int( + lease.get("guardian_pid"), "guardian_pid" + ) + child_pid = _require_int(lease.get("child_pid"), "child_pid") + process_group_id = _require_int( + lease.get("child_process_group_id"), + "child_process_group_id", + ) + guardian_parent_pid, guardian_group_id, _ = _read_proc_stat( + process_procfs_root, guardian_pid + ) + child_parent_pid, child_group_id, _ = _read_proc_stat( + process_procfs_root, child_pid + ) + if ( + guardian_parent_pid != watchdog_pid + or guardian_group_id != process_group_id + or guardian_pid != process_group_id + or child_parent_pid != guardian_pid + or child_group_id != process_group_id + ): + raise LeaseValidationError( + "watchdog, guardian, child, or process group does not match" + ) + command = lease.get("command") + if ( + not isinstance(command, list) + or not command + or not all(isinstance(argument, str) for argument in command) + ): + raise LeaseValidationError("lease field command is invalid") + command_sha256 = _require_string( + lease.get("child_command_sha256"), + "child_command_sha256", + ) + if command_sha256 != _command_sha256(command): + raise LeaseValidationError("monitored command SHA is invalid") + if expected_command is not None and command_sha256 != _command_sha256( + expected_command + ): + raise LeaseValidationError("monitored command SHA does not match") + + process_id = ( + current_process_id + if current_process_id is not None + else os.getpid() + ) + _, current_group_id, _ = _read_proc_stat( + process_procfs_root, process_id + ) + if current_group_id != process_group_id: + raise LeaseValidationError( + "current process is outside the monitored process group" + ) + + heartbeat_path = Path( + _require_string( + lease.get("heartbeat_path"), "heartbeat_path" + ) + ).resolve() + if ( + expected_heartbeat_path is not None + and heartbeat_path + != expected_heartbeat_path.expanduser().resolve() + ): + raise LeaseValidationError("heartbeat path does not match") + heartbeat_max_age = lease.get("max_heartbeat_age_seconds") + if ( + not isinstance(heartbeat_max_age, (int, float)) + or isinstance(heartbeat_max_age, bool) + or not math.isfinite(heartbeat_max_age) + or heartbeat_max_age <= 0 + ): + raise LeaseValidationError( + "lease field max_heartbeat_age_seconds is invalid" + ) + if ( + expected_max_heartbeat_age_seconds is not None + and heartbeat_max_age != expected_max_heartbeat_age_seconds + ): + raise LeaseValidationError("heartbeat max age does not match") + heartbeat = _read_json_object(heartbeat_path) + lease_id = _require_string(lease.get("lease_id"), "lease_id") + if ( + heartbeat.get("format") != HEARTBEAT_FORMAT + or heartbeat.get("version") != HEARTBEAT_VERSION + or heartbeat.get("state") != "active" + or heartbeat.get("lease_id") != lease_id + or heartbeat.get("watchdog_pid") != watchdog_pid + or heartbeat.get("watchdog_start_time_ticks") + != watchdog_start_ticks + or heartbeat.get("child_pid") != child_pid + or heartbeat.get("child_process_group_id") != process_group_id + ): + raise LeaseValidationError("heartbeat identity does not match lease") + updated_monotonic_ns = _require_int( + heartbeat.get("updated_monotonic_ns"), + "heartbeat.updated_monotonic_ns", + ) + _require_int(heartbeat.get("sequence"), "heartbeat.sequence") + _require_string(heartbeat.get("updated_at"), "heartbeat.updated_at") + heartbeat_sample = heartbeat.get("sample") + if not isinstance(heartbeat_sample, dict): + raise LeaseValidationError("heartbeat sample is invalid") + audit_record_sha256 = _require_string( + heartbeat_sample.get("audit_record_sha256"), + "heartbeat.sample.audit_record_sha256", + ) + now_monotonic_ns = (monotonic_ns or time.monotonic_ns)() + age_ns = now_monotonic_ns - updated_monotonic_ns + if age_ns < 0 or age_ns > int(heartbeat_max_age * 1_000_000_000): + raise LeaseValidationError("watchdog heartbeat is stale") + + audit_path = Path( + _require_string(lease.get("audit_path"), "audit_path") + ).resolve() + if ( + expected_audit_path is not None + and audit_path != expected_audit_path.expanduser().resolve() + ): + raise LeaseValidationError("persistent audit path does not match") + audit_fd = _require_int(lease.get("audit_fd"), "audit_fd") + audit_device = _require_int( + lease.get("audit_device"), "audit_device" + ) + audit_inode = _require_int( + lease.get("audit_inode"), "audit_inode" + ) + audit_uid = _require_int(lease.get("audit_uid"), "audit_uid") + audit_mode = _require_int(lease.get("audit_mode"), "audit_mode") + try: + audit_status = audit_path.stat(follow_symlinks=False) + live_audit_status = ( + process_procfs_root + / str(watchdog_pid) + / "fd" + / str(audit_fd) + ).stat() + if ( + not stat.S_ISREG(audit_status.st_mode) + or audit_status.st_dev != audit_device + or audit_status.st_ino != audit_inode + or live_audit_status.st_dev != audit_device + or live_audit_status.st_ino != audit_inode + or audit_status.st_uid != audit_uid + or audit_uid != os.getuid() + or stat.S_IMODE(audit_status.st_mode) != audit_mode + or audit_mode != 0o600 + ): + raise LeaseValidationError( + "persistent audit identity does not match" + ) + audit_descriptor = os.open( + audit_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + try: + try: + fcntl.flock( + audit_descriptor, + fcntl.LOCK_EX | fcntl.LOCK_NB, + ) + except BlockingIOError: + pass + else: + fcntl.flock(audit_descriptor, fcntl.LOCK_UN) + raise LeaseValidationError( + "watchdog does not hold the persistent audit lock" + ) + finally: + os.close(audit_descriptor) + audit_lines = [ + line + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + if line + ] + first_line = next(iter(audit_lines)) + first_record = json.loads(first_line) + if ( + not isinstance(first_record, dict) + or not isinstance(first_record.get("event"), str) + or not isinstance(first_record.get("timestamp"), str) + ): + raise LeaseValidationError( + "persistent audit does not contain watchdog records" + ) + if not any( + _sha256_bytes((line + "\n").encode("utf-8")) + == audit_record_sha256 + for line in audit_lines + ): + raise LeaseValidationError( + "heartbeat audit record does not match persistent audit" + ) + except StopIteration as exc: + raise LeaseValidationError("persistent audit is empty") from exc + except (UnicodeError, json.JSONDecodeError) as exc: + raise LeaseValidationError( + "persistent audit does not contain valid JSONL" + ) from exc + except LeaseValidationError: + raise + except OSError as exc: + raise LeaseValidationError( + f"cannot inspect persistent audit {audit_path}: {exc}" + ) from exc + _, _, final_watchdog_start_ticks = _read_proc_stat( + process_procfs_root, watchdog_pid + ) + if final_watchdog_start_ticks != watchdog_start_ticks: + raise LeaseValidationError( + "watchdog process changed during validation" + ) + if pidfd is not None: + os.close(pidfd) + return lease + + +def start_process_group_lease_guard( + expected_script_path: Path, + *, + startup_timeout_seconds: float = 5.0, + expected_procfs_root: Path = Path("/proc"), + process_procfs_root: Path = Path("/proc"), +) -> threading.Thread: + try: + lease_path = Path( + os.environ["STRIX_MEMORY_WATCHDOG_LEASE_PATH"] + ).resolve() + heartbeat_path = Path( + os.environ["STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH"] + ).resolve() + audit_path = Path( + os.environ["STRIX_MEMORY_WATCHDOG_AUDIT_PATH"] + ).resolve() + max_age_seconds = float( + os.environ[ + "STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS" + ] + ) + except (KeyError, ValueError) as exc: + raise LeaseValidationError( + "watchdog artifact environment is missing or invalid" + ) from exc + current_cmdline = _read_proc_bytes( + process_procfs_root, os.getpid(), "cmdline" + ) + expected_command = tuple( + os.fsdecode(argument) + for argument in current_cmdline.split(b"\0") + if argument + ) + deadline = time.monotonic() + startup_timeout_seconds + while True: + try: + lease = validate_active_lease( + lease_path, + expected_script_path=expected_script_path, + expected_procfs_root=expected_procfs_root, + expected_command=expected_command, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + expected_max_heartbeat_age_seconds=max_age_seconds, + process_procfs_root=process_procfs_root, + ) + break + except Exception: + if time.monotonic() >= deadline: + raise + time.sleep(0.01) + + guardian_pid = _require_int( + lease.get("guardian_pid"), "guardian_pid" + ) + signal.signal(LEASE_GUARD_SIGNAL, _kill_own_process_group) + _set_parent_death_signal(LEASE_GUARD_SIGNAL, guardian_pid) + + def monitor() -> None: + interval = min(1.0, max_age_seconds / 3) + while True: + time.sleep(interval) + try: + validate_active_lease( + lease_path, + expected_script_path=expected_script_path, + expected_procfs_root=expected_procfs_root, + expected_command=expected_command, + expected_heartbeat_path=heartbeat_path, + expected_audit_path=audit_path, + expected_max_heartbeat_age_seconds=max_age_seconds, + process_procfs_root=process_procfs_root, + ) + except Exception: + _kill_own_process_group() + + guard = threading.Thread( + target=monitor, + name="strix-watchdog-lease-guard", + daemon=True, + ) + guard.start() + return guard + + +class AuditLogger: + def __init__( + self, + stream: IO[str], + wall_clock: Callable[[], datetime] | None = None, + ): + self.stream = stream + self.stream_enabled = True + self.wall_clock = wall_clock + self.persistent_stream: IO[str] | None = None + self.lease_manager: LeaseManager | None = None + self.finalized = False + self.final_exit_code = EXIT_INTERNAL_ERROR + self.last_record_sha256: str | None = None + + def open_persistent(self, path: Path) -> None: + resolved_path = path.expanduser().resolve() + try: + descriptor = os.open( + resolved_path, + os.O_CREAT + | os.O_EXCL + | os.O_WRONLY + | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + fcntl.flock( + descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB + ) + self.persistent_stream = os.fdopen( + descriptor, "w", encoding="utf-8" + ) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "audit", + f"cannot create persistent audit {resolved_path}: {detail}", + ) from exc + + def persistent_identity(self) -> dict[str, int]: + if self.persistent_stream is None: + raise ArtifactError( + "audit", "persistent audit is not open" + ) + file_status = os.fstat(self.persistent_stream.fileno()) + return { + "device": file_status.st_dev, + "inode": file_status.st_ino, + "uid": file_status.st_uid, + "mode": stat.S_IMODE(file_status.st_mode), + "fd": self.persistent_stream.fileno(), + } + + def close(self) -> None: + persistent_stream = self.persistent_stream + self.persistent_stream = None + if persistent_stream is not None: + try: + persistent_stream.close() + except (OSError, ValueError): + pass + + def disable_component(self, component: str) -> None: + if component == "audit": + self.close() + elif component == "lease": + self.lease_manager = None + elif component == "stderr": + self.stream_enabled = False + + def emit(self, event: str, **fields: object) -> dict[str, object]: + record = { + "timestamp": _timestamp_utc(self.wall_clock), + "event": event, + **fields, + } + line = ( + json.dumps(record, sort_keys=True, separators=(",", ":")) + + "\n" + ) + if self.stream_enabled: + try: + self.stream.write(line) + self.stream.flush() + except (OSError, ValueError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + raise ArtifactError( + "stderr", + f"cannot write standard error audit: {detail}", + ) from exc + self.last_record_sha256 = _sha256_bytes(line.encode("utf-8")) + if self.persistent_stream is not None: + try: + self.persistent_stream.write(line) + self.persistent_stream.flush() + os.fsync(self.persistent_stream.fileno()) + except OSError as exc: + detail = exc.strerror or str(exc) + raise ArtifactError( + "audit", + f"cannot write persistent audit: {detail}", + ) from exc + return record + + def heartbeat(self, sample: dict[str, object]) -> None: + if self.lease_manager is not None: + self.lease_manager.update_heartbeat( + { + **sample, + "audit_record_sha256": self.last_record_sha256, + } + ) + + def finalize(self, record: dict[str, object]) -> None: + if self.lease_manager is not None: + self.lease_manager.finalize(record) + + def mark_final(self, exit_code: int) -> None: + self.finalized = True + self.final_exit_code = exit_code + + +def _child_status(returncode: int | None, started: bool = True) -> str: + if not started: + return "not_started" + if returncode is None: + return "running" + return "signaled" if returncode < 0 else "exited" + + +def _state_fields( + snapshot: HostSnapshot | None, + peak_used_bytes: int | None, + child: ProcessHandle | None, + child_returncode: int | None, + process_group_status: str, + threshold_reason: str, +) -> dict[str, object]: + return { + "total_bytes": snapshot.total_bytes if snapshot else None, + "available_bytes": snapshot.available_bytes if snapshot else None, + "used_bytes": snapshot.used_bytes if snapshot else None, + "swap_entries": len(snapshot.active_swaps) if snapshot else None, + "peak_used_bytes": peak_used_bytes, + "child_pid": child.pid if child else None, + "child_status": _child_status( + child_returncode, started=child is not None + ), + "child_returncode": child_returncode, + "process_group_id": child.pid if child else None, + "process_group_status": process_group_status, + "threshold_reason": threshold_reason, + } + + +def _emit_final( + audit: AuditLogger, + classification: str, + exit_code: int, + reason: str, + snapshot: HostSnapshot | None, + peak_used_bytes: int | None, + child: ProcessHandle | None = None, + child_returncode: int | None = None, + process_group_status: str = "not_created", + error: str | None = None, + preserve_primary_on_artifact_error: bool = False, + secondary_errors: Sequence[dict[str, str]] | None = None, +) -> int: + fields = _state_fields( + snapshot, + peak_used_bytes, + child, + child_returncode, + process_group_status, + reason, + ) + fields.update(classification=classification, exit_code=exit_code) + if error: + fields["error"] = error + if secondary_errors: + fields["secondary_errors"] = list(secondary_errors) + + def record_artifact_error(exc: ArtifactError) -> None: + nonlocal exit_code + detail = { + "component": exc.component, + "detail": str(exc), + } + if preserve_primary_on_artifact_error: + secondary_errors = fields.setdefault( + "secondary_errors", [] + ) + assert isinstance(secondary_errors, list) + secondary_errors.append(detail) + else: + fields.update( + classification="lease_error", + exit_code=EXIT_LEASE_ERROR, + threshold_reason="watchdog artifact finalization failed", + error=f"{exc.component}: {exc}", + ) + exit_code = EXIT_LEASE_ERROR + + def emit_final_record() -> dict[str, object]: + try: + return audit.emit("final", **fields) + except ArtifactError as exc: + audit.disable_component(exc.component) + record_artifact_error(exc) + try: + return audit.emit("final", **fields) + except ArtifactError as exc: + audit.disable_component(exc.component) + record_artifact_error(exc) + return audit.emit("final", **fields) + + previous_mask = signal.pthread_sigmask( + signal.SIG_BLOCK, PARENT_SIGNALS + ) + try: + record = emit_final_record() + try: + audit.finalize(record) + except ArtifactError as exc: + audit.disable_component(exc.component) + record_artifact_error(exc) + emit_final_record() + audit.mark_final(exit_code) + return exit_code + finally: + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + + +def _signal_process_group(process_group_id: int, signal_number: int) -> str: + try: + os.killpg(process_group_id, signal_number) + except ProcessLookupError: + return "missing" + except OSError as exc: + name = signal.Signals(signal_number).name + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot send {name} to process group {process_group_id}: {detail}" + ) from exc + return f"{signal.Signals(signal_number).name.lower()}_sent" + + +def _process_group_alive(process_group_id: int) -> bool: + if sys.platform.startswith("linux"): + try: + process_paths = Path("/proc").iterdir() + for process_path in process_paths: + if not process_path.name.isdigit(): + continue + try: + content = ( + process_path / "stat" + ).read_text(encoding="utf-8") + close_paren = content.rfind(")") + fields = content[close_paren + 1:].split() + if ( + close_paren >= 0 + and len(fields) >= 3 + and fields[0] != "Z" + and int(fields[2]) == process_group_id + ): + return True + except (OSError, UnicodeError, ValueError): + continue + return False + except OSError: + pass + try: + os.killpg(process_group_id, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + except OSError as exc: + detail = exc.strerror or str(exc) + raise ProcessGroupError( + f"cannot inspect process group {process_group_id}: {detail}" + ) from exc + return True + + +def _raise_parent_signal(signal_number: int, _frame: object) -> None: + raise ParentSignal(signal_number) + + +def _set_parent_signal_handlers( + handler: Any, +) -> dict[int, Any]: + previous: dict[int, Any] = {} + for signal_number in PARENT_SIGNALS: + previous[signal_number] = signal.signal(signal_number, handler) + return previous + + +def _restore_parent_signal_handlers( + previous: dict[int, Any], +) -> None: + for signal_number, handler in previous.items(): + signal.signal(signal_number, handler) + + +def _kill_and_finish( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + signal_group: Callable[[int, int], str], +) -> int: + try: + group_status = signal_group(child.pid, signal.SIGKILL) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + + artifact_error: ArtifactError | None = None + try: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + group_status, + reason, + ), + signal="SIGKILL", + ) + except ArtifactError as exc: + artifact_error = exc + audit.disable_component(exc.component) + try: + child_returncode = child.wait(timeout=5.0) + except subprocess.TimeoutExpired as exc: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "sigkill_timeout", + str(exc), + ) + if artifact_error is not None: + classification = "lease_error" + exit_code = EXIT_LEASE_ERROR + error = f"{artifact_error.component}: {artifact_error}" + else: + error = None + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + group_status, + error, + ) + + +def _graceful_cleanup( + audit: AuditLogger, + child: ProcessHandle, + snapshot: HostSnapshot, + peak_used_bytes: int, + classification: str, + exit_code: int, + reason: str, + graceful_signal: int | None, + grace_seconds: float, + signal_group: Callable[[int, int], str], + group_alive: Callable[[int], bool], + monotonic: Callable[[], float], + sleeper: Callable[[float], None], + process_group_status: str = "active", + escalation_result: tuple[str, int, str] | None = None, + error: str | None = None, +) -> int: + escalated = False + artifact_error: ArtifactError | None = None + guardian_control_error: ProcessGroupError | None = None + try: + if graceful_signal is not None: + process_group_status = signal_group( + child.pid, graceful_signal + ) + try: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + reason, + ), + signal=signal.Signals(graceful_signal).name, + ) + except ArtifactError as exc: + artifact_error = exc + audit.disable_component(exc.component) + if ( + isinstance(child, GuardianProcess) + and child.poll() is None + ): + try: + child.begin_grace() + except ProcessGroupError as exc: + guardian_control_error = exc + deadline = monotonic() + grace_seconds + while ( + guardian_control_error is None + and monotonic() < deadline + ): + child.poll() + if not group_alive(child.pid): + break + if ( + isinstance(child, GuardianProcess) + and child.poll() is None + ): + try: + child.pulse() + except ProcessGroupError as exc: + guardian_control_error = exc + break + sleeper(min(0.05, deadline - monotonic())) + child.poll() + if ( + guardian_control_error is not None + or group_alive(child.pid) + ): + escalated = True + process_group_status = signal_group( + child.pid, signal.SIGKILL + ) + if guardian_control_error is not None: + signal_reason = ( + "guardian control failed during graceful cleanup" + ) + else: + signal_reason = ( + escalation_result[2] + if escalation_result is not None + else reason + ) + try: + audit.emit( + "process_group_signal", + **_state_fields( + snapshot, + peak_used_bytes, + child, + child.poll(), + process_group_status, + signal_reason, + ), + signal="SIGKILL", + ) + except ArtifactError as exc: + if artifact_error is None: + artifact_error = exc + audit.disable_component(exc.component) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + + child_returncode = child.poll() + if child_returncode is None: + try: + child_returncode = child.wait(timeout=5.0) + except subprocess.TimeoutExpired as exc: + return _emit_final( + audit, + "termination_timeout", + EXIT_SIGNAL_ERROR, + reason, + snapshot, + peak_used_bytes, + child, + child.poll(), + "termination_timeout", + str(exc), + ) + + if guardian_control_error is not None: + classification = "signal_error" + exit_code = EXIT_SIGNAL_ERROR + reason = "guardian control failed during graceful cleanup" + error = str(guardian_control_error) + elif escalated and escalation_result is not None: + classification, exit_code, reason = escalation_result + if artifact_error is not None and guardian_control_error is None: + classification = "lease_error" + exit_code = EXIT_LEASE_ERROR + error = f"{artifact_error.component}: {artifact_error}" + return _emit_final( + audit, + classification, + exit_code, + reason, + snapshot, + peak_used_bytes, + child, + child_returncode, + process_group_status, + error, + preserve_primary_on_artifact_error=( + guardian_control_error is not None + ), + secondary_errors=( + [ + { + "component": artifact_error.component, + "detail": str(artifact_error), + } + ] + if ( + guardian_control_error is not None + and artifact_error is not None + ) + else None + ), + ) + + +def _monitor_child( + config: WatchdogConfig, + reader: ProcfsReader, + audit: AuditLogger, + child: ProcessHandle, + state: RuntimeState, + signal_group: Callable[[int, int], str], + group_alive: Callable[[int], bool], + pulse_guardian: Callable[[], None], + monotonic: Callable[[], float], + sleeper: Callable[[float], None], +) -> int: + soft_deadline: float | None = None + + while True: + child_returncode = child.poll() + if child_returncode is not None: + soft_stop = soft_deadline is not None + classification = "soft_limit" if soft_stop else "child_exit" + exit_code = EXIT_SOFT_LIMIT if soft_stop else ( + 128 - child_returncode + if child_returncode < 0 + else child_returncode + ) + reason = ( + "child exited during soft-threshold grace period" + if soft_stop + else "child exited" + ) + if group_alive(child.pid): + grace_seconds = config.grace_seconds + graceful_signal: int | None = signal.SIGTERM + group_status = "active" + if soft_stop: + grace_seconds = max( + 0.0, soft_deadline - monotonic() + ) + graceful_signal = None + group_status = "sigterm_sent" + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + classification, + exit_code, + ( + f"{reason}; process group members still running" + ), + graceful_signal, + grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + group_status, + ( + ( + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired with " + "process group members still running", + ) + if soft_stop + else None + ), + ) + return _emit_final( + audit, + classification, + exit_code, + reason, + state.snapshot, + state.peak_used_bytes, + child, + child_returncode, + "leader_exited", + ) + + now = monotonic() + if soft_deadline is not None and now >= soft_deadline: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "grace_timeout", + EXIT_GRACE_TIMEOUT, + "soft-threshold grace period expired", + signal_group, + ) + + try: + state.snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + signal_group, + ) + + state.peak_used_bytes = max( + state.peak_used_bytes, state.snapshot.used_bytes + ) + if state.snapshot.active_swaps: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "swap_appeared", + EXIT_SWAP_ACTIVE, + "active swap appeared during execution", + signal_group, + ) + if state.snapshot.used_bytes >= config.emergency_bytes: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "emergency_limit", + EXIT_EMERGENCY_LIMIT, + "used_bytes >= emergency_bytes", + signal_group, + ) + soft_signal_fields: dict[str, object] | None = None + if ( + soft_deadline is None + and state.snapshot.used_bytes >= config.soft_bytes + ): + try: + group_status = signal_group(child.pid, signal.SIGTERM) + except ProcessGroupError as exc: + return _emit_final( + audit, + "signal_error", + EXIT_SIGNAL_ERROR, + "used_bytes >= soft_bytes", + state.snapshot, + state.peak_used_bytes, + child, + child.poll(), + "signal_error", + str(exc), + ) + soft_deadline = now + config.grace_seconds + if isinstance(child, GuardianProcess): + try: + child.begin_grace() + except ProcessGroupError as exc: + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "signal_error", + EXIT_SIGNAL_ERROR, + str(exc), + signal_group, + ) + soft_signal_fields = { + **_state_fields( + state.snapshot, + state.peak_used_bytes, + child, + child.poll(), + group_status, + "used_bytes >= soft_bytes", + ), + "signal": "SIGTERM", + "grace_deadline_monotonic": soft_deadline, + } + try: + pulse_guardian() + except ProcessGroupError as exc: + if child.poll() is not None: + continue + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "signal_error", + EXIT_SIGNAL_ERROR, + str(exc), + signal_group, + ) + if soft_signal_fields is not None: + audit.emit( + "process_group_signal", + **soft_signal_fields, + ) + sample_record = audit.emit( + "sample", + **_state_fields( + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", + ) + ) + audit.heartbeat(sample_record) + try: + pulse_guardian() + except ProcessGroupError as exc: + if child.poll() is not None: + continue + return _kill_and_finish( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "signal_error", + EXIT_SIGNAL_ERROR, + str(exc), + signal_group, + ) + + sleep_seconds = config.sample_interval_seconds + if soft_deadline is not None: + sleep_seconds = min( + sleep_seconds, + max(0.0, soft_deadline - monotonic()), + ) + sleeper(sleep_seconds) + + +def run_watchdog( + config: WatchdogConfig, + *, + reader: ProcfsReader | None = None, + audit: AuditLogger | None = None, + launcher: Callable[..., ProcessHandle] | None = None, + signal_group: Callable[[int, int], str] | None = None, + group_alive: Callable[[int], bool] | None = None, + monotonic: Callable[[], float] | None = None, + sleeper: Callable[[float], None] | None = None, +) -> int: + artifact_paths = config.validate() + use_guardian = ( + launcher is None and sys.platform.startswith("linux") + ) + reader = reader or ProcfsReader(config.procfs_root) + audit = audit or AuditLogger(sys.stderr) + launcher = launcher or subprocess.Popen + signal_group = signal_group or _signal_process_group + group_alive = group_alive or _process_group_alive + monotonic = monotonic or time.monotonic + sleeper = sleeper or time.sleep + + if artifact_paths is not None: + try: + audit.open_persistent(artifact_paths.audit) + except ArtifactError as exc: + return _emit_final( + audit, + "lease_error", + EXIT_LEASE_ERROR, + "cannot initialize watchdog artifacts", + None, + None, + error=f"{exc.component}: {exc}", + ) + + try: + snapshot = reader.read_snapshot() + except ProcfsError as exc: + return _emit_final( + audit, + "procfs_error", + EXIT_PROCFS_ERROR, + str(exc), + None, + None, + error=str(exc), + ) + + try: + audit.emit( + "preflight", + **_state_fields( + snapshot, + snapshot.used_bytes, + None, + None, + "not_created", + "none", + ), + soft_bytes=config.soft_bytes, + emergency_bytes=config.emergency_bytes, + strict_ceiling_bytes=STRICT_CEILING_BYTES, + ) + except ArtifactError as exc: + audit.disable_component(exc.component) + return _emit_final( + audit, + "lease_error", + EXIT_LEASE_ERROR, + "cannot write watchdog preflight audit", + snapshot, + snapshot.used_bytes, + error=f"{exc.component}: {exc}", + ) + + if snapshot.active_swaps: + return _emit_final( + audit, + "startup_swap_active", + EXIT_SWAP_ACTIVE, + "active swap present before command launch", + snapshot, + snapshot.used_bytes, + ) + if snapshot.used_bytes >= config.emergency_bytes: + return _emit_final( + audit, + "startup_emergency_limit", + EXIT_EMERGENCY_LIMIT, + "used_bytes >= emergency_bytes before launch", + snapshot, + snapshot.used_bytes, + ) + if snapshot.used_bytes >= config.soft_bytes: + return _emit_final( + audit, + "startup_soft_limit", + EXIT_SOFT_LIMIT, + "used_bytes >= soft_bytes before launch", + snapshot, + snapshot.used_bytes, + ) + + previous_mask = signal.pthread_sigmask( + signal.SIG_BLOCK, PARENT_SIGNALS + ) + mask_restored = False + previous_handlers: dict[int, Any] = {} + child: ProcessHandle | None = None + state = RuntimeState(snapshot, snapshot.used_bytes) + try: + launch_mask = previous_mask + lease_manager = ( + LeaseManager(config, artifact_paths) + if artifact_paths is not None + else None + ) + + def restore_child_signal_mask() -> None: + signal.pthread_sigmask(signal.SIG_SETMASK, launch_mask) + + try: + child_environment = os.environ.copy() + if lease_manager is not None: + child_environment.update( + { + "STRIX_MEMORY_WATCHDOG_LEASE_PATH": str( + lease_manager.lease_path + ), + "STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH": str( + lease_manager.heartbeat_path + ), + "STRIX_MEMORY_WATCHDOG_AUDIT_PATH": str( + lease_manager.audit_path + ), + "STRIX_MEMORY_WATCHDOG_HEARTBEAT_MAX_AGE_SECONDS": ( + str(config.heartbeat_max_age_seconds) + ), + } + ) + if use_guardian: + child = _launch_guardian( + config.command, + child_environment, + config.heartbeat_max_age_seconds, + config.grace_seconds + 1.0, + launch_mask, + ) + elif lease_manager is not None: + child = launcher( + config.command, + start_new_session=True, + preexec_fn=restore_child_signal_mask, + env=child_environment, + ) + else: + child = launcher( + config.command, + start_new_session=True, + preexec_fn=restore_child_signal_mask, + ) + except (OSError, ValueError, subprocess.SubprocessError) as exc: + detail = getattr(exc, "strerror", None) or str(exc) + return _emit_final( + audit, + "launch_error", + EXIT_LAUNCH_ERROR, + "command launch failed", + snapshot, + snapshot.used_bytes, + error=detail, + ) + + previous_handlers = _set_parent_signal_handlers( + _raise_parent_signal + ) + if lease_manager is not None: + lease_manager.start(child, audit) + audit.lease_manager = lease_manager + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + mask_restored = True + audit.emit( + "child_started", + **_state_fields( + state.snapshot, + state.peak_used_bytes, + child, + None, + "active", + "none", + ), + command=list(config.command), + ) + return _monitor_child( + config, + reader, + audit, + child, + state, + signal_group, + group_alive, + child.pulse if isinstance(child, GuardianProcess) else lambda: None, + monotonic, + sleeper, + ) + except ArtifactError as exc: + _set_parent_signal_handlers(signal.SIG_IGN) + audit.disable_component(exc.component) + if child is None: + return _emit_final( + audit, + "lease_error", + EXIT_LEASE_ERROR, + "watchdog artifact initialization failed", + state.snapshot, + state.peak_used_bytes, + error=f"{exc.component}: {exc}", + ) + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "lease_error", + EXIT_LEASE_ERROR, + "watchdog artifact update failed", + signal.SIGTERM, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + error=f"{exc.component}: {exc}", + ) + except ParentSignal as exc: + if audit.finalized: + return audit.final_exit_code + _set_parent_signal_handlers(signal.SIG_IGN) + assert child is not None + signal_name = signal.Signals(exc.signal_number).name + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "parent_signal", + 128 + exc.signal_number, + f"wrapper received {signal_name}", + exc.signal_number, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + ) + except Exception as exc: + if audit.finalized: + return audit.final_exit_code + _set_parent_signal_handlers(signal.SIG_IGN) + if child is None: + return _emit_final( + audit, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected pre-launch exception", + state.snapshot, + state.peak_used_bytes, + error=f"{type(exc).__name__}: {exc}", + ) + return _graceful_cleanup( + audit, + child, + state.snapshot, + state.peak_used_bytes, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected post-launch exception", + signal.SIGTERM, + config.grace_seconds, + signal_group, + group_alive, + monotonic, + sleeper, + error=f"{type(exc).__name__}: {exc}", + ) + finally: + if not mask_restored: + signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask) + if previous_handlers: + _restore_parent_signal_handlers(previous_handlers) + if isinstance(child, GuardianProcess): + child.close() + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be greater than zero") + return parsed + + +def _positive_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError("value must be greater than zero") + return parsed + + +def parse_args(argv: Sequence[str]) -> WatchdogConfig: + parser = argparse.ArgumentParser( + description=( + "Launch a command in a new process group and stop it before " + "host-wide memory use reaches the 120 GiB Strix validation ceiling." + ) + ) + parser.add_argument( + "--procfs-root", + type=Path, + default=Path("/proc"), + help="procfs root containing meminfo and swaps (default: /proc)", + ) + parser.add_argument( + "--soft-gib", + type=_positive_int, + default=116, + help="send SIGTERM at this many GiB used (default: 116)", + ) + parser.add_argument( + "--emergency-gib", + type=_positive_int, + default=118, + help=( + "send SIGKILL at this many GiB used (default: 118, leaving " + "a 2 GiB sampling margin below 120 GiB)" + ), + ) + parser.add_argument( + "--grace-seconds", + type=_positive_float, + default=DEFAULT_GRACE_SECONDS, + help="maximum time after SIGTERM before SIGKILL (default: 30)", + ) + parser.add_argument( + "--sample-interval-seconds", + type=_positive_float, + default=DEFAULT_SAMPLE_INTERVAL_SECONDS, + help="procfs sampling interval (default: 1)", + ) + parser.add_argument( + "--lease-path", + type=Path, + help=( + "atomically publish the watchdog-owned lease JSON; requires " + "--heartbeat-path and --audit-path" + ), + ) + parser.add_argument( + "--heartbeat-path", + type=Path, + help=( + "atomically update watchdog heartbeat JSON on every sample; " + "requires --lease-path and --audit-path" + ), + ) + parser.add_argument( + "--audit-path", + type=Path, + help=( + "create a persistent JSONL audit in addition to standard error; " + "requires --lease-path and --heartbeat-path" + ), + ) + parser.add_argument( + "--heartbeat-max-age-seconds", + type=_positive_float, + default=DEFAULT_HEARTBEAT_MAX_AGE_SECONDS, + help=( + "maximum heartbeat age accepted by a matching harness " + "(default: 5)" + ), + ) + parser.add_argument( + "command", + nargs=argparse.REMAINDER, + help="command and arguments, preceded by --", + ) + args = parser.parse_args(argv) + command = tuple(args.command) + if command and command[0] == "--": + command = command[1:] + return WatchdogConfig( + command=command, + procfs_root=args.procfs_root, + soft_bytes=args.soft_gib * GIB, + emergency_bytes=args.emergency_gib * GIB, + grace_seconds=args.grace_seconds, + sample_interval_seconds=args.sample_interval_seconds, + lease_path=args.lease_path, + heartbeat_path=args.heartbeat_path, + audit_path=args.audit_path, + heartbeat_max_age_seconds=args.heartbeat_max_age_seconds, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = tuple(argv if argv is not None else sys.argv[1:]) + if arguments and arguments[0] == "--internal-guardian": + if len(arguments) < 7 or arguments[5] != "--": + return EXIT_LAUNCH_ERROR + try: + return _guardian_main( + int(arguments[1]), + int(arguments[2]), + _positive_float(arguments[3]), + _positive_float(arguments[4]), + tuple(arguments[6:]), + ) + except (OSError, ValueError): + return EXIT_LAUNCH_ERROR + config = parse_args(arguments) + audit = AuditLogger(sys.stderr) + try: + return run_watchdog(config, audit=audit) + except ValueError as exc: + return _emit_final( + audit, + "configuration_error", + EXIT_PROCFS_ERROR, + "invalid configuration", + None, + None, + error=str(exc), + ) + except Exception as exc: + return _emit_final( + audit, + "internal_error", + EXIT_INTERNAL_ERROR, + "unexpected watchdog error", + None, + None, + error=f"{type(exc).__name__}: {exc}", + ) + finally: + audit.close() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b3d17330ca54..971d75f5706a 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,10 @@ set(LLAMA_CORE_SOURCES llama-chat.cpp llama-context.cpp llama-cparams.cpp + llama-dsv41.cpp + llama-dsv41-engram.cpp + llama-dsv41-expert.cpp + llama-expert-store.cpp llama-grammar.cpp llama-graph.cpp llama-hparams.cpp @@ -28,11 +32,14 @@ set(LLAMA_CORE_SOURCES llama-kv-cache-msa.cpp llama-kv-cache-dsv4.cpp llama-memory.cpp + llama-memory-dsv41.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp llama-memory-hybrid-idx.cpp llama-memory-recurrent.cpp llama-mmap.cpp + llama-bounded-file.cpp + llama-engram.cpp llama-ple-disk.cpp llama-model-loader.cpp llama-model-saver.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index c5e06fe28993..3cfe06514042 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -80,6 +80,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" }, { LLM_ARCH_DEEPSEEK32, "deepseek32" }, { LLM_ARCH_DEEPSEEK4, "deepseek4" }, + { LLM_ARCH_DEEPSEEK41, "deepseek41" }, { LLM_ARCH_CHATGLM, "chatglm" }, { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, @@ -309,6 +310,57 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_PLE_EOS_TOKEN_ID, "%s.ple.eos_token_id" }, { LLM_KV_PLE_IMAGE_TOKEN_ID, "%s.ple.image_token_id" }, + { LLM_KV_DSV41_CONFIG, "%s.config" }, + { LLM_KV_DSV41_VOCAB_SIZE, "%s.vocab_size" }, + { LLM_KV_DSV41_HIDDEN_SIZE, "%s.hidden_size" }, + { LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, "%s.moe_intermediate_size" }, + { LLM_KV_DSV41_NUM_HIDDEN_LAYERS, "%s.num_hidden_layers" }, + { LLM_KV_DSV41_NUM_ATTENTION_HEADS, "%s.num_attention_heads" }, + { LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, "%s.num_key_value_heads" }, + { LLM_KV_DSV41_HEAD_DIM, "%s.head_dim" }, + { LLM_KV_DSV41_QK_ROPE_HEAD_DIM, "%s.qk_rope_head_dim" }, + { LLM_KV_DSV41_Q_LORA_RANK, "%s.q_lora_rank" }, + { LLM_KV_DSV41_O_LORA_RANK, "%s.o_lora_rank" }, + { LLM_KV_DSV41_O_GROUPS, "%s.o_groups" }, + { LLM_KV_DSV41_N_ROUTED_EXPERTS, "%s.n_routed_experts" }, + { LLM_KV_DSV41_N_SHARED_EXPERTS, "%s.n_shared_experts" }, + { LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, "%s.num_experts_per_tok" }, + { LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, "%s.max_position_embeddings" }, + { LLM_KV_DSV41_SLIDING_WINDOW, "%s.sliding_window" }, + { LLM_KV_DSV41_INDEX_N_HEADS, "%s.index_n_heads" }, + { LLM_KV_DSV41_INDEX_HEAD_DIM, "%s.index_head_dim" }, + { LLM_KV_DSV41_INDEX_TOPK, "%s.index_topk" }, + { LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, "%s.candidate_source_layer_id" }, + { LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, "%s.candidate_topk_blocks" }, + { LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, "%s.candidate_block_size" }, + { LLM_KV_DSV41_HC_MULT, "%s.hc_mult" }, + { LLM_KV_DSV41_HC_SINKHORN_ITERS, "%s.hc_sinkhorn_iters" }, + { LLM_KV_DSV41_ROPE_THETA, "%s.rope_theta" }, + { LLM_KV_DSV41_COMPRESS_ROPE_THETA, "%s.compress_rope_theta" }, + { LLM_KV_DSV41_RMS_NORM_EPS, "%s.rms_norm_eps" }, + { LLM_KV_DSV41_HC_EPS, "%s.hc_eps" }, + { LLM_KV_DSV41_SWIGLU_LIMIT, "%s.swiglu_limit" }, + { LLM_KV_DSV41_ROUTED_SCALING_FACTOR, "%s.routed_scaling_factor" }, + { LLM_KV_DSV41_SCORING_FUNC, "%s.scoring_func" }, + { LLM_KV_DSV41_HIDDEN_ACT, "%s.hidden_act" }, + { LLM_KV_DSV41_TOPK_METHOD, "%s.topk_method" }, + { LLM_KV_DSV41_NORM_TOPK_PROB, "%s.norm_topk_prob" }, + { LLM_KV_DSV41_COMPRESS_RATIOS, "%s.compress_ratios" }, + { LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, "%s.kv_source_layer_ids" }, + { LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, "%s.index_source_layer_ids" }, + { LLM_KV_DSV41_ROPE_SCALING_FACTOR, "%s.rope_scaling.factor" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, "%s.rope_scaling.beta_fast" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, "%s.rope_scaling.beta_slow" }, + { LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, "%s.rope_scaling.original_max_position_embeddings" }, + { LLM_KV_DSV41_ENGRAM_ENCODING, "%s.engram.encoding" }, + { LLM_KV_DSV41_ENGRAM_LAYER_IDS, "%s.engram.layer_ids" }, + { LLM_KV_DSV41_ENGRAM_ROWS, "%s.engram.rows" }, + { LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, "%s.engram.compressed_vocab_size" }, + { LLM_KV_DSV41_ENGRAM_PAD_ID, "%s.engram.pad_id" }, + { LLM_KV_DSV41_ENGRAM_TOKEN_MAP, "%s.engram.token_map" }, + { LLM_KV_DSV41_ENGRAM_PRIMES, "%s.engram.primes" }, + { LLM_KV_DSV41_ENGRAM_MULTIPLIERS, "%s.engram.multipliers" }, + { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" }, @@ -545,6 +597,10 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_PLE_NORM_QUERY, "blk.%d.ple_norm_query" }, { LLM_TENSOR_PLE_NORM_CONV, "blk.%d.ple_norm_conv" }, { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, + { LLM_TENSOR_ENGRAM_EMBD, "blk.%d.engram_embd" }, + { LLM_TENSOR_ENGRAM_Q_NORM, "blk.%d.engram_q_norm" }, + { LLM_TENSOR_ENGRAM_K_NORM, "blk.%d.engram_k_norm" }, + { LLM_TENSOR_ENGRAM_KV, "blk.%d.engram_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WKV, "blk.%d.attn_compressor_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "blk.%d.attn_compressor_gate" }, { LLM_TENSOR_ATTN_COMPRESSOR_APE, "blk.%d.attn_compressor_ape" }, @@ -776,6 +832,10 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_PLE_NORM_QUERY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_PLE_NORM_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_PLE_CONV1D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_ENGRAM_EMBD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_ENGRAM_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ENGRAM_K_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ENGRAM_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, @@ -1088,6 +1148,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_MINIMAX_01: return true; default: diff --git a/src/llama-arch.h b/src/llama-arch.h index b8d15e7edd51..29f4308f83ec 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -85,6 +85,7 @@ enum llm_arch { LLM_ARCH_DEEPSEEK2OCR, LLM_ARCH_DEEPSEEK32, LLM_ARCH_DEEPSEEK4, + LLM_ARCH_DEEPSEEK41, LLM_ARCH_CHATGLM, LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, @@ -314,6 +315,57 @@ enum llm_kv { LLM_KV_PLE_EOS_TOKEN_ID, LLM_KV_PLE_IMAGE_TOKEN_ID, + LLM_KV_DSV41_CONFIG, + LLM_KV_DSV41_VOCAB_SIZE, + LLM_KV_DSV41_HIDDEN_SIZE, + LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, + LLM_KV_DSV41_NUM_HIDDEN_LAYERS, + LLM_KV_DSV41_NUM_ATTENTION_HEADS, + LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, + LLM_KV_DSV41_HEAD_DIM, + LLM_KV_DSV41_QK_ROPE_HEAD_DIM, + LLM_KV_DSV41_Q_LORA_RANK, + LLM_KV_DSV41_O_LORA_RANK, + LLM_KV_DSV41_O_GROUPS, + LLM_KV_DSV41_N_ROUTED_EXPERTS, + LLM_KV_DSV41_N_SHARED_EXPERTS, + LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, + LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, + LLM_KV_DSV41_SLIDING_WINDOW, + LLM_KV_DSV41_INDEX_N_HEADS, + LLM_KV_DSV41_INDEX_HEAD_DIM, + LLM_KV_DSV41_INDEX_TOPK, + LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, + LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, + LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, + LLM_KV_DSV41_HC_MULT, + LLM_KV_DSV41_HC_SINKHORN_ITERS, + LLM_KV_DSV41_ROPE_THETA, + LLM_KV_DSV41_COMPRESS_ROPE_THETA, + LLM_KV_DSV41_RMS_NORM_EPS, + LLM_KV_DSV41_HC_EPS, + LLM_KV_DSV41_SWIGLU_LIMIT, + LLM_KV_DSV41_ROUTED_SCALING_FACTOR, + LLM_KV_DSV41_SCORING_FUNC, + LLM_KV_DSV41_HIDDEN_ACT, + LLM_KV_DSV41_TOPK_METHOD, + LLM_KV_DSV41_NORM_TOPK_PROB, + LLM_KV_DSV41_COMPRESS_RATIOS, + LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, + LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, + LLM_KV_DSV41_ROPE_SCALING_FACTOR, + LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, + LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, + LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, + LLM_KV_DSV41_ENGRAM_ENCODING, + LLM_KV_DSV41_ENGRAM_LAYER_IDS, + LLM_KV_DSV41_ENGRAM_ROWS, + LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, + LLM_KV_DSV41_ENGRAM_PAD_ID, + LLM_KV_DSV41_ENGRAM_TOKEN_MAP, + LLM_KV_DSV41_ENGRAM_PRIMES, + LLM_KV_DSV41_ENGRAM_MULTIPLIERS, + LLM_KV_HASH_LAYER_COUNT, LLM_KV_ROPE_DIMENSION_COUNT, @@ -609,6 +661,10 @@ enum llm_tensor { LLM_TENSOR_PLE_NORM_QUERY, // qwen4exp LLM_TENSOR_PLE_NORM_CONV, // qwen4exp LLM_TENSOR_PLE_CONV1D, // qwen4exp + LLM_TENSOR_ENGRAM_EMBD, + LLM_TENSOR_ENGRAM_Q_NORM, + LLM_TENSOR_ENGRAM_K_NORM, + LLM_TENSOR_ENGRAM_KV, LLM_TENSOR_ATTN_COMPRESSOR_WKV, LLM_TENSOR_ATTN_COMPRESSOR_WGATE, LLM_TENSOR_ATTN_COMPRESSOR_APE, diff --git a/src/llama-bounded-file.cpp b/src/llama-bounded-file.cpp new file mode 100644 index 000000000000..3ae6d5c31a18 --- /dev/null +++ b/src/llama-bounded-file.cpp @@ -0,0 +1,244 @@ +#include "llama-bounded-file.h" + +#include "llama-impl.h" + +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#endif + +struct llama_bounded_file::buffer::impl { + void * ptr = nullptr; + size_t bytes = 0; + + ~impl() { + free(ptr); + } +}; + +llama_bounded_file::buffer::buffer() : pimpl(std::make_unique()) {} +llama_bounded_file::buffer::~buffer() = default; +llama_bounded_file::buffer::buffer(buffer && other) noexcept = default; +llama_bounded_file::buffer & llama_bounded_file::buffer::operator=(buffer && other) noexcept = default; + +bool llama_bounded_file::buffer::empty() const { + return pimpl->ptr == nullptr; +} + +struct llama_bounded_file::impl { + std::string fname; + int fd = -1; + uint64_t file_size = 0; + size_t block = 4096; + bool direct = false; + + impl(const std::string & fname, const params & p) : fname(fname) { +#if defined(_WIN32) + GGML_UNUSED(p); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + int direct_error = 0; + if (p.direct_io) { +#if defined(O_DIRECT) + fd = open(fname.c_str(), O_RDONLY | O_DIRECT | O_CLOEXEC); + direct = fd >= 0; + direct_error = direct ? 0 : errno; +#else + fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); + direct_error = fd >= 0 ? ENOTSUP : errno; +#if defined(F_NOCACHE) + if (fd >= 0) { + if (fcntl(fd, F_NOCACHE, 1) == 0) { + direct = true; +#if defined(F_RDAHEAD) + if (fcntl(fd, F_RDAHEAD, 0) != 0) { + direct = false; + direct_error = errno; + } +#endif + } else { + direct_error = errno; + } + } +#endif +#endif + if (!direct && p.direct_io_required) { + const int saved = direct_error ? direct_error : ENOTSUP; + if (fd >= 0) { + close(fd); + fd = -1; + } + throw std::runtime_error(format("llama_bounded_file: uncached open of %s failed: %s", + fname.c_str(), strerror(saved))); + } + if (!direct) { + LLAMA_LOG_WARN("%s: uncached open of %s failed (%s); falling back to buffered reads\n", + __func__, fname.c_str(), strerror(direct_error ? direct_error : ENOTSUP)); + } + if (!direct && fd >= 0) { + close(fd); + fd = -1; + } + } + if (fd < 0) { + fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); + if (fd < 0) { + throw std::runtime_error(format("llama_bounded_file: failed to open %s: %s", + fname.c_str(), strerror(errno))); + } + } + + struct stat st = {}; + if (fstat(fd, &st) != 0) { + const int saved = errno; + close(fd); + fd = -1; + throw std::runtime_error(format("llama_bounded_file: fstat of %s failed: %s", + fname.c_str(), strerror(saved))); + } + if (!S_ISREG(st.st_mode) || st.st_size < 0) { + close(fd); + fd = -1; + throw std::runtime_error(format("llama_bounded_file: %s is not a regular file", fname.c_str())); + } + file_size = (uint64_t) st.st_size; +#endif + } + + ~impl() { +#if !defined(_WIN32) + if (fd >= 0) { + close(fd); + } +#endif + } + + size_t read_size(uint64_t offset, size_t size) const { + if (!direct || size == 0) { + return size; + } + const size_t prefix = (size_t) (offset & (block - 1)); + if (size > SIZE_MAX - prefix || size + prefix > SIZE_MAX - (block - 1)) { + throw std::overflow_error("llama_bounded_file: aligned read size overflow"); + } + return (size + prefix + block - 1) & ~(block - 1); + } + +#if !defined(_WIN32) + void pread_full(void * dst, size_t len, uint64_t offset, size_t need) const { + size_t done = 0; + while (done < len) { + const size_t remaining = len - done; + const ssize_t n = pread(fd, (uint8_t *) dst + done, remaining, (off_t) (offset + done)); + if (n < 0) { + if (errno == EINTR) { + continue; + } + throw std::runtime_error(format("llama_bounded_file: pread of %s failed at %llu: %s", + fname.c_str(), (unsigned long long) (offset + done), strerror(errno))); + } + if (n == 0) { + break; + } + done += (size_t) n; + if (direct && done >= need) { + return; + } + if (direct && (size_t) n < remaining) { + break; + } + } + if (done < need) { + throw std::runtime_error(format("llama_bounded_file: short read in %s: %zu of %zu bytes at %llu", + fname.c_str(), done, need, (unsigned long long) offset)); + } + } +#endif +}; + +llama_bounded_file::llama_bounded_file(const std::string & fname, const params & p) + : pimpl(std::make_unique(fname, p)) {} + +llama_bounded_file::~llama_bounded_file() = default; + +llama_bounded_file::buffer llama_bounded_file::make_buffer(size_t size) const { + buffer result; + if (!pimpl->direct || size == 0) { + return result; + } +#if defined(_WIN32) + GGML_UNUSED(size); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + const size_t bytes = pimpl->read_size(pimpl->block - 1, size); + void * ptr = nullptr; + const int err = posix_memalign(&ptr, pimpl->block, bytes); + if (err != 0) { + throw std::runtime_error(format("llama_bounded_file: posix_memalign failed: %s", strerror(err))); + } + result.pimpl->ptr = ptr; + result.pimpl->bytes = bytes; + return result; +#endif +} + +void llama_bounded_file::read(uint64_t offset, void * dst, size_t size, buffer & scratch) const { +#if defined(_WIN32) + GGML_UNUSED(offset); + GGML_UNUSED(dst); + GGML_UNUSED(size); + GGML_UNUSED(scratch); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + if (size == 0) { + return; + } + if (dst == nullptr || offset > pimpl->file_size || size > pimpl->file_size - offset || + offset > (uint64_t) INT64_MAX || size > (uint64_t) INT64_MAX - offset) { + throw std::invalid_argument(format("llama_bounded_file: invalid read of %zu bytes at %llu in %s", + size, (unsigned long long) offset, pimpl->fname.c_str())); + } + if (!pimpl->direct) { + pimpl->pread_full(dst, size, offset, size); + return; + } + + const uint64_t aligned_offset = offset & ~(uint64_t) (pimpl->block - 1); + const size_t prefix = (size_t) (offset - aligned_offset); + const size_t aligned_size = pimpl->read_size(offset, size); + if (aligned_size > (uint64_t) INT64_MAX - aligned_offset) { + throw std::overflow_error("llama_bounded_file: aligned read offset overflow"); + } + if (scratch.pimpl->ptr == nullptr || scratch.pimpl->bytes < aligned_size) { + throw std::invalid_argument("llama_bounded_file: aligned scratch buffer is too small"); + } + pimpl->pread_full(scratch.pimpl->ptr, aligned_size, aligned_offset, prefix + size); + memcpy(dst, (uint8_t *) scratch.pimpl->ptr + prefix, size); +#endif +} + +uint64_t llama_bounded_file::size() const { + return pimpl->file_size; +} + +size_t llama_bounded_file::alignment() const { + return pimpl->block; +} + +size_t llama_bounded_file::read_size(uint64_t offset, size_t size) const { + return pimpl->read_size(offset, size); +} + +bool llama_bounded_file::direct_io() const { + return pimpl->direct; +} + +const std::string & llama_bounded_file::name() const { + return pimpl->fname; +} diff --git a/src/llama-bounded-file.h b/src/llama-bounded-file.h new file mode 100644 index 000000000000..3af763580a36 --- /dev/null +++ b/src/llama-bounded-file.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include +#include + +struct llama_bounded_file { + struct params { + bool direct_io = true; + bool direct_io_required = false; // fail instead of using the page cache + }; + + struct buffer { + buffer(); + ~buffer(); + buffer(buffer && other) noexcept; + buffer & operator=(buffer && other) noexcept; + + buffer(const buffer &) = delete; + buffer & operator=(const buffer &) = delete; + + bool empty() const; + + struct impl; + std::unique_ptr pimpl; + }; + + llama_bounded_file(const std::string & fname, const params & p); + ~llama_bounded_file(); + + buffer make_buffer(size_t size) const; + // Read only the requested extent. Direct I/O may read its containing aligned blocks into scratch. + void read(uint64_t offset, void * dst, size_t size, buffer & scratch) const; + + uint64_t size() const; + size_t alignment() const; + size_t read_size(uint64_t offset, size_t size) const; + bool direct_io() const; + const std::string & name() const; + + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ccfd52175747..9d9e40d5d728 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,11 +483,15 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + + model.acquire_runtime_context(); } llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + model.release_runtime_work(); + model.release_runtime_context(); // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation if (!model.hparams.no_alloc && !opt_ctx) { @@ -689,18 +693,23 @@ void llama_context::sched_reserve() { } } + size_t graph_workspace_size = 0; for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; ggml_backend_buffer_type_t buft = backend_buft[i]; if (!model.hparams.no_alloc) { backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); } + graph_workspace_size += backend_buf_exp_size[i]; if (backend_buf_exp_size[i] > 1) { LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__, ggml_backend_buft_name(buft), backend_buf_exp_size[i] / 1024.0 / 1024.0); } } + if (memory) { + memory->set_graph_workspace_size(graph_workspace_size); + } if (n_nodes_pp == n_nodes_tg) { LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); @@ -1353,11 +1362,34 @@ bool llama_context::set_adapter_cvec( llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, ggml_status & ret) { if (mctx && !mctx->apply()) { + mctx->rollback(); LLAMA_LOG_ERROR("%s: failed to apply memory context\n", __func__); ret = GGML_STATUS_FAILED; return nullptr; } + struct memory_transaction_guard { + llama_memory_context_i * context; + bool active; + + ~memory_transaction_guard() { + if (active) { + try { + context->rollback(); + } catch (const std::exception & error) { + LLAMA_LOG_ERROR("%s: memory rollback failed: %s\n", __func__, error.what()); + } + } + } + + void commit() { + if (active) { + context->commit(); + active = false; + } + } + } transaction_guard { mctx, mctx != nullptr }; + auto * res = gf_res_prev.get(); auto * gf = res->get_gf(); @@ -1424,11 +1456,23 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); if (status != GGML_STATUS_SUCCESS) { + model.release_runtime_work_after_sync(sched.get()); LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; return nullptr; } + if (model.requires_synchronous_graph()) { + synchronize(); + const std::string error = model.consume_runtime_error(); + if (!error.empty()) { + model.release_runtime_work(); + LLAMA_LOG_ERROR("%s: model runtime failed: %s\n", __func__, error.c_str()); + ret = GGML_STATUS_FAILED; + return nullptr; + } + } + transaction_guard.commit(); ret = GGML_STATUS_SUCCESS; return res; @@ -2351,6 +2395,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_DEEPSEEK41 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || @@ -3731,7 +3776,8 @@ llama_context * llama_init_from_model( } } - if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) { + if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) && + params.type_k != params.type_v) { LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v)); return nullptr; } diff --git a/src/llama-dsv41-engram.cpp b/src/llama-dsv41-engram.cpp new file mode 100644 index 000000000000..b2273a5236f2 --- /dev/null +++ b/src/llama-dsv41-engram.cpp @@ -0,0 +1,400 @@ +#include "llama-dsv41-engram.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +static llama_dsv41_engram_sequence_state dsv41_engram_initial_state() { + llama_dsv41_engram_sequence_state state; + state.history.reset(); + return state; +} + +void llama_dsv41_validate_engram_extent(const llama_dsv41_engram_extent & extent) { + if (extent.fname.empty() || extent.rows == 0 || extent.columns != LLAMA_ENGRAM_ROW_BYTES || + extent.row_count != extent.rows || extent.type != GGML_TYPE_I8) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor must be I8 [264, rows]"); + } + const uint64_t bytes = (uint64_t) extent.rows*LLAMA_ENGRAM_ROW_BYTES; + if (extent.offset > (uint64_t) INT64_MAX || bytes > (uint64_t) INT64_MAX - extent.offset) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor extent overflows"); + } +} + +struct llama_dsv41_engram_transaction::impl { + const llama_dsv41_engram_runtime * owner = nullptr; + uint64_t generation = 0; + bool active = false; + size_t count = 0; + std::vector ids; + std::array, LLAMA_ENGRAM_LAYERS> decoded; + std::vector mask; + std::map next; +}; + +llama_dsv41_engram_transaction::llama_dsv41_engram_transaction() : pimpl(std::make_unique()) {} +llama_dsv41_engram_transaction::~llama_dsv41_engram_transaction() = default; +llama_dsv41_engram_transaction::llama_dsv41_engram_transaction(llama_dsv41_engram_transaction && other) noexcept = default; +llama_dsv41_engram_transaction & llama_dsv41_engram_transaction::operator=(llama_dsv41_engram_transaction && other) noexcept = default; + +size_t llama_dsv41_engram_transaction::token_count() const { + return pimpl->count; +} + +const uint32_t * llama_dsv41_engram_transaction::row_ids(uint32_t layer) const { + if (layer >= LLAMA_ENGRAM_LAYERS || pimpl->count == 0) { + return nullptr; + } + return pimpl->ids.data() + layer*LLAMA_ENGRAM_COLS; +} + +const float * llama_dsv41_engram_transaction::rows(uint32_t layer) const { + if (layer >= LLAMA_ENGRAM_LAYERS || pimpl->decoded[layer].empty()) { + return nullptr; + } + return pimpl->decoded[layer].data(); +} + +const uint8_t * llama_dsv41_engram_transaction::text_mask() const { + return pimpl->mask.empty() ? nullptr : pimpl->mask.data(); +} + +void llama_dsv41_engram_transaction::upload_layer( + uint32_t layer, + size_t token_offset, + size_t token_count, + ggml_tensor * rows_input, + ggml_tensor * text_select_input) const { + if (!pimpl->active) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); + } + if (layer >= LLAMA_ENGRAM_LAYERS || token_offset > pimpl->count || + token_count > pimpl->count - token_offset) { + throw std::invalid_argument("DeepSeek V4.1 Engram upload range is invalid"); + } + const int64_t row_width = LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; + if (rows_input == nullptr || text_select_input == nullptr || + rows_input->type != GGML_TYPE_F32 || rows_input->ne[0] != row_width || + ggml_nelements(rows_input) != row_width*(int64_t) token_count || + text_select_input->type != GGML_TYPE_I32 || + ggml_nelements(text_select_input) != (int64_t) token_count) { + throw std::invalid_argument("DeepSeek V4.1 Engram input tensor shape mismatch"); + } + + const size_t row_offset = token_offset*row_width; + ggml_backend_tensor_set( + rows_input, + pimpl->decoded[layer].data() + row_offset, + 0, + token_count*row_width*sizeof(float)); + std::vector select(token_count); + for (size_t i = 0; i < token_count; ++i) { + select[i] = pimpl->mask[token_offset + i] != 0 ? (int32_t) (token_count + i) : (int32_t) i; + } + ggml_backend_tensor_set( + text_select_input, + select.data(), + 0, + token_count*sizeof(int32_t)); +} + +struct llama_dsv41_engram_runtime::impl { + llama_engram_hasher hasher; + std::array, LLAMA_ENGRAM_LAYERS> tables; + std::map sequences; + size_t max_tokens; + uint64_t generation = 0; + + impl( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens) : + hasher(std::move(layout)), + max_tokens(max_tokens) { + if (max_tokens == 0 || max_tokens > (size_t) INT32_MAX/2) { + throw std::invalid_argument("DeepSeek V4.1 Engram token bound is invalid"); + } + for (size_t i = 0; i < LLAMA_ENGRAM_LAYERS; ++i) { + llama_dsv41_validate_engram_extent(extents[i]); + if (extents[i].rows != hasher.layout().rows[i]) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor rows do not match metadata"); + } + tables[i] = std::make_unique( + extents[i].fname, extents[i].offset, extents[i].rows); + } + } +}; + +llama_dsv41_engram_runtime::llama_dsv41_engram_runtime( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens) : + pimpl(std::make_unique(std::move(layout), extents, max_tokens)) {} + +llama_dsv41_engram_runtime::~llama_dsv41_engram_runtime() = default; + +llama_dsv41_engram_transaction llama_dsv41_engram_runtime::prepare( + const std::vector & tokens) { + if (tokens.size() > pimpl->max_tokens) { + throw std::invalid_argument("DeepSeek V4.1 Engram batch exceeds the configured bound"); + } + + llama_dsv41_engram_transaction result; + result.pimpl->owner = this; + result.pimpl->generation = pimpl->generation; + result.pimpl->count = tokens.size(); + result.pimpl->ids.resize(tokens.size()*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); + result.pimpl->mask.resize(tokens.size()); + result.pimpl->next = pimpl->sequences; + + for (size_t i = 0; i < tokens.size(); ++i) { + const llama_dsv41_engram_token & token = tokens[i]; + if (token.seq_ids.empty()) { + throw std::invalid_argument("DeepSeek V4.1 Engram token has no sequence"); + } + + uint32_t expected[LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS] = {}; + bool have_expected = false; + std::set seen; + for (llama_seq_id seq_id : token.seq_ids) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + if (!seen.insert(seq_id).second) { + throw std::invalid_argument("DeepSeek V4.1 Engram token repeats a sequence ID"); + } + auto inserted = result.pimpl->next.emplace(seq_id, dsv41_engram_initial_state()); + llama_dsv41_engram_sequence_state & state = inserted.first->second; + if (token.pos != state.pos + 1) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence position is not contiguous"); + } + + uint32_t current[LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS] = {}; + pimpl->hasher.hash(state.history, &token.token, &token.text, 1, current); + state.pos = token.pos; + if (have_expected && std::memcmp(expected, current, sizeof(expected)) != 0) { + throw std::invalid_argument("DeepSeek V4.1 coupled sequence histories differ"); + } + std::memcpy(expected, current, sizeof(expected)); + have_expected = true; + } + + std::memcpy( + result.pimpl->ids.data() + i*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + expected, + sizeof(expected)); + result.pimpl->mask[i] = token.text != 0; + } + + const size_t row_values = tokens.size()*LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + result.pimpl->decoded[layer].resize(row_values); + pimpl->tables[layer]->read_batch( + result.pimpl->ids.data() + layer*LLAMA_ENGRAM_COLS, + tokens.size(), + LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + result.pimpl->decoded[layer].data()); + } + result.pimpl->active = true; + return result; +} + +void llama_dsv41_engram_runtime::commit(llama_dsv41_engram_transaction & transaction) { + if (!transaction.pimpl->active || transaction.pimpl->owner != this) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); + } + if (transaction.pimpl->generation != pimpl->generation) { + throw std::runtime_error("DeepSeek V4.1 Engram transaction is stale"); + } + auto next = transaction.pimpl->next; + pimpl->sequences.swap(next); + ++pimpl->generation; + transaction.pimpl->active = false; +} + +void llama_dsv41_engram_runtime::rollback(llama_dsv41_engram_transaction & transaction) { + if (transaction.pimpl->owner != this) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction belongs to another runtime"); + } + transaction.pimpl->active = false; +} + +void llama_dsv41_engram_runtime::seq_reset(llama_seq_id seq_id) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + pimpl->sequences[seq_id] = dsv41_engram_initial_state(); + ++pimpl->generation; +} + +void llama_dsv41_engram_runtime::seq_copy(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) { + if (seq_id_src < 0 || seq_id_dst < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + const auto it = pimpl->sequences.find(seq_id_src); + pimpl->sequences[seq_id_dst] = it == pimpl->sequences.end() ? + dsv41_engram_initial_state() : it->second; + ++pimpl->generation; +} + +void llama_dsv41_engram_runtime::seq_remove(llama_seq_id seq_id) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + pimpl->sequences.erase(seq_id); + ++pimpl->generation; +} + +llama_dsv41_engram_snapshot llama_dsv41_engram_runtime::checkpoint() const { + return { pimpl->sequences }; +} + +void llama_dsv41_engram_runtime::restore(const llama_dsv41_engram_snapshot & snapshot) { + auto restored = snapshot.sequences; + for (auto & item : restored) { + if (item.first < 0 || item.second.pos < -1) { + throw std::invalid_argument("DeepSeek V4.1 Engram snapshot is invalid"); + } + pimpl->hasher.hash(item.second.history, nullptr, nullptr, 0, nullptr); + } + pimpl->sequences.swap(restored); + ++pimpl->generation; +} + +llama_dsv41_engram_sequence_state llama_dsv41_engram_runtime::sequence(llama_seq_id seq_id) const { + const auto it = pimpl->sequences.find(seq_id); + return it == pimpl->sequences.end() ? dsv41_engram_initial_state() : it->second; +} + +size_t llama_dsv41_engram_runtime::max_tokens() const { + return pimpl->max_tokens; +} + +static ggml_tensor * dsv41_bf16_f32(ggml_context * ctx, ggml_tensor * tensor) { + return ggml_cast(ctx, ggml_cast(ctx, tensor, GGML_TYPE_BF16), GGML_TYPE_F32); +} + +static void dsv41_engram_gate_f32( + ggml_tensor * dst, + const ggml_tensor * src, + int ith, + int nth, + void *) { + GGML_ASSERT(dst->type == GGML_TYPE_F32 && src->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src)); + const float * input = static_cast(src->data); + float * output = static_cast(dst->data); + const int64_t count = ggml_nelements(src); + for (int64_t i = ith; i < count; i += nth) { + const float signed_root = std::copysign(std::sqrt(std::max(std::abs(input[i]), 1.0e-6f)), input[i]); + output[i] = 1.0f/(1.0f + std::exp(-signed_root)); + } +} + +ggml_tensor * llama_dsv41_build_engram_gate( + ggml_context * ctx, + ggml_tensor * dot, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || dot == nullptr || sched == nullptr || backend_cpu == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate input is null"); + } + if (!ggml_backend_is_cpu(backend_cpu)) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate backend is not local CPU"); + } + bool found = false; + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + found = found || ggml_backend_sched_get_backend(sched, i) == backend_cpu; + } + if (!found) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate CPU backend is not in the scheduler"); + } + ggml_tensor * gate = ggml_map_custom1(ctx, dot, dsv41_engram_gate_f32, GGML_N_TASKS_MAX, nullptr); + ggml_backend_sched_set_tensor_backend(sched, gate, backend_cpu); + return gate; +} + +ggml_tensor * llama_dsv41_build_engram_add( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * projected, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || residual == nullptr || projected == nullptr || q_norm == nullptr || k_norm == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram graph input is null"); + } + + const int64_t width = residual->ne[0]; + const int64_t streams = residual->ne[1]; + const int64_t tokens = residual->ne[2]; + if (width <= 0 || streams != 4 || tokens <= 0 || tokens > INT32_MAX/2 || + projected->ne[0] != 5*width || projected->ne[1] != tokens || + q_norm->ne[0] != width || q_norm->ne[1] != streams || + k_norm->ne[0] != width || k_norm->ne[1] != streams || + (text_select != nullptr && + (text_select->type != GGML_TYPE_I32 || ggml_nelements(text_select) != tokens))) { + throw std::invalid_argument("DeepSeek V4.1 Engram graph shape mismatch"); + } + + projected = dsv41_bf16_f32(ctx, projected); + ggml_tensor * value = ggml_view_2d( + ctx, projected, width, tokens, projected->nb[1], 4*projected->nb[0]*width); + ggml_tensor * result = nullptr; + for (int64_t stream = 0; stream < streams; ++stream) { + ggml_tensor * hidden = ggml_view_2d( + ctx, residual, width, tokens, residual->nb[2], stream*residual->nb[1]); + ggml_tensor * key = ggml_view_2d( + ctx, projected, width, tokens, projected->nb[1], stream*projected->nb[0]*width); + ggml_tensor * qw = ggml_view_1d(ctx, q_norm, width, stream*q_norm->nb[1]); + ggml_tensor * kw = ggml_view_1d(ctx, k_norm, width, stream*k_norm->nb[1]); + + ggml_tensor * hidden_norm = ggml_rms_norm(ctx, hidden, rms_eps); + ggml_tensor * key_norm = ggml_rms_norm(ctx, key, rms_eps); + ggml_tensor * dot = ggml_mul(ctx, hidden_norm, qw); + dot = ggml_mul(ctx, dot, kw); + dot = ggml_mul(ctx, dot, key_norm); + dot = ggml_scale(ctx, ggml_sum_rows(ctx, dot), 1.0f/std::sqrt((float) width)); + + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot, sched, backend_cpu); + ggml_tensor * updated = dsv41_bf16_f32(ctx, ggml_add(ctx, hidden, ggml_mul(ctx, value, gate))); + if (text_select != nullptr) { + updated = ggml_get_rows(ctx, ggml_concat(ctx, hidden, updated, 1), text_select); + } + updated = ggml_reshape_3d(ctx, updated, width, 1, tokens); + result = result == nullptr ? updated : ggml_concat(ctx, result, updated, 1); + } + return result; +} + +ggml_tensor * llama_dsv41_build_engram( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * rows, + ggml_tensor * engram_kv, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || residual == nullptr || rows == nullptr || engram_kv == nullptr || + rows->ne[0] != LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM || + engram_kv->ne[0] != rows->ne[0] || + engram_kv->ne[1] != 5*residual->ne[0]) { + throw std::invalid_argument("DeepSeek V4.1 Engram projection shape mismatch"); + } + ggml_tensor * projected = ggml_mul_mat(ctx, engram_kv, rows); + return llama_dsv41_build_engram_add( + ctx, residual, projected, q_norm, k_norm, text_select, rms_eps, sched, backend_cpu); +} diff --git a/src/llama-dsv41-engram.h b/src/llama-dsv41-engram.h new file mode 100644 index 000000000000..2edca59c12d0 --- /dev/null +++ b/src/llama-dsv41-engram.h @@ -0,0 +1,123 @@ +#pragma once + +#include "ggml-backend.h" +#include "llama-engram.h" +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +struct llama_dsv41_engram_extent { + std::string fname; + uint64_t offset = 0; + uint32_t rows = 0; + int64_t columns = 0; + int64_t row_count = 0; + int32_t type = 0; +}; + +void llama_dsv41_validate_engram_extent(const llama_dsv41_engram_extent & extent); + +struct llama_dsv41_engram_token { + int32_t token = -1; + llama_pos pos = -1; + std::vector seq_ids; + uint8_t text = 1; +}; + +struct llama_dsv41_engram_sequence_state { + llama_engram_history history = {}; + llama_pos pos = -1; +}; + +struct llama_dsv41_engram_snapshot { + std::map sequences; +}; + +struct llama_dsv41_engram_transaction { + llama_dsv41_engram_transaction(); + ~llama_dsv41_engram_transaction(); + llama_dsv41_engram_transaction(llama_dsv41_engram_transaction && other) noexcept; + llama_dsv41_engram_transaction & operator=(llama_dsv41_engram_transaction && other) noexcept; + + llama_dsv41_engram_transaction(const llama_dsv41_engram_transaction &) = delete; + llama_dsv41_engram_transaction & operator=(const llama_dsv41_engram_transaction &) = delete; + + size_t token_count() const; + const uint32_t * row_ids(uint32_t layer) const; + const float * rows(uint32_t layer) const; + const uint8_t * text_mask() const; + void upload_layer( + uint32_t layer, + size_t token_offset, + size_t token_count, + ggml_tensor * rows_input, + ggml_tensor * text_select_input) const; + + struct impl; + std::unique_ptr pimpl; +}; + +class llama_dsv41_engram_runtime { +public: + llama_dsv41_engram_runtime( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens); + ~llama_dsv41_engram_runtime(); + + llama_dsv41_engram_transaction prepare(const std::vector & tokens); + void commit(llama_dsv41_engram_transaction & transaction); + void rollback(llama_dsv41_engram_transaction & transaction); + + void seq_reset(llama_seq_id seq_id); + void seq_copy(llama_seq_id seq_id_src, llama_seq_id seq_id_dst); + void seq_remove(llama_seq_id seq_id); + llama_dsv41_engram_snapshot checkpoint() const; + void restore(const llama_dsv41_engram_snapshot & snapshot); + llama_dsv41_engram_sequence_state sequence(llama_seq_id seq_id) const; + + size_t max_tokens() const; + +private: + struct impl; + std::unique_ptr pimpl; +}; + +// text_select row i keeps the original residual; row tokens+i selects the BF16-updated residual. +ggml_tensor * llama_dsv41_build_engram_add( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * projected, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_engram_gate( + ggml_context * ctx, + ggml_tensor * dot, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_engram( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * rows, + ggml_tensor * engram_kv, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); diff --git a/src/llama-dsv41-expert.cpp b/src/llama-dsv41-expert.cpp new file mode 100644 index 000000000000..3d82f29b4739 --- /dev/null +++ b/src/llama-dsv41-expert.cpp @@ -0,0 +1,552 @@ +#include "llama-dsv41-expert.h" + +#include "llama-dsv41.h" +#include "llama-impl.h" + +#include "ggml-cpp.h" + +#include +#include +#include +#include +#include +#include +#include + +// The remap node follows ggml-org/llama.cpp#25294 commit 4260e4608. +// Lease publication and release follow ggml-org/llama.cpp#27861 commit bccbacdb8. + +namespace { + +size_t checked_add(size_t a, size_t b, const char * message) { + if (b > std::numeric_limits::max() - a) { + throw std::runtime_error(message); + } + return a + b; +} + +size_t checked_mul(size_t a, size_t b, const char * message) { + if (a != 0 && b > std::numeric_limits::max()/a) { + throw std::runtime_error(message); + } + return a*b; +} + +size_t projection_index(llama_expert_projection projection) { + if (projection < LLAMA_EXPERT_PROJECTION_GATE || projection > LLAMA_EXPERT_PROJECTION_DOWN) { + throw std::invalid_argument("DeepSeek V4.1 expert projection is invalid"); + } + return static_cast(projection); +} + +void require_local_cpu(ggml_backend_sched_t sched, ggml_backend_t backend_cpu) { + if (sched == nullptr || backend_cpu == nullptr || !ggml_backend_is_cpu(backend_cpu)) { + throw std::invalid_argument("DeepSeek V4.1 expert callback requires a local CPU backend"); + } + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + if (ggml_backend_sched_get_backend(sched, i) == backend_cpu) { + return; + } + } + throw std::invalid_argument("DeepSeek V4.1 expert CPU backend is not in the scheduler"); +} + +struct dsv41_expert_callback_state { + llama_dsv41_expert_runtime * runtime = nullptr; + int32_t layer = -1; +}; + +} + +struct llama_dsv41_expert_runtime::impl { + struct logical_slot { + int32_t expert_id = -1; + uint32_t pins = 0; + uint64_t last_use = 0; + }; + + struct layer_state { + std::array tensors = {}; + std::vector slots; + std::unique_ptr lease; + std::vector pinned_slots; + std::vector active_ids; + std::vector active_remap; + }; + + llama_dsv41_expert_runtime_params params; + std::unique_ptr store; + std::vector layers; + std::vector callbacks; + std::vector> contexts; + std::vector buffers; + upload_fn upload; + publish_fn before_publish; + size_t bytes_cache = 0; + size_t bytes_staging = 0; + uint64_t use_clock = 0; + std::string error; + bool context_active = false; + mutable std::mutex mutex; + + impl( + llama_dsv41_expert_runtime * owner, + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload, + publish_fn before_publish) + : params(params), upload(std::move(upload)), before_publish(std::move(before_publish)) { + if (params.cache_bytes == 0 || params.cache_slots == 0) { + throw std::runtime_error("DeepSeek V4.1 expert cache byte and slot capacity must be non-zero"); + } + if (!select_buft) { + throw std::invalid_argument("DeepSeek V4.1 expert cache buffer selector is empty"); + } + if (tensors.size() != LLAMA_DSV41_N_LAYER*3) { + throw std::runtime_error("DeepSeek V4.1 must register 40 gate/up/down routed tensor sets"); + } + + layers.resize(LLAMA_DSV41_N_LAYER); + callbacks.resize(LLAMA_DSV41_N_LAYER); + std::vector layer_plane_bytes(LLAMA_DSV41_N_LAYER, 0); + for (const auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + if (tensor.layer < 0 || tensor.layer >= (int32_t) LLAMA_DSV41_N_LAYER) { + throw std::runtime_error("DeepSeek V4.1 routed tensor layer is invalid"); + } + if (tensor.ne[2] != LLAMA_DSV41_N_EXPERT) { + throw std::runtime_error("DeepSeek V4.1 routed tensor expert count mismatch"); + } + layer_plane_bytes[tensor.layer] = checked_add( + layer_plane_bytes[tensor.layer], + tensor.nb[2], + "DeepSeek V4.1 expert plane byte count overflow"); + bytes_cache = checked_add( + bytes_cache, + checked_mul(tensor.nb[2], params.cache_slots, "DeepSeek V4.1 expert cache byte count overflow"), + "DeepSeek V4.1 expert cache byte count overflow"); + + ggml_backend_buffer_type_t buft = select_buft(tensor); + if (buft == nullptr) { + throw std::runtime_error("DeepSeek V4.1 expert cache has no backend buffer type"); + } + ggml_context * ctx = nullptr; + for (auto & item : contexts) { + if (item.first == buft) { + ctx = item.second.get(); + break; + } + } + if (ctx == nullptr) { + ggml_init_params ctx_params = { + /*.mem_size =*/ ggml_tensor_overhead()*(LLAMA_DSV41_N_LAYER*3 + 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx = ggml_init(ctx_params); + if (ctx == nullptr) { + throw std::runtime_error("DeepSeek V4.1 failed to create expert cache tensor context"); + } + contexts.emplace_back(buft, ctx); + } + + ggml_tensor * cache = ggml_new_tensor_3d( + ctx, tensor.type, tensor.ne[0], tensor.ne[1], params.cache_slots); + ggml_format_name(cache, "%s.cache", tensor.name.c_str()); + layers[tensor.layer].tensors[projection_index(tensor.projection)] = cache; + } + if (bytes_cache > params.cache_bytes) { + throw std::runtime_error(format( + "DeepSeek V4.1 expert cache requires %zu bytes for %zu slots per layer, configured %zu", + bytes_cache, params.cache_slots, params.cache_bytes)); + } + + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (ggml_tensor * tensor : layers[il].tensors) { + if (tensor == nullptr) { + throw std::runtime_error(format("DeepSeek V4.1 layer %d is missing an expert cache tensor", il)); + } + } + layers[il].slots.resize(params.cache_slots); + callbacks[il] = { owner, il }; + bytes_staging = std::max( + bytes_staging, + checked_mul(layer_plane_bytes[il], params.cache_slots, "DeepSeek V4.1 expert staging byte count overflow")); + } + + for (auto & item : contexts) { + ggml_backend_buffer_t buffer = nullptr; + if (params.no_alloc) { + buffer = ggml_backend_buft_alloc_buffer(item.first, 0); + for (ggml_tensor * tensor = ggml_get_first_tensor(item.second.get()); + tensor != nullptr; + tensor = ggml_get_next_tensor(item.second.get(), tensor)) { + tensor->buffer = buffer; + } + } else { + buffer = ggml_backend_alloc_ctx_tensors_from_buft(item.second.get(), item.first); + } + if (buffer == nullptr) { + throw std::runtime_error(format( + "DeepSeek V4.1 failed to allocate %s expert cache buffer", + ggml_backend_buft_name(item.first))); + } + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + buffers.emplace_back(buffer); + } + + if (!params.no_alloc) { + llama_expert_store_params store_params; + store_params.cache_bytes = bytes_staging; + store_params.cache_slots = checked_mul(params.cache_slots, 3, "DeepSeek V4.1 expert staging slot count overflow"); + store_params.direct_io = params.direct_io; + store_params.allow_buffered_io = params.allow_buffered_io; + store = std::make_unique(std::move(tensors), store_params); + } + + if (!this->upload) { + this->upload = [](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + ggml_backend_tensor_set(tensor, data, offset, size); + }; + } + } + + const llama_expert_store::payload & payload( + const std::vector & payloads, + int32_t layer, + llama_expert_projection projection, + int32_t expert_id) const { + for (const auto & payload : payloads) { + if (payload.layer == layer && payload.projection == projection && payload.expert_id == expert_id) { + return payload; + } + } + throw std::runtime_error("DeepSeek V4.1 expert store did not return a requested payload"); + } + + std::vector remap(int32_t layer, const std::vector & expert_ids) { + std::lock_guard lock(mutex); + if (store == nullptr) { + throw std::runtime_error("DeepSeek V4.1 expert cache is metadata-only"); + } + if (layer < 0 || layer >= (int32_t) layers.size()) { + throw std::invalid_argument("DeepSeek V4.1 expert layer is invalid"); + } + layer_state & state = layers[layer]; + if (state.lease) { + if (state.active_ids == expert_ids) { + return state.active_remap; + } + throw std::runtime_error(format("DeepSeek V4.1 expert layer %d still has an in-flight lease", layer)); + } + + std::vector unique = expert_ids; + for (int32_t expert_id : unique) { + if (expert_id < 0 || expert_id >= (int32_t) LLAMA_DSV41_N_EXPERT) { + throw std::runtime_error("DeepSeek V4.1 selected expert ID is out of range"); + } + } + std::sort(unique.begin(), unique.end()); + unique.erase(std::unique(unique.begin(), unique.end()), unique.end()); + if (unique.size() > params.cache_slots) { + throw std::runtime_error(format( + "DeepSeek V4.1 selected expert union has %zu entries, cache has %zu slots", + unique.size(), params.cache_slots)); + } + + std::map resident; + std::vector empty; + std::vector victims; + for (uint32_t slot = 0; slot < state.slots.size(); ++slot) { + const logical_slot & entry = state.slots[slot]; + if (entry.expert_id >= 0) { + resident.emplace(entry.expert_id, slot); + if (entry.pins == 0 && !std::binary_search(unique.begin(), unique.end(), entry.expert_id)) { + victims.push_back(slot); + } + } else { + empty.push_back(slot); + } + } + std::sort(victims.begin(), victims.end(), [&](uint32_t a, uint32_t b) { + if (state.slots[a].last_use != state.slots[b].last_use) { + return state.slots[a].last_use < state.slots[b].last_use; + } + return a < b; + }); + + std::vector misses; + for (int32_t expert_id : unique) { + if (resident.count(expert_id) == 0) { + misses.push_back(expert_id); + } + } + if (misses.size() > empty.size() + victims.size()) { + throw std::runtime_error("DeepSeek V4.1 expert cache capacity is pinned"); + } + + std::vector targets = empty; + targets.insert(targets.end(), victims.begin(), victims.end()); + targets.resize(misses.size()); + std::sort(targets.begin(), targets.end()); + + auto lease = std::make_unique(); + if (!misses.empty()) { + *lease = store->acquire({ + { layer, LLAMA_EXPERT_PROJECTION_GATE, misses }, + { layer, LLAMA_EXPERT_PROJECTION_UP, misses }, + { layer, LLAMA_EXPERT_PROJECTION_DOWN, misses }, + }); + const auto payloads = lease->payloads(); + try { + for (size_t i = 0; i < misses.size(); ++i) { + const uint32_t slot = targets[i]; + for (llama_expert_projection projection : { + LLAMA_EXPERT_PROJECTION_GATE, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN }) { + const auto & item = payload(payloads, layer, projection, misses[i]); + ggml_tensor * tensor = state.tensors[projection_index(projection)]; + upload(tensor, slot*tensor->nb[2], item.data, item.size); + } + } + } catch (...) { + for (uint32_t slot : targets) { + state.slots[slot] = {}; + } + throw; + } + + for (size_t i = 0; i < misses.size(); ++i) { + const uint32_t slot = targets[i]; + state.slots[slot].expert_id = misses[i]; + resident[misses[i]] = slot; + } + } + + std::vector pinned_slots; + pinned_slots.reserve(unique.size()); + for (int32_t expert_id : unique) { + const uint32_t slot = resident.at(expert_id); + pinned_slots.push_back(slot); + } + + std::vector result; + result.reserve(expert_ids.size()); + for (int32_t expert_id : expert_ids) { + result.push_back((int32_t) resident.at(expert_id)); + } + std::vector active_ids = expert_ids; + std::vector active_remap = result; + + if (before_publish) { + before_publish(); + } + for (uint32_t slot : pinned_slots) { + logical_slot & entry = state.slots[slot]; + entry.last_use = ++use_clock; + entry.pins++; + } + state.pinned_slots = std::move(pinned_slots); + state.lease = std::move(lease); + state.active_ids = std::move(active_ids); + state.active_remap = std::move(active_remap); + return result; + } + + void release(int32_t layer) { + std::lock_guard lock(mutex); + if (layer < 0 || layer >= (int32_t) layers.size()) { + return; + } + layer_state & state = layers[layer]; + for (uint32_t slot : state.pinned_slots) { + if (state.slots[slot].pins == 0) { + GGML_ABORT("DeepSeek V4.1 expert slot pin underflow"); + } + state.slots[slot].pins--; + } + state.pinned_slots.clear(); + state.lease.reset(); + state.active_ids.clear(); + state.active_remap.clear(); + } +}; + +llama_dsv41_expert_runtime::llama_dsv41_expert_runtime( + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload, + publish_fn before_publish) + : pimpl(std::make_unique( + this, std::move(tensors), params, std::move(select_buft), std::move(upload), std::move(before_publish))) { +} + +llama_dsv41_expert_runtime::~llama_dsv41_expert_runtime() = default; + +std::vector llama_dsv41_expert_runtime::remap( + int32_t layer, const std::vector & expert_ids) { + return pimpl->remap(layer, expert_ids); +} + +void llama_dsv41_expert_runtime::release(int32_t layer) { + pimpl->release(layer); +} + +void llama_dsv41_expert_runtime::release_all() { + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + pimpl->release(il); + } +} + +void llama_dsv41_expert_runtime::release_all_after_sync(ggml_backend_sched_t sched) { + if (sched != nullptr) { + ggml_backend_sched_synchronize(sched); + } + release_all(); +} + +ggml_tensor * llama_dsv41_expert_runtime::cache_tensor( + int32_t layer, llama_expert_projection projection) const { + if (layer < 0 || layer >= (int32_t) pimpl->layers.size()) { + throw std::invalid_argument("DeepSeek V4.1 expert layer is invalid"); + } + return pimpl->layers[layer].tensors[projection_index(projection)]; +} + +size_t llama_dsv41_expert_runtime::cache_slots() const { + return pimpl->params.cache_slots; +} + +size_t llama_dsv41_expert_runtime::cache_bytes() const { + return pimpl->bytes_cache; +} + +size_t llama_dsv41_expert_runtime::staging_bytes() const { + return pimpl->bytes_staging; +} + +void llama_dsv41_expert_runtime::set_error(const std::string & error) { + std::lock_guard lock(pimpl->mutex); + if (pimpl->error.empty()) { + pimpl->error = error; + } +} + +std::string llama_dsv41_expert_runtime::consume_error() { + std::lock_guard lock(pimpl->mutex); + std::string result; + result.swap(pimpl->error); + return result; +} + +void llama_dsv41_expert_runtime::acquire_context() { + std::lock_guard lock(pimpl->mutex); + if (pimpl->context_active) { + throw std::runtime_error("DeepSeek V4.1 bounded expert runtime supports one context per model"); + } + pimpl->context_active = true; +} + +void llama_dsv41_expert_runtime::release_context() { + std::lock_guard lock(pimpl->mutex); + pimpl->context_active = false; +} + +std::vector llama_dsv41_register_expert_tensors( + const std::function & ne)> & register_tensor) { + if (!register_tensor) { + throw std::invalid_argument("DeepSeek V4.1 expert tensor registrar is empty"); + } + std::vector result; + result.reserve(LLAMA_DSV41_N_LAYER*3); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_gate_exps.weight", + il, LLAMA_EXPERT_PROJECTION_GATE, + { LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EXPERT })); + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_up_exps.weight", + il, LLAMA_EXPERT_PROJECTION_UP, + { LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EXPERT })); + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_down_exps.weight", + il, LLAMA_EXPERT_PROJECTION_DOWN, + { LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_EXPERT })); + } + return result; +} + +static void dsv41_expert_remap_callback( + ggml_tensor * dst, + const ggml_tensor * src, + int, + int, + void * userdata) { + auto * state = static_cast(userdata); + GGML_ASSERT(dst->type == GGML_TYPE_I32 && src->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src)); + const int32_t * ids = static_cast(src->data); + const size_t count = ggml_nelements(src); + try { + const std::vector remapped = state->runtime->remap( + state->layer, std::vector(ids, ids + count)); + memcpy(dst->data, remapped.data(), remapped.size()*sizeof(int32_t)); + } catch (const std::exception & error) { + std::fill_n(static_cast(dst->data), count, 0); + state->runtime->set_error(error.what()); + } +} + +static void dsv41_expert_release_callback( + ggml_tensor * dst, + const ggml_tensor * src, + int, + int, + void * userdata) { + auto * state = static_cast(userdata); + memcpy(dst->data, src->data, ggml_nbytes(src)); + state->runtime->release(state->layer); +} + +ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context * ctx, + ggml_tensor * selected_experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || selected_experts == nullptr || selected_experts->type != GGML_TYPE_I32) { + throw std::invalid_argument("DeepSeek V4.1 expert remap input is invalid"); + } + require_local_cpu(sched, backend_cpu); + ggml_tensor * original = ggml_cont(ctx, selected_experts); + ggml_tensor * remapped = ggml_map_custom1( + ctx, original, dsv41_expert_remap_callback, 1, &runtime.pimpl->callbacks.at(layer)); + ggml_backend_sched_set_tensor_backend(sched, remapped, backend_cpu); + return remapped; +} + +ggml_tensor * llama_dsv41_build_expert_release( + ggml_context * ctx, + ggml_tensor * experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || experts == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 expert release input is invalid"); + } + require_local_cpu(sched, backend_cpu); + ggml_tensor * completion = ggml_sum(ctx, experts); + completion = ggml_map_custom1( + ctx, completion, dsv41_expert_release_callback, 1, &runtime.pimpl->callbacks.at(layer)); + ggml_backend_sched_set_tensor_backend(sched, completion, backend_cpu); + return completion; +} diff --git a/src/llama-dsv41-expert.h b/src/llama-dsv41-expert.h new file mode 100644 index 000000000000..82ce5037e8fc --- /dev/null +++ b/src/llama-dsv41-expert.h @@ -0,0 +1,89 @@ +#pragma once + +#include "llama-expert-store.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_cgraph; +struct ggml_context; +struct ggml_tensor; + +struct llama_dsv41_expert_runtime_params { + size_t cache_bytes = 0; + size_t cache_slots = 0; + bool direct_io = true; + bool allow_buffered_io = false; + bool no_alloc = false; +}; + +struct llama_dsv41_expert_runtime { + using buft_selector = std::function; + using upload_fn = std::function; + using publish_fn = std::function; + + llama_dsv41_expert_runtime( + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload = {}, + publish_fn before_publish = {}); + ~llama_dsv41_expert_runtime(); + + llama_dsv41_expert_runtime(const llama_dsv41_expert_runtime &) = delete; + llama_dsv41_expert_runtime & operator=(const llama_dsv41_expert_runtime &) = delete; + + std::vector remap(int32_t layer, const std::vector & expert_ids); + void acquire_context(); + void release_context(); + void release(int32_t layer); + void release_all(); + void release_all_after_sync(ggml_backend_sched_t sched); + + ggml_tensor * cache_tensor(int32_t layer, llama_expert_projection projection) const; + size_t cache_slots() const; + size_t cache_bytes() const; + size_t staging_bytes() const; + + void set_error(const std::string & error); + std::string consume_error(); + +private: + friend ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context *, ggml_tensor *, llama_dsv41_expert_runtime &, int32_t, ggml_backend_sched_t, ggml_backend_t); + friend ggml_tensor * llama_dsv41_build_expert_release( + ggml_context *, ggml_tensor *, llama_dsv41_expert_runtime &, int32_t, ggml_backend_sched_t, ggml_backend_t); + + struct impl; + std::unique_ptr pimpl; +}; + +std::vector llama_dsv41_register_expert_tensors( + const std::function & ne)> & register_tensor); + +ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context * ctx, + ggml_tensor * selected_experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_expert_release( + ggml_context * ctx, + ggml_tensor * experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); diff --git a/src/llama-dsv41.cpp b/src/llama-dsv41.cpp new file mode 100644 index 000000000000..2b578afa443c --- /dev/null +++ b/src/llama-dsv41.cpp @@ -0,0 +1,634 @@ +#include "llama-dsv41.h" + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +static constexpr int32_t DSV41_KV_SOURCES[] = { 2, 8, 14, 20 }; +static constexpr int32_t DSV41_INDEX_SOURCES[] = { 2, 8, 14, 20, 24, 28, 32, 36 }; +static constexpr uint32_t DSV41_ENGRAM_LAYERS[] = { 1, 14 }; +static constexpr uint32_t DSV41_ENGRAM_ROWS[] = { 384006168, 384016682 }; + +static void dsv41_require(bool condition, const char * message) { + if (!condition) { + throw std::runtime_error(std::string("DeepSeek V4.1 metadata: ") + message); + } +} + +void llama_dsv41_validate_config(const llama_dsv41_config & config) { + dsv41_require(config.n_ctx_train == LLAMA_DSV41_N_CTX, "max_position_embeddings must be 1048576"); + dsv41_require(config.n_embd == LLAMA_DSV41_N_EMBD, "hidden_size must be 5120"); + dsv41_require(config.n_layer == LLAMA_DSV41_N_LAYER, "num_hidden_layers must be 40"); + dsv41_require(config.n_vocab == LLAMA_DSV41_N_VOCAB, "vocab_size must be 129280"); + dsv41_require(config.n_head == LLAMA_DSV41_N_HEAD, "num_attention_heads must be 64"); + dsv41_require(config.n_head_kv == LLAMA_DSV41_N_HEAD_KV, "num_key_value_heads must be 1"); + dsv41_require(config.n_head_dim == LLAMA_DSV41_N_HEAD_DIM, "head_dim must be 512"); + dsv41_require(config.n_rot == LLAMA_DSV41_N_ROT, "qk_rope_head_dim must be 64"); + dsv41_require(config.n_lora_q == LLAMA_DSV41_N_LORA_Q, "q_lora_rank must be 1280"); + dsv41_require(config.n_lora_o == LLAMA_DSV41_N_LORA_O, "o_lora_rank must be 1024"); + dsv41_require(config.n_o_group == LLAMA_DSV41_N_O_GROUP, "o_group_num must be 8"); + dsv41_require(config.n_ff_dense == LLAMA_DSV41_N_FF_DENSE, "intermediate_size must be 18432"); + dsv41_require(config.n_ff_expert == LLAMA_DSV41_N_FF_EXP, "moe_intermediate_size must be 2304"); + dsv41_require(config.n_expert == LLAMA_DSV41_N_EXPERT, "n_routed_experts must be 384"); + dsv41_require(config.n_expert_used == LLAMA_DSV41_N_EXPERT_USED, "num_experts_per_tok must be 6"); + dsv41_require(config.n_expert_shared == LLAMA_DSV41_N_EXPERT_SHARED, "n_shared_experts must be 1"); + dsv41_require(config.indexer_n_head == LLAMA_DSV41_N_INDEX_HEAD, "index_n_heads must be 32"); + dsv41_require(config.indexer_head_size == LLAMA_DSV41_N_INDEX_HEAD_DIM, "index_head_dim must be 128"); + dsv41_require(config.indexer_top_k == LLAMA_DSV41_N_INDEX_TOP_K, "index_topk must be 512"); + dsv41_require(config.hc_count == LLAMA_DSV41_HC_MULT, "hc_num_streams must be 4"); + dsv41_require(config.hc_sinkhorn_iters == LLAMA_DSV41_HC_SINKHORN_ITERS, "hc_sinkhorn_iters must be 20"); + dsv41_require(config.raw_window == LLAMA_DSV41_N_SWA, "sliding_window must be 128"); + dsv41_require(config.candidate_source_layer == LLAMA_DSV41_CANDIDATE_SOURCE_LAYER, "candidate_source_layer must be 20"); + dsv41_require(config.candidate_topk_blocks == LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS, "candidate_topk_blocks must be 2048"); + dsv41_require(config.candidate_block_size == LLAMA_DSV41_CANDIDATE_BLOCK_SIZE, "candidate_block_size must be 8"); + dsv41_require(config.f_norm_rms_eps == 1.0e-20f, "rms_norm_eps must be 1e-20"); + dsv41_require(config.hc_eps == 1.0e-6f, "hc_eps must be 1e-6"); + dsv41_require(config.swiglu_clamp == 10.0f, "swiglu_clamp_limit must be 10"); + dsv41_require(config.routed_scale == 1.5f, "routed_scaling_factor must be 1.5"); + dsv41_require(config.rope_theta == 10000, "rope_theta must be 10000"); + dsv41_require(config.compress_rope_theta == 160000, "compress_rope_theta must be 160000"); + dsv41_require(config.yarn_factor == 16.0f, "rope_scaling.factor must be 16"); + dsv41_require(config.yarn_beta_fast == 32.0f, "rope_scaling.beta_fast must be 32"); + dsv41_require(config.yarn_beta_slow == 1.0f, "rope_scaling.beta_slow must be 1"); + dsv41_require(config.yarn_original_context == 65536.0f, "rope_scaling.original_max_position_embeddings must be 65536"); + dsv41_require(config.expert_weights_norm, "norm_topk_prob must be true"); + dsv41_require(config.hidden_act == "silu", "hidden_act must be silu"); + dsv41_require(config.scoring_func == "sqrtsoftplus", "scoring_func must be sqrtsoftplus"); + dsv41_require(config.topk_method == "noaux_tc", "topk_method must be noaux_tc"); + + dsv41_require(config.compress_ratios.size() == LLAMA_DSV41_N_LAYER, "compress_ratios must have 40 entries"); + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + dsv41_require(config.compress_ratios[il] == llama_dsv41_compress_ratio(il), "compress_ratios has an invalid main-layer value"); + } + dsv41_require(config.kv_sources == std::vector(std::begin(DSV41_KV_SOURCES), std::end(DSV41_KV_SOURCES)), "kv_source_layers must be [2,8,14,20]"); + dsv41_require(config.index_sources == std::vector(std::begin(DSV41_INDEX_SOURCES), std::end(DSV41_INDEX_SOURCES)), "index_source_layers must be [2,8,14,20,24,28,32,36]"); + dsv41_require(config.engram_layers == std::vector(std::begin(DSV41_ENGRAM_LAYERS), std::end(DSV41_ENGRAM_LAYERS)), "engram.layer_ids must be [1,14]"); + dsv41_require(config.engram_rows == std::vector(std::begin(DSV41_ENGRAM_ROWS), std::end(DSV41_ENGRAM_ROWS)), "engram.rows must be [384006168,384016682]"); + dsv41_require(config.engram_encoding == LLAMA_DSV41_ENGRAM_ENCODING, "engram.encoding must be e4m3_e8m0_32_row264"); + dsv41_require(config.engram_compressed_vocab_size == LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB, "engram.compressed_vocab_size must be 99092"); + dsv41_require(config.engram_pad_id == LLAMA_DSV41_ENGRAM_PAD_ID, "engram.pad_id must be 2"); + dsv41_require(config.engram_token_map_size == LLAMA_DSV41_N_VOCAB, "engram.token_map must contain 129280 entries"); + dsv41_require(config.engram_primes_size == LLAMA_DSV41_ENGRAM_PRIMES_COUNT, "engram.primes must contain 48 entries"); + dsv41_require(config.engram_multipliers_size == LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT, "engram.multipliers must contain 8 entries"); + if (!config.engram_token_map.empty() || !config.engram_primes.empty() || !config.engram_multipliers.empty()) { + llama_dsv41_make_engram_layout(config); + } +} + +llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config) { + dsv41_require(config.engram_token_map.size() == LLAMA_DSV41_N_VOCAB, "engram.token_map data must contain 129280 entries"); + dsv41_require(config.engram_primes.size() == LLAMA_DSV41_ENGRAM_PRIMES_COUNT, "engram.primes data must contain 48 entries"); + dsv41_require(config.engram_multipliers.size() == LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT, "engram.multipliers data must contain 8 entries"); + + llama_engram_layout layout; + layout.encoding = config.engram_encoding; + std::copy(config.engram_layers.begin(), config.engram_layers.end(), layout.layer_ids.begin()); + layout.token_map = config.engram_token_map; + layout.compressed_vocab_size = config.engram_compressed_vocab_size; + layout.pad_id = config.engram_pad_id; + std::copy(config.engram_rows.begin(), config.engram_rows.end(), layout.rows.begin()); + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + std::copy_n( + config.engram_primes.begin() + layer*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS, + layout.primes[layer].begin()); + std::copy_n( + config.engram_multipliers.begin() + layer*LLAMA_ENGRAM_NGRAM, + LLAMA_ENGRAM_NGRAM, + layout.multipliers[layer].begin()); + } + llama_engram_hasher validate(layout); + return layout; +} + + + +static int32_t dsv41_source_layer(const int32_t * sources, size_t n, uint32_t il) { + int32_t result = -1; + for (size_t i = 0; i < n; ++i) { + if ((uint32_t) sources[i] > il) { + break; + } + result = sources[i]; + } + return result; +} + +llama_dsv41_cache_state::llama_dsv41_cache_state(uint32_t compressed_cache_size) : + compressed_cache_size(compressed_cache_size), + raw(LLAMA_DSV41_N_SWA, -1) { + if (compressed_cache_size == 0) { + throw std::runtime_error("DeepSeek V4.1 compressed cache is empty"); + } + for (int32_t source : DSV41_KV_SOURCES) { + const uint32_t ratio = llama_dsv41_compress_ratio(source); + sources.emplace(source, source_state { + ratio, + std::vector(compressed_cache_size, -1), + std::vector(ratio, -1), + }); + } +} + +void llama_dsv41_cache_state::clear() { + pos = -1; + std::fill(raw.begin(), raw.end(), -1); + for (auto & item : sources) { + std::fill(item.second.compressed.begin(), item.second.compressed.end(), -1); + std::fill(item.second.pending.begin(), item.second.pending.end(), -1); + } + candidates.clear(); +} + +void llama_dsv41_cache_state::append(llama_pos next_pos) { + if (next_pos != pos + 1) { + throw std::runtime_error("DeepSeek V4.1 cache requires contiguous single-sequence positions"); + } + + for (const auto & item : sources) { + const source_state & state = item.second; + if ((next_pos + 1)%state.ratio == 0 && (uint64_t) (next_pos/state.ratio) >= compressed_cache_size) { + throw std::runtime_error("DeepSeek V4.1 compressed cache overflow"); + } + } + + raw[next_pos%raw.size()] = next_pos; + for (auto & item : sources) { + source_state & state = item.second; + state.pending[next_pos%state.ratio] = next_pos; + if ((next_pos + 1)%state.ratio != 0) { + continue; + } + + const uint64_t dst = next_pos/state.ratio; + state.compressed[dst] = next_pos + 1 - state.ratio; + } + pos = next_pos; +} + +void llama_dsv41_cache_state::set_candidate_blocks(const std::vector & blocks) { + candidates = blocks; +} + +llama_pos llama_dsv41_cache_state::position() const { + return pos; +} + +const std::vector & llama_dsv41_cache_state::raw_slots() const { + return raw; +} + +const std::vector & llama_dsv41_cache_state::compressed_slots(uint32_t source_layer) const { + const auto it = sources.find(source_layer); + if (it == sources.end()) { + throw std::runtime_error("DeepSeek V4.1 compressed cache source layer is invalid"); + } + return it->second.compressed; +} + +const std::vector & llama_dsv41_cache_state::pending_slots(uint32_t source_layer) const { + const auto it = sources.find(source_layer); + if (it == sources.end()) { + throw std::runtime_error("DeepSeek V4.1 compressor state source layer is invalid"); + } + return it->second.pending; +} + +const std::vector & llama_dsv41_cache_state::candidate_blocks() const { + return candidates; +} + +uint64_t llama_dsv41_memory_accounting::total() const { + return raw_kv + compressed_kv + index_keys + compressor_carry + + candidate_scores + candidate_ids + position_state + graph_workspace; +} + +llama_dsv41_memory_accounting llama_dsv41_account_memory( + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_tokens, + uint32_t kv_element_size, + uint32_t index_element_size, + uint64_t graph_workspace) { + if (n_ctx == 0 || n_seq == 0 || n_tokens == 0 || kv_element_size == 0 || index_element_size == 0) { + throw std::runtime_error("DeepSeek V4.1 memory accounting dimensions must be non-zero"); + } + + llama_dsv41_memory_accounting result; + const uint64_t raw_rows = (uint64_t) LLAMA_DSV41_N_LAYER*LLAMA_DSV41_N_SWA*n_seq; + const uint64_t ratio_2_rows = ((uint64_t) n_ctx + 1)/2; + const uint64_t ratio_1_rows = n_ctx; + const uint64_t compressed_rows = (3*ratio_2_rows + ratio_1_rows)*n_seq; + const uint64_t candidate_blocks = ((uint64_t) n_ctx + LLAMA_DSV41_CANDIDATE_BLOCK_SIZE - 1)/ + LLAMA_DSV41_CANDIDATE_BLOCK_SIZE; + uint64_t pending_rows = 0; + uint64_t gated_pending_rows = 0; + for (int32_t source : DSV41_KV_SOURCES) { + const uint32_t ratio = llama_dsv41_compress_ratio(source); + pending_rows += ratio; + if (ratio == 2) { + gated_pending_rows += ratio; + } + } + + result.raw_kv = raw_rows*LLAMA_DSV41_N_HEAD_DIM*kv_element_size; + result.compressed_kv = compressed_rows*LLAMA_DSV41_N_HEAD_DIM*kv_element_size; + result.index_keys = compressed_rows*LLAMA_DSV41_N_INDEX_HEAD_DIM*index_element_size; + result.compressor_carry = + (pending_rows + gated_pending_rows)*LLAMA_DSV41_N_HEAD_DIM*kv_element_size*n_seq; + result.candidate_scores = candidate_blocks*sizeof(float)*n_tokens; + result.candidate_ids = std::min(candidate_blocks, LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS)* + sizeof(int32_t)*n_tokens; + const uint64_t position_rows = + ((uint64_t) LLAMA_DSV41_N_SWA + compressed_rows/n_seq + pending_rows)*n_seq; + result.position_state = position_rows*sizeof(llama_pos); + result.graph_workspace = graph_workspace; + return result; +} + +int32_t llama_dsv41_kv_source_layer(uint32_t il) { + return dsv41_source_layer(DSV41_KV_SOURCES, sizeof(DSV41_KV_SOURCES)/sizeof(DSV41_KV_SOURCES[0]), il); +} + +int32_t llama_dsv41_index_source_layer(uint32_t il) { + return dsv41_source_layer(DSV41_INDEX_SOURCES, sizeof(DSV41_INDEX_SOURCES)/sizeof(DSV41_INDEX_SOURCES[0]), il); +} + +uint32_t llama_dsv41_compress_ratio(uint32_t il) { + if (il < 2) { + return 0; + } + if (il < 20) { + return 2; + } + if (il < LLAMA_DSV41_N_LAYER) { + return 1; + } + throw std::runtime_error("DeepSeek V4.1 layer index is out of range"); +} + +llama_dsv41_layer_plan llama_dsv41_build_layer_plan( + uint32_t il, + const std::vector & positions, + uint32_t compressed_cache_size) { + if (positions.empty()) { + throw std::runtime_error("DeepSeek V4.1 graph plan requires at least one token"); + } + if (positions.front() < 0) { + throw std::runtime_error("DeepSeek V4.1 graph plan requires non-negative positions"); + } + for (size_t i = 1; i < positions.size(); ++i) { + if (positions[i] != positions[i - 1] + 1) { + throw std::runtime_error("DeepSeek V4.1 graph plan requires one contiguous sequence"); + } + } + + llama_dsv41_layer_plan plan = {}; + plan.layer = il; + plan.ratio = llama_dsv41_compress_ratio(il); + plan.kv_source_layer = llama_dsv41_kv_source_layer(il); + plan.index_source_layer = llama_dsv41_index_source_layer(il); + plan.owns_kv_source = plan.kv_source_layer == (int32_t) il; + plan.owns_index_source = plan.index_source_layer == (int32_t) il; + plan.builds_candidates = il == LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + plan.uses_candidates = plan.owns_index_source && il > LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + plan.reuses_index_selection = !plan.owns_index_source && plan.index_source_layer >= 0; + plan.collapses_output = il + 1 == LLAMA_DSV41_N_LAYER; + plan.raw_ring_order = llama_dsv41_raw_ring_order(positions.back(), LLAMA_DSV41_N_SWA); + if (plan.owns_kv_source && plan.ratio != 0) { + plan.compression = llama_dsv41_build_compression_plan(positions, plan.ratio, compressed_cache_size); + } + return plan; +} + +llama_dsv41_compression_plan llama_dsv41_build_compression_plan( + const std::vector & positions, + uint32_t ratio, + uint32_t cache_size) { + if (ratio != 1 && ratio != 2) { + throw std::runtime_error("DeepSeek V4.1 compression ratio must be 1 or 2"); + } + if (cache_size == 0) { + throw std::runtime_error("DeepSeek V4.1 compressed cache is empty"); + } + + llama_dsv41_compression_plan plan; + plan.n_visible.resize(positions.size(), 0); + + std::vector latest_state_src(ratio, -1); + std::vector latest_state_pos(ratio, -1); + + const int32_t scratch_offset = ratio; + for (size_t i = 0; i < positions.size(); ++i) { + const llama_pos pos = positions[i]; + if (pos < 0) { + continue; + } + + plan.n_visible[i] = (int32_t) ((pos + 1)/ratio); + plan.n_kv = std::max(plan.n_kv, plan.n_visible[i]); + + if (ratio == 1) { + if ((uint64_t) pos >= cache_size) { + throw std::runtime_error("DeepSeek V4.1 compressed cache overflow"); + } + plan.write_idxs.push_back(pos); + plan.write_pos.push_back((int32_t) pos); + continue; + } + + const int32_t row = (int32_t) (pos%ratio); + plan.state_pos.push_back(row); + if (latest_state_src[row] < 0 || pos > latest_state_pos[row]) { + latest_state_src[row] = (int32_t) i; + latest_state_pos[row] = pos; + } + + if ((pos + 1)%ratio != 0) { + continue; + } + + const llama_pos source_start = pos + 1 - ratio; + const int64_t write_idx = pos/ratio; + if ((uint64_t) write_idx >= cache_size) { + throw std::runtime_error("DeepSeek V4.1 compressed cache overflow"); + } + + for (uint32_t j = 0; j < ratio; ++j) { + const llama_pos source_pos = source_start + j; + int32_t source_idx = (int32_t) (source_pos%ratio); + for (size_t k = 0; k < positions.size(); ++k) { + if (positions[k] == source_pos) { + source_idx = scratch_offset + (int32_t) k; + break; + } + } + plan.state_read_idxs.push_back(source_idx); + } + + plan.write_idxs.push_back(write_idx); + plan.write_pos.push_back((int32_t) source_start); + } + + if (ratio == 2) { + for (uint32_t row = 0; row < ratio; ++row) { + if (latest_state_src[row] >= 0) { + plan.state_persist_src_idxs.push_back(latest_state_src[row]); + plan.state_persist_dst_idxs.push_back((int32_t) row); + } + } + } + + plan.n_kv = std::min(cache_size, std::max(1, plan.n_kv)); + return plan; +} + +std::vector llama_dsv41_select_candidate_blocks( + const std::vector & scores, + uint32_t n_visible, + uint32_t block_size, + uint32_t top_k_blocks) { + if (block_size == 0) { + throw std::runtime_error("DeepSeek V4.1 candidate block size must be non-zero"); + } + + n_visible = std::min(n_visible, scores.size()); + if (n_visible == 0 || top_k_blocks == 0) { + return {}; + } + + const uint32_t n_blocks = (n_visible + block_size - 1)/block_size; + std::vector block_scores(n_blocks, -std::numeric_limits::infinity()); + for (uint32_t block = 0; block < n_blocks; ++block) { + const uint32_t i0 = block*block_size; + const uint32_t i1 = std::min(n_visible, i0 + block_size); + for (uint32_t i = i0; i < i1; ++i) { + if (std::isnan(scores[i])) { + throw std::runtime_error("DeepSeek V4.1 candidate score is NaN"); + } + block_scores[block] = std::max(block_scores[block], scores[i]); + } + } + + const int32_t final_block = (int32_t) n_blocks - 1; + std::vector blocks(n_blocks - 1); + std::iota(blocks.begin(), blocks.end(), 0); + std::stable_sort(blocks.begin(), blocks.end(), [&](int32_t a, int32_t b) { + if (block_scores[a] != block_scores[b]) { + return block_scores[a] > block_scores[b]; + } + return a < b; + }); + + const uint32_t n_selected = std::min(top_k_blocks, n_blocks); + std::vector selected; + selected.reserve(n_selected); + selected.push_back(final_block); + selected.insert(selected.end(), blocks.begin(), blocks.begin() + n_selected - 1); + return selected; +} + +std::vector llama_dsv41_candidate_rows( + const std::vector & blocks, + uint32_t n_visible, + uint32_t block_size) { + if (block_size == 0) { + throw std::runtime_error("DeepSeek V4.1 candidate block size must be non-zero"); + } + + std::vector rows; + for (int32_t block : blocks) { + if (block < 0) { + throw std::runtime_error("DeepSeek V4.1 candidate block index must be non-negative"); + } + const uint64_t i0 = (uint64_t) block*block_size; + const uint64_t i1 = std::min(n_visible, i0 + block_size); + for (uint64_t i = i0; i < i1; ++i) { + rows.push_back((int32_t) i); + } + } + return rows; +} + +std::vector llama_dsv41_raw_ring_order(llama_pos pos, uint32_t window) { + if (window == 0 || pos < 0) { + return {}; + } + + const uint32_t n_raw = std::min((uint64_t) pos + 1, window); + const uint32_t start = (uint32_t) ((pos + 1 - n_raw)%window); + + std::vector result(n_raw); + for (uint32_t i = 0; i < n_raw; ++i) { + result[i] = (start + i)%window; + } + return result; +} + +std::vector llama_dsv41_output_collapse( + const std::vector & residual, + const std::vector & pre, + uint32_t n_embd, + uint32_t hc_mult) { + if (pre.size() != hc_mult || residual.size() != (size_t) n_embd*hc_mult) { + throw std::runtime_error("DeepSeek V4.1 output collapse shape mismatch"); + } + + std::vector result(n_embd, 0.0f); + for (uint32_t h = 0; h < hc_mult; ++h) { + for (uint32_t d = 0; d < n_embd; ++d) { + result[d] += residual[(size_t) h*n_embd + d]*pre[h]; + } + } + return result; +} + +ggml_tensor * llama_dsv41_build_ratio_pool( + ggml_context * ctx, + ggml_tensor * kv, + ggml_tensor * gate, + uint32_t ratio) { + if (ratio != 1 && ratio != 2) { + throw std::runtime_error("DeepSeek V4.1 graph compression ratio must be 1 or 2"); + } + if (kv->ne[1] != ratio) { + throw std::runtime_error("DeepSeek V4.1 graph compressor input shape mismatch"); + } + if (ratio == 1) { + return ggml_reshape_2d(ctx, kv, kv->ne[0], kv->ne[2]); + } + if (gate == nullptr || !ggml_are_same_shape(kv, gate)) { + throw std::runtime_error("DeepSeek V4.1 ratio-2 graph requires matching KV and gate tensors"); + } + + ggml_tensor * kv_t = ggml_cont(ctx, ggml_permute(ctx, kv, 1, 0, 2, 3)); + ggml_tensor * gate_t = ggml_cont(ctx, ggml_permute(ctx, gate, 1, 0, 2, 3)); + ggml_tensor * weights = ggml_soft_max(ctx, gate_t); + ggml_tensor * pooled = ggml_sum_rows(ctx, ggml_mul(ctx, kv_t, weights)); + return ggml_reshape_2d(ctx, pooled, kv->ne[0], kv->ne[2]); +} + +ggml_tensor * llama_dsv41_build_shared_softmax( + ggml_context * ctx, + ggml_tensor * raw_scores, + ggml_tensor * compressed_scores) { + if (raw_scores == nullptr && compressed_scores == nullptr) { + throw std::runtime_error("DeepSeek V4.1 attention requires raw or compressed scores"); + } + ggml_tensor * scores = raw_scores; + if (scores == nullptr) { + scores = compressed_scores; + } else if (compressed_scores != nullptr) { + if (raw_scores->ne[1] != compressed_scores->ne[1] || + raw_scores->ne[2] != compressed_scores->ne[2] || + raw_scores->ne[3] != compressed_scores->ne[3]) { + throw std::runtime_error("DeepSeek V4.1 raw and compressed score shapes are incompatible"); + } + scores = ggml_concat(ctx, raw_scores, compressed_scores, 0); + } + return ggml_soft_max(ctx, scores); +} + +ggml_tensor * llama_dsv41_build_candidate_blocks( + ggml_context * ctx, + ggml_tensor * block_scores, + ggml_tensor * final_blocks, + uint32_t n_candidate) { + if (block_scores == nullptr || final_blocks == nullptr || + block_scores->type != GGML_TYPE_F32 || + final_blocks->type != GGML_TYPE_I32 || + final_blocks->ne[0] != 1 || + final_blocks->ne[1] != block_scores->ne[1]) { + throw std::runtime_error("DeepSeek V4.1 candidate graph shape mismatch"); + } + + const int64_t n_blocks = block_scores->ne[0]; + const int64_t n_tokens = block_scores->ne[1]; + if (n_candidate == 0 || n_candidate > (uint32_t) n_blocks) { + throw std::runtime_error("DeepSeek V4.1 candidate graph width is invalid"); + } + if (n_candidate == 1) { + return ggml_cont(ctx, final_blocks); + } + + ggml_tensor * local = ggml_reshape_2d( + ctx, ggml_arange(ctx, 0.0f, (float) (n_blocks - 1), 1.0f), + n_blocks - 1, 1); + local = ggml_repeat_4d( + ctx, local, n_blocks - 1, n_tokens, 1, 1); + ggml_tensor * final_f32 = ggml_cast(ctx, final_blocks, GGML_TYPE_F32); + ggml_tensor * shift = ggml_step( + ctx, ggml_scale_bias( + ctx, ggml_sub(ctx, local, final_f32), 1.0f, 0.5f)); + ggml_tensor * ordinary_ids = ggml_cast( + ctx, ggml_add(ctx, local, shift), GGML_TYPE_I32); + + ggml_tensor * ordinary_scores = ggml_get_rows( + ctx, + ggml_reshape_3d(ctx, block_scores, 1, n_blocks, n_tokens), + ordinary_ids); + ordinary_scores = ggml_reshape_2d( + ctx, ordinary_scores, n_blocks - 1, n_tokens); + ggml_tensor * selected_local = ggml_argsort_top_k( + ctx, ordinary_scores, n_candidate - 1); + ggml_tensor * selected = ggml_get_rows( + ctx, + ggml_reshape_3d(ctx, ordinary_ids, 1, n_blocks - 1, n_tokens), + selected_local); + selected = ggml_reshape_2d( + ctx, selected, n_candidate - 1, n_tokens); + return ggml_cont(ctx, ggml_concat(ctx, final_blocks, selected, 0)); +} + +ggml_tensor * llama_dsv41_build_output_collapse( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * pre, + uint32_t n_embd, + uint32_t hc_mult, + uint32_t n_tokens) { + if (residual->ne[0] != n_embd || residual->ne[1] != hc_mult || residual->ne[2] != n_tokens || + pre->ne[0] != hc_mult || pre->ne[1] != n_tokens) { + throw std::runtime_error("DeepSeek V4.1 graph output collapse shape mismatch"); + } + + ggml_tensor * residual_t = ggml_cont(ctx, ggml_permute(ctx, residual, 1, 0, 2, 3)); + ggml_tensor * pre_t = ggml_reshape_3d(ctx, pre, hc_mult, 1, n_tokens); + ggml_tensor * collapsed = ggml_sum_rows(ctx, ggml_mul(ctx, residual_t, pre_t)); + collapsed = ggml_reshape_2d(ctx, collapsed, n_embd, n_tokens); + return ggml_cast(ctx, collapsed, GGML_TYPE_BF16); +} + +ggml_tensor * llama_dsv41_build_output_norm_input( + ggml_context * ctx, + ggml_tensor * collapsed) { + if (collapsed == nullptr || collapsed->type != GGML_TYPE_BF16) { + throw std::runtime_error("DeepSeek V4.1 output collapse must be BF16"); + } + return ggml_cast(ctx, collapsed, GGML_TYPE_F32); +} + +ggml_tensor * llama_dsv41_build_output( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * pre, + ggml_tensor * output_norm, + ggml_tensor * output, + float rms_eps, + uint32_t hc_mult) { + if (output_norm->ne[0] != residual->ne[0] || output->ne[0] != residual->ne[0]) { + throw std::runtime_error("DeepSeek V4.1 output tensor shape mismatch"); + } + + ggml_tensor * collapsed = llama_dsv41_build_output_collapse( + ctx, residual, pre, residual->ne[0], hc_mult, residual->ne[2]); + ggml_tensor * normalized = ggml_rms_norm( + ctx, llama_dsv41_build_output_norm_input(ctx, collapsed), rms_eps); + normalized = ggml_mul(ctx, normalized, output_norm); + return ggml_mul_mat(ctx, output, normalized); +} diff --git a/src/llama-dsv41.h b/src/llama-dsv41.h new file mode 100644 index 000000000000..86fca484d8de --- /dev/null +++ b/src/llama-dsv41.h @@ -0,0 +1,249 @@ +#pragma once + +#include "llama-engram.h" +#include "llama.h" + +#include +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +static constexpr uint32_t LLAMA_DSV41_N_LAYER = 40; +static constexpr uint32_t LLAMA_DSV41_N_EMBD = 5120; +static constexpr uint32_t LLAMA_DSV41_N_VOCAB = 129280; +static constexpr uint32_t LLAMA_DSV41_N_HEAD = 64; +static constexpr uint32_t LLAMA_DSV41_N_HEAD_KV = 1; +static constexpr uint32_t LLAMA_DSV41_N_HEAD_DIM = 512; +static constexpr uint32_t LLAMA_DSV41_N_ROT = 64; +static constexpr uint32_t LLAMA_DSV41_N_LORA_Q = 1280; +static constexpr uint32_t LLAMA_DSV41_N_LORA_O = 1024; +static constexpr uint32_t LLAMA_DSV41_N_O_GROUP = 8; +static constexpr uint32_t LLAMA_DSV41_N_FF_DENSE = 18432; +static constexpr uint32_t LLAMA_DSV41_N_EXPERT = 384; +static constexpr uint32_t LLAMA_DSV41_N_EXPERT_USED = 6; +static constexpr uint32_t LLAMA_DSV41_N_EXPERT_SHARED = 1; +static constexpr uint32_t LLAMA_DSV41_N_FF_EXP = 2304; +static constexpr uint32_t LLAMA_DSV41_N_INDEX_HEAD = 32; +static constexpr uint32_t LLAMA_DSV41_N_INDEX_HEAD_DIM = 128; +static constexpr uint32_t LLAMA_DSV41_N_INDEX_TOP_K = 512; +static constexpr uint32_t LLAMA_DSV41_N_SWA = 128; +static constexpr uint32_t LLAMA_DSV41_N_CTX = 1048576; +static constexpr uint32_t LLAMA_DSV41_HC_MULT = 4; +static constexpr uint32_t LLAMA_DSV41_HC_SINKHORN_ITERS = 20; +static constexpr uint32_t LLAMA_DSV41_CANDIDATE_SOURCE_LAYER = 20; +static constexpr uint32_t LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS = 2048; +static constexpr uint32_t LLAMA_DSV41_CANDIDATE_BLOCK_SIZE = 8; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB = 99092; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_PAD_ID = 2; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_PRIMES_COUNT = 48; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT = 8; +static constexpr const char * LLAMA_DSV41_ENGRAM_ENCODING = "e4m3_e8m0_32_row264"; + +struct llama_dsv41_config { + uint32_t n_ctx_train; + uint32_t n_embd; + uint32_t n_layer; + uint32_t n_vocab; + uint32_t n_head; + uint32_t n_head_kv; + uint32_t n_head_dim; + uint32_t n_rot; + uint32_t n_lora_q; + uint32_t n_lora_o; + uint32_t n_o_group; + uint32_t n_ff_dense; + uint32_t n_ff_expert; + uint32_t n_expert; + uint32_t n_expert_used; + uint32_t n_expert_shared; + uint32_t indexer_n_head; + uint32_t indexer_head_size; + uint32_t indexer_top_k; + uint32_t hc_count; + uint32_t hc_sinkhorn_iters; + uint32_t raw_window; + uint32_t candidate_source_layer; + uint32_t candidate_topk_blocks; + uint32_t candidate_block_size; + float f_norm_rms_eps; + float hc_eps; + float swiglu_clamp; + float routed_scale; + uint32_t rope_theta; + uint32_t compress_rope_theta; + float yarn_factor; + float yarn_beta_fast; + float yarn_beta_slow; + float yarn_original_context; + bool expert_weights_norm; + std::string hidden_act; + std::string scoring_func; + std::string topk_method; + std::vector compress_ratios; + std::vector kv_sources; + std::vector index_sources; + std::vector engram_layers; + std::vector engram_rows; + std::string engram_encoding; + uint32_t engram_compressed_vocab_size; + uint32_t engram_pad_id; + uint32_t engram_token_map_size; + uint32_t engram_primes_size; + uint32_t engram_multipliers_size; + std::vector engram_token_map; + std::vector engram_primes; + std::vector engram_multipliers; +}; + +void llama_dsv41_validate_config(const llama_dsv41_config & config); +llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config); + +struct llama_dsv41_compression_plan { + std::vector state_pos; + std::vector state_persist_src_idxs; + std::vector state_persist_dst_idxs; + std::vector state_read_idxs; + std::vector write_idxs; + std::vector write_pos; + std::vector n_visible; + int64_t n_kv = 0; +}; + +struct llama_dsv41_layer_plan { + uint32_t layer; + uint32_t ratio; + int32_t kv_source_layer; + int32_t index_source_layer; + bool owns_kv_source; + bool owns_index_source; + bool builds_candidates; + bool uses_candidates; + bool reuses_index_selection; + bool collapses_output; + std::vector raw_ring_order; + llama_dsv41_compression_plan compression; +}; + +class llama_dsv41_cache_state { +public: + explicit llama_dsv41_cache_state(uint32_t compressed_cache_size); + + void clear(); + void append(llama_pos pos); + void set_candidate_blocks(const std::vector & blocks); + + llama_pos position() const; + const std::vector & raw_slots() const; + const std::vector & compressed_slots(uint32_t source_layer) const; + const std::vector & pending_slots(uint32_t source_layer) const; + const std::vector & candidate_blocks() const; + +private: + struct source_state { + uint32_t ratio; + std::vector compressed; + std::vector pending; + }; + + uint32_t compressed_cache_size; + llama_pos pos = -1; + std::vector raw; + std::map sources; + std::vector candidates; +}; + +struct llama_dsv41_memory_accounting { + uint64_t raw_kv = 0; + uint64_t compressed_kv = 0; + uint64_t index_keys = 0; + uint64_t compressor_carry = 0; + uint64_t candidate_scores = 0; + uint64_t candidate_ids = 0; + uint64_t position_state = 0; + uint64_t graph_workspace = 0; + + uint64_t total() const; +}; + +llama_dsv41_memory_accounting llama_dsv41_account_memory( + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_tokens, + uint32_t kv_element_size, + uint32_t index_element_size, + uint64_t graph_workspace); + +int32_t llama_dsv41_kv_source_layer(uint32_t il); +int32_t llama_dsv41_index_source_layer(uint32_t il); +uint32_t llama_dsv41_compress_ratio(uint32_t il); + +llama_dsv41_layer_plan llama_dsv41_build_layer_plan( + uint32_t il, + const std::vector & positions, + uint32_t compressed_cache_size); + +llama_dsv41_compression_plan llama_dsv41_build_compression_plan( + const std::vector & positions, + uint32_t ratio, + uint32_t cache_size); + +std::vector llama_dsv41_select_candidate_blocks( + const std::vector & scores, + uint32_t n_visible, + uint32_t block_size, + uint32_t top_k_blocks); + +std::vector llama_dsv41_candidate_rows( + const std::vector & blocks, + uint32_t n_visible, + uint32_t block_size); + +std::vector llama_dsv41_raw_ring_order(llama_pos pos, uint32_t window); + +std::vector llama_dsv41_output_collapse( + const std::vector & residual, + const std::vector & pre, + uint32_t n_embd, + uint32_t hc_mult); + +ggml_tensor * llama_dsv41_build_ratio_pool( + ggml_context * ctx, + ggml_tensor * kv, + ggml_tensor * gate, + uint32_t ratio); + +ggml_tensor * llama_dsv41_build_shared_softmax( + ggml_context * ctx, + ggml_tensor * raw_scores, + ggml_tensor * compressed_scores); + +ggml_tensor * llama_dsv41_build_candidate_blocks( + ggml_context * ctx, + ggml_tensor * block_scores, + ggml_tensor * final_blocks, + uint32_t n_candidate); + +ggml_tensor * llama_dsv41_build_output_collapse( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * pre, + uint32_t n_embd, + uint32_t hc_mult, + uint32_t n_tokens); + +ggml_tensor * llama_dsv41_build_output_norm_input( + ggml_context * ctx, + ggml_tensor * collapsed); + +ggml_tensor * llama_dsv41_build_output( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * pre, + ggml_tensor * output_norm, + ggml_tensor * output, + float rms_eps, + uint32_t hc_mult); diff --git a/src/llama-engram.cpp b/src/llama-engram.cpp new file mode 100644 index 000000000000..336aae39f2d7 --- /dev/null +++ b/src/llama-engram.cpp @@ -0,0 +1,305 @@ +#include "llama-engram.h" + +#include "llama-bounded-file.h" +#include "llama-impl.h" + +#include +#include +#include +#include +#include +#include + +static void llama_engram_validate_layout(const llama_engram_layout & layout) { + if (layout.encoding != "e4m3_e8m0_32_row264") { + throw std::invalid_argument("llama_engram: unsupported row encoding"); + } + if (layout.layer_ids[0] == layout.layer_ids[1]) { + throw std::invalid_argument("llama_engram: layer IDs must be distinct"); + } + if (layout.token_map.empty() || layout.token_map.size() > (size_t) INT32_MAX) { + throw std::invalid_argument("llama_engram: invalid token map size"); + } + if (layout.compressed_vocab_size == 0 || layout.compressed_vocab_size > (uint32_t) INT32_MAX || + layout.pad_id >= layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: invalid compressed vocabulary"); + } + for (uint32_t token : layout.token_map) { + if (token >= layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: token map entry is out of range"); + } + } + + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + const uint64_t multiplier = layout.multipliers[layer][i]; + if ((multiplier & 1) == 0 || + multiplier > (uint64_t) INT64_MAX / layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: invalid hash multiplier"); + } + } + + uint64_t row_count = 0; + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + const uint32_t prime = layout.primes[layer][col]; + if (prime < 2) { + throw std::invalid_argument("llama_engram: invalid hash prime"); + } + row_count += prime; + } + if (row_count > UINT32_MAX || row_count != layout.rows[layer]) { + throw std::invalid_argument("llama_engram: row extent does not match hash buckets"); + } + } +} + +void llama_engram_history::reset() { + tail.fill(LLAMA_ENGRAM_DEAD); +} + +struct llama_engram_hasher::impl { + llama_engram_layout layout; + + explicit impl(llama_engram_layout layout) : layout(std::move(layout)) { + llama_engram_validate_layout(this->layout); + } + + void validate_history(const llama_engram_history & history) const { + for (int32_t token : history.tail) { + if (token < LLAMA_ENGRAM_DEAD || + (token >= 0 && (uint32_t) token >= layout.compressed_vocab_size)) { + throw std::invalid_argument("llama_engram: invalid token history"); + } + } + } + + void hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const { + if (count > SIZE_MAX / (LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS * sizeof(*rows)) || + (count != 0 && (tokens == nullptr || rows == nullptr))) { + throw std::invalid_argument("llama_engram: invalid hash buffers"); + } + validate_history(history); + for (size_t i = 0; i < count; ++i) { + if (tokens[i] < 0 || (size_t) tokens[i] >= layout.token_map.size()) { + throw std::invalid_argument("llama_engram: token ID is out of range"); + } + } + + llama_engram_history next = history; + uint32_t * output = rows; + for (size_t pos = 0; pos < count; ++pos) { + const int32_t current = mask != nullptr && mask[pos] == 0 ? + LLAMA_ENGRAM_DEAD : (int32_t) layout.token_map[tokens[pos]]; + uint32_t ids[LLAMA_ENGRAM_NGRAM]; + bool blocked = false; + for (size_t depth = 0; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + const int32_t id = depth == 0 ? current : next.tail[depth - 1]; + blocked = blocked || id == LLAMA_ENGRAM_DEAD; + ids[depth] = blocked ? layout.pad_id : (uint32_t) id; + } + + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + uint64_t hash = (uint64_t) ids[0] * layout.multipliers[layer][0]; + uint64_t offset = 0; + for (size_t depth = 1; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + hash ^= (uint64_t) ids[depth] * layout.multipliers[layer][depth]; + for (size_t head = 0; head < LLAMA_ENGRAM_HEADS; ++head) { + const size_t col = (depth - 1) * LLAMA_ENGRAM_HEADS + head; + const uint32_t prime = layout.primes[layer][col]; + *output++ = (uint32_t) (hash % prime + offset); + offset += prime; + } + } + } + + for (size_t i = next.tail.size() - 1; i > 0; --i) { + next.tail[i] = next.tail[i - 1]; + } + next.tail[0] = current; + } + history = next; + } +}; + +llama_engram_hasher::llama_engram_hasher(llama_engram_layout layout) + : pimpl(std::make_unique(std::move(layout))) {} + +llama_engram_hasher::~llama_engram_hasher() = default; + +void llama_engram_hasher::hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const { + pimpl->hash(history, tokens, mask, count, rows); +} + +const llama_engram_layout & llama_engram_hasher::layout() const { + return pimpl->layout; +} + +static float llama_engram_e4m3(uint8_t code) { + const int exponent = (code >> 3) & 15; + const int mantissa = code & 7; + const float value = exponent != 0 ? + std::ldexp((float) (8 + mantissa), exponent - 10) : + std::ldexp((float) mantissa, -9); + return (code & 128) != 0 ? -value : value; +} + +void llama_engram_decode_row(const uint8_t row[LLAMA_ENGRAM_ROW_BYTES], float output[LLAMA_ENGRAM_DIM]) { + if (row == nullptr || output == nullptr) { + throw std::invalid_argument("llama_engram: invalid row decode buffers"); + } + + float decoded[LLAMA_ENGRAM_DIM]; + for (size_t i = 0; i < LLAMA_ENGRAM_DIM; ++i) { + const uint8_t code = row[i]; + const uint8_t scale = row[LLAMA_ENGRAM_DIM + i / 32]; + if ((code & 127) == 127) { + throw std::domain_error("llama_engram: E4M3 NaN encoding"); + } + if (scale == 255) { + throw std::domain_error("llama_engram: E8M0 scale 255"); + } + + float value = std::ldexp(llama_engram_e4m3(code), (int) scale - 127); + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u; + memcpy(&value, &bits, sizeof(value)); + if (!std::isfinite(value)) { + throw std::domain_error("llama_engram: decoded value is not finite"); + } + decoded[i] = value; + } + memcpy(output, decoded, sizeof(decoded)); +} + +struct llama_engram_table::impl { + struct request { + uint32_t row; + uint32_t output; + }; + + static constexpr size_t BATCH_TOKENS = 2048; + + llama_bounded_file file; + llama_bounded_file::buffer scratch; + uint64_t offset; + uint32_t rows; + std::mutex mutex; + std::vector requests; + + impl(const std::string & fname, uint64_t offset, uint32_t rows) + : file(fname, { true, true }), + scratch(file.make_buffer(LLAMA_ENGRAM_ROW_BYTES)), + offset(offset), + rows(rows) { + const uint64_t bytes = (uint64_t) rows * LLAMA_ENGRAM_ROW_BYTES; + if (rows == 0 || offset > (uint64_t) INT64_MAX || bytes > (uint64_t) INT64_MAX - offset || + offset > file.size() || bytes > file.size() - offset) { + throw std::invalid_argument("llama_engram: invalid table extent"); + } + requests.reserve(BATCH_TOKENS * LLAMA_ENGRAM_COLS); + } + + void validate_rows(const uint32_t * row_ids, size_t count) const { + if (count > SIZE_MAX / (LLAMA_ENGRAM_DIM * sizeof(float)) || + (count != 0 && row_ids == nullptr)) { + throw std::invalid_argument("llama_engram: invalid row list"); + } + for (size_t i = 0; i < count; ++i) { + if (row_ids[i] >= rows) { + throw std::invalid_argument("llama_engram: row ID is out of range"); + } + } + } + + void read_one(uint32_t row, float * output) { + uint8_t raw[LLAMA_ENGRAM_ROW_BYTES]; + file.read(offset + (uint64_t) row * LLAMA_ENGRAM_ROW_BYTES, raw, sizeof(raw), scratch); + llama_engram_decode_row(raw, output); + } + + void read(const uint32_t * row_ids, size_t count, float * output) { + validate_rows(row_ids, count); + if (count != 0 && output == nullptr) { + throw std::invalid_argument("llama_engram: invalid row output"); + } + + std::lock_guard lock(mutex); + for (size_t i = 0; i < count; ++i) { + read_one(row_ids[i], output + i * LLAMA_ENGRAM_DIM); + } + } + + void read_batch(const uint32_t * row_ids, size_t tokens, size_t stride, float * output) { + if (tokens > SIZE_MAX / (LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM * sizeof(float)) || + (tokens != 0 && (row_ids == nullptr || output == nullptr || stride < LLAMA_ENGRAM_COLS)) || + (tokens != 0 && tokens - 1 > (SIZE_MAX / sizeof(*row_ids) - LLAMA_ENGRAM_COLS) / stride)) { + throw std::invalid_argument("llama_engram: invalid batch buffers"); + } + for (size_t token = 0; token < tokens; ++token) { + validate_rows(row_ids + token * stride, LLAMA_ENGRAM_COLS); + } + if (tokens == 0) { + return; + } + + std::lock_guard lock(mutex); + for (size_t start = 0; start < tokens; start += BATCH_TOKENS) { + const size_t count_tokens = std::min(tokens - start, BATCH_TOKENS); + const size_t count_rows = count_tokens * LLAMA_ENGRAM_COLS; + requests.resize(count_rows); + for (size_t i = 0; i < count_rows; ++i) { + requests[i].row = row_ids[(start + i / LLAMA_ENGRAM_COLS) * stride + i % LLAMA_ENGRAM_COLS]; + requests[i].output = (uint32_t) i; + } + std::sort(requests.begin(), requests.end(), [](const request & a, const request & b) { + return a.row < b.row; + }); + + float * chunk = output + start * LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM; + const float * previous = nullptr; + for (size_t i = 0; i < count_rows; ++i) { + float * dst = chunk + (size_t) requests[i].output * LLAMA_ENGRAM_DIM; + if (i != 0 && requests[i].row == requests[i - 1].row) { + memcpy(dst, previous, LLAMA_ENGRAM_DIM * sizeof(*dst)); + } else { + read_one(requests[i].row, dst); + previous = dst; + } + } + } + } +}; + +llama_engram_table::llama_engram_table(const std::string & fname, uint64_t offset, uint32_t rows) + : pimpl(std::make_unique(fname, offset, rows)) {} + +llama_engram_table::~llama_engram_table() = default; + +void llama_engram_table::read(const uint32_t * rows, size_t count, float * output) { + pimpl->read(rows, count, output); +} + +void llama_engram_table::read_batch(const uint32_t * rows, size_t tokens, size_t stride, float * output) { + pimpl->read_batch(rows, tokens, stride, output); +} + +uint32_t llama_engram_table::n_rows() const { + return pimpl->rows; +} + +std::string llama_engram_table::describe() const { + return format("%s @ %llu: %u rows x %u bytes, uncached aligned reads", + pimpl->file.name().c_str(), (unsigned long long) pimpl->offset, + pimpl->rows, LLAMA_ENGRAM_ROW_BYTES); +} diff --git a/src/llama-engram.h b/src/llama-engram.h new file mode 100644 index 000000000000..6c1c0f78b34c --- /dev/null +++ b/src/llama-engram.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +enum { + LLAMA_ENGRAM_LAYERS = 2, + LLAMA_ENGRAM_NGRAM = 4, + LLAMA_ENGRAM_HEADS = 8, + LLAMA_ENGRAM_COLS = 24, + LLAMA_ENGRAM_DIM = 256, + LLAMA_ENGRAM_ROW_BYTES = 264, + LLAMA_ENGRAM_DEAD = -1, +}; + +struct llama_engram_layout { + std::string encoding; + std::array layer_ids = {}; + std::vector token_map; + uint32_t compressed_vocab_size = 0; + uint32_t pad_id = 0; + std::array rows = {}; + std::array, LLAMA_ENGRAM_LAYERS> multipliers = {}; + std::array, LLAMA_ENGRAM_LAYERS> primes = {}; +}; + +struct llama_engram_history { + // Newest compressed token first. LLAMA_ENGRAM_DEAD breaks all n-grams that cross it. + std::array tail = {}; + + void reset(); +}; + +struct llama_engram_hasher { + explicit llama_engram_hasher(llama_engram_layout layout); + ~llama_engram_hasher(); + + // Output is [token][layer][column]. Invalid input does not change history or output. + void hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const; + + const llama_engram_layout & layout() const; + + struct impl; + std::unique_ptr pimpl; +}; + +// Decode 256 E4M3 values and eight E8M0 scales to BF16-rounded values in F32 storage. +void llama_engram_decode_row(const uint8_t row[LLAMA_ENGRAM_ROW_BYTES], float output[LLAMA_ENGRAM_DIM]); + +struct llama_engram_table { + // The table is never mapped or cached. Construction fails unless uncached reads are available. + llama_engram_table(const std::string & fname, uint64_t offset, uint32_t rows); + ~llama_engram_table(); + + void read(const uint32_t * rows, size_t count, float * output); + // Read 24 rows per token. Input uses row-ID stride; output is packed [token][column][dimension]. + void read_batch(const uint32_t * rows, size_t tokens, size_t stride, float * output); + + uint32_t n_rows() const; + std::string describe() const; + + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-expert-store.cpp b/src/llama-expert-store.cpp new file mode 100644 index 000000000000..c8597e52845f --- /dev/null +++ b/src/llama-expert-store.cpp @@ -0,0 +1,646 @@ +#include "llama-expert-store.h" + +#include "llama-impl.h" +#include "llama-mmap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#endif + +// The positional I/O and reservation model is adapted from ggml-org/llama.cpp#25294. +// Leases add the in-flight publication safety described in ggml-org/llama.cpp#27861. + +namespace { + +bool checked_add_u64(uint64_t a, uint64_t b, uint64_t * result) { + if (b > std::numeric_limits::max() - a) { + return false; + } + *result = a + b; + return true; +} + +bool checked_mul_u64(uint64_t a, uint64_t b, uint64_t * result) { + if (a != 0 && b > std::numeric_limits::max() / a) { + return false; + } + *result = a * b; + return true; +} + +bool is_power_of_two(size_t value) { + return value != 0 && (value & (value - 1)) == 0; +} + +struct expert_key { + int32_t layer; + llama_expert_projection projection; + int32_t expert_id; + + bool operator<(const expert_key & other) const { + if (layer != other.layer) { + return layer < other.layer; + } + if (projection != other.projection) { + return projection < other.projection; + } + return expert_id < other.expert_id; + } + + bool operator==(const expert_key & other) const { + return layer == other.layer && projection == other.projection && expert_id == other.expert_id; + } +}; + +struct tensor_key { + int32_t layer; + llama_expert_projection projection; + + bool operator<(const tensor_key & other) const { + if (layer != other.layer) { + return layer < other.layer; + } + return projection < other.projection; + } +}; + +struct aligned_buffer { + uint8_t * data = nullptr; + size_t size = 0; + + aligned_buffer() = default; + + aligned_buffer(size_t size, size_t alignment) { + reset(size, alignment); + } + + aligned_buffer(aligned_buffer && other) noexcept : data(other.data), size(other.size) { + other.data = nullptr; + other.size = 0; + } + + aligned_buffer & operator=(aligned_buffer && other) noexcept { + if (this != &other) { + clear(); + data = other.data; + size = other.size; + other.data = nullptr; + other.size = 0; + } + return *this; + } + + ~aligned_buffer() { + clear(); + } + + aligned_buffer(const aligned_buffer &) = delete; + aligned_buffer & operator=(const aligned_buffer &) = delete; + + void reset(size_t new_size, size_t alignment) { + clear(); + if (new_size == 0) { + return; + } + alignment = std::max(alignment, alignof(void *)); +#if defined(_WIN32) + data = static_cast(_aligned_malloc(new_size, alignment)); + if (data == nullptr) { + throw std::bad_alloc(); + } +#else + void * ptr = nullptr; + if (posix_memalign(&ptr, alignment, new_size) != 0) { + throw std::bad_alloc(); + } + data = static_cast(ptr); +#endif + size = new_size; + } + + void clear() { +#if defined(_WIN32) + _aligned_free(data); +#else + free(data); +#endif + data = nullptr; + size = 0; + } +}; + +struct expert_file { + std::string fname; + uint64_t size = 0; + bool direct = false; + std::unique_ptr file; + + expert_file(const std::string & fname, bool direct_io, bool allow_buffered_io) : fname(fname) { + reopen(direct_io); + if (direct_io && !direct && !allow_buffered_io) { + throw std::runtime_error(format("llama_expert_store: direct I/O is required but unavailable for %s", fname.c_str())); + } + if (direct_io && !direct) { + LLAMA_LOG_WARN("%s: direct I/O is unavailable for %s; using explicitly allowed buffered reads\n", + __func__, fname.c_str()); + } + } + + expert_file(const expert_file &) = delete; + expert_file & operator=(const expert_file &) = delete; + + void reopen(bool direct_io) { + file = std::make_unique(fname.c_str(), "rb", direct_io); + size = file->size(); + direct = direct_io && file->has_direct_io(); + } + + size_t pread_at_least(void * dst, size_t len, uint64_t offset, size_t need) const { + if (need > len) { + throw std::runtime_error("llama_expert_store: invalid read requirement"); + } + size_t total = 0; + while (total < need) { +#if defined(_WIN32) + file->seek(offset + total, SEEK_SET); + file->read_raw(static_cast(dst) + total, len - total); + const size_t n = len - total; +#else + if (offset + total > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("llama_expert_store: file offset exceeds off_t"); + } + const ssize_t result = pread(file->file_id(), static_cast(dst) + total, len - total, + static_cast(offset + total)); + if (result < 0) { + if (errno == EINTR) { + continue; + } + throw std::runtime_error(format("llama_expert_store: pread failed for %s at %llu: %s", + fname.c_str(), (unsigned long long) (offset + total), strerror(errno))); + } + const size_t n = static_cast(result); +#endif + if (n == 0) { + break; + } + total += n; + } + if (total < need) { + throw std::runtime_error(format("llama_expert_store: short read for %s: %zu bytes, need %zu at %llu", + fname.c_str(), total, need, (unsigned long long) offset)); + } + return total; + } +}; + +} + +llama_expert_store_aligned_read llama_expert_store_align_read( + uint64_t offset, size_t size, size_t alignment, uint64_t file_size) { + if (!is_power_of_two(alignment)) { + throw std::runtime_error("llama_expert_store: I/O alignment must be a power of two"); + } + + uint64_t end; + if (!checked_add_u64(offset, size, &end) || end > file_size) { + throw std::runtime_error("llama_expert_store: read is outside the source file"); + } + + const uint64_t aligned_offset = offset & ~static_cast(alignment - 1); + const size_t prefix = static_cast(offset - aligned_offset); + uint64_t needed; + if (!checked_add_u64(prefix, size, &needed)) { + throw std::runtime_error("llama_expert_store: aligned read size overflow"); + } + uint64_t rounded; + if (!checked_add_u64(needed, alignment - 1, &rounded)) { + throw std::runtime_error("llama_expert_store: aligned read size overflow"); + } + rounded &= ~static_cast(alignment - 1); + llama_expert_store_aligned_read result; + result.offset = aligned_offset; + if (rounded > std::numeric_limits::max()) { + throw std::runtime_error("llama_expert_store: aligned read exceeds addressable memory"); + } + result.size = static_cast(rounded); + result.prefix = prefix; + return result; +} + +void llama_expert_store_validate_tensor(const llama_expert_store_tensor & tensor) { + if (tensor.name.empty() || tensor.fname.empty()) { + throw std::runtime_error("llama_expert_store: tensor name and source file are required"); + } + if (tensor.layer < 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s has an invalid layer", tensor.name.c_str())); + } + if (tensor.projection < LLAMA_EXPERT_PROJECTION_GATE || tensor.projection > LLAMA_EXPERT_PROJECTION_DOWN) { + throw std::runtime_error(format("llama_expert_store: tensor %s has an invalid projection", tensor.name.c_str())); + } + const ggml_type expected_type = tensor.projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + if (tensor.type != expected_type) { + throw std::runtime_error(format("llama_expert_store: tensor %s must be %s, got %s", + tensor.name.c_str(), ggml_type_name(expected_type), ggml_type_name(tensor.type))); + } + if (tensor.ne[0] <= 0 || tensor.ne[1] <= 0 || tensor.ne[2] <= 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s has invalid dimensions", tensor.name.c_str())); + } + if (tensor.ne[2] > std::numeric_limits::max()) { + throw std::runtime_error(format("llama_expert_store: tensor %s has too many experts", tensor.name.c_str())); + } + if (tensor.ne[0] % ggml_blck_size(tensor.type) != 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s rows are not whole quantization blocks", tensor.name.c_str())); + } + + const size_t row_size = ggml_row_size(tensor.type, tensor.ne[0]); + uint64_t plane_size; + uint64_t tensor_size; + if (!checked_mul_u64(row_size, static_cast(tensor.ne[1]), &plane_size) || + !checked_mul_u64(plane_size, static_cast(tensor.ne[2]), &tensor_size)) { + throw std::runtime_error(format("llama_expert_store: tensor %s size overflows", tensor.name.c_str())); + } + if (tensor.nb[0] != ggml_type_size(tensor.type) || tensor.nb[1] != row_size || tensor.nb[2] != plane_size) { + throw std::runtime_error(format("llama_expert_store: tensor %s is not a contiguous merged-expert tensor", tensor.name.c_str())); + } + uint64_t tensor_end; + if (!checked_add_u64(tensor.file_offset, tensor_size, &tensor_end) || tensor_end > tensor.file_size) { + throw std::runtime_error(format("llama_expert_store: tensor %s is outside the source file", tensor.name.c_str())); + } +} + +struct llama_expert_store::impl { + struct slot { + bool occupied = false; + expert_key key = {}; + const llama_expert_store_tensor * tensor = nullptr; + aligned_buffer bytes; + uint64_t last_use = 0; + uint32_t pins = 0; + }; + + llama_expert_store_params params; + std::map tensors; + std::map> files; + std::vector slots; + size_t bytes_resident = 0; + uint64_t use_clock = 0; + llama_expert_store_stats counters; + mutable std::mutex mutex; + + impl(std::vector tensors, const llama_expert_store_params & params) : params(params) { + if (params.cache_bytes == 0 || params.cache_slots == 0) { + throw std::runtime_error("llama_expert_store: cache byte and slot budgets must be non-zero"); + } + if (params.cache_slots > std::numeric_limits::max()) { + throw std::runtime_error("llama_expert_store: cache slot budget exceeds the slot ID range"); + } + if (!is_power_of_two(params.io_alignment)) { + throw std::runtime_error("llama_expert_store: I/O alignment must be a power of two"); + } + + for (auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + const tensor_key key = { tensor.layer, tensor.projection }; + if (this->tensors.count(key) != 0) { + throw std::runtime_error(format("llama_expert_store: duplicate tensor for layer %d projection %d", + tensor.layer, static_cast(tensor.projection))); + } + if (tensor.nb[2] > params.cache_bytes) { + throw std::runtime_error(format("llama_expert_store: tensor %s expert plane exceeds the cache byte budget", + tensor.name.c_str())); + } + auto file_it = files.find(tensor.fname); + if (file_it == files.end()) { + file_it = files.emplace(tensor.fname, + std::make_unique(tensor.fname, params.direct_io, params.allow_buffered_io)).first; + } + if (file_it->second->size != tensor.file_size) { + throw std::runtime_error(format("llama_expert_store: source file size changed for %s", tensor.fname.c_str())); + } + this->tensors.emplace(key, std::move(tensor)); + } + + if (this->tensors.empty()) { + throw std::runtime_error("llama_expert_store: no tensors registered"); + } + for (auto it = this->tensors.begin(); it != this->tensors.end();) { + const int32_t layer = it->first.layer; + const auto gate = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_GATE }); + const auto up = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_UP }); + const auto down = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_DOWN }); + if (gate == this->tensors.end() || up == this->tensors.end() || down == this->tensors.end()) { + throw std::runtime_error(format("llama_expert_store: layer %d must register gate, up, and down tensors", layer)); + } + if (gate->second.ne[0] != up->second.ne[0] || + gate->second.ne[1] != up->second.ne[1] || + gate->second.ne[2] != up->second.ne[2] || + down->second.ne[0] != gate->second.ne[1] || + down->second.ne[1] != gate->second.ne[0] || + down->second.ne[2] != gate->second.ne[2]) { + throw std::runtime_error(format("llama_expert_store: layer %d expert tensor dimensions do not match", layer)); + } + it = this->tensors.upper_bound({ layer, LLAMA_EXPERT_PROJECTION_DOWN }); + } + slots.resize(params.cache_slots); + } + + const llama_expert_store_tensor & get_tensor(const expert_key & key) const { + const auto it = tensors.find({ key.layer, key.projection }); + if (it == tensors.end()) { + throw std::runtime_error(format("llama_expert_store: no tensor for layer %d projection %d", + key.layer, static_cast(key.projection))); + } + if (key.expert_id < 0 || key.expert_id >= it->second.ne[2]) { + throw std::runtime_error(format("llama_expert_store: expert ID %d is outside [0, %lld)", + key.expert_id, (long long) it->second.ne[2])); + } + return it->second; + } + + aligned_buffer read_expert(const llama_expert_store_tensor & tensor, int32_t expert_id, uint64_t * bytes_read) const { + uint64_t expert_delta; + uint64_t expert_offset; + if (!checked_mul_u64(static_cast(expert_id), tensor.nb[2], &expert_delta) || + !checked_add_u64(tensor.file_offset, expert_delta, &expert_offset)) { + throw std::runtime_error(format("llama_expert_store: expert offset overflow for %s", tensor.name.c_str())); + } + + aligned_buffer payload(tensor.nb[2], params.io_alignment); + auto & file = *files.at(tensor.fname); + if (file.direct) { + const llama_expert_store_aligned_read read = + llama_expert_store_align_read(expert_offset, tensor.nb[2], params.io_alignment, tensor.file_size); + aligned_buffer bounce(read.size, params.io_alignment); + try { + *bytes_read += file.pread_at_least(bounce.data, read.size, read.offset, read.prefix + tensor.nb[2]); + memcpy(payload.data, bounce.data + read.prefix, tensor.nb[2]); + } catch (const std::runtime_error & e) { + if (!params.allow_buffered_io) { + throw std::runtime_error(format("llama_expert_store: direct I/O failed for %s and buffered fallback is disabled: %s", + tensor.fname.c_str(), e.what())); + } + LLAMA_LOG_WARN("%s: direct read failed for %s; retrying with buffered I/O: %s\n", + __func__, tensor.fname.c_str(), e.what()); + file.reopen(false); + *bytes_read += file.pread_at_least(payload.data, tensor.nb[2], expert_offset, tensor.nb[2]); + } + } else { + *bytes_read += file.pread_at_least(payload.data, tensor.nb[2], expert_offset, tensor.nb[2]); + } + + if (!ggml_validate_row_data(tensor.type, payload.data, tensor.nb[2])) { + throw std::runtime_error(format("llama_expert_store: tensor %s expert %d has invalid payload", + tensor.name.c_str(), expert_id)); + } + return payload; + } + + void unpin(const std::vector & slot_ids) { + std::lock_guard lock(mutex); + for (uint32_t slot_id : slot_ids) { + if (slot_id >= slots.size() || slots[slot_id].pins == 0) { + GGML_ABORT("llama_expert_store: invalid lease slot"); + } + slots[slot_id].pins--; + } + } + + std::vector get_payloads(const std::vector & slot_ids) const { + std::lock_guard lock(mutex); + std::vector result; + result.reserve(slot_ids.size()); + for (uint32_t slot_id : slot_ids) { + const slot & entry = slots.at(slot_id); + GGML_ASSERT(entry.occupied && entry.pins > 0); + result.push_back({ + entry.key.layer, + entry.key.projection, + entry.key.expert_id, + slot_id, + entry.tensor->type, + entry.bytes.data, + entry.bytes.size, + }); + } + return result; + } +}; + +struct llama_expert_store::lease::impl { + std::shared_ptr store; + std::vector pinned_slots; + std::vector> remapped_slots; + + ~impl() { + if (store) { + store->unpin(pinned_slots); + } + } +}; + +llama_expert_store::lease::lease() = default; +llama_expert_store::lease::lease(lease && other) noexcept = default; +llama_expert_store::lease & llama_expert_store::lease::operator=(lease && other) noexcept = default; +llama_expert_store::lease::~lease() = default; + +const std::vector> & llama_expert_store::lease::slot_ids() const { + static const std::vector> empty; + return pimpl ? pimpl->remapped_slots : empty; +} + +std::vector llama_expert_store::lease::payloads() const { + return pimpl ? pimpl->store->get_payloads(pimpl->pinned_slots) : std::vector(); +} + +llama_expert_store::llama_expert_store( + std::vector tensors, const llama_expert_store_params & params) + : pimpl(std::make_shared(std::move(tensors), params)) { +} + +llama_expert_store::~llama_expert_store() = default; + +llama_expert_store::lease llama_expert_store::acquire(const std::vector & requests) { + std::lock_guard lock(pimpl->mutex); + + std::vector> request_keys; + std::vector unique_keys; + request_keys.reserve(requests.size()); + for (const auto & request : requests) { + std::vector keys; + keys.reserve(request.expert_ids.size()); + for (int32_t expert_id : request.expert_ids) { + const expert_key key = { request.layer, request.projection, expert_id }; + pimpl->get_tensor(key); + keys.push_back(key); + unique_keys.push_back(key); + } + request_keys.push_back(std::move(keys)); + } + std::sort(unique_keys.begin(), unique_keys.end()); + unique_keys.erase(std::unique(unique_keys.begin(), unique_keys.end()), unique_keys.end()); + + std::map resident; + for (uint32_t i = 0; i < pimpl->slots.size(); ++i) { + if (pimpl->slots[i].occupied) { + resident.emplace(pimpl->slots[i].key, i); + } + } + + std::vector misses; + std::vector hit_slots; + size_t miss_bytes = 0; + for (const expert_key & key : unique_keys) { + const auto hit = resident.find(key); + if (hit != resident.end()) { + hit_slots.push_back(hit->second); + continue; + } + const auto & tensor = pimpl->get_tensor(key); + if (miss_bytes > pimpl->params.cache_bytes || tensor.nb[2] > pimpl->params.cache_bytes - miss_bytes) { + throw std::runtime_error("llama_expert_store: requested expert union exceeds the cache byte budget"); + } + miss_bytes += tensor.nb[2]; + misses.push_back(key); + } + + std::vector empty_slots; + std::vector candidates; + std::sort(hit_slots.begin(), hit_slots.end()); + for (uint32_t i = 0; i < pimpl->slots.size(); ++i) { + const auto & entry = pimpl->slots[i]; + if (!entry.occupied) { + empty_slots.push_back(i); + } else if (entry.pins == 0 && !std::binary_search(hit_slots.begin(), hit_slots.end(), i)) { + candidates.push_back(i); + } + } + std::sort(candidates.begin(), candidates.end(), [&](uint32_t a, uint32_t b) { + const auto & lhs = pimpl->slots[a]; + const auto & rhs = pimpl->slots[b]; + if (lhs.last_use != rhs.last_use) { + return lhs.last_use < rhs.last_use; + } + return a < b; + }); + + const size_t min_victims = misses.size() > empty_slots.size() ? misses.size() - empty_slots.size() : 0; + std::vector victims; + if (miss_bytes > std::numeric_limits::max() - pimpl->bytes_resident) { + throw std::runtime_error("llama_expert_store: cache byte accounting overflow"); + } + size_t bytes_after = pimpl->bytes_resident + miss_bytes; + for (uint32_t candidate : candidates) { + if (victims.size() >= min_victims && bytes_after <= pimpl->params.cache_bytes) { + break; + } + victims.push_back(candidate); + bytes_after -= pimpl->slots[candidate].bytes.size; + } + if (victims.size() < min_victims || bytes_after > pimpl->params.cache_bytes) { + throw std::runtime_error("llama_expert_store: cache capacity is pinned or too small for the requested expert union"); + } + + std::vector target_slots = empty_slots; + target_slots.insert(target_slots.end(), victims.begin(), victims.end()); + std::sort(target_slots.begin(), target_slots.end()); + + uint64_t bytes_read = 0; + std::vector staged; + staged.reserve(misses.size()); + for (const expert_key & key : misses) { + const auto & tensor = pimpl->get_tensor(key); + staged.push_back(pimpl->read_expert(tensor, key.expert_id, &bytes_read)); + } + + for (uint32_t victim : victims) { + auto & entry = pimpl->slots[victim]; + resident.erase(entry.key); + pimpl->bytes_resident -= entry.bytes.size; + entry.bytes.clear(); + entry.occupied = false; + entry.tensor = nullptr; + entry.last_use = 0; + pimpl->counters.evictions++; + } + + for (size_t i = 0; i < misses.size(); ++i) { + const expert_key & key = misses[i]; + const auto & tensor = pimpl->get_tensor(key); + auto & entry = pimpl->slots[target_slots[i]]; + entry.bytes = std::move(staged[i]); + entry.occupied = true; + entry.key = key; + entry.tensor = &tensor; + entry.pins = 0; + pimpl->bytes_resident += entry.bytes.size; + resident[entry.key] = target_slots[i]; + } + + auto lease_impl = std::make_unique(); + lease_impl->remapped_slots.reserve(request_keys.size()); + for (const auto & keys : request_keys) { + std::vector remapped; + remapped.reserve(keys.size()); + for (const expert_key & key : keys) { + remapped.push_back(resident.at(key)); + } + lease_impl->remapped_slots.push_back(std::move(remapped)); + } + lease_impl->pinned_slots.reserve(unique_keys.size()); + for (const expert_key & key : unique_keys) { + lease_impl->pinned_slots.push_back(resident.at(key)); + } + for (uint32_t slot_id : lease_impl->pinned_slots) { + auto & entry = pimpl->slots[slot_id]; + entry.last_use = ++pimpl->use_clock; + entry.pins++; + } + lease_impl->store = pimpl; + + pimpl->counters.hits += unique_keys.size() - misses.size(); + pimpl->counters.misses += misses.size(); + pimpl->counters.bytes_read += bytes_read; + + lease result; + result.pimpl = std::move(lease_impl); + return result; +} + +llama_expert_store_stats llama_expert_store::stats() const { + std::lock_guard lock(pimpl->mutex); + return pimpl->counters; +} + +size_t llama_expert_store::resident_bytes() const { + std::lock_guard lock(pimpl->mutex); + return pimpl->bytes_resident; +} + +size_t llama_expert_store::resident_entries() const { + std::lock_guard lock(pimpl->mutex); + size_t result = 0; + for (const auto & slot : pimpl->slots) { + result += slot.occupied ? 1 : 0; + } + return result; +} + +bool llama_expert_store::direct_io_active() const { + std::lock_guard lock(pimpl->mutex); + for (const auto & item : pimpl->files) { + if (!item.second->direct) { + return false; + } + } + return true; +} diff --git a/src/llama-expert-store.h b/src/llama-expert-store.h new file mode 100644 index 000000000000..10ce3a1a19d7 --- /dev/null +++ b/src/llama-expert-store.h @@ -0,0 +1,106 @@ +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include +#include + +enum llama_expert_projection { + LLAMA_EXPERT_PROJECTION_GATE = 0, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN, +}; + +struct llama_expert_store_tensor { + std::string name; + std::string fname; + size_t file_index = 0; + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + ggml_type type = GGML_TYPE_COUNT; + int64_t ne[3] = {}; + size_t nb[3] = {}; + uint64_t file_offset = 0; + uint64_t file_size = 0; +}; + +struct llama_expert_store_params { + size_t cache_bytes = 0; + size_t cache_slots = 0; + size_t io_alignment = 4096; + bool direct_io = true; + bool allow_buffered_io = false; // opt-in only; page-cache bytes are outside cache_bytes +}; + +struct llama_expert_store_request { + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + std::vector expert_ids; +}; + +struct llama_expert_store_stats { + uint64_t hits = 0; + uint64_t misses = 0; + uint64_t bytes_read = 0; + uint64_t evictions = 0; +}; + +struct llama_expert_store_aligned_read { + uint64_t offset = 0; + size_t size = 0; + size_t prefix = 0; +}; + +llama_expert_store_aligned_read llama_expert_store_align_read( + uint64_t offset, size_t size, size_t alignment, uint64_t file_size); + +void llama_expert_store_validate_tensor(const llama_expert_store_tensor & tensor); + +struct llama_expert_store { + struct payload { + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + int32_t expert_id = -1; + uint32_t slot_id = 0; + ggml_type type = GGML_TYPE_COUNT; + const uint8_t * data = nullptr; + size_t size = 0; + }; + + struct lease { + lease(); + lease(lease && other) noexcept; + lease & operator=(lease && other) noexcept; + ~lease(); + + lease(const lease &) = delete; + lease & operator=(const lease &) = delete; + + const std::vector> & slot_ids() const; + std::vector payloads() const; + + private: + friend struct llama_expert_store; + + struct impl; + std::unique_ptr pimpl; + }; + + llama_expert_store(std::vector tensors, const llama_expert_store_params & params); + ~llama_expert_store(); + + // The lease pins every unique returned slot. Keep it until the backend upload completes. + lease acquire(const std::vector & requests); + + llama_expert_store_stats stats() const; + size_t resident_bytes() const; + size_t resident_entries() const; + bool direct_io_active() const; + +private: + struct impl; + std::shared_ptr pimpl; +}; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 85ed7f345afb..10549704054f 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1576,15 +1576,17 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s) const { + ggml_tensor * w_s, + ggml_tensor * ids_scale) const { ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids); if (w_s) { + ids_scale = ids_scale ? ids_scale : ids; const int64_t n_expert = w_s->ne[0]; const int64_t n_tokens = cur->ne[2]; ggml_tensor * s = ggml_reshape_3d(ctx0, w_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, ids); + s = ggml_get_rows(ctx0, s, ids_scale); res = ggml_mul(ctx0, res, s); } for (const auto & lora : *loras) { @@ -1861,7 +1863,8 @@ ggml_tensor * llm_graph_context::build_ffn( const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || + (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit); } else { tmp = ggml_clamp(ctx0, tmp, -limit, limit); @@ -1976,6 +1979,26 @@ ggml_tensor * llm_graph_context::build_ffn( return cur; } +llm_moe_expert_ids llm_build_moe_expert_ids( + ggml_context * ctx, + llm_arch arch, + ggml_tensor * selected, + ggml_tensor * lookup, + int64_t n_expert, + int64_t n_expert_total, + uint32_t n_group_experts) { + if (lookup != nullptr) { + return { selected, lookup }; + } + if (arch == LLM_ARCH_GROVEMOE && n_expert != n_expert_total) { + GGML_ASSERT(n_group_experts > 0); + // TODO: Use scalar div instead when/if implemented + ggml_tensor * f_sel = ggml_cast(ctx, selected, GGML_TYPE_F32); + lookup = ggml_cast(ctx, ggml_scale(ctx, f_sel, 1.0f / float(n_group_experts)), GGML_TYPE_I32); + } + return { selected, lookup ? lookup : selected }; +} + ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * cur, ggml_tensor * gate_inp, @@ -1995,7 +2018,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * selected_experts_lookup_in) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -2016,7 +2040,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up_exps_s, gate_exps_s, down_exps_s, - selected_experts_in + selected_experts_in, + selected_experts_lookup_in ); } @@ -2044,7 +2069,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * selected_experts_lookup_in) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN @@ -2141,10 +2167,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } cb(selected_experts, "ffn_moe_topk", il); + const llm_moe_expert_ids expert_ids = llm_build_moe_expert_ids( + ctx0, arch, selected_experts, selected_experts_lookup_in, + n_expert, hparams.n_expert, hparams.n_group_experts); + selected_experts = expert_ids.routing; + ggml_tensor * selected_experts_lookup = expert_ids.lookup; + if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) { - // TODO: Use scalar div instead when/if implemented - ggml_tensor * f_sel = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32); - selected_experts = ggml_cast(ctx0, ggml_scale(ctx0, f_sel, 1.0f / float(hparams.n_group_experts)), GGML_TYPE_I32); probs = ggml_reshape_3d(ctx0, probs, 1, hparams.n_expert, n_tokens); } else { probs = ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens); @@ -2198,7 +2227,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views - ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens] + ggml_tensor * gate_up = build_lora_mm_id( + gate_up_exps, cur, selected_experts_lookup, up_exps_s, selected_experts); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il); if (up_exps_s) { @@ -2217,7 +2247,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(up, "ffn_moe_up", il); } else { // separate gate and up path - up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens] + up = build_lora_mm_id( + up_exps, cur, selected_experts_lookup, up_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il); if (up_exps_s) { @@ -2230,7 +2261,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } if (gate_exps) { - cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens] + cur = build_lora_mm_id( + gate_exps, cur, selected_experts_lookup, gate_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); } else { cur = up; @@ -2255,7 +2287,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || + (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { cur = ggml_swiglu_clamp(ctx0, cur, up, limit); } else { up = ggml_clamp(ctx0, up, -limit, limit); @@ -2331,7 +2364,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( GGML_ABORT("fatal error"); } - experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] + experts = build_lora_mm_id( + down_exps, cur, selected_experts_lookup, down_exps_s, selected_experts); // [n_embd, n_expert_used, n_tokens] cb(experts, "ffn_moe_down", il); if (down_exps_s) { diff --git a/src/llama-graph.h b/src/llama-graph.h index ce68832ce47a..e885bf4a8513 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -75,6 +75,20 @@ enum llm_norm_type { LLM_NORM_GROUP, }; +struct llm_moe_expert_ids { + ggml_tensor * routing; + ggml_tensor * lookup; +}; + +llm_moe_expert_ids llm_build_moe_expert_ids( + ggml_context * ctx, + llm_arch arch, + ggml_tensor * selected, + ggml_tensor * lookup, + int64_t n_expert, + int64_t n_expert_total, + uint32_t n_group_experts); + // TODO: tmp - need something better to pass the data from the encoder to the decoder struct llama_cross { // the output embeddings from the encoder as a ggml tensor @@ -1060,7 +1074,8 @@ struct llm_graph_context { ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s = nullptr) const; + ggml_tensor * w_s = nullptr, + ggml_tensor * ids_scale = nullptr) const; ggml_tensor * build_norm( ggml_tensor * cur, @@ -1129,7 +1144,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * selected_experts_lookup_in = nullptr) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -1155,7 +1171,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * selected_experts_lookup_in = nullptr) const; // // inputs diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 34b3c688019a..b3d49ac482f7 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1,5 +1,6 @@ #include "llama-hparams.h" +#include "llama-dsv41.h" #include "ggml.h" #include @@ -47,6 +48,28 @@ bool llama_hparams::is_swa_any() const { return false; } +int32_t llama_hparams::dsv41_kv_source(uint32_t il) const { + if (il >= n_layer()) { + GGML_ABORT("fatal error"); + } + return dsv41_kv_source_layer[il]; +} + +int32_t llama_hparams::dsv41_index_source(uint32_t il) const { + if (il >= n_layer()) { + GGML_ABORT("fatal error"); + } + return dsv41_index_source_layer[il]; +} + +bool llama_hparams::dsv41_is_kv_source(uint32_t il) const { + return dsv41_kv_source(il) == (int32_t) il; +} + +bool llama_hparams::dsv41_is_index_source(uint32_t il) const { + return dsv41_index_source(il) == (int32_t) il; +} + uint32_t llama_hparams::n_head(uint32_t il) const { if (il < n_layer_all) { return n_head_arr[il]; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 3afa49ebe861..268e28809494 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -294,6 +294,19 @@ struct llama_hparams { float dsv4_hc_eps = 0.0f; std::array dsv4_compress_ratios; + // DeepSeek-V4.1 + uint32_t dsv41_candidate_source_layer = 0; + uint32_t dsv41_candidate_topk_blocks = 0; + uint32_t dsv41_candidate_block_size = 0; + std::array dsv41_kv_source_layer; + std::array dsv41_index_source_layer; + std::bitset dsv41_engram_layers; + + int32_t dsv41_kv_source(uint32_t il) const; + int32_t dsv41_index_source(uint32_t il) const; + bool dsv41_is_kv_source(uint32_t il) const; + bool dsv41_is_index_source(uint32_t il) const; + // 0 = full rank (DeepSeek-V4) uint32_t hc_low_rank = 0; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a5c0970e9bb3..1a284eac0440 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -341,6 +341,7 @@ llama_kv_cache::llama_kv_cache( // always create Hadamard rotation tensors for DeepSeek lightning indexers if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_DEEPSEEK41 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_DOTS3NOTE) && hparams.n_embd_head_k_full == hparams.indexer_head_size) { attn_rot_k = true; diff --git a/src/llama-memory-dsv41.cpp b/src/llama-memory-dsv41.cpp new file mode 100644 index 000000000000..68bcc1161b98 --- /dev/null +++ b/src/llama-memory-dsv41.cpp @@ -0,0 +1,1562 @@ +#include "llama-memory-dsv41.h" + +#include "ggml-backend.h" +#include "ggml-cpp.h" +#include "llama-batch.h" +#include "llama-impl.h" +#include "llama-io.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr uint64_t DSV41_STATE_MAGIC = 0x314d454d31345644ULL; +constexpr uint32_t DSV41_STATE_VERSION = 1; + +uint64_t hash_mix(uint64_t hash, uint64_t value) { + hash ^= value; + return hash*1099511628211ULL; +} + +uint64_t hash_string(uint64_t hash, const char * value) { + while (*value != '\0') { + hash = hash_mix(hash, (uint8_t) *value++); + } + return hash; +} + +uint32_t source_capacity(const llama_dsv41_memory_config & config, uint32_t source) { + const uint32_t ratio = config.ratios.at(source); + return (config.n_ctx + ratio - 1)/ratio; +} + +llama_dsv41_memory_config make_model_config( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch, + std::unique_ptr engram) { + llama_dsv41_memory_config config; + config.n_ctx = n_ctx; + config.n_seq = n_seq; + config.n_ubatch = n_ubatch; + config.n_layer = model.hparams.n_layer(); + config.raw_window = model.hparams.n_swa; + config.kv_width = model.hparams.n_embd_head_k(); + config.index_width = model.hparams.indexer_head_size; + config.candidate_topk_blocks = model.hparams.dsv41_candidate_topk_blocks; + config.candidate_block_size = model.hparams.dsv41_candidate_block_size; + config.candidate_source_layer = model.hparams.dsv41_candidate_source_layer; + config.type_k = type_k; + config.type_index = type_k; + config.no_alloc = model.hparams.no_alloc; + config.attach_no_alloc_buffers = config.no_alloc; + config.engram_enabled = model.hparams.dsv41_engram_layers.any(); + config.expert_enabled = model.requires_synchronous_graph(); + config.ratios.resize(config.n_layer); + config.kv_sources.clear(); + config.index_sources.clear(); + for (uint32_t il = 0; il < config.n_layer; ++il) { + config.ratios[il] = model.hparams.dsv4_compress_ratios[il]; + if (model.hparams.dsv41_is_kv_source(il)) { + config.kv_sources.push_back(il); + } + if (model.hparams.dsv41_is_index_source(il)) { + config.index_sources.push_back(il); + } + } + config.buft_for_layer = [&model, offload](int32_t il) { + return offload ? model.select_buft(il) : ggml_backend_cpu_buffer_type(); + }; + config.engram = std::move(engram); + return config; +} + +} + +bool llama_dsv41_graph_topology::same_topology(const llama_dsv41_graph_topology & other) const { + return n_tokens == other.n_tokens && + n_seqs == other.n_seqs && + n_outputs == other.n_outputs && + visible_raw_widths == other.visible_raw_widths && + visible_compressed_widths == other.visible_compressed_widths && + source_ratios == other.source_ratios && + source_carry_counts == other.source_carry_counts && + candidate_width == other.candidate_width && + backend_layout == other.backend_layout && + engram_enabled == other.engram_enabled && + expert_enabled == other.expert_enabled; +} + +struct llama_memory_dsv41::impl { + struct sequence_state { + llama_pos pos = -1; + std::vector candidates; + }; + + struct tensor_snapshot { + ggml_tensor * tensor = nullptr; + size_t offset = 0; + std::vector data; + }; + + struct rollback_state { + llama_pos start_pos = -1; + sequence_state sequence; + std::vector snapshots; + llama_dsv41_engram_sequence_state engram; + bool has_engram = false; + }; + + struct source_storage { + uint32_t layer = 0; + uint32_t ratio = 0; + uint32_t capacity = 0; + ggml_tensor * kv = nullptr; + ggml_tensor * index = nullptr; + ggml_tensor * carry_kv = nullptr; + ggml_tensor * carry_score = nullptr; + }; + + struct buffer_group { + ggml_backend_buffer_type_t buft = nullptr; + ggml_context_ptr ctx; + ggml_backend_buffer_ptr buffer; + }; + + llama_dsv41_memory_config config; + std::vector sequences; + std::vector raw; + std::map sources; + std::vector groups; + ggml_tensor * candidate_scores = nullptr; + ggml_tensor * candidate_ids = nullptr; + ggml_tensor * committed_candidate_ids = nullptr; + ggml_tensor * positions = nullptr; + uint32_t position_rows = 0; + uint32_t candidate_blocks = 0; + uint32_t candidate_width = 0; + uint64_t generation = 0; + uint64_t backend_layout = 1469598103934665603ULL; + uint64_t graph_workspace = 0; + bool transaction_active = false; + std::map rollback_states; + + explicit impl(llama_dsv41_memory_config config) : config(std::move(config)) { + this->config.engram_enabled = + this->config.engram_enabled || this->config.engram != nullptr; + if (this->config.n_ctx == 0 || this->config.n_seq == 0 || this->config.n_ubatch == 0 || + this->config.n_layer == 0 || this->config.raw_window == 0 || + this->config.kv_width == 0 || this->config.index_width == 0 || + this->config.candidate_block_size == 0 || this->config.candidate_topk_blocks == 0) { + throw std::invalid_argument("DeepSeek V4.1 memory dimensions must be non-zero"); + } + if (this->config.ratios.empty()) { + this->config.ratios.resize(this->config.n_layer); + for (uint32_t il = 0; il < this->config.n_layer; ++il) { + this->config.ratios[il] = llama_dsv41_compress_ratio(il); + } + } + if (this->config.ratios.size() != this->config.n_layer) { + throw std::invalid_argument("DeepSeek V4.1 memory ratio map has the wrong size"); + } + if (this->config.candidate_source_layer >= this->config.n_layer) { + throw std::invalid_argument("DeepSeek V4.1 candidate source layer is out of range"); + } + if (!this->config.buft_for_layer) { + this->config.buft_for_layer = [](int32_t) { return ggml_backend_cpu_buffer_type(); }; + } + for (uint32_t source : this->config.kv_sources) { + if (source >= this->config.n_layer || (this->config.ratios[source] != 1 && this->config.ratios[source] != 2)) { + throw std::invalid_argument("DeepSeek V4.1 memory KV source is invalid"); + } + } + for (uint32_t source : this->config.index_sources) { + if (source >= this->config.n_layer) { + throw std::invalid_argument("DeepSeek V4.1 memory index source is invalid"); + } + } + if (this->config.engram && this->config.engram->max_tokens() < this->config.n_ubatch) { + throw std::invalid_argument("DeepSeek V4.1 Engram runtime is smaller than n_ubatch"); + } + + sequences.resize(this->config.n_seq); + raw.resize(this->config.n_layer); + candidate_blocks = (this->config.n_ctx + this->config.candidate_block_size - 1)/ + this->config.candidate_block_size; + candidate_width = std::min(candidate_blocks, this->config.candidate_topk_blocks); + + std::map tensor_counts; + for (uint32_t il = 0; il < this->config.n_layer; ++il) { + tensor_counts[this->config.buft_for_layer(il)]++; + } + for (uint32_t source : this->config.kv_sources) { + tensor_counts[this->config.buft_for_layer(source)] += this->config.ratios[source] == 2 ? 4 : 2; + } + ggml_backend_buffer_type_t candidate_buft = + this->config.buft_for_layer(this->config.candidate_source_layer); + tensor_counts[candidate_buft] += 3; + tensor_counts[ggml_backend_cpu_buffer_type()]++; + + for (const auto & entry : tensor_counts) { + ggml_init_params params = { + /*.mem_size =*/ (entry.second + 4)*ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + throw std::runtime_error("failed to create DeepSeek V4.1 memory tensor context"); + } + groups.push_back({ entry.first, ggml_context_ptr(ctx), nullptr }); + } + + auto context_for = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { + for (auto & group : groups) { + if (group.buft == buft) { + return group.ctx.get(); + } + } + throw std::runtime_error("DeepSeek V4.1 memory buffer type is missing"); + }; + + for (uint32_t il = 0; il < this->config.n_layer; ++il) { + ggml_context * ctx = context_for(this->config.buft_for_layer(il)); + raw[il] = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, this->config.raw_window, this->config.n_seq); + ggml_format_name(raw[il], "dsv41_raw_k_l%u", il); + } + + for (uint32_t source : this->config.kv_sources) { + source_storage storage; + storage.layer = source; + storage.ratio = this->config.ratios[source]; + storage.capacity = source_capacity(this->config, source); + ggml_context * ctx = context_for(this->config.buft_for_layer(source)); + storage.kv = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, storage.capacity, this->config.n_seq); + storage.index = ggml_new_tensor_3d( + ctx, this->config.type_index, this->config.index_width, storage.capacity, this->config.n_seq); + ggml_format_name(storage.kv, "dsv41_comp_kv_l%u", source); + ggml_format_name(storage.index, "dsv41_index_k_l%u", source); + if (storage.ratio == 2) { + storage.carry_kv = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, storage.ratio, this->config.n_seq); + storage.carry_score = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, storage.ratio, this->config.n_seq); + ggml_format_name(storage.carry_kv, "dsv41_carry_kv_l%u", source); + ggml_format_name(storage.carry_score, "dsv41_carry_score_l%u", source); + } + sources.emplace(source, storage); + } + + { + ggml_context * ctx = context_for(candidate_buft); + candidate_scores = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, candidate_blocks, this->config.n_ubatch); + candidate_ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, candidate_width, this->config.n_ubatch); + committed_candidate_ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, candidate_width, this->config.n_seq); + ggml_set_name(candidate_scores, "dsv41_candidate_scores"); + ggml_set_name(candidate_ids, "dsv41_candidate_ids"); + ggml_set_name(committed_candidate_ids, "dsv41_committed_candidate_ids"); + } + + position_rows = 1 + this->config.raw_window; + for (uint32_t source : this->config.kv_sources) { + position_rows += source_capacity(this->config, source); + if (this->config.ratios[source] == 2) { + position_rows += 2; + } + } + { + ggml_context * ctx = context_for(ggml_backend_cpu_buffer_type()); + positions = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, position_rows, this->config.n_seq); + ggml_set_name(positions, "dsv41_position_state"); + } + + for (auto & group : groups) { + const size_t buffer_size = + ggml_backend_alloc_ctx_tensors_from_buft_size(group.ctx.get(), group.buft); + ggml_backend_buffer_t buffer = nullptr; + if (this->config.no_alloc && this->config.attach_no_alloc_buffers) { + buffer = ggml_backend_buft_alloc_buffer(group.buft, 0); + for (ggml_tensor * tensor = ggml_get_first_tensor(group.ctx.get()); + tensor != nullptr; + tensor = ggml_get_next_tensor(group.ctx.get(), tensor)) { + tensor->buffer = buffer; + } + } else if (!this->config.no_alloc) { + buffer = + ggml_backend_alloc_ctx_tensors_from_buft(group.ctx.get(), group.buft); + } + if (buffer == nullptr && (!this->config.no_alloc || this->config.attach_no_alloc_buffers)) { + throw std::runtime_error("failed to allocate DeepSeek V4.1 memory buffer"); + } + if (buffer != nullptr) { + group.buffer.reset(buffer); + } + if (buffer != nullptr && !this->config.no_alloc) { + ggml_backend_buffer_clear(buffer, 0); + } + backend_layout = hash_string(backend_layout, ggml_backend_buft_name(group.buft)); + backend_layout = hash_mix(backend_layout, buffer_size); + } + for (ggml_tensor * tensor : all_tensors()) { + backend_layout = hash_mix(backend_layout, tensor->type); + backend_layout = hash_mix(backend_layout, tensor->ne[0]); + backend_layout = hash_mix(backend_layout, tensor->ne[1]); + backend_layout = hash_mix(backend_layout, tensor->ne[2]); + } + for (uint32_t seq = 0; seq < this->config.n_seq; ++seq) { + update_position_state(seq); + } + } + + std::vector all_tensors() const { + std::vector result; + result.insert(result.end(), raw.begin(), raw.end()); + for (const auto & entry : sources) { + result.push_back(entry.second.kv); + result.push_back(entry.second.index); + if (entry.second.carry_kv != nullptr) { + result.push_back(entry.second.carry_kv); + result.push_back(entry.second.carry_score); + } + } + result.push_back(candidate_scores); + result.push_back(candidate_ids); + result.push_back(committed_candidate_ids); + result.push_back(positions); + return result; + } + + std::vector sequence_tensors() const { + std::vector result; + result.insert(result.end(), raw.begin(), raw.end()); + for (const auto & entry : sources) { + result.push_back(entry.second.kv); + result.push_back(entry.second.index); + if (entry.second.carry_kv != nullptr) { + result.push_back(entry.second.carry_kv); + result.push_back(entry.second.carry_score); + } + } + result.push_back(committed_candidate_ids); + result.push_back(positions); + return result; + } + + size_t sequence_plane(const ggml_tensor * tensor) const { + if (tensor->ne[2] == (int64_t) config.n_seq) { + return tensor->nb[2]; + } + if (tensor->ne[1] == (int64_t) config.n_seq) { + return tensor->nb[1]; + } + throw std::runtime_error("DeepSeek V4.1 tensor has no sequence plane"); + } + + bool valid_seq(llama_seq_id seq_id) const { + return seq_id >= 0 && (uint32_t) seq_id < config.n_seq; + } + + void require_idle() const { + if (transaction_active) { + throw std::runtime_error("DeepSeek V4.1 memory operation is not allowed during a transaction"); + } + } + + void update_position_state(llama_seq_id seq_id) { + const std::vector values = position_state_values(sequences[seq_id]); + if (!config.no_alloc) { + ggml_backend_tensor_set( + positions, values.data(), (size_t) seq_id*positions->nb[1], values.size()*sizeof(int32_t)); + } + } + + std::vector position_state_values(const sequence_state & sequence) const { + std::vector values(position_rows, -1); + const llama_pos pos = sequence.pos; + values[0] = pos; + if (pos >= 0) { + const llama_pos first_raw = std::max(0, pos + 1 - config.raw_window); + for (llama_pos current = first_raw; current <= pos; ++current) { + values[1 + current%config.raw_window] = current; + } + } + uint32_t offset = 1 + config.raw_window; + for (uint32_t source : config.kv_sources) { + const uint32_t ratio = config.ratios[source]; + const uint32_t capacity = source_capacity(config, source); + const uint32_t visible = pos < 0 ? 0 : (uint32_t) (pos + 1)/ratio; + for (uint32_t row = 0; row < visible; ++row) { + values[offset + row] = row*ratio; + } + offset += capacity; + if (ratio == 2) { + for (uint32_t row = 0; row < 2; ++row) { + if (pos >= (llama_pos) row) { + values[offset + row] = pos - ((pos - row)%2); + } + } + offset += 2; + } + } + return values; + } + + void update_candidate_state(llama_seq_id seq_id) { + const std::vector values = candidate_state_values(sequences[seq_id]); + if (!config.no_alloc) { + ggml_backend_tensor_set( + committed_candidate_ids, + values.data(), + (size_t) seq_id*committed_candidate_ids->nb[1], + values.size()*sizeof(int32_t)); + } + } + + std::vector candidate_state_values(const sequence_state & sequence) const { + std::vector values(candidate_width); + const auto & candidates = sequence.candidates; + std::copy_n(candidates.begin(), std::min(candidates.size(), values.size()), values.begin()); + return values; + } + + void clear_sequence_data(llama_seq_id seq_id) { + if (config.no_alloc) { + return; + } + for (ggml_tensor * tensor : sequence_tensors()) { + const size_t plane = sequence_plane(tensor); + std::vector zeros(plane); + ggml_backend_tensor_set(tensor, zeros.data(), (size_t) seq_id*plane, plane); + } + } + + void copy_sequence_data(llama_seq_id src, llama_seq_id dst) { + if (config.no_alloc || src == dst) { + return; + } + for (ggml_tensor * tensor : sequence_tensors()) { + const size_t plane = sequence_plane(tensor); + std::vector data(plane); + ggml_backend_tensor_get(tensor, data.data(), (size_t) src*plane, plane); + ggml_backend_tensor_set(tensor, data.data(), (size_t) dst*plane, plane); + } + } +}; + +struct llama_memory_dsv41_context::transaction_state { + llama_dsv41_memory_plan plan; + std::vector next_sequences; + std::vector snapshots; + std::unique_ptr engram; + std::map engram_before; + std::vector seq_ids; + std::vector start_positions; +}; + +llama_memory_dsv41::llama_memory_dsv41(llama_dsv41_memory_config config) : + pimpl(std::make_unique(std::move(config))) { +} + +llama_memory_dsv41::llama_memory_dsv41( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch, + std::unique_ptr engram) : + llama_memory_dsv41(make_model_config( + model, type_k, offload, n_ctx, n_seq, n_ubatch, std::move(engram))) { +} + +llama_memory_dsv41::~llama_memory_dsv41() = default; + +llama_memory_context_ptr llama_memory_dsv41::init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) { + GGML_UNUSED(embd_all); + if (n_ubatch == 0 || n_ubatch > pimpl->config.n_ubatch || pimpl->transaction_active) { + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + balloc.split_reset(); + std::vector ubatches; + while (true) { + llama_ubatch ubatch = balloc.split_seq(n_ubatch); + if (ubatch.n_tokens == 0) { + break; + } + ubatches.push_back(std::move(ubatch)); + } + if (balloc.get_n_used() != balloc.get_n_tokens() || ubatches.empty()) { + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + return std::make_unique(this, std::move(ubatches)); +} + +llama_memory_context_ptr llama_memory_dsv41::init_full() { + return std::make_unique(this, true); +} + +llama_memory_context_ptr llama_memory_dsv41::init_update(llama_context * lctx, bool optimize) { + GGML_UNUSED(lctx); + GGML_UNUSED(optimize); + return std::make_unique(LLAMA_MEMORY_STATUS_NO_UPDATE); +} + +bool llama_memory_dsv41::get_can_shift() const { + return false; +} + +void llama_memory_dsv41::clear(bool data) { + pimpl->require_idle(); + for (auto & sequence : pimpl->sequences) { + sequence = {}; + } + pimpl->rollback_states.clear(); + if (pimpl->config.engram) { + for (uint32_t seq = 0; seq < pimpl->config.n_seq; ++seq) { + pimpl->config.engram->seq_remove(seq); + } + } + if (data && !pimpl->config.no_alloc) { + for (auto & group : pimpl->groups) { + ggml_backend_buffer_clear(group.buffer.get(), 0); + } + } + for (uint32_t seq = 0; seq < pimpl->config.n_seq; ++seq) { + pimpl->update_position_state(seq); + pimpl->update_candidate_state(seq); + } + ++pimpl->generation; +} + +bool llama_memory_dsv41::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + try { + pimpl->require_idle(); + } catch (const std::exception & error) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, error.what()); + return false; + } + if (seq_id < 0) { + if (p0 == -1 && p1 == -1) { + clear(true); + return true; + } + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 only supports wildcard removal for the full memory\n", __func__); + return false; + } + if (!pimpl->valid_seq(seq_id)) { + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 sequence ID %d is out of range\n", __func__, seq_id); + return false; + } + + const llama_pos max_pos = pimpl->sequences[seq_id].pos; + const llama_pos begin = p0 < 0 ? 0 : p0; + const llama_pos end = p1 < 0 ? std::numeric_limits::max() : p1; + if (begin >= end || begin > max_pos) { + return true; + } + if (end <= max_pos) { + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 only supports full removal or suffix rollback\n", __func__); + return false; + } + + auto rollback = pimpl->rollback_states.end(); + if (begin > 0) { + rollback = pimpl->rollback_states.find(seq_id); + if (rollback == pimpl->rollback_states.end() || rollback->second.start_pos != begin) { + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 only supports rollback of the most recent committed ubatch\n", __func__); + return false; + } + } + + if (begin == 0) { + pimpl->sequences[seq_id] = {}; + pimpl->clear_sequence_data(seq_id); + pimpl->update_position_state(seq_id); + pimpl->update_candidate_state(seq_id); + if (pimpl->config.engram) { + pimpl->config.engram->seq_remove(seq_id); + } + pimpl->rollback_states.erase(seq_id); + } else { + for (const auto & snapshot : rollback->second.snapshots) { + ggml_backend_tensor_set( + snapshot.tensor, snapshot.data.data(), + snapshot.offset, snapshot.data.size()); + } + pimpl->sequences[seq_id] = rollback->second.sequence; + if (rollback->second.has_engram) { + llama_dsv41_engram_snapshot snapshot = pimpl->config.engram->checkpoint(); + snapshot.sequences[seq_id] = rollback->second.engram; + pimpl->config.engram->restore(snapshot); + } + pimpl->update_position_state(seq_id); + pimpl->update_candidate_state(seq_id); + pimpl->rollback_states.erase(rollback); + } + ++pimpl->generation; + return true; +} + +void llama_memory_dsv41::seq_cp( + llama_seq_id seq_id_src, + llama_seq_id seq_id_dst, + llama_pos p0, + llama_pos p1) { + pimpl->require_idle(); + if (!pimpl->valid_seq(seq_id_src) || !pimpl->valid_seq(seq_id_dst)) { + throw std::invalid_argument("DeepSeek V4.1 sequence copy ID is out of range"); + } + const llama_pos max_pos = pimpl->sequences[seq_id_src].pos; + if ((p0 > 0) || (p1 >= 0 && p1 <= max_pos)) { + throw std::invalid_argument("DeepSeek V4.1 only supports full sequence copy"); + } + pimpl->copy_sequence_data(seq_id_src, seq_id_dst); + pimpl->sequences[seq_id_dst] = pimpl->sequences[seq_id_src]; + pimpl->rollback_states.erase(seq_id_dst); + if (pimpl->config.engram) { + pimpl->config.engram->seq_copy(seq_id_src, seq_id_dst); + } + ++pimpl->generation; +} + +void llama_memory_dsv41::seq_keep(llama_seq_id seq_id) { + pimpl->require_idle(); + if (!pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 sequence keep ID is out of range"); + } + for (uint32_t current = 0; current < pimpl->config.n_seq; ++current) { + if ((llama_seq_id) current == seq_id) { + continue; + } + pimpl->sequences[current] = {}; + pimpl->clear_sequence_data(current); + pimpl->update_position_state(current); + pimpl->update_candidate_state(current); + if (pimpl->config.engram) { + pimpl->config.engram->seq_remove(current); + } + pimpl->rollback_states.erase(current); + } + ++pimpl->generation; +} + +[[noreturn]] void llama_memory_dsv41::seq_add( + llama_seq_id, + llama_pos, + llama_pos, + llama_pos) { + throw std::invalid_argument("DeepSeek V4.1 memory does not support position shifts"); +} + +[[noreturn]] void llama_memory_dsv41::seq_div( + llama_seq_id, + llama_pos, + llama_pos, + int) { + throw std::invalid_argument("DeepSeek V4.1 memory does not support position division"); +} + +llama_pos llama_memory_dsv41::seq_pos_min(llama_seq_id seq_id) const { + if (!pimpl->valid_seq(seq_id)) { + return -1; + } + return pimpl->sequences[seq_id].pos < 0 ? -1 : 0; +} + +llama_pos llama_memory_dsv41::seq_pos_max(llama_seq_id seq_id) const { + if (!pimpl->valid_seq(seq_id)) { + return -1; + } + return pimpl->sequences[seq_id].pos; +} + +std::map llama_memory_dsv41::memory_breakdown() const { + std::map result; + for (const auto & group : pimpl->groups) { + const size_t size = pimpl->config.no_alloc ? + ggml_backend_alloc_ctx_tensors_from_buft_size(group.ctx.get(), group.buft) : + ggml_backend_buffer_get_size(group.buffer.get()); + result[group.buft] += size; + } + return result; +} + +void llama_memory_dsv41::set_graph_workspace_size(size_t size) { + pimpl->graph_workspace = size; +} + +void llama_memory_dsv41::state_write( + llama_io_write_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags) const { + GGML_UNUSED(flags); + pimpl->require_idle(); + if (seq_id != -1 && !pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 state sequence ID is out of range"); + } + + io.write(&DSV41_STATE_MAGIC, sizeof(DSV41_STATE_MAGIC)); + io.write(&DSV41_STATE_VERSION, sizeof(DSV41_STATE_VERSION)); + const uint32_t count = seq_id == -1 ? pimpl->config.n_seq : 1; + io.write(&count, sizeof(count)); + const auto tensors = pimpl->sequence_tensors(); + const uint32_t tensor_count = tensors.size(); + io.write(&tensor_count, sizeof(tensor_count)); + + const llama_dsv41_engram_snapshot engram_snapshot = + pimpl->config.engram ? pimpl->config.engram->checkpoint() : llama_dsv41_engram_snapshot {}; + for (uint32_t i = 0; i < count; ++i) { + const llama_seq_id current = seq_id == -1 ? (llama_seq_id) i : seq_id; + io.write(¤t, sizeof(current)); + const auto & state = pimpl->sequences[current]; + io.write(&state.pos, sizeof(state.pos)); + const uint32_t n_candidates = state.candidates.size(); + io.write(&n_candidates, sizeof(n_candidates)); + if (n_candidates > 0) { + io.write(state.candidates.data(), n_candidates*sizeof(int32_t)); + } + const uint8_t has_engram = pimpl->config.engram ? 1 : 0; + io.write(&has_engram, sizeof(has_engram)); + if (has_engram) { + const auto found = engram_snapshot.sequences.find(current); + llama_dsv41_engram_sequence_state engram_state; + engram_state.history.reset(); + if (found != engram_snapshot.sequences.end()) { + engram_state = found->second; + } + io.write(&engram_state.pos, sizeof(engram_state.pos)); + io.write(engram_state.history.tail.data(), sizeof(engram_state.history.tail)); + } + for (ggml_tensor * tensor : tensors) { + const uint64_t plane = pimpl->sequence_plane(tensor); + io.write(&plane, sizeof(plane)); + io.write_tensor(tensor, (size_t) current*plane, plane); + } + } +} + +void llama_memory_dsv41::state_read( + llama_io_read_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags) { + GGML_UNUSED(flags); + pimpl->require_idle(); + if (seq_id != -1 && !pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 state destination sequence ID is out of range"); + } + + uint64_t magic = 0; + uint32_t version = 0; + uint32_t count = 0; + uint32_t tensor_count = 0; + io.read(&magic, sizeof(magic)); + io.read(&version, sizeof(version)); + io.read(&count, sizeof(count)); + io.read(&tensor_count, sizeof(tensor_count)); + const auto tensors = pimpl->sequence_tensors(); + if (magic != DSV41_STATE_MAGIC || version != DSV41_STATE_VERSION || + tensor_count != tensors.size() || count == 0 || + (seq_id != -1 && count != 1) || (seq_id == -1 && count != pimpl->config.n_seq)) { + throw std::runtime_error("DeepSeek V4.1 state header is incompatible"); + } + + const auto sequence_is_empty = [&](llama_seq_id current) { + if (pimpl->sequences[current].pos >= 0 || + !pimpl->sequences[current].candidates.empty() || + pimpl->rollback_states.count(current) != 0) { + return false; + } + return !pimpl->config.engram || pimpl->config.engram->sequence(current).pos < 0; + }; + if (seq_id == -1) { + for (uint32_t current = 0; current < pimpl->config.n_seq; ++current) { + if (!sequence_is_empty(current)) { + throw std::runtime_error("DeepSeek V4.1 full state restore requires empty memory"); + } + } + } else if (!sequence_is_empty(seq_id)) { + throw std::runtime_error("DeepSeek V4.1 sequence state restore requires an empty destination"); + } + + llama_dsv41_engram_snapshot engram_snapshot = + pimpl->config.engram && seq_id != -1 ? + pimpl->config.engram->checkpoint() : llama_dsv41_engram_snapshot {}; + std::set restored; + std::vector tensor_data; + struct staged_sequence { + llama_seq_id target = -1; + llama_pos pos = -1; + std::vector candidates; + std::vector planes; + }; + std::vector staged; + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + staged.reserve(count); + } + + try { + for (uint32_t i = 0; i < count; ++i) { + llama_seq_id stored = -1; + llama_pos pos = -1; + uint32_t n_candidates = 0; + uint8_t has_engram = 0; + io.read(&stored, sizeof(stored)); + io.read(&pos, sizeof(pos)); + io.read(&n_candidates, sizeof(n_candidates)); + const llama_seq_id target = seq_id == -1 ? stored : seq_id; + if (!pimpl->valid_seq(stored) || !pimpl->valid_seq(target) || + (seq_id == -1 && !restored.insert(target).second) || + pos < -1 || pos >= (llama_pos) pimpl->config.n_ctx || + n_candidates > pimpl->candidate_width) { + throw std::runtime_error("DeepSeek V4.1 state metadata is invalid"); + } + std::vector candidates(n_candidates); + if (n_candidates > 0) { + io.read(candidates.data(), n_candidates*sizeof(int32_t)); + } + io.read(&has_engram, sizeof(has_engram)); + if ((has_engram != 0) != (pimpl->config.engram != nullptr)) { + throw std::runtime_error("DeepSeek V4.1 Engram state availability differs"); + } + if (has_engram) { + llama_dsv41_engram_sequence_state engram_state; + io.read(&engram_state.pos, sizeof(engram_state.pos)); + io.read(engram_state.history.tail.data(), sizeof(engram_state.history.tail)); + engram_snapshot.sequences[target] = engram_state; + } + std::vector planes; + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + planes.reserve(tensors.size()); + } + for (ggml_tensor * tensor : tensors) { + uint64_t plane = 0; + io.read(&plane, sizeof(plane)); + if (plane != pimpl->sequence_plane(tensor)) { + throw std::runtime_error("DeepSeek V4.1 state tensor layout differs"); + } + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + planes.push_back(plane); + } else { + tensor_data.resize(plane); + io.read(tensor_data.data(), plane); + ggml_backend_tensor_set(tensor, tensor_data.data(), (size_t) target*plane, plane); + } + } + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + staged.push_back({ target, pos, std::move(candidates), std::move(planes) }); + } else { + pimpl->sequences[target].pos = pos; + pimpl->sequences[target].candidates = std::move(candidates); + pimpl->update_position_state(target); + pimpl->update_candidate_state(target); + } + } + if (pimpl->config.engram) { + pimpl->config.engram->restore(engram_snapshot); + } + for (auto & sequence : staged) { + for (size_t i = 0; i < tensors.size(); ++i) { + io.read_tensor( + tensors[i], + (size_t) sequence.target*sequence.planes[i], + sequence.planes[i]); + } + pimpl->sequences[sequence.target].pos = sequence.pos; + pimpl->sequences[sequence.target].candidates = std::move(sequence.candidates); + pimpl->update_position_state(sequence.target); + pimpl->update_candidate_state(sequence.target); + } + if (seq_id == -1) { + pimpl->rollback_states.clear(); + } else { + pimpl->rollback_states.erase(seq_id); + } + ++pimpl->generation; + } catch (...) { + if (seq_id == -1) { + clear(true); + } else { + seq_rm(seq_id, -1, -1); + } + throw; + } +} + +ggml_tensor * llama_memory_dsv41::raw_k(uint32_t layer) const { + return layer < pimpl->raw.size() ? pimpl->raw[layer] : nullptr; +} + +ggml_tensor * llama_memory_dsv41::compressed_kv(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.kv; +} + +ggml_tensor * llama_memory_dsv41::index_keys(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.index; +} + +ggml_tensor * llama_memory_dsv41::compressor_carry_kv(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.carry_kv; +} + +ggml_tensor * llama_memory_dsv41::compressor_carry_score(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.carry_score; +} + +ggml_tensor * llama_memory_dsv41::candidate_scores() const { + return pimpl->candidate_scores; +} + +ggml_tensor * llama_memory_dsv41::candidate_ids() const { + return pimpl->candidate_ids; +} + +ggml_tensor * llama_memory_dsv41::committed_candidate_ids() const { + return pimpl->committed_candidate_ids; +} + +ggml_tensor * llama_memory_dsv41::position_state() const { + return pimpl->positions; +} + +const llama_dsv41_memory_config & llama_memory_dsv41::config() const { + return pimpl->config; +} + +llama_dsv41_memory_accounting llama_memory_dsv41::accounting() const { + llama_dsv41_memory_accounting result; + for (ggml_tensor * tensor : pimpl->raw) { + result.raw_kv += ggml_nbytes(tensor); + } + for (const auto & entry : pimpl->sources) { + result.compressed_kv += ggml_nbytes(entry.second.kv); + result.index_keys += ggml_nbytes(entry.second.index); + if (entry.second.carry_kv != nullptr) { + result.compressor_carry += ggml_nbytes(entry.second.carry_kv); + result.compressor_carry += ggml_nbytes(entry.second.carry_score); + } + } + result.candidate_scores = ggml_nbytes(pimpl->candidate_scores); + result.candidate_ids = ggml_nbytes(pimpl->candidate_ids) + ggml_nbytes(pimpl->committed_candidate_ids); + result.position_state = ggml_nbytes(pimpl->positions); + result.graph_workspace = pimpl->graph_workspace; + return result; +} + +std::vector llama_memory_dsv41::sequence_candidate_ids(llama_seq_id seq_id) const { + if (!pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 candidate sequence ID is out of range"); + } + return pimpl->sequences[seq_id].candidates; +} + +size_t llama_memory_dsv41::retained_rollback_count() const { + return pimpl->rollback_states.size(); +} + +bool llama_memory_dsv41::engram_enabled() const { + return pimpl->config.engram_enabled; +} + +llama_memory_dsv41_context::llama_memory_dsv41_context(llama_memory_status status) : + status(status) { +} + +llama_memory_dsv41_context::llama_memory_dsv41_context( + llama_memory_dsv41 * memory, + bool full) : + status(LLAMA_MEMORY_STATUS_SUCCESS), + mem(memory), + full(full) { + if (!memory || !full) { + return; + } + + llama_batch_allocr allocator(1); + llama_ubatch ubatch = allocator.ubatch_reserve(memory->pimpl->config.n_ubatch, 1); + const uint32_t n_tokens = ubatch.n_tokens; + const uint32_t n_streams = 1; + ubatch.data->seq_id_data.resize((size_t) n_tokens*n_streams); + ubatch.data->seq_id_unq.resize(n_streams); + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.n_seqs_unq = n_streams; + for (uint32_t seq = 0; seq < n_streams; ++seq) { + ubatch.data->seq_id_unq[seq] = seq; + ubatch.seq_idx[seq] = seq; + } + for (uint32_t token = 0; token < n_tokens; ++token) { + ubatch.token[token] = 0; + ubatch.pos[token] = token; + ubatch.n_seq_id[token] = n_streams; + ubatch.seq_id[token] = ubatch.data->seq_id_data.data() + (size_t) token*n_streams; + for (uint32_t seq = 0; seq < n_streams; ++seq) { + ubatch.seq_id[token][seq] = seq; + } + } + ubatches.push_back(std::move(ubatch)); + + transaction = std::make_unique(); + transaction->plan.generation = memory->pimpl->generation; + transaction->next_sequences = memory->pimpl->sequences; + for (uint32_t seq = 0; seq < n_streams; ++seq) { + transaction->seq_ids.push_back(seq); + transaction->start_positions.push_back(0); + transaction->next_sequences[seq].pos = n_tokens - 1; + } + const uint32_t raw_width = memory->pimpl->config.raw_window; + const uint32_t persist_first = n_tokens > raw_width ? n_tokens - raw_width : 0; + for (uint32_t token = 0; token < n_tokens; ++token) { + std::vector token_sequences(n_streams); + for (uint32_t seq = 0; seq < n_streams; ++seq) { + token_sequences[seq] = seq; + if (token >= persist_first) { + transaction->plan.raw.persist_src_idxs.push_back(token); + transaction->plan.raw.write_idxs.push_back( + (int64_t) seq*raw_width + token%raw_width); + } + } + transaction->plan.token_seq_ids.push_back(std::move(token_sequences)); + transaction->plan.positions.push_back(token); + const uint32_t visible = std::min(raw_width, token + 1); + transaction->plan.raw.n_visible.push_back(visible); + const auto order = llama_dsv41_raw_ring_order(token, raw_width); + for (uint32_t row = 0; row < raw_width; ++row) { + transaction->plan.raw.read_idxs.push_back(row < visible ? order[row] : -1); + transaction->plan.raw.mask.push_back( + row < visible ? 0.0f : -std::numeric_limits::infinity()); + } + } + for (uint32_t source : memory->pimpl->config.kv_sources) { + llama_dsv41_source_plan source_plan; + source_plan.source_layer = source; + source_plan.ratio = memory->pimpl->config.ratios[source]; + source_plan.capacity = source_capacity(memory->pimpl->config, source); + source_plan.compression = llama_dsv41_build_compression_plan( + transaction->plan.positions, source_plan.ratio, source_plan.capacity); + const uint32_t max_visible = source_plan.compression.n_visible.empty() ? 0 : + *std::max_element( + source_plan.compression.n_visible.begin(), + source_plan.compression.n_visible.end()); + for (int32_t visible : source_plan.compression.n_visible) { + for (uint32_t row = 0; row < max_visible; ++row) { + source_plan.read_idxs.push_back(row < (uint32_t) visible ? row : -1); + } + } + transaction->plan.sources.push_back(std::move(source_plan)); + } + const uint32_t visible_blocks = + (n_tokens + memory->pimpl->config.candidate_block_size - 1)/ + memory->pimpl->config.candidate_block_size; + transaction->plan.candidate_width = std::min( + visible_blocks, memory->pimpl->config.candidate_topk_blocks); +} + +llama_memory_dsv41_context::llama_memory_dsv41_context( + llama_memory_dsv41 * memory, + std::vector ubatches) : + status(LLAMA_MEMORY_STATUS_SUCCESS), + mem(memory), + ubatches(std::move(ubatches)) { +} + +llama_memory_dsv41_context::~llama_memory_dsv41_context() { + rollback(); +} + +bool llama_memory_dsv41_context::next() { + if (status != LLAMA_MEMORY_STATUS_SUCCESS || transaction) { + return false; + } + if (++i_next >= ubatches.size()) { + return false; + } + return true; +} + +bool llama_memory_dsv41_context::apply() { + if (status == LLAMA_MEMORY_STATUS_NO_UPDATE) { + return true; + } + if (status != LLAMA_MEMORY_STATUS_SUCCESS || full || mem == nullptr || transaction || + i_next >= ubatches.size() || mem->pimpl->transaction_active) { + status = LLAMA_MEMORY_STATUS_FAILED_PREPARE; + return false; + } + + const llama_ubatch & ubatch = ubatches[i_next]; + auto next = std::make_unique(); + try { + if (ubatch.n_tokens == 0 || ubatch.n_tokens > mem->pimpl->config.n_ubatch || + ubatch.n_pos != 1 || ubatch.pos == nullptr || ubatch.n_seq_id == nullptr || + ubatch.seq_id == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 ubatch shape is invalid"); + } + + next->next_sequences = mem->pimpl->sequences; + std::vector topology; + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + if (ubatch.n_seq_id[token] <= 0 || ubatch.seq_id[token] == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 token has no sequence"); + } + std::vector token_sequences( + ubatch.seq_id[token], ubatch.seq_id[token] + ubatch.n_seq_id[token]); + std::sort(token_sequences.begin(), token_sequences.end()); + if (std::adjacent_find(token_sequences.begin(), token_sequences.end()) != token_sequences.end()) { + throw std::invalid_argument("DeepSeek V4.1 token repeats a sequence ID"); + } + if (token == 0) { + topology = token_sequences; + } else if (topology != token_sequences) { + throw std::invalid_argument("DeepSeek V4.1 ubatch changes sequence topology"); + } + next->plan.token_seq_ids.push_back(token_sequences); + next->plan.positions.push_back(ubatch.pos[token]); + } + if (topology.empty()) { + throw std::invalid_argument("DeepSeek V4.1 ubatch has no sequence topology"); + } + if (topology.size() != 1 || ubatch.n_seqs != 1) { + throw std::invalid_argument("DeepSeek V4.1 memory supports one sequence per ubatch"); + } + next->seq_ids = topology; + for (llama_seq_id seq_id : topology) { + if (!mem->pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 sequence ID is out of range"); + } + next->start_positions.push_back(mem->pimpl->sequences[seq_id].pos + 1); + } + if (!std::all_of(next->start_positions.begin(), next->start_positions.end(), + [&](llama_pos pos) { return pos == next->start_positions.front(); })) { + throw std::invalid_argument("DeepSeek V4.1 coupled sequences have different positions"); + } + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + const llama_pos expected = next->start_positions.front() + token; + if (ubatch.pos[token] != expected || expected < 0 || expected >= (llama_pos) mem->pimpl->config.n_ctx) { + throw std::invalid_argument("DeepSeek V4.1 positions must be contiguous and within n_ctx"); + } + } + + next->plan.generation = mem->pimpl->generation + 1; + const uint32_t raw_width = mem->pimpl->config.raw_window; + const uint32_t persist_first = + ubatch.n_tokens > raw_width ? ubatch.n_tokens - raw_width : 0; + next->plan.raw.n_visible.resize(ubatch.n_tokens); + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + const llama_pos pos = ubatch.pos[token]; + for (llama_seq_id seq_id : topology) { + if (token >= persist_first) { + next->plan.raw.persist_src_idxs.push_back(token); + next->plan.raw.write_idxs.push_back( + (int64_t) seq_id*raw_width + pos%raw_width); + } + next->next_sequences[seq_id].pos = pos; + } + const uint32_t visible = std::min(raw_width, pos + 1); + next->plan.raw.n_visible[token] = visible; + const auto order = llama_dsv41_raw_ring_order(pos, raw_width); + for (uint32_t row = 0; row < raw_width; ++row) { + next->plan.raw.read_idxs.push_back( + row < visible ? (int32_t) ((uint32_t) topology.front()*raw_width + order[row]) : -1); + next->plan.raw.mask.push_back(row < visible ? 0.0f : -std::numeric_limits::infinity()); + } + } + + for (uint32_t source : mem->pimpl->config.kv_sources) { + llama_dsv41_source_plan source_plan; + source_plan.source_layer = source; + source_plan.ratio = mem->pimpl->config.ratios[source]; + source_plan.capacity = source_capacity(mem->pimpl->config, source); + source_plan.compression = llama_dsv41_build_compression_plan( + next->plan.positions, source_plan.ratio, source_plan.capacity); + const uint32_t max_visible = source_plan.compression.n_visible.empty() ? 0 : + *std::max_element(source_plan.compression.n_visible.begin(), source_plan.compression.n_visible.end()); + for (int32_t visible : source_plan.compression.n_visible) { + for (uint32_t row = 0; row < max_visible; ++row) { + source_plan.read_idxs.push_back( + row < (uint32_t) visible ? + (int32_t) ((uint32_t) topology.front()*source_plan.capacity + row) : -1); + } + } + next->plan.sources.push_back(std::move(source_plan)); + } + const llama_pos final_pos = next->plan.positions.back(); + const uint32_t candidate_visible = final_pos + 1; + const uint32_t visible_blocks = + (candidate_visible + mem->pimpl->config.candidate_block_size - 1)/ + mem->pimpl->config.candidate_block_size; + next->plan.candidate_width = std::min( + visible_blocks, mem->pimpl->config.candidate_topk_blocks); + + if (mem->pimpl->config.engram) { + if (ubatch.token == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram requires token inputs"); + } + std::vector tokens; + tokens.reserve(ubatch.n_tokens); + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + tokens.push_back({ + ubatch.token[token], + ubatch.pos[token], + next->plan.token_seq_ids[token], + 1, + }); + } + for (llama_seq_id seq_id : topology) { + next->engram_before[seq_id] = mem->pimpl->config.engram->sequence(seq_id); + } + next->engram = std::make_unique( + mem->pimpl->config.engram->prepare(tokens)); + } + + std::set> snapshot_keys; + auto snapshot = [&](ggml_tensor * tensor, size_t offset, size_t size) { + if (tensor == nullptr || mem->pimpl->config.no_alloc || + !snapshot_keys.emplace(tensor, offset, size).second) { + return; + } + llama_memory_dsv41::impl::tensor_snapshot entry; + entry.tensor = tensor; + entry.offset = offset; + entry.data.resize(size); + ggml_backend_tensor_get(tensor, entry.data.data(), offset, size); + next->snapshots.push_back(std::move(entry)); + }; + for (llama_seq_id seq_id : topology) { + for (ggml_tensor * raw : mem->pimpl->raw) { + const size_t row_bytes = raw->nb[1]; + for (llama_pos pos : next->plan.positions) { + const size_t row = (size_t) seq_id*mem->pimpl->config.raw_window + + (uint32_t) pos%mem->pimpl->config.raw_window; + snapshot(raw, row*row_bytes, row_bytes); + } + } + for (const auto & source_plan : next->plan.sources) { + const auto & storage = mem->pimpl->sources.at(source_plan.source_layer); + for (int64_t row : source_plan.compression.write_idxs) { + const size_t physical = (size_t) seq_id*storage.capacity + row; + snapshot(storage.kv, physical*storage.kv->nb[1], storage.kv->nb[1]); + snapshot(storage.index, physical*storage.index->nb[1], storage.index->nb[1]); + } + if (storage.carry_kv != nullptr) { + snapshot( + storage.carry_kv, + (size_t) seq_id*storage.carry_kv->nb[2], + storage.carry_kv->nb[2]); + snapshot( + storage.carry_score, + (size_t) seq_id*storage.carry_score->nb[2], + storage.carry_score->nb[2]); + } + } + } + + mem->pimpl->transaction_active = true; + transaction = std::move(next); + return true; + } catch (const std::exception & error) { + if (next->engram) { + mem->pimpl->config.engram->rollback(*next->engram); + } + LLAMA_LOG_ERROR("%s: %s\n", __func__, error.what()); + status = LLAMA_MEMORY_STATUS_FAILED_PREPARE; + return false; + } +} + +void llama_memory_dsv41_context::commit() { + if (!transaction || mem == nullptr) { + return; + } + GGML_ASSERT(transaction->seq_ids.size() == 1); + const llama_seq_id seq_id = transaction->seq_ids.front(); + llama_memory_dsv41::impl::rollback_state rollback; + rollback.start_pos = transaction->start_positions.front(); + rollback.sequence = mem->pimpl->sequences[seq_id]; + if (transaction->engram) { + rollback.engram = transaction->engram_before.at(seq_id); + rollback.has_engram = true; + } + + std::vector committed_candidates; + if (!mem->pimpl->config.no_alloc && transaction->plan.candidate_width > 0) { + const size_t row_bytes = mem->pimpl->candidate_ids->nb[1]; + std::vector ids(transaction->plan.candidate_width); + uint32_t token = 0; + for (uint32_t current = 0; current < transaction->plan.token_seq_ids.size(); ++current) { + if (std::find( + transaction->plan.token_seq_ids[current].begin(), + transaction->plan.token_seq_ids[current].end(), + seq_id) != transaction->plan.token_seq_ids[current].end()) { + token = current; + } + } + ggml_backend_tensor_get( + mem->pimpl->candidate_ids, + ids.data(), + (size_t) token*row_bytes, + ids.size()*sizeof(int32_t)); + transaction->next_sequences[seq_id].candidates = std::move(ids); + } + + const std::vector position_values = + mem->pimpl->position_state_values(transaction->next_sequences[seq_id]); + committed_candidates = + mem->pimpl->candidate_state_values(transaction->next_sequences[seq_id]); + + auto rollback_slot = mem->pimpl->rollback_states.end(); + bool inserted_rollback = false; + try { + std::tie(rollback_slot, inserted_rollback) = + mem->pimpl->rollback_states.emplace(seq_id, llama_memory_dsv41::impl::rollback_state {}); + if (transaction->engram) { + mem->pimpl->config.engram->commit(*transaction->engram); + } + } catch (...) { + if (inserted_rollback) { + mem->pimpl->rollback_states.erase(rollback_slot); + } + throw; + } + + if (!mem->pimpl->config.no_alloc) { + ggml_backend_tensor_set( + mem->pimpl->committed_candidate_ids, + committed_candidates.data(), + (size_t) seq_id*mem->pimpl->committed_candidate_ids->nb[1], + committed_candidates.size()*sizeof(int32_t)); + ggml_backend_tensor_set( + mem->pimpl->positions, + position_values.data(), + (size_t) seq_id*mem->pimpl->positions->nb[1], + position_values.size()*sizeof(int32_t)); + } + mem->pimpl->sequences.swap(transaction->next_sequences); + rollback.snapshots.swap(transaction->snapshots); + rollback_slot->second = std::move(rollback); + mem->pimpl->generation = transaction->plan.generation; + mem->pimpl->transaction_active = false; + transaction.reset(); +} + +void llama_memory_dsv41_context::rollback() { + if (!transaction || mem == nullptr) { + return; + } + if (full) { + transaction.reset(); + return; + } + if (transaction->engram) { + mem->pimpl->config.engram->rollback(*transaction->engram); + } + if (!mem->pimpl->config.no_alloc) { + for (const auto & snapshot : transaction->snapshots) { + ggml_backend_tensor_set( + snapshot.tensor, snapshot.data.data(), snapshot.offset, snapshot.data.size()); + } + } + mem->pimpl->transaction_active = false; + transaction.reset(); +} + +llama_memory_status llama_memory_dsv41_context::get_status() const { + return status; +} + +const llama_ubatch & llama_memory_dsv41_context::get_ubatch() const { + if (status != LLAMA_MEMORY_STATUS_SUCCESS || i_next >= ubatches.size()) { + throw std::runtime_error("DeepSeek V4.1 memory context has no current ubatch"); + } + return ubatches[i_next]; +} + +const llama_dsv41_memory_plan & llama_memory_dsv41_context::plan() const { + if (!transaction) { + throw std::runtime_error("DeepSeek V4.1 memory transaction is not prepared"); + } + return transaction->plan; +} + +const llama_dsv41_memory_plan & llama_memory_dsv41_context::graph_plan( + const llama_ubatch & ubatch) const { + if (!full) { + return plan(); + } + if (mem == nullptr || ubatch.n_tokens == 0 || + ubatch.n_tokens > mem->pimpl->config.n_ubatch) { + throw std::invalid_argument("DeepSeek V4.1 full graph ubatch is invalid"); + } + if (transaction && transaction->plan.positions.size() == ubatch.n_tokens) { + return transaction->plan; + } + + auto next = std::make_unique(); + next->plan.generation = mem->pimpl->generation; + next->next_sequences = mem->pimpl->sequences; + next->seq_ids = { 0 }; + next->start_positions = { 0 }; + next->next_sequences[0].pos = ubatch.n_tokens - 1; + + const uint32_t raw_width = mem->pimpl->config.raw_window; + const uint32_t persist_first = + ubatch.n_tokens > raw_width ? ubatch.n_tokens - raw_width : 0; + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + next->plan.token_seq_ids.push_back({ 0 }); + next->plan.positions.push_back(token); + if (token >= persist_first) { + next->plan.raw.persist_src_idxs.push_back(token); + next->plan.raw.write_idxs.push_back(token%raw_width); + } + const uint32_t visible = std::min(raw_width, token + 1); + next->plan.raw.n_visible.push_back(visible); + const auto order = llama_dsv41_raw_ring_order(token, raw_width); + for (uint32_t row = 0; row < raw_width; ++row) { + next->plan.raw.read_idxs.push_back(row < visible ? order[row] : -1); + next->plan.raw.mask.push_back( + row < visible ? 0.0f : -std::numeric_limits::infinity()); + } + } + for (uint32_t source : mem->pimpl->config.kv_sources) { + llama_dsv41_source_plan source_plan; + source_plan.source_layer = source; + source_plan.ratio = mem->pimpl->config.ratios[source]; + source_plan.capacity = source_capacity(mem->pimpl->config, source); + source_plan.compression = llama_dsv41_build_compression_plan( + next->plan.positions, source_plan.ratio, source_plan.capacity); + const uint32_t max_visible = source_plan.compression.n_visible.empty() ? 0 : + *std::max_element( + source_plan.compression.n_visible.begin(), + source_plan.compression.n_visible.end()); + for (int32_t visible : source_plan.compression.n_visible) { + for (uint32_t row = 0; row < max_visible; ++row) { + source_plan.read_idxs.push_back(row < (uint32_t) visible ? row : -1); + } + } + next->plan.sources.push_back(std::move(source_plan)); + } + const uint32_t visible_blocks = + (ubatch.n_tokens + mem->pimpl->config.candidate_block_size - 1)/ + mem->pimpl->config.candidate_block_size; + next->plan.candidate_width = std::min( + visible_blocks, mem->pimpl->config.candidate_topk_blocks); + transaction = std::move(next); + return transaction->plan; +} + +llama_dsv41_graph_topology llama_memory_dsv41_context::topology( + const llama_ubatch & ubatch, + uint32_t n_outputs) const { + if (mem == nullptr) { + throw std::runtime_error("DeepSeek V4.1 memory context has no memory"); + } + if (full) { + graph_plan(ubatch); + } + llama_dsv41_graph_topology result; + result.n_tokens = ubatch.n_tokens; + result.n_seqs = transaction ? transaction->seq_ids.size() : ubatch.n_seqs; + result.n_outputs = n_outputs; + result.backend_layout = mem->pimpl->backend_layout; + result.engram_enabled = mem->pimpl->config.engram_enabled; + result.expert_enabled = mem->pimpl->config.expert_enabled; + result.transaction_generation = transaction ? transaction->plan.generation : mem->pimpl->generation; + if (transaction) { + result.seq_ids = transaction->seq_ids; + result.start_positions = transaction->start_positions; + for (llama_seq_id seq_id : transaction->seq_ids) { + const llama_pos final_pos = transaction->next_sequences[seq_id].pos; + result.visible_raw_widths.push_back(std::min( + mem->pimpl->config.raw_window, final_pos + 1)); + } + for (const auto & source : transaction->plan.sources) { + result.source_ratios.push_back(source.ratio); + result.visible_compressed_widths.push_back( + source.compression.n_visible.empty() ? 0 : + *std::max_element(source.compression.n_visible.begin(), source.compression.n_visible.end())); + result.source_carry_counts.push_back( + source.ratio == 2 ? (transaction->plan.positions.back() + 1)%source.ratio : 0); + } + result.candidate_width = transaction->plan.candidate_width; + } else { + const uint32_t n_seqs = std::min(ubatch.n_seqs, mem->pimpl->config.n_seq); + for (uint32_t seq = 0; seq < n_seqs; ++seq) { + result.seq_ids.push_back(seq); + result.start_positions.push_back(0); + result.visible_raw_widths.push_back(mem->pimpl->config.raw_window); + } + for (uint32_t source : mem->pimpl->config.kv_sources) { + result.source_ratios.push_back(mem->pimpl->config.ratios[source]); + result.visible_compressed_widths.push_back(source_capacity(mem->pimpl->config, source)); + result.source_carry_counts.push_back(mem->pimpl->config.ratios[source] == 2 ? 1 : 0); + } + result.candidate_width = mem->pimpl->candidate_width; + } + return result; +} + +const llama_dsv41_engram_transaction * llama_memory_dsv41_context::engram_transaction() const { + return transaction && transaction->engram ? transaction->engram.get() : nullptr; +} + +std::vector llama_memory_dsv41_context::engram_row_ids(uint32_t engram_layer) const { + const auto * prepared = engram_transaction(); + if (prepared == nullptr) { + return {}; + } + if (engram_layer >= LLAMA_ENGRAM_LAYERS) { + throw std::invalid_argument("DeepSeek V4.1 Engram layer index is out of range"); + } + + const uint32_t * source = prepared->row_ids(engram_layer); + std::vector result(prepared->token_count()*LLAMA_ENGRAM_COLS); + for (size_t token = 0; token < prepared->token_count(); ++token) { + for (uint32_t column = 0; column < LLAMA_ENGRAM_COLS; ++column) { + const uint32_t row = + source[token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS + column]; + if (row > (uint32_t) INT32_MAX) { + throw std::runtime_error("DeepSeek V4.1 Engram row ID exceeds I32"); + } + result[token*LLAMA_ENGRAM_COLS + column] = row; + } + } + return result; +} + +void llama_memory_dsv41_context::stage_candidate_ids( + uint32_t token, + const std::vector & ids) { + if (!transaction || mem == nullptr || token >= transaction->plan.positions.size() || + ids.size() != transaction->plan.candidate_width) { + throw std::invalid_argument("DeepSeek V4.1 candidate staging shape is invalid"); + } + if (!mem->pimpl->config.no_alloc && !ids.empty()) { + ggml_backend_tensor_set( + mem->pimpl->candidate_ids, + ids.data(), + (size_t) token*mem->pimpl->candidate_ids->nb[1], + ids.size()*sizeof(int32_t)); + } +} diff --git a/src/llama-memory-dsv41.h b/src/llama-memory-dsv41.h new file mode 100644 index 000000000000..a66a09ba31ea --- /dev/null +++ b/src/llama-memory-dsv41.h @@ -0,0 +1,176 @@ +#pragma once + +#include "llama-dsv41.h" +#include "llama-dsv41-engram.h" +#include "llama-memory.h" + +#include +#include +#include +#include +#include +#include + +struct ggml_tensor; +struct llama_model; + +struct llama_dsv41_memory_config { + uint32_t n_ctx = 0; + uint32_t n_seq = 0; + uint32_t n_ubatch = 0; + uint32_t n_layer = LLAMA_DSV41_N_LAYER; + uint32_t raw_window = LLAMA_DSV41_N_SWA; + uint32_t kv_width = LLAMA_DSV41_N_HEAD_DIM; + uint32_t index_width = LLAMA_DSV41_N_INDEX_HEAD_DIM; + uint32_t candidate_topk_blocks = LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS; + uint32_t candidate_block_size = LLAMA_DSV41_CANDIDATE_BLOCK_SIZE; + uint32_t candidate_source_layer = LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + ggml_type type_k = GGML_TYPE_F16; + ggml_type type_index = GGML_TYPE_F16; + bool no_alloc = false; + bool attach_no_alloc_buffers = false; + bool engram_enabled = false; + bool expert_enabled = false; + std::vector kv_sources = { 2, 8, 14, 20 }; + std::vector index_sources = { 2, 8, 14, 20, 24, 28, 32, 36 }; + std::vector ratios; + std::function buft_for_layer; + std::unique_ptr engram; +}; + +struct llama_dsv41_raw_plan { + std::vector persist_src_idxs; + std::vector write_idxs; + std::vector read_idxs; + std::vector mask; + std::vector n_visible; +}; + +struct llama_dsv41_source_plan { + uint32_t source_layer = 0; + uint32_t ratio = 0; + uint32_t capacity = 0; + llama_dsv41_compression_plan compression; + std::vector read_idxs; +}; + +struct llama_dsv41_graph_topology { + uint32_t n_tokens = 0; + uint32_t n_seqs = 0; + uint32_t n_outputs = 0; + std::vector seq_ids; + std::vector start_positions; + std::vector visible_raw_widths; + std::vector visible_compressed_widths; + std::vector source_ratios; + std::vector source_carry_counts; + uint32_t candidate_width = 0; + uint64_t backend_layout = 0; + uint64_t transaction_generation = 0; + bool engram_enabled = false; + bool expert_enabled = false; + + bool same_topology(const llama_dsv41_graph_topology & other) const; +}; + +struct llama_dsv41_memory_plan { + uint64_t generation = 0; + std::vector positions; + std::vector> token_seq_ids; + llama_dsv41_raw_plan raw; + std::vector sources; + uint32_t candidate_width = 0; +}; + +class llama_memory_dsv41; + +class llama_memory_dsv41_context : public llama_memory_context_i { +public: + llama_memory_dsv41_context(llama_memory_status status); + llama_memory_dsv41_context(llama_memory_dsv41 * memory, bool full); + llama_memory_dsv41_context(llama_memory_dsv41 * memory, std::vector ubatches); + ~llama_memory_dsv41_context() override; + + bool next() override; + bool apply() override; + void commit() override; + void rollback() override; + + llama_memory_status get_status() const override; + const llama_ubatch & get_ubatch() const override; + + const llama_dsv41_memory_plan & plan() const; + const llama_dsv41_memory_plan & graph_plan(const llama_ubatch & ubatch) const; + llama_dsv41_graph_topology topology(const llama_ubatch & ubatch, uint32_t n_outputs) const; + const llama_dsv41_engram_transaction * engram_transaction() const; + std::vector engram_row_ids(uint32_t engram_layer) const; + const llama_memory_dsv41 * memory() const { return mem; } + + void stage_candidate_ids(uint32_t token, const std::vector & ids); + +private: + struct transaction_state; + + llama_memory_status status; + llama_memory_dsv41 * mem = nullptr; + bool full = false; + size_t i_next = 0; + std::vector ubatches; + mutable std::unique_ptr transaction; +}; + +class llama_memory_dsv41 : public llama_memory_i { +public: + explicit llama_memory_dsv41(llama_dsv41_memory_config config); + llama_memory_dsv41( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch, + std::unique_ptr engram); + ~llama_memory_dsv41() override; + + llama_memory_context_ptr init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) override; + llama_memory_context_ptr init_full() override; + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + bool get_can_shift() const override; + void clear(bool data) override; + bool seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; + void seq_keep(llama_seq_id seq_id) override; + [[noreturn]] void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + [[noreturn]] void seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override; + llama_pos seq_pos_min(llama_seq_id seq_id) const override; + llama_pos seq_pos_max(llama_seq_id seq_id) const override; + std::map memory_breakdown() const override; + void set_graph_workspace_size(size_t size) override; + void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; + void state_read(llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + + ggml_tensor * raw_k(uint32_t layer) const; + ggml_tensor * compressed_kv(uint32_t source_layer) const; + ggml_tensor * index_keys(uint32_t source_layer) const; + ggml_tensor * compressor_carry_kv(uint32_t source_layer) const; + ggml_tensor * compressor_carry_score(uint32_t source_layer) const; + ggml_tensor * candidate_scores() const; + ggml_tensor * candidate_ids() const; + ggml_tensor * committed_candidate_ids() const; + ggml_tensor * position_state() const; + + const llama_dsv41_memory_config & config() const; + llama_dsv41_memory_accounting accounting() const; + std::vector sequence_candidate_ids(llama_seq_id seq_id) const; + size_t retained_rollback_count() const; + bool engram_enabled() const; + +private: + friend class llama_memory_dsv41_context; + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..9ab61a480355 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -59,6 +59,10 @@ struct llama_memory_context_i { // return false on failure virtual bool apply() = 0; + // publish or discard state prepared by apply() + virtual void commit() {} + virtual void rollback() {} + // get the current ubatch virtual const llama_ubatch & get_ubatch() const = 0; @@ -117,6 +121,7 @@ struct llama_memory_i { virtual llama_pos seq_pos_max(llama_seq_id seq_id) const = 0; virtual std::map memory_breakdown() const = 0; + virtual void set_graph_workspace_size(size_t size) { GGML_UNUSED(size); } // // state write/read diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 715a6e3548e6..141640f9d192 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -172,7 +173,13 @@ struct llama_file::impl { } bool has_direct_io() const { - return true; + // Windows uses cached CRT I/O until FILE_FLAG_NO_BUFFERING support is added. + return false; + } + + void discard_cache(size_t offset, size_t length) const { + GGML_UNUSED(offset); + GGML_UNUSED(length); } ~impl() { @@ -374,6 +381,19 @@ struct llama_file::impl { return fd != -1 && alignment > 1; } + void discard_cache(size_t offset, size_t length) const { +#if defined(POSIX_FADV_DONTNEED) + const int file_id = fd == -1 ? fileno(fp) : fd; + const int result = posix_fadvise(file_id, offset, length, POSIX_FADV_DONTNEED); + if (result != 0) { + LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_DONTNEED) failed: %s\n", strerror(result)); + } +#else + GGML_UNUSED(offset); + GGML_UNUSED(length); +#endif + } + ~impl() { if (fd != -1) { close(fd); @@ -408,6 +428,7 @@ size_t llama_file::size() const { return pimpl->size; } size_t llama_file::read_alignment() const { return pimpl->read_alignment(); } bool llama_file::has_direct_io() const { return pimpl->has_direct_io(); } +void llama_file::discard_cache(size_t offset, size_t length) const { pimpl->discard_cache(offset, length); } int llama_file::file_id() const { #ifdef _WIN32 @@ -439,7 +460,6 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); } // llama_mmap -#if defined(_POSIX_MAPPED_FILES) || defined(_WIN32) // merge `ranges` and return their complement within [0, limit) static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t limit) { llama_mmap::ranges res; @@ -457,27 +477,38 @@ static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t li if (pos < limit) { res.emplace_back(pos, limit); } - return res; } -#endif struct llama_mmap::impl { #ifdef _POSIX_MAPPED_FILES std::vector> mapped_fragments; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; if (numa) { prefetch = 0; } #ifdef __linux__ - if (posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL)) { - LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", - strerror(errno)); + const bool sequential = llama_mmap::use_sequential_file_advice(strict_exclusion); + const int file_advice = sequential ? POSIX_FADV_SEQUENTIAL : POSIX_FADV_RANDOM; + const int advice_error = file_advice_override ? + file_advice_override(fd, file_advice) : posix_fadvise(fd, 0, 0, file_advice); + if (advice_error) { + if (strict_exclusion) { + throw std::runtime_error(format( + "posix_fadvise(.., POSIX_FADV_RANDOM) failed for external tensor mapping: %s", + strerror(advice_error))); + } + LLAMA_LOG_WARN("warning: posix_fadvise(.., %s) failed: %s\n", + sequential ? "POSIX_FADV_SEQUENTIAL" : "POSIX_FADV_RANDOM", strerror(advice_error)); } - // MAP_POPULATE would fault in the lazy ranges too - if (prefetch && lazy_ranges.empty()) { flags |= MAP_POPULATE; } + // MAP_POPULATE would fault in excluded ranges too + if (prefetch && excluded_ranges.empty()) { flags |= MAP_POPULATE; } +#else + GGML_UNUSED(file_advice_override); #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { @@ -497,12 +528,11 @@ struct llama_mmap::impl { } }; - if (prefetch > 0) { - for (const auto & range : ranges_complement(lazy_ranges, std::min(file->size(), prefetch))) { - advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); - } + for (const auto & range : + llama_mmap::planned_prefetch_ranges(file->size(), prefetch, excluded_ranges, strict_exclusion)) { + advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); } - for (const auto & range : lazy_ranges) { + for (const auto & range : excluded_ranges) { advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); } if (numa) { @@ -573,8 +603,11 @@ struct llama_mmap::impl { #elif defined(_WIN32) HANDLE hMapping = nullptr; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(numa); + GGML_UNUSED(file_advice_override); size = file->size(); @@ -604,7 +637,8 @@ struct llama_mmap::impl { if (pPrefetchVirtualMemory) { std::vector entries; - for (const auto & range : ranges_complement(lazy_ranges, std::min(size, prefetch))) { + for (const auto & range : + llama_mmap::planned_prefetch_ranges(size, prefetch, excluded_ranges, strict_exclusion)) { WIN32_MEMORY_RANGE_ENTRY entry; entry.VirtualAddress = (char *) addr + range.first; entry.NumberOfBytes = (SIZE_T) (range.second - range.first); @@ -642,11 +676,15 @@ struct llama_mmap::impl { } } #else - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); - GGML_UNUSED(lazy_ranges); + GGML_UNUSED(excluded_ranges); + GGML_UNUSED(strict_exclusion); + GGML_UNUSED(file_advice_override); throw std::runtime_error("mmap not supported"); } @@ -664,9 +702,22 @@ struct llama_mmap::impl { }; llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, - const ranges & lazy_ranges) : pimpl(std::make_unique(file, prefetch, numa, lazy_ranges)) {} + const ranges & excluded_ranges, bool strict_exclusion, file_advice_override file_advice) : + pimpl(std::make_unique(file, prefetch, numa, excluded_ranges, strict_exclusion, file_advice)) {} llama_mmap::~llama_mmap() = default; +bool llama_mmap::use_sequential_file_advice(bool strict_exclusion) { + return !strict_exclusion; +} + +llama_mmap::ranges llama_mmap::planned_prefetch_ranges( + size_t file_size, size_t prefetch, const ranges & excluded_ranges, bool strict_exclusion) { + if (strict_exclusion || prefetch == 0) { + return {}; + } + return ranges_complement(excluded_ranges, std::min(file_size, prefetch)); +} + size_t llama_mmap::size() const { return pimpl->size; } void * llama_mmap::addr() const { return pimpl->addr; } @@ -782,6 +833,18 @@ struct llama_mlock::impl { impl() : addr(NULL), size(0), failed_already(false) {} + static void align_range(size_t * first, size_t * last) { + const size_t granularity = lock_granularity(); + *first &= ~(granularity - 1); + const size_t remainder = *last & (granularity - 1); + if (remainder != 0) { + if (*last > std::numeric_limits::max() - (granularity - remainder)) { + throw std::runtime_error("mlock range overflow"); + } + *last += granularity - remainder; + } + } + void init(void * ptr) { GGML_ASSERT(addr == NULL && size == 0); addr = ptr; @@ -814,6 +877,7 @@ llama_mlock::~llama_mlock() = default; void llama_mlock::init(void * ptr) { pimpl->init(ptr); } void llama_mlock::grow_to(size_t target_size) { pimpl->grow_to(target_size); } +void llama_mlock::align_range(size_t * first, size_t * last) { impl::align_range(first, last); } #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32) const bool llama_mlock::SUPPORTED = true; diff --git a/src/llama-mmap.h b/src/llama-mmap.h index cc28c8a73fa5..c64c5acfb6a1 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -31,6 +31,8 @@ struct llama_file { void read_aligned_chunk(void * dest, size_t size); uint32_t read_u32(); + void discard_cache(size_t offset, size_t length) const; + void write_raw(const void * ptr, size_t len) const; void write_u32(uint32_t val) const; @@ -44,10 +46,12 @@ struct llama_file { struct llama_mmap { // list of [first, last) byte ranges within a file using ranges = std::vector>; + using file_advice_override = int (*)(int fd, int advice); llama_mmap(const llama_mmap &) = delete; llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, - const ranges & lazy_ranges = {}); + const ranges & excluded_ranges = {}, bool strict_exclusion = false, + file_advice_override file_advice = nullptr); ~llama_mmap(); size_t size() const; @@ -56,6 +60,9 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); static const bool SUPPORTED; + static bool use_sequential_file_advice(bool strict_exclusion); + static ranges planned_prefetch_ranges( + size_t file_size, size_t prefetch, const ranges & excluded_ranges, bool strict_exclusion); private: struct impl; @@ -69,6 +76,7 @@ struct llama_mlock { void init(void * ptr); void grow_to(size_t target_size); + static void align_range(size_t * first, size_t * last); static const bool SUPPORTED; private: diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index fde86b19541f..420656116cf1 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include static const size_t kiB = 1024; @@ -429,6 +430,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); @@ -724,6 +726,41 @@ llama_model_loader::llama_model_loader( } else { get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false); llm_kv = LLM_KV(llm_arch_from_string(arch_name)); + if (no_alloc && gguf_get_n_tensors(metadata) > 0) { + const int64_t tensor_count = gguf_get_n_tensors(metadata); + const size_t overhead = ggml_tensor_overhead(); + if ((uint64_t) tensor_count > std::numeric_limits::max()/overhead - 1) { + throw std::runtime_error("no-allocation metadata tensor count overflows the context size"); + } + ggml_init_params params = { + /*.mem_size =*/ overhead*((size_t) tensor_count + 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + throw std::runtime_error("failed to create no-allocation metadata tensor context"); + } + contexts.emplace_back(ctx); + for (int64_t i = 0; i < tensor_count; ++i) { + ggml_tensor * tensor = ggml_new_tensor( + ctx, + gguf_get_tensor_type(metadata, i), + GGML_MAX_DIMS, + gguf_get_tensor_ne(metadata, i)); + ggml_set_name(tensor, gguf_get_tensor_name(metadata, i)); + n_elements += ggml_nelements(tensor); + n_bytes += ggml_nbytes(tensor); + const auto inserted = weights_map.emplace( + ggml_get_name(tensor), + llama_tensor_weight(0, metadata, tensor)); + if (!inserted.second) { + throw std::runtime_error(format( + "invalid model: tensor '%s' is duplicated", + ggml_get_name(tensor))); + } + } + } } n_kv = gguf_get_n_kv(metadata); @@ -1131,6 +1168,14 @@ bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_ten return true; } +void llama_model_loader::external_read::add(const llama_tensor_weight & w) { + const std::string name = ggml_get_name(w.tensor); + if (!tensors.insert(name).second) { + throw std::runtime_error(format("external tensor '%s' is already registered", name.c_str())); + } + ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(w.tensor)); +} + struct ggml_tensor * llama_model_loader::create_tensor( const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { @@ -1197,7 +1242,9 @@ struct ggml_tensor * llama_model_loader::create_tensor( const size_t nbytes = ggml_nbytes(t_meta); LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", tn.str().c_str(), nbytes); - size_data -= nbytes; + if (!files.empty()) { + size_data -= nbytes; + } n_created++; return nullptr; @@ -1311,12 +1358,6 @@ struct ggml_tensor * llama_model_loader::create_tensor( if (flags & TENSOR_SKIP_IF_VIRTUAL) { return nullptr; } - ggml_type type = GGML_TYPE_F32; - const int64_t tid = gguf_find_tensor(metadata, tn.str().c_str()); - if (tid != -1) { - type = gguf_get_tensor_type(metadata, tid); - } - // for tensors that are not required some of the dimensions can be invalid: if (flags & TENSOR_NOT_REQUIRED) { for (size_t dim = 0; dim < ne.size(); dim++) { @@ -1326,6 +1367,15 @@ struct ggml_tensor * llama_model_loader::create_tensor( } } + ggml_type type = GGML_TYPE_F32; + const int64_t tid = gguf_find_tensor(metadata, tn.str().c_str()); + if (tid != -1) { + const ggml_tensor * declared = check_tensor_dims( + tn.str(), ne, true, flags & TENSOR_ALLOW_RESHAPE); + GGML_ASSERT(declared != nullptr); + type = declared->type; + } + ggml_tensor t_meta; memset(&t_meta, 0, sizeof(ggml_tensor)); t_meta.type = type; @@ -1344,10 +1394,15 @@ struct ggml_tensor * llama_model_loader::create_tensor( ggml_set_name(&t_meta, tn.str().c_str()); ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta); - GGML_ASSERT(buft != nullptr); + if (buft == nullptr) { + return nullptr; + } ggml_context * ctx = ctx_for_buft(buft); ggml_tensor * ret = ggml_dup_tensor(ctx, &t_meta); ggml_set_name(ret, tn.str().c_str()); + if (tid != -1 && !(flags & TENSOR_DUPLICATED)) { + n_created++; + } return ret; } @@ -1407,6 +1462,49 @@ struct ggml_tensor * llama_model_loader::create_tensor( return tensor; } +llama_expert_store_tensor llama_model_loader::register_external_tensor( + const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + const ggml_tensor * tensor = check_tensor_dims(name, ne, true, false); + GGML_ASSERT(tensor != nullptr); + if (tensor->ne[3] != 1) { + throw std::runtime_error(format("external tensor '%s' must have three dimensions", name.c_str())); + } + + const llama_tensor_weight & weight = require_weight(name.c_str()); + if (!no_alloc && (fnames.at(weight.idx).empty() || fnames.at(weight.idx) == "(file*)")) { + throw std::runtime_error(format("external tensor '%s' requires a reopenable source file", name.c_str())); + } + llama_expert_store_tensor result; + result.name = name; + result.fname = no_alloc ? "(no_alloc)" : fnames.at(weight.idx); + result.file_index = weight.idx; + result.layer = layer; + result.projection = projection; + result.type = tensor->type; + for (size_t i = 0; i < 3; ++i) { + result.ne[i] = tensor->ne[i]; + result.nb[i] = tensor->nb[i]; + } + result.file_offset = weight.offs; + if (no_alloc) { + const size_t tensor_size = ggml_nbytes(tensor); + if (result.file_offset > std::numeric_limits::max() - tensor_size) { + throw std::runtime_error(format("external tensor '%s' extent overflows the file size", name.c_str())); + } + result.file_size = result.file_offset + tensor_size; + } else { + result.file_size = files.at(weight.idx)->size(); + } + + llama_expert_store_validate_tensor(result); + external.add(weight); + n_created++; + return result; +} + void llama_model_loader::done_getting_tensors(bool partial) const { if (n_created > n_tensors) { throw std::runtime_error(format("%s: too many tensors created; expected %d, got %d", __func__, n_tensors, n_created)); @@ -1415,7 +1513,7 @@ void llama_model_loader::done_getting_tensors(bool partial) const { if (!partial) { throw std::runtime_error(format("%s: wrong number of tensors; expected %d, got %d", __func__, n_tensors, n_created)); } - LLAMA_LOG_INFO("%s: partial load — used %d of %d tensors in the file (rest belong to a sibling model on the same .gguf)\n", + LLAMA_LOG_INFO("%s: partial load - used %d of %d tensors in the file (rest belong to a sibling model on the same .gguf)\n", __func__, n_created, n_tensors); } if (n_tensors_moved > 0) { @@ -1446,10 +1544,17 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps const size_t prefetch_size = prefetch && use_mmap ? -1 : 0; - std::unique_ptr mapping = std::make_unique(file.get(), prefetch_size, is_numa, - lazy.for_file(idx)); + llama_mmap::ranges excluded = lazy.for_file(idx); + const auto & external_ranges = external.for_file(idx); + excluded.insert(excluded.end(), external_ranges.begin(), external_ranges.end()); + + std::unique_ptr mapping = std::make_unique( + file.get(), prefetch_size, is_numa, excluded, !external_ranges.empty()); + for (const auto & range : external_ranges) { + mapping->unmap_fragment(range.first, range.second); + } mmaps_used.emplace_back(mapping->size(), 0); - if (mlock_mmaps) { + if (mlock_mmaps && external_ranges.empty()) { std::unique_ptr mlock_mmap(new llama_mlock()); mlock_mmap->init(mapping->addr()); mlock_mmaps->emplace_back(std::move(mlock_mmap)); @@ -1460,7 +1565,9 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps // compute the total size of all tensors for progress reporting for (const auto & it : weights_map) { - size_data += ggml_nbytes(it.second.tensor); + if (!external.has(it.second.tensor)) { + size_data += ggml_nbytes(it.second.tensor); + } } } @@ -1511,6 +1618,8 @@ const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, bool llama_model_loader::load_all_data( struct ggml_context * ctx, llama_buf_map & bufs, + bool load_from_mmap, + bool discard_file_cache, llama_mlocks * lmlocks, llama_progress_callback progress_callback, void * progress_callback_user_data) { @@ -1537,12 +1646,44 @@ bool llama_model_loader::load_all_data( // 64MB works well for NVMe drives const size_t buffer_size = alignment != 1 ? 64 * 1024 * 1024 + 2 * alignment : 1 * 1024 * 1024; - std::vector host_buffers; - std::vector events; - std::vector host_ptrs; + struct async_upload_resources { + std::vector host_buffers; + std::vector events; + std::vector host_ptrs; + ggml_backend_t backend = nullptr; + + void reset() { + for (auto * event : events) { + if (backend != nullptr) { + ggml_backend_event_synchronize(event); + } + ggml_backend_event_free(event); + } + events.clear(); + for (auto * buffer : host_buffers) { + ggml_backend_buffer_free(buffer); + } + host_buffers.clear(); + host_ptrs.clear(); + ggml_backend_free(backend); + backend = nullptr; + } + + ~async_upload_resources() { + reset(); + } + } async_upload; + async_upload.host_buffers.reserve(n_buffers); + async_upload.events.reserve(n_buffers); + async_upload.host_ptrs.reserve(n_buffers); + + auto & host_buffers = async_upload.host_buffers; + auto & events = async_upload.events; + auto & host_ptrs = async_upload.host_ptrs; + auto & upload_backend = async_upload.backend; size_t buffer_idx = 0; // buffer to use for async loads - ggml_backend_t upload_backend = [&](const char * func) -> ggml_backend_t { - if (use_mmap || check_tensors) { + upload_backend = [&](const char * func) -> ggml_backend_t { + if (load_from_mmap || check_tensors) { return nullptr; } // When not using mmaped io use async uploads from pinned memory to GPU memory. @@ -1589,6 +1730,7 @@ bool llama_model_loader::load_all_data( if (!buf) { LLAMA_LOG_DEBUG("%s: failed to allocate host buffer for async uploads for device %s\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1599,6 +1741,7 @@ bool llama_model_loader::load_all_data( if (!event) { LLAMA_LOG_DEBUG("%s: failed to create event for async uploads for device %s\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1609,6 +1752,7 @@ bool llama_model_loader::load_all_data( if (!backend) { LLAMA_LOG_DEBUG("%s: failed to initialize backend for device %s for async uploads\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1629,7 +1773,7 @@ bool llama_model_loader::load_all_data( // without mmap, tensors in non-host buffers are staged through a temporary buffer sized like the tensor // load them biggest-first so the largest staging buffer is allocated while the fewest weights are resident - if (!use_mmap) { + if (!load_from_mmap) { std::stable_sort(tensors.begin(), tensors.end(), [](const ggml_tensor * a, const ggml_tensor * b) { const bool staged_a = a->buffer && !ggml_backend_buffer_is_host(a->buffer); const bool staged_b = b->buffer && !ggml_backend_buffer_is_host(b->buffer); @@ -1655,7 +1799,7 @@ bool llama_model_loader::load_all_data( size_t n_size = ggml_nbytes(cur); - const bool from_mapping = use_mmap || lazy.has(cur); + const bool from_mapping = load_from_mmap || lazy.has(cur); if (from_mapping) { const auto & mapping = mappings.at(weight->idx); @@ -1763,20 +1907,15 @@ bool llama_model_loader::load_all_data( } } } + if (discard_file_cache) { + file->discard_cache(weight->offs, n_size); + } } size_done += n_size; } - // free temporary resources used for async uploads - for (auto * event : events) { - ggml_backend_event_synchronize(event); - ggml_backend_event_free(event); - } - for (auto * buf : host_buffers) { - ggml_backend_buffer_free(buf); - } - ggml_backend_free(upload_backend); + async_upload.reset(); // check validation results bool validation_failed = false; diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 72bbd53e7d03..e3e7899b917f 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -4,6 +4,7 @@ #include "llama-impl.h" #include "llama-arch.h" +#include "llama-expert-store.h" #include "llama-hparams.h" #include "llama-mmap.h" @@ -11,6 +12,7 @@ #include #include +#include #include #include #include @@ -48,6 +50,19 @@ struct llama_model_loader { throw std::runtime_error(format("tensor '%s' data is not within the file bounds, model is corrupted or incomplete", ggml_get_name(tensor))); } } + + llama_tensor_weight(uint16_t idx, const struct gguf_context * gguf_ctx, ggml_tensor * tensor) : idx(idx), tensor(tensor) { + const int tensor_idx = gguf_find_tensor(gguf_ctx, ggml_get_name(tensor)); + if (tensor_idx < 0) { + throw std::runtime_error(format("tensor '%s' not found in the model", ggml_get_name(tensor))); + } + const size_t data_offset = gguf_get_data_offset(gguf_ctx); + const size_t tensor_offset = gguf_get_tensor_offset(gguf_ctx, tensor_idx); + if (data_offset > std::numeric_limits::max() - tensor_offset) { + throw std::runtime_error(format("tensor '%s' offset overflows", ggml_get_name(tensor))); + } + offs = data_offset + tensor_offset; + } }; // custom comparator to sort weights more nicely by layer @@ -117,6 +132,38 @@ struct llama_model_loader { std::set tensors; } lazy; + struct external_read { + void add(const llama_tensor_weight & w); + + bool any() const { + return !ranges.empty(); + } + + bool has(const ggml_tensor * t) const { + return tensors.count(ggml_get_name(t)) > 0; + } + + const llama_mmap::ranges & for_file(uint32_t idx) const { + static const llama_mmap::ranges none; + + const auto it = ranges.find(idx); + return it == ranges.end() ? none : it->second; + } + + bool intersects(uint32_t idx, size_t first, size_t last) const { + for (const auto & range : for_file(idx)) { + if (range.first < last && first < range.second) { + return true; + } + } + return false; + } + + private: + std::map ranges; + std::set tensors; + } external; + llama_files files; std::vector fnames; // one per entry of files, for readers that outlive the loader llama_ftype ftype; @@ -239,6 +286,12 @@ struct llama_model_loader { const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags); + llama_expert_store_tensor register_external_tensor( + const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne); + void done_getting_tensors(bool partial = false) const; void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr); @@ -256,6 +309,8 @@ struct llama_model_loader { bool load_all_data( struct ggml_context * ctx, llama_buf_map & bufs, + bool load_from_mmap, + bool discard_file_cache, llama_mlocks * lmlocks, llama_progress_callback progress_callback, void * progress_callback_user_data); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c6d3e5dc6a7f..e6f17d8c9ddb 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -6,6 +6,7 @@ #include "llama-impl.h" #include "llama-mmap.h" #include "llama-cparams.h" +#include "llama-dsv41.h" #include "llama-model-loader.h" #include "llama-kv-cache.h" @@ -17,6 +18,7 @@ #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-hybrid-idx.h" +#include "llama-memory-dsv41.h" #include "llama-memory-recurrent.h" #include "llama.h" @@ -200,6 +202,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_dots3note(params); case LLM_ARCH_DEEPSEEK4: return new llama_model_deepseek4(params); + case LLM_ARCH_DEEPSEEK41: + return new llama_model_deepseek41(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); case LLM_ARCH_MISTRAL4: @@ -373,6 +377,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const llama_hparams & hparams = ud->model->hparams; const std::string tensor_name = tensor->name; const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || + ud->model->arch == LLM_ARCH_DEEPSEEK41 || (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0); static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight"); @@ -384,7 +389,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); - static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); + static const std::regex pattern_dsv4_state ("dsv4(1)?_(csa|hca|lid|comp|index)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); @@ -1224,6 +1229,33 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { return; } + if (arch == LLM_ARCH_DEEPSEEK41) { + std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0); + std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0); + std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0); + std::fill(hparams.n_ff_exp_arr.begin(), hparams.n_ff_exp_arr.end(), 0); + std::fill(hparams.n_expert_used_arr.begin(), hparams.n_expert_used_arr.end(), 0); + std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0); + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), 1); + std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0); + std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), 0); + std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0); + std::fill(hparams.dsv41_kv_source_layer.begin(), hparams.dsv41_kv_source_layer.end(), -1); + std::fill(hparams.dsv41_index_source_layer.begin(), hparams.dsv41_index_source_layer.end(), -1); + std::fill(hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.end(), 0); + std::fill(hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.end(), 0.0f); + std::fill(hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.end(), 0.0f); + hparams.dsv41_engram_layers.reset(); + + load_arch_hparams(ml); + + pimpl->n_bytes = ml.n_bytes; + pimpl->desc_str = arch_name() + " " + type_name() + " " + ml.ftype_name(); + pimpl->ftype = ml.ftype; + hparams.rope_type = llama_model_rope_type(this); + return; + } + ml.get_key(LLM_KV_CONTEXT_LENGTH, hparams.n_ctx_train); ml.get_key(LLM_KV_EMBEDDING_LENGTH, hparams.n_embd); ml.get_key(LLM_KV_EMBEDDING_LENGTH_OUT, hparams.n_embd_out_impl, false); @@ -1706,12 +1738,22 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // With the n-gram table left on disk, a populated mapping would pull the table's // third of the file resident for nothing; readahead alone carries the sequential load. - ml.init_mappings(!(params.ple_on_disk || params.lazy_mode == LLAMA_LAZY_MODE_DIRECT), use_mlock ? &pimpl->mlock_mmaps : nullptr); + // Do not prefetch files with disk-owned tensor holes either: unsafe contexts load their + // resident tensors through bounded staging and discard copied source pages. + llama_mlocks * mmap_locks = use_mlock && !ml.external.any() ? &pimpl->mlock_mmaps : nullptr; + ml.init_mappings(!(params.ple_on_disk || params.lazy_mode == LLAMA_LAZY_MODE_DIRECT || ml.external.any()), mmap_locks); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers - std::vector> ctx_buf_maps; + struct ctx_buf_map { + ggml_context * ctx; + llama_buf_map bufs; + bool load_from_mmap; + bool discard_file_cache; + }; + std::vector ctx_buf_maps; ctx_buf_maps.reserve(ml.ctx_map.size()); + bool keep_mappings = false; // Ensure we have enough capacity for the maximum backend buffer we will potentially create const size_t n_max_backend_buffer = ml.ctx_map.size() * ml.files.size(); @@ -1747,9 +1789,24 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // a lazy context is mapped whatever the load mode, but the memory-fit pass maps nothing const bool is_lazy_mapped = ctx_key.lazy && !ml.no_alloc; + bool context_mmap_safe = true; + if (ml.use_mmap && ml.external.any() && !is_lazy_mapped) { + for (uint32_t idx = 0; idx < ml.files.size(); ++idx) { + void * addr = nullptr; + size_t first; + size_t last; + ml.get_mapping_range(&first, &last, &addr, idx, ctx); + if (first < last && ml.external.intersects(idx, first, last)) { + context_mmap_safe = false; + break; + } + } + } - if ((ml.use_mmap || is_lazy_mapped) && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft) { + if ((ml.use_mmap || is_lazy_mapped) && (use_mmap_buffer || is_lazy_mapped) && + context_mmap_safe && buffer_from_host_ptr_supported && is_default_buft) { GGML_ASSERT(!ml.no_alloc); + keep_mappings = true; for (uint32_t idx = 0; idx < ml.files.size(); idx++) { // only the mmap region containing the tensors in the model is mapped to the backend buffer // this is important for metal with apple silicon: if the entire model could be mapped to a metal buffer, @@ -1766,6 +1823,15 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { if (buf == nullptr) { throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); } + if (use_mlock && ml.external.any() && !is_lazy_mapped) { + size_t lock_first = first; + size_t lock_last = last; + llama_mlock::align_range(&lock_first, &lock_last); + pimpl->mlock_mmaps.emplace_back(new llama_mlock); + auto & mlock_mmap = pimpl->mlock_mmaps.back(); + mlock_mmap->init((char *) addr + lock_first); + mlock_mmap->grow_to(lock_last - lock_first); + } bufs.emplace_back(buf); buf_map.emplace(idx, buf); } @@ -1802,7 +1868,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs)); - ctx_buf_maps.emplace_back(ctx, buf_map); + const bool load_from_mmap = context_mmap_safe && (ml.use_mmap || is_lazy_mapped); + ctx_buf_maps.push_back({ ctx, std::move(buf_map), load_from_mmap, ml.use_mmap && !load_from_mmap }); } if (llama_supports_gpu_offload()) { @@ -1834,21 +1901,24 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } // without mmap, load non-host buffers first: their tensors go through a staging buffer, which is cheapest while the fewest weights are resident - if (!ml.use_mmap) { + if (!ml.use_mmap || ml.external.any()) { std::stable_partition(ctx_buf_maps.begin(), ctx_buf_maps.end(), [](const auto & ctx_buf_map) { - const auto & buf_map = ctx_buf_map.second; - return !buf_map.empty() && !ggml_backend_buffer_is_host(buf_map.begin()->second); + const auto & buf_map = ctx_buf_map.bufs; + return !ctx_buf_map.load_from_mmap && !buf_map.empty() && + !ggml_backend_buffer_is_host(buf_map.begin()->second); }); } // load tensor data - for (auto & [ctx, buf_map] : ctx_buf_maps) { - if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) { + for (auto & ctx_buf_map : ctx_buf_maps) { + if (!ml.load_all_data(ctx_buf_map.ctx, ctx_buf_map.bufs, ctx_buf_map.load_from_mmap, + ctx_buf_map.discard_file_cache, mmap_locks, + params.progress_callback, params.progress_callback_user_data)) { return false; } } - if (use_mmap_buffer) { + if (keep_mappings) { for (auto & mapping : ml.mappings) { pimpl->mappings.emplace_back(std::move(mapping)); } @@ -2207,6 +2277,20 @@ ggml_backend_buffer_type_t llama_model::select_buft(int il) const { }); } +ggml_backend_buffer_type_t llama_model::select_moe_buft( + int il, enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) const { + return ::select_buft( + *pimpl->dev_layer.at(il).buft_list, + [&](ggml_context * ctx) { + const int64_t n_expert_used = hparams.n_expert_used_max(); + GGML_ASSERT(n_expert_used > 0); + ggml_tensor * weight = ggml_new_tensor_3d(ctx, type, ne0, ne1, ne2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, ne0, n_expert_used, 512); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512); + return ggml_mul_mat_id(ctx, weight, input, ids); + }); +} + bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } @@ -2480,6 +2564,18 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } break; + case LLM_ARCH_DEEPSEEK41: + { + const auto & model_dsv41 = static_cast(*this); + res = new llama_memory_dsv41( + *this, + params.type_k, + cparams.offload_kqv, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + model_dsv41.create_memory_engram_runtime(cparams.n_ubatch)); + } break; case LLM_ARCH_DFLASH: { // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring) @@ -2803,6 +2899,8 @@ llama_model_params llama_model_default_params() { /*.main_gpu =*/ 0, /*.ple_io_threads =*/ 64, /*.ple_cache_mb =*/ 256, + /*.expert_cache_bytes =*/ 0, + /*.expert_cache_slots =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, @@ -2871,7 +2969,7 @@ int32_t llama_model_n_head_kv(const llama_model * model) { int32_t llama_model_n_swa(const llama_model * model) { // dsv4 kv-cache has SWA but it cannot be used as a rollback because of // other compression ratios, so we return 0 here - if (model->arch == LLM_ARCH_DEEPSEEK4) { + if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) { return 0; } return model->hparams.n_swa; @@ -2957,6 +3055,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK2OCR: case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_PLM: case LLM_ARCH_CHATGLM: diff --git a/src/llama-model.h b/src/llama-model.h index 75643a8523d4..cd4d9fb5833d 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -564,6 +564,10 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + struct ggml_tensor * engram_kv = nullptr; + struct ggml_tensor * engram_q_norm = nullptr; + struct ggml_tensor * engram_k_norm = nullptr; + // MSA struct ggml_tensor * index_q_proj = nullptr; struct ggml_tensor * index_k_proj = nullptr; @@ -748,6 +752,8 @@ struct llama_model { ggml_backend_dev_t dev_output() const; ggml_backend_buffer_type_t select_buft(int il) const; + ggml_backend_buffer_type_t select_moe_buft( + int il, enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) const; bool has_tensor_overrides() const; @@ -767,6 +773,13 @@ struct llama_model { virtual void load_vocab (llama_model_loader & ml) = 0; virtual bool load_tensors(llama_model_loader & ml) = 0; // returns false if cancelled by progress_callback + virtual bool requires_synchronous_graph() const { return false; } + virtual std::string consume_runtime_error() const { return {}; } + virtual void release_runtime_work() const {} + virtual void release_runtime_work_after_sync(ggml_backend_sched_t) const { release_runtime_work(); } + virtual void acquire_runtime_context() const {} + virtual void release_runtime_context() const {} + // model must define these virtual void load_arch_hparams(llama_model_loader & ml) = 0; virtual void load_arch_tensors(llama_model_loader & ml) = 0; diff --git a/src/llama-ple-disk.cpp b/src/llama-ple-disk.cpp index e21ab0bd1f2b..66957a9d4543 100644 --- a/src/llama-ple-disk.cpp +++ b/src/llama-ple-disk.cpp @@ -1,5 +1,6 @@ #include "llama-ple-disk.h" +#include "llama-bounded-file.h" #include "llama-impl.h" #include @@ -7,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -14,22 +16,21 @@ #include #if !defined(_WIN32) -#include #include #include #endif struct llama_ple_disk::impl { std::string fname; - int fd = -1; bool direct = false; size_t offs = 0; ggml_type type = GGML_TYPE_COUNT; int64_t ne0 = 0; int64_t nrows = 0; size_t rs = 0; // bytes per row - size_t block = 4096; // O_DIRECT alignment for offset, length and buffer ggml_to_float_t to_float = nullptr; + std::unique_ptr file; + int fd = -1; // buffered-mode fd, used only for prefetch() fadvise hints // direct-mapped cache of raw rows: slot = row & mask, tag = row size_t n_slots = 0; @@ -43,7 +44,7 @@ struct llama_ple_disk::impl { std::vector raw; // [distinct rows, rs] std::vector> misses; // (row, index into raw) float * dst = nullptr; // destination of the gather in progress - uint8_t * bounce0 = nullptr; // calling-thread bounce buffer + llama_bounded_file::buffer bounce0; // calling-thread bounce buffer std::mutex mtx; // one gather at a time; a model is shared by every context built on it @@ -60,6 +61,7 @@ struct llama_ple_disk::impl { uint64_t gen = 0; size_t pending = 0; std::atomic next{0}; + std::exception_ptr worker_error; // first error seen by a worker, rethrown on the calling thread bool stop = false; uint64_t st_calls = 0, st_rows = 0, st_uniq = 0, st_hits = 0, st_reads = 0, st_bytes = 0; @@ -70,6 +72,9 @@ struct llama_ple_disk::impl { #if defined(_WIN32) throw std::runtime_error("llama_ple_disk: not supported on Windows"); #else + if (ne0 <= 0 || nrows <= 0) { + throw std::runtime_error("llama_ple_disk: invalid table shape"); + } if (ggml_is_quantized(type) && ne0 % ggml_blck_size(type) != 0) { throw std::runtime_error(format("llama_ple_disk: row of %lld %s elements is not a whole number of blocks", (long long) ne0, ggml_type_name(type))); @@ -82,35 +87,20 @@ struct llama_ple_disk::impl { } } - direct = p.direct_io; - if (direct) { - // only Linux has O_DIRECT; Darwin turns the page cache off per descriptor - // with F_NOCACHE after the open -#if defined(O_DIRECT) - fd = open(fname.c_str(), O_RDONLY | O_DIRECT | O_CLOEXEC); -#else - fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); -#if defined(F_NOCACHE) - if (fd >= 0 && fcntl(fd, F_NOCACHE, 1) < 0) { - LLAMA_LOG_WARN("%s: F_NOCACHE on %s failed (%s); reads go through the page cache\n", - __func__, fname.c_str(), strerror(errno)); - direct = false; - } -#else - direct = false; -#endif -#endif - if (fd < 0) { - LLAMA_LOG_WARN("%s: direct open of %s failed (%s); falling back to buffered reads\n", - __func__, fname.c_str(), strerror(errno)); - direct = false; - } + llama_bounded_file::params fp; + fp.direct_io = p.direct_io; + file = std::make_unique(fname, fp); + direct = file->direct_io(); + if (rs == 0 || (uint64_t) nrows > UINT64_MAX / rs) { + throw std::runtime_error("llama_ple_disk: table extent overflow"); + } + const uint64_t bytes = (uint64_t) nrows * rs; + if (offs > file->size() || bytes > file->size() - offs) { + throw std::runtime_error("llama_ple_disk: table extent is outside the file"); } - if (fd < 0) { + if (!direct) { + // best-effort fd used only to hint the page cache in prefetch(); failure is not fatal fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); - if (fd < 0) { - throw std::runtime_error(format("llama_ple_disk: failed to open %s: %s", fname.c_str(), strerror(errno))); - } } n_threads = std::max(1, p.n_threads); @@ -138,7 +128,6 @@ struct llama_ple_disk::impl { for (auto & w : workers) { w.join(); } - free(bounce0); if (fd >= 0) { close(fd); } @@ -146,52 +135,9 @@ struct llama_ple_disk::impl { } #if !defined(_WIN32) - size_t bounce_size() const { - return ((rs + block - 1) / block) * block + 2 * block; - } - - uint8_t * alloc_bounce() const { - void * ptr = nullptr; - if (posix_memalign(&ptr, block, bounce_size()) != 0) { - throw std::runtime_error("llama_ple_disk: posix_memalign failed"); - } - return (uint8_t *) ptr; - } - - // read `len` bytes at `off` into `dst`; a short read is only tolerated past `need` - void pread_full(uint8_t * dst, size_t len, off_t off, size_t need) const { - size_t got = 0; - while (got < len) { - const ssize_t r = pread(fd, dst + got, len - got, off + (off_t) got); - if (r < 0) { - if (errno == EINTR) { - continue; - } - GGML_ABORT("llama_ple_disk: pread(%s, %zu @ %lld) failed: %s", - fname.c_str(), len, (long long) off, strerror(errno)); - } - if (r == 0) { - break; // EOF - } - got += (size_t) r; - } - if (got < need) { - GGML_ABORT("llama_ple_disk: short read in %s: %zu of %zu bytes at %lld", - fname.c_str(), got, need, (long long) off); - } - } - - void read_row(int64_t row, uint8_t * dst, uint8_t * bounce) const { - const off_t off = (off_t) offs + (off_t) row * (off_t) rs; - if (!direct) { - pread_full(dst, rs, off, rs); - return; - } - const off_t a0 = off & ~(off_t) (block - 1); - const size_t need = (size_t) (off - a0) + rs; - const size_t len = ((need + block - 1) / block) * block; - pread_full(bounce, len, a0, need); - memcpy(dst, bounce + (off - a0), rs); + void read_row(int64_t row, uint8_t * dst, llama_bounded_file::buffer & bounce) const { + const uint64_t off = (uint64_t) offs + (uint64_t) row * rs; + file->read(off, dst, rs, bounce); } // dequantize distinct row u into the first position that asked for it, then copy it to the others @@ -208,7 +154,7 @@ struct llama_ple_disk::impl { } } - void run_items(uint8_t * bounce) { + void run_items(llama_bounded_file::buffer & bounce) { for (;;) { const size_t i0 = next.fetch_add(grain); if (i0 >= n_items) { @@ -226,8 +172,8 @@ struct llama_ple_disk::impl { } void worker() { - uint8_t * bounce = direct ? alloc_bounce() : nullptr; - uint64_t seen = 0; + llama_bounded_file::buffer bounce; + uint64_t seen = 0; for (;;) { { std::unique_lock lk(pm); @@ -237,15 +183,25 @@ struct llama_ple_disk::impl { } seen = gen; } - run_items(bounce); + std::exception_ptr error; + try { + if (direct && bounce.empty()) { + bounce = file->make_buffer(rs); + } + run_items(bounce); + } catch (...) { + error = std::current_exception(); + } { std::lock_guard lk(pm); + if (error && !worker_error) { + worker_error = error; + } if (--pending == 0) { cv_done.notify_one(); } } } - free(bounce); } // run `kind` over n items; up to serial_max of them stay on the calling thread, @@ -254,8 +210,8 @@ struct llama_ple_disk::impl { job = kind; n_items = n; grain = g; - if (direct && bounce0 == nullptr) { - bounce0 = alloc_bounce(); + if (direct && bounce0.empty()) { + bounce0 = file->make_buffer(rs); } if (n_threads <= 1 || n <= serial_max) { next = 0; @@ -270,14 +226,29 @@ struct llama_ple_disk::impl { } { std::lock_guard lk(pm); - next = 0; - pending = workers.size(); + next = 0; + pending = workers.size(); + worker_error = nullptr; ++gen; } cv_work.notify_all(); - run_items(bounce0); + // the calling thread also helps drain items, so its own read errors join the same + // worker_error slot instead of throwing early and abandoning still-running workers + std::exception_ptr error; + try { + run_items(bounce0); + } catch (...) { + error = std::current_exception(); + } std::unique_lock lk(pm); cv_done.wait(lk, [&] { return pending == 0; }); + if (!error) { + error = worker_error; + } + lk.unlock(); + if (error) { + std::rethrow_exception(error); + } } void gather(const int32_t * idx, size_t n, float * out) { @@ -338,7 +309,9 @@ struct llama_ple_disk::impl { st_uniq += n_uniq; st_hits += n_uniq - misses.size(); st_reads += misses.size(); - st_bytes += misses.size() * (direct ? block : rs); + for (const auto & m : misses) { + st_bytes += file->read_size((uint64_t) offs + (uint64_t) m.first * rs, rs); + } st_ms += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); } #endif @@ -381,11 +354,13 @@ void llama_ple_disk::prefetch(const int32_t * idx, size_t n) const { } std::sort(uniq.begin(), uniq.end()); uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end()); +#if defined(POSIX_FADV_WILLNEED) for (const int32_t row : uniq) { const off_t off = (off_t) pimpl->offs + (off_t) row * (off_t) pimpl->rs; posix_fadvise(pimpl->fd, off, (off_t) pimpl->rs, POSIX_FADV_WILLNEED); } #endif +#endif } bool llama_ple_disk::page_cached() const { return !pimpl->direct; } diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp new file mode 100644 index 000000000000..8f5946ec7900 --- /dev/null +++ b/src/models/deepseek41.cpp @@ -0,0 +1,1582 @@ +#include "llama-dsv41.h" +#include "llama-dsv41-engram.h" +#include "llama-dsv41-expert.h" +#include "llama-hparams.h" +#include "llama-memory-dsv41.h" +#include "models.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +static float dsv41_rope_attn_factor(float freq_scale) { + return 1.0f/(1.0f + 0.1f*logf(1.0f/freq_scale)); +} + +std::string llama_dsv41_graph_trace_name(const char * trace, uint32_t layer); +ggml_tensor * llama_dsv41_graph_append_zero_row(ggml_context * ctx, ggml_tensor * tensor); +ggml_tensor * llama_dsv41_graph_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type); + +std::string llama_dsv41_graph_trace_name( + const char * trace, + uint32_t layer) { + return "dsv41.trace." + std::string(trace) + ".l" + + std::to_string(layer); +} + +struct llama_model_deepseek41::engram_model { + llama_engram_layout layout; + std::array extents; +}; + +std::unique_ptr llama_model_deepseek41::create_memory_engram_runtime( + size_t max_tokens) const { + if (hparams.no_alloc) { + return nullptr; + } + return engram ? + std::make_unique(engram->layout, engram->extents, max_tokens) : + nullptr; +} + +namespace { + +struct dsv41_graph_source_input { + uint32_t layer = 0; + uint32_t ratio = 0; + uint32_t capacity = 0; + uint32_t read_width = 0; + + ggml_tensor * read_idxs = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * carry_read_idxs = nullptr; + ggml_tensor * state_read_idxs = nullptr; + ggml_tensor * state_persist_src_idxs = nullptr; + ggml_tensor * state_persist_dst_idxs = nullptr; + ggml_tensor * write_idxs = nullptr; + ggml_tensor * write_pos = nullptr; + ggml_tensor * candidate_pad_mask = nullptr; + ggml_tensor * candidate_block_bias = nullptr; + ggml_tensor * candidate_final_block = nullptr; + ggml_tensor * row_blocks = nullptr; +}; + +class dsv41_graph_input final : public llm_graph_input_i { +public: + dsv41_graph_input( + ggml_context * ctx, + const llama_cparams & cparams, + const llama_hparams & hparams, + const llama_memory_dsv41_context * mctx, + const llama_ubatch & ubatch, + uint32_t n_outputs) : + cparams(cparams), + mctx(mctx), + topology(mctx->topology(ubatch, n_outputs)) { + const llama_memory_dsv41 * memory = mctx->memory(); + const uint32_t n_tokens = ubatch.n_tokens; + const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + const llama_dsv41_memory_plan & plan = mctx->graph_plan(ubatch); + + initial_pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.dsv4_hc_mult, n_tokens); + ggml_set_input(initial_pre); + ggml_set_name(initial_pre, "dsv41.inp.carried_pre"); + + raw_persist_src_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, plan.raw.persist_src_idxs.size()); + raw_write_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, plan.raw.write_idxs.size()); + raw_read_idxs = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, hparams.n_swa, n_tokens); + raw_mask = ggml_new_tensor_4d(ctx, type_mask, hparams.n_swa, 1, 1, n_tokens); + ggml_set_input(raw_persist_src_idxs); + ggml_set_input(raw_write_idxs); + ggml_set_input(raw_read_idxs); + ggml_set_input(raw_mask); + ggml_set_name(raw_persist_src_idxs, "dsv41.inp.raw.persist_src_idxs"); + ggml_set_name(raw_write_idxs, "dsv41.inp.raw.write_idxs"); + ggml_set_name(raw_read_idxs, "dsv41.inp.raw.read_idxs"); + ggml_set_name(raw_mask, "dsv41.inp.raw.mask"); + + for (const llama_dsv41_source_plan & source_plan : plan.sources) { + dsv41_graph_source_input source; + source.layer = source_plan.source_layer; + source.ratio = source_plan.ratio; + source.capacity = source_plan.capacity; + source.read_width = source_plan.read_idxs.empty() ? 0 : + source_plan.read_idxs.size()/n_tokens; + if (source.read_width > 0) { + source.read_idxs = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, source.read_width, n_tokens); + source.mask = ggml_new_tensor_4d( + ctx, type_mask, source.read_width, 1, 1, n_tokens); + ggml_set_input(source.read_idxs); + ggml_set_input(source.mask); + } + if (!source_plan.compression.state_read_idxs.empty()) { + source.carry_read_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source.ratio); + source.state_read_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source_plan.compression.state_read_idxs.size()); + ggml_set_input(source.carry_read_idxs); + ggml_set_input(source.state_read_idxs); + } + if (!source_plan.compression.state_persist_src_idxs.empty()) { + source.state_persist_src_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source_plan.compression.state_persist_src_idxs.size()); + source.state_persist_dst_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, + source_plan.compression.state_persist_dst_idxs.size()*topology.n_seqs); + ggml_set_input(source.state_persist_src_idxs); + ggml_set_input(source.state_persist_dst_idxs); + } + if (!source_plan.compression.write_idxs.empty()) { + source.write_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, + source_plan.compression.write_idxs.size()*topology.n_seqs); + source.write_pos = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source_plan.compression.write_pos.size()); + ggml_set_input(source.write_idxs); + ggml_set_input(source.write_pos); + } + if (source.layer == hparams.dsv41_candidate_source_layer && source.read_width > 0) { + const uint32_t block_size = hparams.dsv41_candidate_block_size; + const uint32_t n_blocks = (source.read_width + block_size - 1)/block_size; + const uint32_t padded = n_blocks*block_size; + source.candidate_pad_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, padded, n_tokens); + source.candidate_block_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_blocks, n_tokens); + source.candidate_final_block = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, n_tokens); + source.row_blocks = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, source.read_width); + ggml_set_input(source.candidate_pad_mask); + ggml_set_input(source.candidate_block_bias); + ggml_set_input(source.candidate_final_block); + ggml_set_input(source.row_blocks); + } + sources.emplace(source.layer, source); + } + + if (topology.engram_enabled) { + for (uint32_t index = 0; index < LLAMA_ENGRAM_LAYERS; ++index) { + engram_rows[index] = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, n_tokens); + engram_select[index] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + engram_row_ids[index] = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, LLAMA_ENGRAM_COLS, n_tokens); + ggml_set_input(engram_rows[index]); + ggml_set_input(engram_select[index]); + ggml_set_input(engram_row_ids[index]); + const uint32_t layer = index == 0 ? 1 : 14; + ggml_format_name( + engram_row_ids[index], + "dsv41.inp.engram.row_ids.l%u", layer); + } + } + + GGML_UNUSED(memory); + } + + void set_input(const llama_ubatch * ubatch) override { + const llama_dsv41_memory_plan & plan = mctx->plan(); + const llama_memory_dsv41 * memory = mctx->memory(); + const uint32_t n_tokens = ubatch->n_tokens; + const uint32_t n_streams = plan.token_seq_ids.empty() ? 1 : plan.token_seq_ids.front().size(); + if (n_streams != 1) { + throw std::runtime_error( + "DeepSeek V4.1 graph supports one active sequence per ubatch"); + } + + std::vector pre((size_t) LLAMA_DSV41_HC_MULT*n_tokens, 0.0f); + for (uint32_t token = 0; token < n_tokens; ++token) { + pre[(size_t) token*LLAMA_DSV41_HC_MULT] = 1.0f; + } + set_tensor(initial_pre, pre); + set_tensor(raw_persist_src_idxs, plan.raw.persist_src_idxs); + set_tensor(raw_write_idxs, plan.raw.write_idxs); + + std::vector raw_idxs = plan.raw.read_idxs; + const int32_t raw_cache_rows = + (int32_t) (memory->config().raw_window*memory->config().n_seq); + const int32_t raw_sentinel = raw_cache_rows + n_tokens; + const llama_pos batch_start = plan.positions.front(); + for (uint32_t token = 0; token < n_tokens; ++token) { + const uint32_t visible = plan.raw.n_visible[token]; + const llama_pos first = plan.positions[token] + 1 - visible; + for (uint32_t row = 0; row < visible; ++row) { + const llama_pos pos = first + row; + if (pos >= batch_start) { + raw_idxs[(size_t) token*memory->config().raw_window + row] = + raw_cache_rows + pos - batch_start; + } + } + } + std::replace(raw_idxs.begin(), raw_idxs.end(), -1, raw_sentinel); + set_tensor(raw_read_idxs, raw_idxs); + set_mask(raw_mask, plan.raw.mask); + + for (const llama_dsv41_source_plan & source_plan : plan.sources) { + dsv41_graph_source_input & source = sources.at(source_plan.source_layer); + if (source.read_idxs != nullptr) { + std::vector read_idxs = source_plan.read_idxs; + const int32_t sentinel = + (int32_t) (source_plan.capacity*memory->config().n_seq); + std::replace(read_idxs.begin(), read_idxs.end(), -1, sentinel); + set_tensor(source.read_idxs, read_idxs); + std::vector mask(read_idxs.size()); + std::transform( + source_plan.read_idxs.begin(), + source_plan.read_idxs.end(), + mask.begin(), + [](int32_t idx) { + return idx >= 0 ? 0.0f : -std::numeric_limits::infinity(); + }); + set_mask(source.mask, mask); + } + if (source.state_read_idxs != nullptr) { + std::vector carry_read_idxs(source.ratio); + const llama_seq_id seq_id = plan.token_seq_ids.front().front(); + for (uint32_t row = 0; row < source.ratio; ++row) { + carry_read_idxs[row] = seq_id*source.ratio + row; + } + set_tensor(source.carry_read_idxs, carry_read_idxs); + set_tensor(source.state_read_idxs, source_plan.compression.state_read_idxs); + } + if (source.state_persist_src_idxs != nullptr) { + set_tensor(source.state_persist_src_idxs, source_plan.compression.state_persist_src_idxs); + std::vector dst; + dst.reserve(source_plan.compression.state_persist_dst_idxs.size()*n_streams); + for (int32_t row : source_plan.compression.state_persist_dst_idxs) { + for (llama_seq_id seq_id : plan.token_seq_ids.front()) { + dst.push_back((int64_t) seq_id*source.ratio + row); + } + } + set_tensor(source.state_persist_dst_idxs, dst); + } + if (source.write_idxs != nullptr) { + std::vector write_idxs; + write_idxs.reserve(source_plan.compression.write_idxs.size()*n_streams); + for (int64_t row : source_plan.compression.write_idxs) { + for (llama_seq_id seq_id : plan.token_seq_ids.front()) { + write_idxs.push_back((int64_t) seq_id*source.capacity + row); + } + } + set_tensor(source.write_idxs, write_idxs); + set_tensor(source.write_pos, source_plan.compression.write_pos); + } + if (source.candidate_pad_mask != nullptr) { + const uint32_t block_size = memory->config().candidate_block_size; + const uint32_t n_blocks = + (source.read_width + block_size - 1)/block_size; + const uint32_t padded = n_blocks*block_size; + std::vector pad((size_t) padded*n_tokens, 0.0f); + for (uint32_t token = 0; token < n_tokens; ++token) { + for (uint32_t row = source.read_width; row < padded; ++row) { + pad[(size_t) token*padded + row] = + -std::numeric_limits::infinity(); + } + } + set_tensor(source.candidate_pad_mask, pad); + + std::vector bias((size_t) n_blocks*n_tokens, 0.0f); + std::vector final_blocks(n_tokens); + for (uint32_t token = 0; token < n_tokens; ++token) { + const uint32_t visible = + source_plan.compression.n_visible[token]; + const uint32_t visible_blocks = + (visible + block_size - 1)/block_size; + if (visible_blocks == 0) { + throw std::runtime_error( + "DeepSeek V4.1 candidate source has no visible block"); + } + for (uint32_t block = visible_blocks; block < n_blocks; ++block) { + bias[(size_t) token*n_blocks + block] = + -std::numeric_limits::infinity(); + } + final_blocks[token] = (int32_t) visible_blocks - 1; + } + set_tensor(source.candidate_block_bias, bias); + set_tensor(source.candidate_final_block, final_blocks); + + std::vector row_blocks(source.read_width); + for (uint32_t row = 0; row < source.read_width; ++row) { + row_blocks[row] = row/block_size; + } + set_tensor(source.row_blocks, row_blocks); + } + } + + const llama_dsv41_engram_transaction * transaction = mctx->engram_transaction(); + for (uint32_t index = 0; index < LLAMA_ENGRAM_LAYERS; ++index) { + if (engram_rows[index] == nullptr) { + continue; + } + if (transaction == nullptr) { + std::vector zeros( + (size_t) LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM*n_tokens); + std::vector select(n_tokens); + std::vector ids((size_t) LLAMA_ENGRAM_COLS*n_tokens); + set_tensor(engram_rows[index], zeros); + set_tensor(engram_select[index], select); + set_tensor(engram_row_ids[index], ids); + continue; + } + transaction->upload_layer( + index, 0, n_tokens, engram_rows[index], engram_select[index]); + set_tensor( + engram_row_ids[index], + mctx->engram_row_ids(index)); + } + } + + bool can_reuse(const llm_graph_params & params) override { + const auto * next = static_cast(params.mctx); + const llama_dsv41_graph_topology next_topology = + next->topology(params.ubatch, params.n_outputs); + if (!topology.same_topology(next_topology)) { + return false; + } + mctx = next; + topology = next_topology; + return true; + } + + dsv41_graph_source_input * source(uint32_t layer) { + const auto found = sources.find(layer); + return found == sources.end() ? nullptr : &found->second; + } + + const llama_cparams cparams; + const llama_memory_dsv41_context * mctx; + llama_dsv41_graph_topology topology; + ggml_tensor * initial_pre = nullptr; + ggml_tensor * raw_persist_src_idxs = nullptr; + ggml_tensor * raw_write_idxs = nullptr; + ggml_tensor * raw_read_idxs = nullptr; + ggml_tensor * raw_mask = nullptr; + std::map sources; + std::array engram_rows = {}; + std::array engram_select = {}; + std::array engram_row_ids = {}; + +private: + template + static void set_tensor(ggml_tensor * tensor, const std::vector & values) { + if (tensor == nullptr || tensor->buffer == nullptr) { + return; + } + if (ggml_nelements(tensor) != (int64_t) values.size()) { + throw std::runtime_error("DeepSeek V4.1 graph input shape changed"); + } + ggml_backend_tensor_set(tensor, values.data(), 0, values.size()*sizeof(T)); + } + + static void set_mask(ggml_tensor * tensor, const std::vector & values) { + if (tensor == nullptr || tensor->buffer == nullptr) { + return; + } + if (ggml_nelements(tensor) != (int64_t) values.size()) { + throw std::runtime_error("DeepSeek V4.1 graph mask shape changed"); + } + if (tensor->type == GGML_TYPE_F16) { + std::vector converted(values.size()); + ggml_fp32_to_fp16_row(values.data(), converted.data(), values.size()); + ggml_backend_tensor_set( + tensor, converted.data(), 0, converted.size()*sizeof(ggml_fp16_t)); + } else { + ggml_backend_tensor_set( + tensor, values.data(), 0, values.size()*sizeof(float)); + } + } +}; + +static ggml_tensor * dsv41_flatten_memory( + ggml_context * ctx, + ggml_tensor * tensor) { + return ggml_reshape_2d( + ctx, tensor, tensor->ne[0], tensor->ne[1]*tensor->ne[2]); +} + +static ggml_tensor * dsv41_append_zero_row( + ggml_context * ctx, + ggml_tensor * tensor) { + const ggml_type type = tensor->type; + ggml_tensor * row = ggml_view_2d( + ctx, tensor, tensor->ne[0], 1, tensor->nb[1], 0); + if (row->type != GGML_TYPE_F32) { + row = ggml_cast(ctx, row, GGML_TYPE_F32); + } + row = ggml_scale(ctx, row, 0.0f); + if (type != GGML_TYPE_F32) { + row = ggml_cast(ctx, row, type); + } + return ggml_concat(ctx, tensor, row, 1); +} + +static ggml_tensor * dsv41_hc_mean( + ggml_context * ctx, + ggml_tensor * streams) { + ggml_tensor * result = ggml_view_2d( + ctx, streams, streams->ne[0], streams->ne[2], streams->nb[2], 0); + for (int64_t stream = 1; stream < streams->ne[1]; ++stream) { + result = ggml_add(ctx, result, ggml_view_2d( + ctx, streams, streams->ne[0], streams->ne[2], + streams->nb[2], stream*streams->nb[1])); + } + return ggml_scale(ctx, result, 1.0f/streams->ne[1]); +} + +static ggml_tensor * dsv41_view_1d( + ggml_context * ctx, + ggml_tensor * tensor, + int64_t ne0, + int64_t offset) { + return ggml_view_1d(ctx, tensor, ne0, ggml_row_size(tensor->type, offset)); +} + +static ggml_tensor * dsv41_view_2d( + ggml_context * ctx, + ggml_tensor * tensor, + int64_t ne0, + int64_t ne1, + int64_t offset) { + return ggml_view_2d( + ctx, tensor, ne0, ne1, tensor->nb[1], + ggml_row_size(tensor->type, offset)); +} + +static ggml_tensor * dsv41_sort_row_ids( + ggml_context * ctx, + ggml_tensor * ids) { + ggml_tensor * order = ggml_argsort( + ctx, ggml_cast(ctx, ids, GGML_TYPE_F32), + GGML_SORT_ORDER_ASC); + ggml_tensor * sorted = ggml_get_rows( + ctx, + ggml_reshape_3d(ctx, ids, 1, ids->ne[0], ids->ne[1]), + order); + return ggml_cont( + ctx, + ggml_reshape_2d(ctx, sorted, ids->ne[0], ids->ne[1])); +} + +} + +void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { + llama_dsv41_config config = {}; + std::string raw_config; + ml.get_key(LLM_KV_DSV41_CONFIG, raw_config); + ml.get_key(LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, config.n_ctx_train); + ml.get_key(LLM_KV_DSV41_HIDDEN_SIZE, config.n_embd); + ml.get_key(LLM_KV_DSV41_NUM_HIDDEN_LAYERS, config.n_layer); + ml.get_key(LLM_KV_DSV41_VOCAB_SIZE, config.n_vocab); + ml.get_key(LLM_KV_DSV41_NUM_ATTENTION_HEADS, config.n_head); + ml.get_key(LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, config.n_head_kv); + ml.get_key(LLM_KV_DSV41_HEAD_DIM, config.n_head_dim); + ml.get_key(LLM_KV_DSV41_QK_ROPE_HEAD_DIM, config.n_rot); + ml.get_key(LLM_KV_DSV41_Q_LORA_RANK, config.n_lora_q); + ml.get_key(LLM_KV_DSV41_O_LORA_RANK, config.n_lora_o); + ml.get_key(LLM_KV_DSV41_O_GROUPS, config.n_o_group); + ml.get_key(LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, config.n_ff_expert); + ml.get_key(LLM_KV_DSV41_N_ROUTED_EXPERTS, config.n_expert); + ml.get_key(LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, config.n_expert_used); + ml.get_key(LLM_KV_DSV41_N_SHARED_EXPERTS, config.n_expert_shared); + ml.get_key(LLM_KV_DSV41_INDEX_N_HEADS, config.indexer_n_head); + ml.get_key(LLM_KV_DSV41_INDEX_HEAD_DIM, config.indexer_head_size); + ml.get_key(LLM_KV_DSV41_INDEX_TOPK, config.indexer_top_k); + ml.get_key(LLM_KV_DSV41_HC_MULT, config.hc_count); + ml.get_key(LLM_KV_DSV41_HC_SINKHORN_ITERS, config.hc_sinkhorn_iters); + ml.get_key(LLM_KV_DSV41_SLIDING_WINDOW, config.raw_window); + ml.get_key(LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, config.candidate_source_layer); + ml.get_key(LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, config.candidate_topk_blocks); + ml.get_key(LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, config.candidate_block_size); + ml.get_key(LLM_KV_DSV41_RMS_NORM_EPS, config.f_norm_rms_eps); + ml.get_key(LLM_KV_DSV41_HC_EPS, config.hc_eps); + ml.get_key(LLM_KV_DSV41_SWIGLU_LIMIT, config.swiglu_clamp); + ml.get_key(LLM_KV_DSV41_ROUTED_SCALING_FACTOR, config.routed_scale); + ml.get_key(LLM_KV_DSV41_ROPE_THETA, config.rope_theta); + ml.get_key(LLM_KV_DSV41_COMPRESS_ROPE_THETA, config.compress_rope_theta); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_FACTOR, config.yarn_factor); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, config.yarn_beta_fast); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, config.yarn_beta_slow); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, config.yarn_original_context); + ml.get_key(LLM_KV_DSV41_NORM_TOPK_PROB, config.expert_weights_norm); + ml.get_key(LLM_KV_DSV41_HIDDEN_ACT, config.hidden_act); + ml.get_key(LLM_KV_DSV41_SCORING_FUNC, config.scoring_func); + ml.get_key(LLM_KV_DSV41_TOPK_METHOD, config.topk_method); + ml.get_arr(LLM_KV_DSV41_COMPRESS_RATIOS, config.compress_ratios); + ml.get_arr(LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, config.kv_sources); + ml.get_arr(LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, config.index_sources); + ml.get_key(LLM_KV_DSV41_ENGRAM_ENCODING, config.engram_encoding); + ml.get_arr(LLM_KV_DSV41_ENGRAM_LAYER_IDS, config.engram_layers); + ml.get_arr(LLM_KV_DSV41_ENGRAM_ROWS, config.engram_rows); + ml.get_key(LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, config.engram_compressed_vocab_size); + ml.get_key(LLM_KV_DSV41_ENGRAM_PAD_ID, config.engram_pad_id); + ml.get_arr_n(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, config.engram_token_map_size); + ml.get_arr_n(LLM_KV_DSV41_ENGRAM_PRIMES, config.engram_primes_size); + ml.get_arr_n(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, config.engram_multipliers_size); + ml.get_arr(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, config.engram_token_map); + ml.get_arr(LLM_KV_DSV41_ENGRAM_PRIMES, config.engram_primes); + ml.get_arr(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, config.engram_multipliers); + + config.n_ff_dense = LLAMA_DSV41_N_FF_DENSE; + llama_dsv41_validate_config(config); + engram = std::make_shared(); + engram->layout = llama_dsv41_make_engram_layout(config); + + if (raw_config.empty()) { + throw std::runtime_error("DeepSeek V4.1 metadata: config must not be empty"); + } + hparams.n_ctx_train = config.n_ctx_train; + hparams.n_embd = config.n_embd; + hparams.n_embd_out_impl = config.n_embd; + hparams.n_layer_all = config.n_layer; + hparams.n_layer_nextn = 0; + hparams.n_expert = config.n_expert; + hparams.n_expert_shared = config.n_expert_shared; + hparams.n_lora_q = config.n_lora_q; + hparams.n_ff_shexp = config.n_ff_expert; + hparams.n_embd_head_k_full = config.n_head_dim; + hparams.n_embd_head_v_full = config.n_head_dim; + hparams.n_embd_head_k_swa = config.n_head_dim; + hparams.n_embd_head_v_swa = config.n_head_dim; + hparams.n_rot_full = config.n_rot; + hparams.n_rot_swa = config.n_rot; + hparams.n_swa = config.raw_window; + hparams.indexer_n_head = config.indexer_n_head; + hparams.indexer_head_size = config.indexer_head_size; + hparams.indexer_top_k = config.indexer_top_k; + hparams.dsv4_o_group_count = config.n_o_group; + hparams.dsv4_o_lora_rank = config.n_lora_o; + hparams.dsv4_hc_mult = config.hc_count; + hparams.dsv4_hc_sinkhorn_iters = config.hc_sinkhorn_iters; + hparams.dsv4_compress_rope_base = (float) config.compress_rope_theta; + hparams.dsv4_hc_eps = config.hc_eps; + hparams.dsv41_candidate_source_layer = config.candidate_source_layer; + hparams.dsv41_candidate_topk_blocks = config.candidate_topk_blocks; + hparams.dsv41_candidate_block_size = config.candidate_block_size; + hparams.f_norm_rms_eps = config.f_norm_rms_eps; + hparams.expert_weights_scale = config.routed_scale; + hparams.expert_weights_norm = config.expert_weights_norm; + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS; + hparams.rope_freq_base_train = (float) config.rope_theta; + hparams.rope_freq_base_train_swa = (float) config.rope_theta; + hparams.rope_freq_scale_train = 1.0f/config.yarn_factor; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + hparams.n_ctx_orig_yarn = (uint32_t) config.yarn_original_context; + hparams.yarn_beta_fast = config.yarn_beta_fast; + hparams.yarn_beta_slow = config.yarn_beta_slow; + hparams.yarn_ext_factor = 1.0f; + hparams.rope_attn_factor = dsv41_rope_attn_factor(hparams.rope_freq_scale_train); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + hparams.causal_attn = true; + + for (uint32_t il = 0; il < config.n_layer; ++il) { + hparams.n_head_arr[il] = config.n_head; + hparams.n_head_kv_arr[il] = config.n_head_kv; + hparams.n_ff_arr[il] = config.n_ff_dense; + hparams.n_ff_exp_arr[il] = config.n_ff_expert; + hparams.n_expert_used_arr[il] = config.n_expert_used; + hparams.swiglu_clamp_exp[il] = config.swiglu_clamp; + hparams.swiglu_clamp_shexp[il] = config.swiglu_clamp; + hparams.dsv4_compress_ratios[il] = config.compress_ratios[il]; + hparams.dsv41_kv_source_layer[il] = llama_dsv41_kv_source_layer(il); + hparams.dsv41_index_source_layer[il] = llama_dsv41_index_source_layer(il); + hparams.is_swa_impl[il] = 1; + } + for (uint32_t il : config.engram_layers) { + hparams.dsv41_engram_layers.set(il); + } + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t n_ff_exp = hparams.n_ff_exp(); + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_embd_head = hparams.n_embd_head_k(); + const int64_t o_groups = hparams.dsv4_o_group_count; + const int64_t o_lora_rank = hparams.dsv4_o_lora_rank; + const int64_t hc_mult = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc_mult*n_embd; + const int64_t hc_mix_dim = (2 + hc_mult)*hc_mult; + + const std::vector expert_tensors = + llama_dsv41_register_expert_tensors([&](const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + return ml.register_external_tensor(name, layer, projection, ne); + }); + if (params.expert_cache_bytes == 0 || params.expert_cache_slots <= 0) { + throw std::runtime_error( + "DeepSeek V4.1 requires non-zero expert_cache_bytes and expert_cache_slots before tensor allocation"); + } + llama_dsv41_expert_runtime_params expert_params; + expert_params.cache_bytes = params.expert_cache_bytes; + expert_params.cache_slots = params.expert_cache_slots; + expert_params.direct_io = true; + expert_params.allow_buffered_io = false; + expert_params.no_alloc = ml.no_alloc; + experts = std::make_shared( + expert_tensors, + expert_params, + [this](const llama_expert_store_tensor & tensor) { + return select_moe_buft( + tensor.layer, tensor.type, tensor.ne[0], tensor.ne[1], params.expert_cache_slots); + }); + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, 0); + + for (int32_t il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0); + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", il), { n_head }, 0); + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, 0); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head*n_embd_head }, 0); + layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", il), { n_embd, n_embd_head }, 0); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { n_embd_head }, 0); + layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", il), { n_head*n_embd_head/o_groups, o_lora_rank, o_groups }, TENSOR_ALLOW_RESHAPE); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", il), { o_groups*o_lora_rank, n_embd }, 0); + + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", il), { hc_dim, hc_mix_dim }, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", il), { hc_mix_dim }, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", il), { 3 }, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", il), { hc_dim, hc_mix_dim }, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", il), { hc_mix_dim }, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", il), { 3 }, 0); + + if (hparams.dsv41_is_kv_source(il)) { + layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", il), { n_embd, n_embd_head }, 0); + layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", il), { n_embd_head }, 0); + if (hparams.dsv4_compress_ratios[il] == 2) { + layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", il), { n_embd, n_embd_head }, 0); + } + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", il), { n_embd_head, hparams.indexer_head_size }, 0); + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { hparams.indexer_head_size }, 0); + } + if (hparams.dsv41_is_index_source(il)) { + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", il), { n_embd, hparams.indexer_n_head }, 0); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", il), { q_lora_rank, hparams.indexer_n_head*hparams.indexer_head_size }, 0); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, 0); + layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", il), { n_expert }, TENSOR_NOT_REQUIRED); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, 0); + layer.ffn_gate_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_GATE); + layer.ffn_down_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_DOWN); + layer.ffn_up_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_UP); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_exp*n_expert_shared, n_embd }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); + + if (hparams.dsv41_engram_layers.test(il)) { + const size_t index = il == (int32_t) engram->layout.layer_ids[0] ? 0 : 1; + const std::string table_name = tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(); + const auto * table = ml.get_weight(table_name.c_str()); + if (table == nullptr) { + throw std::runtime_error("DeepSeek V4.1 is missing required Engram tensor " + table_name); + } + llama_dsv41_engram_extent & extent = engram->extents[index]; + extent.fname = ml.no_alloc ? "(no_alloc)" : ml.fnames.at(table->idx); + extent.offset = table->offs; + extent.rows = engram->layout.rows[index]; + extent.columns = table->tensor->ne[0]; + extent.row_count = table->tensor->ne[1]; + extent.type = table->tensor->type; + llama_dsv41_validate_engram_extent(extent); + + create_tensor( + tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il), + { LLAMA_ENGRAM_ROW_BYTES, (int64_t) extent.rows }, + TENSOR_SKIP); + layer.engram_q_norm = create_tensor( + tn(LLM_TENSOR_ENGRAM_Q_NORM, "weight", il), + { n_embd, hc_mult }, + 0); + layer.engram_k_norm = create_tensor( + tn(LLM_TENSOR_ENGRAM_K_NORM, "weight", il), + { n_embd, hc_mult }, + 0); + layer.engram_kv = create_tensor( + tn(LLM_TENSOR_ENGRAM_KV, "weight", il), + { LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, (hc_mult + 1)*n_embd }, + 0); + } + } + +} + +bool llama_model_deepseek41::requires_synchronous_graph() const { + return experts != nullptr; +} + +std::string llama_model_deepseek41::consume_runtime_error() const { + return experts ? experts->consume_error() : std::string(); +} + +void llama_model_deepseek41::release_runtime_work() const { + if (experts) { + experts->release_all(); + } +} + +void llama_model_deepseek41::release_runtime_work_after_sync(ggml_backend_sched_t sched) const { + if (experts) { + experts->release_all_after_sync(sched); + } +} + +void llama_model_deepseek41::acquire_runtime_context() const { + if (experts) { + experts->acquire_context(); + } +} + +void llama_model_deepseek41::release_runtime_context() const { + if (experts) { + experts->release_context(); + } +} + +namespace { + +struct dsv41_hc_mix { + ggml_tensor * pre; + ggml_tensor * post; + ggml_tensor * comb; +}; + +static dsv41_hc_mix dsv41_build_hc_mix( + const llama_model_deepseek41::graph & graph, + ggml_tensor * streams, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + int il) { + const int64_t hc = graph.hparams.dsv4_hc_mult; + const int64_t nt = streams->ne[2]; + ggml_tensor * flat = ggml_reshape_2d(graph.ctx0, streams, graph.n_embd*hc, nt); + flat = ggml_rms_norm(graph.ctx0, flat, graph.norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(graph.ctx0, hc_fn, flat); + graph.cb(mixes, "hc_mixes", il); + + ggml_tensor * scale_pre = dsv41_view_1d(graph.ctx0, hc_scale, 1, 0); + ggml_tensor * scale_post = dsv41_view_1d(graph.ctx0, hc_scale, 1, 1); + ggml_tensor * scale_comb = dsv41_view_1d(graph.ctx0, hc_scale, 1, 2); + ggml_tensor * base_pre = dsv41_view_1d(graph.ctx0, hc_base, hc, 0); + ggml_tensor * base_post = dsv41_view_1d(graph.ctx0, hc_base, hc, hc); + ggml_tensor * base_comb = dsv41_view_1d(graph.ctx0, hc_base, hc*hc, 2*hc); + + ggml_tensor * pre = dsv41_view_2d(graph.ctx0, mixes, hc, nt, 0); + pre = ggml_add(graph.ctx0, ggml_mul(graph.ctx0, pre, scale_pre), base_pre); + pre = ggml_scale_bias( + graph.ctx0, ggml_sigmoid(graph.ctx0, pre), 1.0f, + graph.hparams.dsv4_hc_eps); + graph.cb(pre, "hc_pre", il); + + ggml_tensor * post = dsv41_view_2d(graph.ctx0, mixes, hc, nt, hc); + post = ggml_add(graph.ctx0, ggml_mul(graph.ctx0, post, scale_post), base_post); + post = ggml_scale(graph.ctx0, ggml_sigmoid(graph.ctx0, post), 2.0f); + graph.cb(post, "hc_post", il); + + ggml_tensor * comb = nullptr; + if (graph.cparams.fused_dsv4_hc_comb) { + comb = ggml_dsv4_hc_comb( + graph.ctx0, mixes, hc_scale, hc_base, + graph.hparams.dsv4_hc_eps, + (int32_t) graph.hparams.dsv4_hc_sinkhorn_iters); + graph.res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, comb, il}); + } else { + comb = dsv41_view_2d(graph.ctx0, mixes, hc*hc, nt, 2*hc); + comb = ggml_add(graph.ctx0, ggml_mul(graph.ctx0, comb, scale_comb), base_comb); + comb = ggml_reshape_3d(graph.ctx0, comb, hc, hc, nt); + comb = graph.build_hc_sinkhorn(comb, il); + } + graph.cb(comb, "hc_comb", il); + return { pre, post, comb }; +} + +static ggml_tensor * dsv41_cast_for_store( + ggml_context * ctx, + ggml_tensor * source, + const ggml_tensor * destination) { + return source->type == destination->type ? + source : ggml_cast(ctx, source, destination->type); +} + +static ggml_tensor * dsv41_concat( + ggml_context * ctx, + ggml_tensor * a, + ggml_tensor * b, + int dim, + const char * label) { + if (a == nullptr || b == nullptr || a->type != b->type) { + throw std::runtime_error(format( + "DeepSeek V4.1 %s concat type mismatch: %s and %s", + label, + a == nullptr ? "null" : ggml_type_name(a->type), + b == nullptr ? "null" : ggml_type_name(b->type))); + } + return ggml_concat(ctx, a, b, dim); +} + +static ggml_tensor * dsv41_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type) { + ggml_tensor * marker = ggml_view_1d(ctx, dependency, 1, 0); + marker = ggml_argsort_top_k(ctx, marker, 1); + marker = ggml_cast(ctx, marker, GGML_TYPE_F32); + marker = ggml_scale(ctx, marker, 0.0f); + return type == GGML_TYPE_F32 ? marker : ggml_cast(ctx, marker, type); +} + +static ggml_tensor * dsv41_build_candidate_mask( + const llama_model_deepseek41::graph & graph, + ggml_tensor * candidate_blocks, + const dsv41_graph_source_input & source) { + const int64_t n_blocks = source.candidate_block_bias->ne[0]; + const int64_t n_tokens = candidate_blocks->ne[1]; + ggml_tensor * all = ggml_new_tensor_3d( + graph.ctx0, GGML_TYPE_F32, 1, n_blocks, n_tokens); + all = ggml_fill(graph.ctx0, all, -INFINITY); + ggml_tensor * zeros = ggml_scale( + graph.ctx0, ggml_cast(graph.ctx0, candidate_blocks, GGML_TYPE_F32), 0.0f); + zeros = ggml_reshape_3d( + graph.ctx0, zeros, 1, candidate_blocks->ne[0], n_tokens); + ggml_tensor * blocks = ggml_set_rows(graph.ctx0, all, zeros, candidate_blocks); + blocks = ggml_reshape_2d(graph.ctx0, blocks, n_blocks, n_tokens); + blocks = ggml_cont(graph.ctx0, ggml_transpose(graph.ctx0, blocks)); + ggml_tensor * rows = ggml_get_rows(graph.ctx0, blocks, source.row_blocks); + return ggml_cont(graph.ctx0, ggml_transpose(graph.ctx0, rows)); +} + +static ggml_tensor * dsv41_build_index_selection( + const llama_model_deepseek41::graph & graph, + const llama_model & model, + dsv41_graph_input & input, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * & candidate_blocks, + ggml_tensor * & selected_local, + int il) { + const auto & layer = model.layers[il]; + const uint32_t source_layer = graph.hparams.dsv41_kv_source_layer[il]; + dsv41_graph_source_input * source = input.source(source_layer); + if (source == nullptr || source->read_width == 0) { + return nullptr; + } + + const llama_memory_dsv41 * memory = input.mctx->memory(); + ggml_tensor * index_cache = dsv41_flatten_memory( + graph.ctx0, memory->index_keys(source_layer)); + index_cache = dsv41_append_zero_row(graph.ctx0, index_cache); + ggml_tensor * index_k = ggml_get_rows( + graph.ctx0, index_cache, + ggml_reshape_1d( + graph.ctx0, source->read_idxs, + source->read_width*graph.n_tokens)); + index_k = ggml_reshape_3d( + graph.ctx0, index_k, + graph.hparams.indexer_head_size, + source->read_width, graph.n_tokens); + graph.cb(index_k, "dsv41_index_k", il); + + ggml_tensor * index_q = graph.build_lora_mm(layer.indexer_attn_q_b, qr); + index_q = ggml_reshape_3d( + graph.ctx0, index_q, + graph.hparams.indexer_head_size, + graph.hparams.indexer_n_head, + graph.n_tokens); + index_q = ggml_rope_ext( + graph.ctx0, index_q, inp_pos, nullptr, graph.hparams.n_rot(), + graph.rope_type, graph.n_ctx_orig, + graph.hparams.dsv4_compress_rope_base, + graph.freq_scale, graph.ext_factor, + dsv41_rope_attn_factor(graph.freq_scale), + graph.beta_fast, graph.beta_slow); + index_q = ggml_rope_set_offset( + index_q, graph.hparams.indexer_head_size - graph.hparams.n_rot()); + graph.cb(index_q, "dsv41_index_q", il); + + ggml_tensor * weights = graph.build_lora_mm(layer.indexer_proj, cur); + weights = ggml_scale( + graph.ctx0, weights, + 1.0f/std::sqrt( + (float) (graph.hparams.indexer_head_size*graph.hparams.indexer_n_head))); + weights = ggml_reshape_3d( + graph.ctx0, weights, graph.hparams.indexer_n_head, 1, graph.n_tokens); + + ggml_tensor * scores = ggml_mul_mat(graph.ctx0, index_k, index_q); + ggml_prec_set_acc(scores, GGML_PREC_F32); + scores = ggml_relu(graph.ctx0, scores); + scores = ggml_cont(graph.ctx0, ggml_permute(graph.ctx0, scores, 1, 0, 2, 3)); + scores = ggml_mul(graph.ctx0, scores, weights); + scores = ggml_sum_rows(graph.ctx0, scores); + scores = ggml_reshape_2d(graph.ctx0, scores, source->read_width, graph.n_tokens); + scores = ggml_add( + graph.ctx0, scores, + ggml_reshape_2d( + graph.ctx0, source->mask, source->read_width, graph.n_tokens)); + + if (candidate_blocks != nullptr) { + scores = ggml_add( + graph.ctx0, scores, + dsv41_build_candidate_mask(graph, candidate_blocks, *source)); + } + graph.cb(scores, "dsv41_index_scores", il); + + if ((uint32_t) il == graph.hparams.dsv41_candidate_source_layer) { + const uint32_t block_size = graph.hparams.dsv41_candidate_block_size; + const int64_t padded = source->candidate_pad_mask->ne[0]; + ggml_tensor * padded_scores = ggml_pad( + graph.ctx0, scores, padded - source->read_width, 0, 0, 0); + padded_scores = ggml_add( + graph.ctx0, padded_scores, source->candidate_pad_mask); + ggml_tensor * block_scores = ggml_pool_1d( + graph.ctx0, padded_scores, GGML_OP_POOL_MAX, + block_size, block_size, 0); + block_scores = ggml_add( + graph.ctx0, block_scores, source->candidate_block_bias); + graph.cb(block_scores, "dsv41_candidate_scores", il); + + const uint32_t n_candidate = input.topology.candidate_width; + candidate_blocks = llama_dsv41_build_candidate_blocks( + graph.ctx0, block_scores, + source->candidate_final_block, n_candidate); + ggml_set_name( + candidate_blocks, + llama_dsv41_graph_trace_name( + "attn.candidate_blocks", il).c_str()); + ggml_build_forward_expand(graph.gf, candidate_blocks); + + ggml_tensor * score_store = memory->candidate_scores(); + ggml_tensor * score_view = ggml_view_2d( + graph.ctx0, score_store, block_scores->ne[0], graph.n_tokens, + score_store->nb[1], 0); + ggml_build_forward_expand( + graph.gf, ggml_cpy(graph.ctx0, block_scores, score_view)); + ggml_tensor * id_store = memory->candidate_ids(); + ggml_tensor * id_view = ggml_view_2d( + graph.ctx0, id_store, candidate_blocks->ne[0], graph.n_tokens, + id_store->nb[1], 0); + ggml_build_forward_expand( + graph.gf, ggml_cpy(graph.ctx0, candidate_blocks, id_view)); + } + + const uint32_t n_top_k = std::min( + source->read_width, graph.hparams.indexer_top_k); + selected_local = ggml_cont( + graph.ctx0, + ggml_argsort_top_k(graph.ctx0, scores, n_top_k)); + selected_local = dsv41_sort_row_ids( + graph.ctx0, selected_local); + ggml_tensor * selected = ggml_get_rows( + graph.ctx0, + ggml_reshape_3d( + graph.ctx0, source->read_idxs, 1, + source->read_width, graph.n_tokens), + selected_local); + selected = ggml_cont( + graph.ctx0, + ggml_reshape_2d( + graph.ctx0, selected, + n_top_k, graph.n_tokens)); + if ((uint32_t) il > graph.hparams.dsv41_candidate_source_layer) { + ggml_set_name( + selected, + llama_dsv41_graph_trace_name( + "attn.candidates", il).c_str()); + ggml_build_forward_expand(graph.gf, selected); + } + return selected; +} + +static ggml_tensor * dsv41_build_attention( + const llama_model_deepseek41::graph & graph, + const llama_model & model, + dsv41_graph_input & input, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * & selected, + ggml_tensor * & selected_local, + ggml_tensor * & candidate_blocks, + int il) { + const auto & layer = model.layers[il]; + const int64_t n_embd_head = graph.hparams.n_embd_head_k(); + const int64_t n_rot = graph.hparams.n_rot(); + const int64_t n_nope = n_embd_head - n_rot; + const int64_t ratio = graph.hparams.dsv4_compress_ratios[il]; + const int64_t n_groups = graph.hparams.dsv4_o_group_count; + const int64_t heads_per_group = graph.n_head/n_groups; + const int64_t group_width = heads_per_group*n_embd_head; + const int64_t o_lora_rank = graph.hparams.dsv4_o_lora_rank; + const float rope_base = ratio == 0 ? + graph.freq_base : graph.hparams.dsv4_compress_rope_base; + const float rope_scale = ratio == 0 ? 1.0f : graph.freq_scale; + const float rope_ext = ratio == 0 ? 0.0f : graph.ext_factor; + const float rope_attn = ratio == 0 ? + 1.0f : dsv41_rope_attn_factor(graph.freq_scale); + const float rope_beta_fast = ratio == 0 ? 0.0f : graph.beta_fast; + const float rope_beta_slow = ratio == 0 ? 0.0f : graph.beta_slow; + const int32_t rope_ctx = ratio == 0 ? 0 : graph.n_ctx_orig; + + ggml_tensor * qr = graph.build_lora_mm(layer.wq_a, cur); + qr = graph.build_norm( + qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + graph.cb(qr, "dsv41_qr", il); + + ggml_tensor * q = graph.build_lora_mm(layer.wq_b, qr); + q = ggml_reshape_3d( + graph.ctx0, q, n_embd_head, graph.n_head, graph.n_tokens); + q = ggml_rms_norm(graph.ctx0, q, graph.norm_rms_eps); + q = ggml_rope_ext( + graph.ctx0, q, inp_pos, nullptr, n_rot, graph.rope_type, + rope_ctx, rope_base, rope_scale, rope_ext, rope_attn, + rope_beta_fast, rope_beta_slow); + q = ggml_rope_set_offset(q, n_nope); + graph.cb(q, "dsv41_q", il); + + ggml_tensor * kv = graph.build_lora_mm(layer.wkv, cur); + kv = graph.build_norm( + kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv = ggml_reshape_3d( + graph.ctx0, kv, n_embd_head, 1, graph.n_tokens); + ggml_tensor * raw_kv = ggml_rope_ext( + graph.ctx0, kv, inp_pos, nullptr, n_rot, graph.rope_type, + rope_ctx, rope_base, rope_scale, rope_ext, rope_attn, + rope_beta_fast, rope_beta_slow); + raw_kv = ggml_rope_set_offset(raw_kv, n_nope); + graph.cb(raw_kv, "dsv41_raw_kv", il); + + const llama_memory_dsv41 * memory = input.mctx->memory(); + ggml_tensor * raw_store = dsv41_flatten_memory( + graph.ctx0, memory->raw_k(il)); + ggml_tensor * raw_write = ggml_reshape_2d( + graph.ctx0, raw_kv, n_embd_head, graph.n_tokens); + ggml_tensor * raw_persist = ggml_get_rows( + graph.ctx0, raw_write, input.raw_persist_src_idxs); + raw_write = dsv41_cast_for_store( + graph.ctx0, raw_write, raw_store); + + ggml_tensor * raw_read = dsv41_concat( + graph.ctx0, raw_store, raw_write, 1, "raw state"); + raw_read = dsv41_append_zero_row(graph.ctx0, raw_read); + raw_read = ggml_get_rows( + graph.ctx0, raw_read, + ggml_reshape_1d( + graph.ctx0, input.raw_read_idxs, + graph.hparams.n_swa*graph.n_tokens)); + raw_read = ggml_reshape_4d( + graph.ctx0, raw_read, n_embd_head, 1, + graph.hparams.n_swa, graph.n_tokens); + + const int32_t kv_source = graph.hparams.dsv41_kv_source_layer[il]; + dsv41_graph_source_input * source = + kv_source >= 0 ? input.source(kv_source) : nullptr; + if (source != nullptr && source->layer == (uint32_t) il) { + ggml_tensor * compressed = graph.build_lora_mm( + layer.attn_comp_wkv, cur); + ggml_tensor * gate = ratio == 2 ? + graph.build_lora_mm(layer.attn_comp_wgate, cur) : nullptr; + ggml_tensor * compressed_current = compressed; + ggml_tensor * gate_current = gate; + ggml_tensor * carry_kv = nullptr; + ggml_tensor * carry_gate = nullptr; + ggml_tensor * persist_kv = nullptr; + ggml_tensor * persist_gate = nullptr; + + if (ratio == 2) { + carry_kv = dsv41_flatten_memory( + graph.ctx0, memory->compressor_carry_kv(source->layer)); + carry_gate = dsv41_flatten_memory( + graph.ctx0, memory->compressor_carry_score(source->layer)); + if (source->write_idxs != nullptr) { + ggml_tensor * carry_kv_first = ggml_get_rows( + graph.ctx0, carry_kv, source->carry_read_idxs); + ggml_tensor * carry_gate_first = ggml_get_rows( + graph.ctx0, carry_gate, source->carry_read_idxs); + carry_kv_first = dsv41_cast_for_store( + graph.ctx0, carry_kv_first, compressed); + carry_gate_first = dsv41_cast_for_store( + graph.ctx0, carry_gate_first, gate); + ggml_tensor * source_kv = dsv41_concat( + graph.ctx0, carry_kv_first, compressed, 1, "compressor carry KV"); + ggml_tensor * source_gate = dsv41_concat( + graph.ctx0, carry_gate_first, gate, 1, "compressor carry score"); + source_kv = ggml_get_rows( + graph.ctx0, source_kv, source->state_read_idxs); + source_gate = ggml_get_rows( + graph.ctx0, source_gate, source->state_read_idxs); + const int64_t n_write = source->write_pos->ne[0]; + source_kv = ggml_reshape_3d( + graph.ctx0, source_kv, n_embd_head, ratio, n_write); + source_gate = ggml_reshape_3d( + graph.ctx0, source_gate, n_embd_head, ratio, n_write); + compressed = llama_dsv41_build_ratio_pool( + graph.ctx0, source_kv, source_gate, ratio); + } else { + compressed = nullptr; + } + + persist_kv = ggml_get_rows( + graph.ctx0, compressed_current, + source->state_persist_src_idxs); + persist_gate = ggml_get_rows( + graph.ctx0, gate_current, + source->state_persist_src_idxs); + } else { + compressed = llama_dsv41_build_ratio_pool( + graph.ctx0, + ggml_reshape_3d( + graph.ctx0, compressed, n_embd_head, 1, + graph.n_tokens), + nullptr, 1); + } + + if (compressed != nullptr) { + compressed = graph.build_norm( + compressed, layer.attn_comp_norm, nullptr, + LLM_NORM_RMS, il); + graph.cb(compressed, "dsv41_compressed_unrotated", il); + + ggml_tensor * index_k = graph.build_lora_mm( + layer.indexer_attn_k, compressed); + index_k = graph.build_norm( + index_k, layer.indexer_k_norm, nullptr, + LLM_NORM_RMS, il); + index_k = ggml_reshape_3d( + graph.ctx0, index_k, + graph.hparams.indexer_head_size, 1, index_k->ne[1]); + index_k = ggml_rope_ext( + graph.ctx0, index_k, source->write_pos, nullptr, + graph.hparams.n_rot(), graph.rope_type, graph.n_ctx_orig, + graph.hparams.dsv4_compress_rope_base, + graph.freq_scale, graph.ext_factor, + dsv41_rope_attn_factor(graph.freq_scale), + graph.beta_fast, graph.beta_slow); + index_k = ggml_rope_set_offset( + index_k, + graph.hparams.indexer_head_size - graph.hparams.n_rot()); + + ggml_tensor * compressed_rope = ggml_reshape_3d( + graph.ctx0, compressed, n_embd_head, 1, compressed->ne[1]); + compressed_rope = ggml_rope_ext( + graph.ctx0, compressed_rope, source->write_pos, nullptr, + n_rot, graph.rope_type, graph.n_ctx_orig, + graph.hparams.dsv4_compress_rope_base, + graph.freq_scale, graph.ext_factor, + dsv41_rope_attn_factor(graph.freq_scale), + graph.beta_fast, graph.beta_slow); + compressed_rope = ggml_rope_set_offset(compressed_rope, n_nope); + + ggml_tensor * comp_store = dsv41_flatten_memory( + graph.ctx0, memory->compressed_kv(source->layer)); + ggml_tensor * index_store = dsv41_flatten_memory( + graph.ctx0, memory->index_keys(source->layer)); + ggml_tensor * comp_write = ggml_reshape_2d( + graph.ctx0, compressed_rope, n_embd_head, + compressed_rope->ne[2]); + ggml_tensor * index_write = ggml_reshape_2d( + graph.ctx0, index_k, graph.hparams.indexer_head_size, + index_k->ne[2]); + comp_write = dsv41_cast_for_store( + graph.ctx0, comp_write, comp_store); + index_write = dsv41_cast_for_store( + graph.ctx0, index_write, index_store); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, comp_store, comp_write, + source->write_idxs)); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, index_store, index_write, + source->write_idxs)); + } + + if (persist_kv != nullptr) { + if (compressed != nullptr) { + ggml_tensor * completion = dsv41_completion_zero( + graph.ctx0, compressed, persist_kv->type); + persist_kv = ggml_add( + graph.ctx0, persist_kv, completion); + persist_gate = ggml_add( + graph.ctx0, persist_gate, completion); + } + persist_kv = dsv41_cast_for_store( + graph.ctx0, persist_kv, carry_kv); + persist_gate = dsv41_cast_for_store( + graph.ctx0, persist_gate, carry_gate); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, carry_kv, persist_kv, + source->state_persist_dst_idxs)); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, carry_gate, persist_gate, + source->state_persist_dst_idxs)); + } + } + + if (graph.hparams.dsv41_is_index_source(il)) { + selected = dsv41_build_index_selection( + graph, model, input, qr, cur, inp_pos, + candidate_blocks, selected_local, + il); + } + + ggml_tensor * source_trace = nullptr; + if (selected != nullptr) { + source_trace = ggml_cont(graph.ctx0, selected); + ggml_set_name( + source_trace, + llama_dsv41_graph_trace_name( + "attn.source", il).c_str()); + } else { + source_trace = ggml_cont(graph.ctx0, input.raw_read_idxs); + ggml_set_name( + source_trace, + llama_dsv41_graph_trace_name( + "attn.source", il).c_str()); + } + ggml_build_forward_expand(graph.gf, source_trace); + + ggml_tensor * k_all = raw_read; + ggml_tensor * mask_all = input.raw_mask; + int64_t n_kv_max = graph.hparams.n_swa; + if (selected != nullptr) { + ggml_tensor * comp_store = dsv41_flatten_memory( + graph.ctx0, memory->compressed_kv(kv_source)); + comp_store = dsv41_append_zero_row( + graph.ctx0, comp_store); + ggml_tensor * compressed = ggml_get_rows( + graph.ctx0, comp_store, + ggml_reshape_1d( + graph.ctx0, selected, + selected->ne[0]*graph.n_tokens)); + compressed = ggml_reshape_4d( + graph.ctx0, compressed, n_embd_head, 1, + selected->ne[0], graph.n_tokens); + k_all = dsv41_concat( + graph.ctx0, raw_read, compressed, 2, "raw and compressed attention"); + + ggml_tensor * source_mask = ggml_reshape_3d( + graph.ctx0, source->mask, 1, + source->read_width, graph.n_tokens); + ggml_tensor * compressed_mask = ggml_get_rows( + graph.ctx0, source_mask, selected_local); + compressed_mask = ggml_cont( + graph.ctx0, + ggml_permute( + graph.ctx0, compressed_mask, 1, 0, 2, 3)); + compressed_mask = ggml_reshape_4d( + graph.ctx0, compressed_mask, + selected->ne[0], 1, 1, graph.n_tokens); + compressed_mask = dsv41_cast_for_store( + graph.ctx0, compressed_mask, input.raw_mask); + mask_all = dsv41_concat( + graph.ctx0, input.raw_mask, compressed_mask, 0, "attention mask"); + n_kv_max += selected->ne[0]; + } + + ggml_tensor * out = graph.build_attn_mha( + q, k_all, k_all, nullptr, mask_all, + layer.attn_sinks, nullptr, n_kv_max, + 1.0f/std::sqrt((float) n_embd_head), il); + + // Keep the prior ring intact until attention has consumed it. + ggml_tensor * completion = dsv41_completion_zero( + graph.ctx0, out, raw_persist->type); + raw_persist = ggml_add( + graph.ctx0, raw_persist, completion); + raw_persist = dsv41_cast_for_store( + graph.ctx0, raw_persist, raw_store); + ggml_tensor * raw_update = ggml_set_rows( + graph.ctx0, raw_store, raw_persist, + input.raw_write_idxs); + ggml_build_forward_expand(graph.gf, raw_update); + out = ggml_reshape_3d( + graph.ctx0, out, n_embd_head, graph.n_head, + graph.n_tokens); + out = ggml_rope_ext_back( + graph.ctx0, out, inp_pos, nullptr, n_rot, + graph.rope_type, rope_ctx, rope_base, rope_scale, + rope_ext, rope_attn, rope_beta_fast, rope_beta_slow); + out = ggml_rope_set_offset(out, n_nope); + graph.cb(out, "dsv41_attn_derope", il); + + out = ggml_reshape_3d( + graph.ctx0, out, group_width, n_groups, + graph.n_tokens); + out = ggml_permute(graph.ctx0, out, 0, 2, 1, 3); + if (graph.n_tokens > 1 && graph.n_tokens <= 8) { + out = ggml_cont(graph.ctx0, out); + } + ggml_tensor * oa = ggml_mul_mat(graph.ctx0, layer.wo_a, out); + oa = ggml_permute(graph.ctx0, oa, 0, 2, 1, 3); + oa = ggml_cont_2d( + graph.ctx0, oa, o_lora_rank*n_groups, + graph.n_tokens); + out = graph.build_lora_mm(layer.wo_b, oa); + graph.cb(out, "dsv41_attn_out", il); + return out; +} + +static std::pair dsv41_build_router( + const llama_model_deepseek41::graph & graph, + const llama_layer & layer, + ggml_tensor * cur, + int il) { + ggml_tensor * logits = graph.build_lora_mm(layer.ffn_gate_inp, cur); + ggml_prec_set_acc(logits, GGML_PREC_F32); + ggml_tensor * probs = ggml_sqrt( + graph.ctx0, ggml_softplus(graph.ctx0, logits)); + ggml_tensor * selection = ggml_add( + graph.ctx0, probs, layer.ffn_exp_probs_b); + ggml_tensor * ids = ggml_cont( + graph.ctx0, + ggml_argsort_top_k( + graph.ctx0, selection, + graph.hparams.n_expert_used())); + ids = dsv41_sort_row_ids(graph.ctx0, ids); + ggml_set_name( + ids, + llama_dsv41_graph_trace_name( + "expert.ids", il).c_str()); + + ggml_tensor * weights = ggml_get_rows( + graph.ctx0, + ggml_reshape_3d( + graph.ctx0, probs, 1, graph.n_expert, + graph.n_tokens), + ids); + weights = ggml_reshape_2d( + graph.ctx0, weights, + graph.hparams.n_expert_used(), graph.n_tokens); + ggml_tensor * sum = ggml_clamp( + graph.ctx0, ggml_sum_rows(graph.ctx0, weights), + 6.103515625e-5, INFINITY); + weights = ggml_div(graph.ctx0, weights, sum); + weights = ggml_scale( + graph.ctx0, weights, + graph.hparams.expert_weights_scale); + weights = ggml_cont(graph.ctx0, weights); + ggml_set_name( + weights, + llama_dsv41_graph_trace_name( + "expert.weights", il).c_str()); + ggml_build_forward_expand(graph.gf, weights); + return { logits, ids }; +} + +} + +ggml_tensor * llama_dsv41_graph_append_zero_row( + ggml_context * ctx, + ggml_tensor * tensor) { + return dsv41_append_zero_row(ctx, tensor); +} + +ggml_tensor * llama_dsv41_graph_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type) { + return dsv41_completion_zero(ctx, dependency, type); +} + +llama_model_deepseek41::graph::graph( + const llama_model & model, + const llm_graph_params & params) : + llama_model_deepseek4::graph(params) { + GGML_ASSERT(n_layer == LLAMA_DSV41_N_LAYER); + GGML_ASSERT(hparams.dsv4_hc_mult == LLAMA_DSV41_HC_MULT); + GGML_ASSERT(mctx != nullptr); + + const auto * dsv41_mctx = + static_cast(mctx); + auto input_owner = std::make_unique( + ctx0, cparams, hparams, dsv41_mctx, + ubatch, n_outputs); + auto * input = static_cast( + res->add_input(std::move(input_owner))); + if (input->topology.n_seqs != 1) { + throw std::invalid_argument( + "DeepSeek V4.1 graph supports one active sequence per ubatch"); + } + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + ggml_tensor * streams = ggml_reshape_3d( + ctx0, inp, n_embd, 1, n_tokens); + streams = ggml_repeat_4d( + ctx0, streams, n_embd, + hparams.dsv4_hc_mult, n_tokens, 1); + ggml_tensor * carried_pre = input->initial_pre; + ggml_tensor * selected = nullptr; + ggml_tensor * selected_local = nullptr; + ggml_tensor * candidate_blocks = nullptr; + cb(streams, "dsv41_hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + if ((size_t) il < cparams.embeddings_layer_inp.size() && + cparams.embeddings_layer_inp[il]) { + res->t_layer_inp[il] = dsv41_hc_mean(ctx0, streams); + cb(res->t_layer_inp[il], "layer_inp", il); + ggml_build_forward_expand(gf, res->t_layer_inp[il]); + } + + if (hparams.dsv41_engram_layers.test(il)) { + const uint32_t index = il == 1 ? 0 : 1; + ggml_tensor * row_ids = ggml_cont( + ctx0, input->engram_row_ids[index]); + ggml_set_name( + row_ids, + llama_dsv41_graph_trace_name( + "engram.row_ids", il).c_str()); + ggml_build_forward_expand(gf, row_ids); + streams = llama_dsv41_build_engram( + ctx0, streams, + input->engram_rows[index], + model.layers[il].engram_kv, + model.layers[il].engram_q_norm, + model.layers[il].engram_k_norm, + input->engram_select[index], + norm_rms_eps, sched, backend_cpu); + cb(streams, "dsv41_engram", il); + } + + ggml_tensor * residual = streams; + const dsv41_hc_mix attn_mix = dsv41_build_hc_mix( + *this, streams, + model.layers[il].hc_attn_fn, + model.layers[il].hc_attn_scale, + model.layers[il].hc_attn_base, + il); + ggml_tensor * cur = build_hc_pre( + streams, carried_pre, il); + cb(cur, "dsv41_hc_attn_carried_pre", il); + cur = build_norm( + cur, model.layers[il].attn_norm, + nullptr, LLM_NORM_RMS, il); + cur = dsv41_build_attention( + *this, model, *input, cur, inp_pos, + selected, selected_local, candidate_blocks, il); + streams = build_hc_post( + cur, residual, attn_mix.post, attn_mix.comb, il); + cb(streams, "dsv41_hc_attn_post", il); + + residual = streams; + const dsv41_hc_mix ffn_mix = dsv41_build_hc_mix( + *this, streams, + model.layers[il].hc_ffn_fn, + model.layers[il].hc_ffn_scale, + model.layers[il].hc_ffn_base, + il); + cur = build_hc_pre(streams, attn_mix.pre, il); + cb(cur, "dsv41_hc_ffn_attn_pre", il); + cur = build_norm( + cur, model.layers[il].ffn_norm, + nullptr, LLM_NORM_RMS, il); + + const auto [router_logits, original_ids] = + dsv41_build_router(*this, model.layers[il], cur, il); + ggml_tensor * slot_ids = llama_dsv41_build_expert_remap( + ctx0, original_ids, + *static_cast(model).experts, + il, sched, backend_cpu); + ggml_tensor * moe_out = build_moe_ffn( + cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, hparams.n_expert_used(), + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + router_logits, + nullptr, nullptr, nullptr, nullptr, + original_ids, + slot_ids); + cb(moe_out, "dsv41_ffn_moe", il); + + ggml_tensor * shared = build_ffn( + cur, + model.layers[il].ffn_up_shexp, nullptr, nullptr, + model.layers[il].ffn_gate_shexp, nullptr, nullptr, + model.layers[il].ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shared, "dsv41_ffn_shared", il); + cur = ggml_add(ctx0, moe_out, shared); + cb(cur, "dsv41_ffn_out", il); + + ggml_tensor * release = llama_dsv41_build_expert_release( + ctx0, moe_out, + *static_cast(model).experts, + il, sched, backend_cpu); + ggml_build_forward_expand(gf, release); + + streams = build_hc_post( + cur, residual, ffn_mix.post, ffn_mix.comb, il); + streams = build_cvec(streams, il); + carried_pre = ffn_mix.pre; + cb(streams, "dsv41_layer_out", il); + } + + if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && + cparams.embeddings_layer_inp[n_layer]) { + res->t_layer_inp[n_layer] = dsv41_hc_mean(ctx0, streams); + cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); + ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); + } + + if (inp_out_ids != nullptr) { + ggml_tensor * flat = ggml_reshape_2d( + ctx0, streams, + n_embd*hparams.dsv4_hc_mult, n_tokens); + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + streams = ggml_reshape_3d( + ctx0, flat, n_embd, + hparams.dsv4_hc_mult, n_outputs); + carried_pre = ggml_get_rows(ctx0, carried_pre, inp_out_ids); + } + + ggml_tensor * cur = llama_dsv41_build_output_collapse( + ctx0, streams, carried_pre, n_embd, + hparams.dsv4_hc_mult, + inp_out_ids ? n_outputs : n_tokens); + cb(cur, "dsv41_output_collapse", -1); + cur = llama_dsv41_build_output_norm_input(ctx0, cur); + cb(cur, "dsv41_output_collapse_f32", -1); + cur = build_norm( + cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + +std::unique_ptr llama_model_deepseek41::build_arch_graph( + const llm_graph_params & params) const { + return std::make_unique(*this, params); +} diff --git a/src/models/models.h b/src/models/models.h index 3f4e2eeae6e8..4acf4d90e1ad 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -9,6 +9,8 @@ #include class llama_memory_hybrid_idx_context; +class llama_dsv41_engram_runtime; +struct llama_dsv41_expert_runtime; // ref: https://github.com/ggml-org/llama.cpp/pull/28068 static inline ggml_tensor * build_gdn_l2_norm(ggml_context * ctx, ggml_tensor * x, float eps) { @@ -1313,6 +1315,31 @@ struct llama_model_deepseek4 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct llama_model_deepseek41 : public llama_model_deepseek4 { + llama_model_deepseek41(const struct llama_model_params & params) : llama_model_deepseek4(params) {} + + struct graph : public llama_model_deepseek4::graph { + graph(const llama_model & model, const llm_graph_params & params); + }; + + struct engram_model; + std::shared_ptr engram; + std::shared_ptr experts; + + std::unique_ptr create_memory_engram_runtime(size_t max_tokens) const; + + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + bool requires_synchronous_graph() const override; + std::string consume_runtime_error() const override; + void release_runtime_work() const override; + void release_runtime_work_after_sync(ggml_backend_sched_t sched) const override; + void acquire_runtime_context() const override; + void release_runtime_context() const override; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_deepseek2ocr : public llama_model_base { llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e013c9a984c3..dd58dd04fb03 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -198,7 +198,22 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW + llama_build_and_test(test-deepseek41-schema.cpp) + llama_build_and_test(test-deepseek41-engram.cpp) + llama_build_and_test(test-deepseek41-expert.cpp) + llama_build_and_test(test-expert-store.cpp) + llama_build_and_test(test-deepseek41-memory.cpp) + # test-deepseek41-runtime.cpp needs tools/deepseek-v41-trace/trace-components.h + # to validate the trace tensor naming contract; that header lands with the + # trace harness, so its registration here is deferred to that layer. + llama_build_and_test(test-engram.cpp) llama_build(test-llama-archs.cpp) + llama_test( + test-llama-archs + NAME test-deepseek41-no-alloc + LABEL main + ARGS -a deepseek41 -s 1 + ) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") @@ -269,6 +284,17 @@ llama_build_and_test(test-chat-template.cpp) # debug tool for chat template differential analysis (not registered as a test, run it manually) llama_build(test-chat-analysis.cpp) llama_build_and_test(test-log.cpp) + +find_package(Python3 3.10 COMPONENTS Interpreter QUIET) +if (CMAKE_SYSTEM_NAME STREQUAL "Linux" AND Python3_Interpreter_FOUND) + llama_test_cmd( + ${Python3_EXECUTABLE} + NAME test-strix-memory-watchdog + LABEL python + ARGS ${CMAKE_CURRENT_SOURCE_DIR}/test_strix_memory_watchdog.py + ) +endif() + llama_build_and_test( test-peg-parser.cpp peg-parser/simple-tokenize.cpp diff --git a/tests/test-deepseek41-engram.cpp b/tests/test-deepseek41-engram.cpp new file mode 100644 index 000000000000..ee73c48628f2 --- /dev/null +++ b/tests/test-deepseek41-engram.cpp @@ -0,0 +1,529 @@ +#include "../src/llama-dsv41-engram.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#endif + +static void check(bool condition, const char * message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message); + std::exit(1); + } +} + +static void expect_invalid(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::invalid_argument &) { + return; + } + check(false, message); +} + +static void expect_runtime(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_engram_layout make_layout() { + llama_engram_layout layout; + layout.encoding = "e4m3_e8m0_32_row264"; + layout.layer_ids = { 1, 14 }; + layout.token_map.resize(32); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = (uint32_t) i; + } + layout.compressed_vocab_size = 32; + layout.pad_id = 2; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + layout.multipliers[layer][i] = 101 + 8*layer + 2*i; + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 2; + layout.rows[layer] += 2; + } + } + return layout; +} + +static void fill_row(uint8_t row[LLAMA_ENGRAM_ROW_BYTES], uint32_t id) { + const uint8_t code = (uint8_t) (8 + id%100); + std::fill(row, row + LLAMA_ENGRAM_DIM, code); + std::fill(row + LLAMA_ENGRAM_DIM, row + LLAMA_ENGRAM_ROW_BYTES, 127); +} + +#if !defined(_WIN32) +static void write_full(int fd, const void * data, size_t size, uint64_t offset) { + const ssize_t written = pwrite(fd, data, size, (off_t) offset); + check(written == (ssize_t) size, "failed to write DeepSeek V4.1 Engram test data"); +} + +struct test_file { + std::string path; + int fd = -1; + std::array extents; + + test_file(const llama_engram_layout & layout) { + char name[] = "/tmp/llama-dsv41-engram-XXXXXX"; + fd = mkstemp(name); + check(fd >= 0, "failed to create DeepSeek V4.1 Engram test file"); + path = name; + + uint64_t offset = 4096; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + extents[layer] = { + path, + offset, + layout.rows[layer], + LLAMA_ENGRAM_ROW_BYTES, + layout.rows[layer], + GGML_TYPE_I8, + }; + for (uint32_t row_id = 0; row_id < layout.rows[layer]; ++row_id) { + uint8_t row[LLAMA_ENGRAM_ROW_BYTES]; + fill_row(row, row_id + 7*layer); + write_full(fd, row, sizeof(row), offset + (uint64_t) row_id*sizeof(row)); + } + offset += (uint64_t) layout.rows[layer]*LLAMA_ENGRAM_ROW_BYTES + 4096; + } + } + + ~test_file() { + if (fd >= 0) { + close(fd); + } + if (!path.empty()) { + unlink(path.c_str()); + } + } +}; + +static void test_extent_validation() { + llama_dsv41_engram_extent extent = { + "/tmp/model.gguf", 4096, 48, LLAMA_ENGRAM_ROW_BYTES, 48, GGML_TYPE_I8, + }; + llama_dsv41_validate_engram_extent(extent); + extent.type = GGML_TYPE_F32; + expect_invalid([&] { llama_dsv41_validate_engram_extent(extent); }, "non-I8 Engram extent was accepted"); + extent.type = GGML_TYPE_I8; + extent.columns = LLAMA_ENGRAM_ROW_BYTES - 1; + expect_invalid([&] { llama_dsv41_validate_engram_extent(extent); }, "short Engram row was accepted"); +} + +static void test_transactions_and_sequences() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + + std::vector tokens = { + { 3, 0, { 7, 9 }, 1 }, + { 5, 1, { 7, 9 }, 0 }, + { 11, 0, { 12 }, 1 }, + }; + llama_dsv41_engram_transaction transaction = runtime.prepare(tokens); + check(transaction.token_count() == tokens.size(), "Engram transaction token count mismatch"); + check(transaction.row_ids(1) == transaction.row_ids(0) + LLAMA_ENGRAM_COLS, + "Engram layer row selection is not the second 24-ID half"); + llama_engram_hasher hasher(layout); + llama_engram_history expected_history_7; + llama_engram_history expected_history_12; + expected_history_7.reset(); + expected_history_12.reset(); + uint32_t expected[3*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS]; + const int32_t coupled_tokens[] = { 3, 5 }; + const uint8_t coupled_mask[] = { 1, 0 }; + hasher.hash(expected_history_7, coupled_tokens, coupled_mask, 2, expected); + hasher.hash( + expected_history_12, + &tokens[2].token, + &tokens[2].text, + 1, + expected + 2*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); + for (size_t token = 0; token < tokens.size(); ++token) { + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + transaction.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + expected + (token*LLAMA_ENGRAM_LAYERS + layer)*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "DeepSeek V4.1 Engram IDs or 48-ID token stride differ from the core hasher"); + } + } + check(transaction.text_mask()[0] == 1 && transaction.text_mask()[1] == 0, + "Engram text mask changed"); + check(runtime.sequence(7).pos == -1, "Engram prepare committed sequence state early"); + + runtime.commit(transaction); + check(runtime.sequence(7).pos == 1 && runtime.sequence(9).pos == 1, + "Engram commit did not advance coupled sequences"); + check(runtime.sequence(7).history.tail[0] == LLAMA_ENGRAM_DEAD, + "masked Engram token did not break sequence history"); + + runtime.seq_copy(7, 13); + check(runtime.sequence(13).history.tail == runtime.sequence(7).history.tail, + "Engram sequence copy changed history"); + const llama_dsv41_engram_snapshot snapshot = runtime.checkpoint(); + runtime.seq_reset(7); + check(runtime.sequence(7).pos == -1, "Engram sequence reset did not clear position"); + runtime.restore(snapshot); + check(runtime.sequence(7).pos == 1, "Engram checkpoint restore lost position"); + runtime.seq_remove(13); + check(runtime.sequence(13).pos == -1, "Engram sequence remove retained state"); + + llama_dsv41_engram_transaction stale = runtime.prepare({ { 7, 2, { 7 }, 1 } }); + runtime.seq_copy(7, 14); + expect_runtime([&] { runtime.commit(stale); }, "stale Engram transaction was committed"); + runtime.rollback(stale); + + llama_dsv41_engram_snapshot invalid = runtime.checkpoint(); + invalid.sequences[7].history.tail[0] = (int32_t) layout.compressed_vocab_size; + expect_invalid([&] { runtime.restore(invalid); }, "invalid Engram snapshot was restored"); + expect_invalid( + [&] { runtime.prepare({ { 8, 2, { 7, 7 }, 1 } }); }, + "duplicate Engram sequence ID was accepted"); +} + +static void test_scheduler_upload() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + llama_dsv41_engram_transaction transaction = runtime.prepare({ + { 3, 0, { 0 }, 1 }, + { 5, 1, { 0 }, 0 }, + { 7, 2, { 0 }, 1 }, + }); + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram upload context"); + ggml_tensor * rows = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, transaction.token_count()); + ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, transaction.token_count()); + + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 Engram upload backend"); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + check(buffer != nullptr, "failed to allocate DeepSeek V4.1 Engram upload tensors"); + + std::vector actual_rows(ggml_nelements(rows)); + std::vector actual_select(ggml_nelements(select)); + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + transaction.upload_layer(layer, 0, transaction.token_count(), rows, select); + ggml_backend_tensor_get(rows, actual_rows.data(), 0, ggml_nbytes(rows)); + ggml_backend_tensor_get(select, actual_select.data(), 0, ggml_nbytes(select)); + check(std::memcmp( + actual_rows.data(), + transaction.rows(layer), + ggml_nbytes(rows)) == 0, + "scheduler-backed Engram row upload changed the bounded pack"); + check(actual_select == std::vector({ 3, 1, 5 }), + "scheduler-backed Engram text selection upload changed"); + } + + runtime.commit(transaction); + expect_invalid( + [&] { transaction.upload_layer(0, 0, 3, rows, select); }, + "committed Engram transaction uploaded stale input"); + + ggml_backend_buffer_free(buffer); + ggml_backend_free(backend); + ggml_free(ctx); +} + +static void test_chunked_prefill() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime whole_runtime(layout, file.extents, 8); + llama_dsv41_engram_runtime chunked_runtime(layout, file.extents, 8); + const std::vector tokens = { + { 3, 0, { 0 }, 1 }, + { 5, 1, { 0 }, 1 }, + { 7, 2, { 0 }, 0 }, + { 9, 3, { 0 }, 1 }, + }; + + llama_dsv41_engram_transaction whole = whole_runtime.prepare(tokens); + llama_dsv41_engram_transaction first = chunked_runtime.prepare({ tokens[0], tokens[1] }); + for (size_t token = 0; token < 2; ++token) { + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + whole.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + first.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "first Engram prefill chunk differs from whole-chunk hashing"); + } + } + chunked_runtime.commit(first); + + llama_dsv41_engram_transaction second = chunked_runtime.prepare({ tokens[2], tokens[3] }); + for (size_t token = 0; token < 2; ++token) { + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + whole.row_ids(layer) + (token + 2)*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + second.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "later Engram prefill chunk lost committed history"); + } + } +} + +static void test_transactional_read_failure() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + + llama_dsv41_engram_transaction first = runtime.prepare({ { 1, 0, { 0 }, 1 } }); + runtime.commit(first); + const llama_dsv41_engram_sequence_state before = runtime.sequence(0); + + const uint64_t truncated = file.extents[1].offset + LLAMA_ENGRAM_ROW_BYTES; + check(ftruncate(file.fd, (off_t) truncated) == 0, "failed to truncate Engram transaction test file"); + expect_runtime( + [&] { runtime.prepare({ { 2, 1, { 0 }, 1 } }); }, + "Engram read failure was not surfaced"); + const llama_dsv41_engram_sequence_state after = runtime.sequence(0); + check(after.pos == before.pos && after.history.tail == before.history.tail, + "failed Engram transaction advanced sequence state"); +} +#endif + +static float bf16(float value) { + return ggml_bf16_to_fp32(ggml_fp32_to_bf16(value)); +} + +static void test_graph_gate() { + constexpr int64_t width = 8; + constexpr int64_t streams = 4; + constexpr int64_t tokens = 2; + + ggml_init_params params = { + /*.mem_size =*/ 8*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram graph context"); + + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, width, streams, tokens); + ggml_tensor * rows = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, tokens); + ggml_tensor * engram_kv = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, 5*width); + ggml_tensor * q_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); + ggml_tensor * k_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); + ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); + + ggml_set_input(residual); + ggml_set_input(rows); + ggml_set_input(engram_kv); + ggml_set_input(q_norm); + ggml_set_input(k_norm); + ggml_set_input(select); + + std::vector residual_data(ggml_nelements(residual)); + std::vector rows_data(ggml_nelements(rows)); + std::vector engram_kv_data(ggml_nelements(engram_kv)); + std::vector q_data(ggml_nelements(q_norm)); + std::vector k_data(ggml_nelements(k_norm)); + for (int64_t i = 0; i < width*streams*tokens; ++i) { + residual_data[i] = i < width*streams ? + bf16(0.125f + (float) (i%13)/16.0f) : + 0.1234567f + (float) (i%13)/17.0f; + } + residual_data[width*streams] = -0.0f; + residual_data[width*streams + 1] = 0.0f; + std::vector projected_data(5*width*tokens); + for (int64_t token = 0; token < tokens; ++token) { + rows_data[token*rows->ne[0]] = 1.0f; + for (int64_t i = 0; i < 5*width; ++i) { + projected_data[token*5*width + i] = 0.0625f + (float) ((token*5*width + i)%11)/32.0f; + engram_kv_data[i*engram_kv->ne[0] + token] = projected_data[token*5*width + i]; + } + } + for (int64_t i = 0; i < width*streams; ++i) { + q_data[i] = 0.5f + (float) (i%5)/8.0f; + k_data[i] = 0.75f - (float) (i%3)/16.0f; + } + const int32_t select_data[] = { (int32_t) tokens, 1 }; + + const std::vector original = residual_data; + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 Engram graph backend"); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 512, false, true); + check(sched != nullptr, "failed to create DeepSeek V4.1 Engram graph scheduler"); + ggml_tensor * output = llama_dsv41_build_engram( + ctx, residual, rows, engram_kv, q_norm, k_norm, select, 1.0e-20f, sched, backend); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 Engram graph"); + ggml_backend_tensor_set(residual, residual_data.data(), 0, ggml_nbytes(residual)); + ggml_backend_tensor_set(rows, rows_data.data(), 0, ggml_nbytes(rows)); + ggml_backend_tensor_set(engram_kv, engram_kv_data.data(), 0, ggml_nbytes(engram_kv)); + ggml_backend_tensor_set(q_norm, q_data.data(), 0, ggml_nbytes(q_norm)); + ggml_backend_tensor_set(k_norm, k_data.data(), 0, ggml_nbytes(k_norm)); + ggml_backend_tensor_set(select, select_data, 0, sizeof(select_data)); + check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, + "DeepSeek V4.1 Engram graph execution failed"); + + std::vector actual(ggml_nelements(output)); + ggml_backend_tensor_get(output, actual.data(), 0, ggml_nbytes(output)); + for (int64_t stream = 0; stream < streams; ++stream) { + double hidden_sq = 0.0; + double key_sq = 0.0; + double dot = 0.0; + for (int64_t i = 0; i < width; ++i) { + const float hidden = original[stream*width + i]; + const float key = bf16(projected_data[stream*width + i]); + hidden_sq += hidden*hidden; + key_sq += key*key; + dot += hidden*q_data[stream*width + i]*k_data[stream*width + i]*key; + } + dot /= std::sqrt(hidden_sq/width + 1.0e-20); + dot /= std::sqrt(key_sq/width + 1.0e-20); + dot /= std::sqrt((double) width); + const double gate = 1.0/(1.0 + std::exp(-std::copysign(std::sqrt(std::max(std::abs(dot), 1.0e-6)), dot))); + for (int64_t i = 0; i < width; ++i) { + const float value = bf16(projected_data[4*width + i]); + const float expected = bf16(original[stream*width + i] + (float) gate*value); + check(std::abs(actual[stream*width + i] - expected) <= std::max(1.0e-6f, std::abs(expected)/128.0f), + "DeepSeek V4.1 Engram gate differs from scalar reference"); + const size_t masked = width*streams + stream*width + i; + check(std::memcmp(actual.data() + masked, original.data() + masked, sizeof(float)) == 0, + "masked DeepSeek V4.1 Engram row changed"); + } + } + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +static void test_signed_zero_gate() { + constexpr int64_t matrix_size = 32; + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 signed-zero graph context"); + + ggml_tensor * dot = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); + ggml_tensor * matrix_a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, matrix_size, matrix_size); + ggml_tensor * matrix_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, matrix_size, matrix_size); + ggml_set_input(dot); + ggml_set_input(matrix_a); + ggml_set_input(matrix_b); + ggml_tensor * accelerated = ggml_mul_mat(ctx, matrix_a, matrix_b); + + ggml_backend_load_all(); + ggml_backend_t backend_accel = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t device = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_ACCEL && + ggml_backend_dev_supports_op(device, accelerated)) { + backend_accel = ggml_backend_dev_init(device, nullptr); + if (backend_accel != nullptr) { + break; + } + } + } + + ggml_backend_t backend_cpu = ggml_backend_cpu_init(); + check(backend_cpu != nullptr, "failed to create DeepSeek V4.1 signed-zero CPU backend"); + ggml_backend_t backends[] = { + backend_accel != nullptr ? backend_accel : backend_cpu, + backend_cpu, + }; + ggml_backend_buffer_type_t bufts[] = { + ggml_backend_get_default_buffer_type(backends[0]), + ggml_backend_cpu_buffer_type(), + }; + const int n_backends = backend_accel != nullptr ? 2 : 1; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, n_backends, 32, false, true); + check(sched != nullptr, "failed to create DeepSeek V4.1 signed-zero scheduler"); + if (backend_accel != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, accelerated, backend_accel); + expect_invalid( + [&] { llama_dsv41_build_engram_gate(ctx, dot, sched, backend_accel); }, + "DeepSeek V4.1 gate accepted a non-CPU backend"); + } + + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot, sched, backend_cpu); + ggml_tensor * output = ggml_add(ctx, gate, ggml_repeat(ctx, ggml_sum(ctx, accelerated), gate)); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 signed-zero graph"); + if (backend_accel != nullptr) { + check(ggml_backend_sched_get_tensor_backend(sched, accelerated) == backend_accel, + "DeepSeek V4.1 gate moved surrounding work off its accelerator"); + } + check(ggml_backend_sched_get_tensor_backend(sched, gate) == backend_cpu, + "DeepSeek V4.1 gate was not assigned to the local CPU backend"); + + const float input[] = { 0.0f, -0.0f }; + std::vector zeros(matrix_size*matrix_size); + ggml_backend_tensor_set(dot, input, 0, sizeof(input)); + ggml_backend_tensor_set(matrix_a, zeros.data(), 0, ggml_nbytes(matrix_a)); + ggml_backend_tensor_set(matrix_b, zeros.data(), 0, ggml_nbytes(matrix_b)); + check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, + "DeepSeek V4.1 signed-zero scheduler execution failed"); + const float positive = 1.0f/(1.0f + std::exp(-0.001f)); + const float negative = 1.0f/(1.0f + std::exp(0.001f)); + float actual[2]; + ggml_backend_tensor_get(output, actual, 0, sizeof(actual)); + check(std::abs(actual[0] - positive) < 1.0e-7f && actual[0] > 0.5f, + "DeepSeek V4.1 positive-zero gate lost copysign semantics"); + check(std::abs(actual[1] - negative) < 1.0e-7f && actual[1] < 0.5f, + "DeepSeek V4.1 negative-zero gate lost copysign semantics"); + + ggml_backend_sched_free(sched); + if (backend_accel != nullptr) { + ggml_backend_free(backend_accel); + } + ggml_backend_free(backend_cpu); + ggml_free(ctx); +} + +int main() { +#if !defined(_WIN32) + test_extent_validation(); + test_transactions_and_sequences(); + test_scheduler_upload(); + test_chunked_prefill(); + test_transactional_read_failure(); +#endif + test_graph_gate(); + test_signed_zero_gate(); + std::puts("DeepSeek V4.1 Engram runtime and graph: PASS"); + return 0; +} diff --git a/tests/test-deepseek41-expert.cpp b/tests/test-deepseek41-expert.cpp new file mode 100644 index 000000000000..269922196301 --- /dev/null +++ b/tests/test-deepseek41-expert.cpp @@ -0,0 +1,583 @@ +#include "../src/llama-dsv41-expert.h" +#include "../src/llama-dsv41.h" +#include "../src/llama-graph.h" +#include "../src/llama-model-loader.h" + +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" +#include "gguf.h" +#include "../ggml/src/ggml-backend-impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(cond) do { if (!(cond)) { throw std::runtime_error("requirement failed: " #cond); } } while (0) + +namespace { + +struct temp_file { + std::filesystem::path path; + + temp_file() { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-dsv41-expert-" + std::to_string(stamp) + "-" + std::to_string(++sequence) + ".gguf"); + } + + ~temp_file() { + std::error_code ec; + std::filesystem::remove(path, ec); + } +}; + +template +void require_throws(F && fn) { + bool threw = false; + try { + fn(); + } catch (const std::exception &) { + threw = true; + } + REQUIRE(threw); +} + +struct fixture { + static constexpr int64_t n_embd = 256; + static constexpr int64_t n_ff = 256; + static constexpr int64_t n_expert = LLAMA_DSV41_N_EXPERT; + + temp_file file; + std::vector tensors; + + fixture() { + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd)*n_ff; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff)*n_embd; + const size_t data_size = (2*gate_plane + down_plane)*n_expert; + ggml_init_params params = { + /*.mem_size =*/ data_size + 8*ggml_tensor_overhead() + 4096, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * down = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_ff, n_embd, n_expert); + ggml_set_name(gate, "blk.0.ffn_gate_exps.weight"); + ggml_set_name(up, "blk.0.ffn_up_exps.weight"); + ggml_set_name(down, "blk.0.ffn_down_exps.weight"); + fill(gate, 0x10); + fill(up, 0x20); + fill(down, 0x30); + + gguf_context * gguf = gguf_init_empty(); + REQUIRE(gguf != nullptr); + gguf_set_val_str(gguf, "general.architecture", "deepseek41"); + gguf_add_tensor(gguf, gate); + gguf_add_tensor(gguf, up); + gguf_add_tensor(gguf, down); + REQUIRE(gguf_write_to_file(gguf, file.path.string().c_str(), false)); + gguf_free(gguf); + ggml_free(ctx); + + std::vector splits; + llama_model_loader loader( + nullptr, nullptr, nullptr, file.path.string(), splits, nullptr, + LLAMA_LOAD_MODE_MMAP, false, false, false, nullptr, nullptr); + const auto gate_extent = loader.register_external_tensor( + "blk.0.ffn_gate_exps.weight", 0, LLAMA_EXPERT_PROJECTION_GATE, { n_embd, n_ff, n_expert }); + const auto up_extent = loader.register_external_tensor( + "blk.0.ffn_up_exps.weight", 0, LLAMA_EXPERT_PROJECTION_UP, { n_embd, n_ff, n_expert }); + const auto down_extent = loader.register_external_tensor( + "blk.0.ffn_down_exps.weight", 0, LLAMA_EXPERT_PROJECTION_DOWN, { n_ff, n_embd, n_expert }); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (auto extent : { gate_extent, up_extent, down_extent }) { + extent.layer = il; + extent.name = "blk." + std::to_string(il) + extent.name.substr(5); + tensors.push_back(std::move(extent)); + } + } + } + + static void fill(ggml_tensor * tensor, uint8_t tag) { + memset(tensor->data, 0, ggml_nbytes(tensor)); + for (int64_t expert = 0; expert < tensor->ne[2]; ++expert) { + uint8_t * plane = static_cast(tensor->data) + expert*tensor->nb[2]; + plane[tensor->nb[2] - 1] = tag + expert%32; + } + } + + size_t cache_bytes(size_t slots) const { + size_t result = 0; + for (const auto & tensor : tensors) { + result += tensor.nb[2]*slots; + } + return result; + } + + llama_dsv41_expert_runtime make_runtime( + size_t slots, + llama_dsv41_expert_runtime::upload_fn upload = {}, + llama_dsv41_expert_runtime::publish_fn before_publish = {}) const { + llama_dsv41_expert_runtime_params params; + params.cache_slots = slots; + params.cache_bytes = cache_bytes(slots); + params.direct_io = false; + return llama_dsv41_expert_runtime( + tensors, + params, + [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }, + std::move(upload), + std::move(before_publish)); + } +}; + +void test_registration() { + std::vector names; + const auto tensors = llama_dsv41_register_expert_tensors( + [&](const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + names.push_back(name); + llama_expert_store_tensor tensor; + tensor.name = name; + tensor.fname = "unused"; + tensor.file_index = layer % 3; + tensor.layer = layer; + tensor.projection = projection; + tensor.type = projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + std::copy(ne.begin(), ne.end(), tensor.ne); + tensor.nb[0] = ggml_type_size(tensor.type); + tensor.nb[1] = ggml_row_size(tensor.type, tensor.ne[0]); + tensor.nb[2] = tensor.nb[1]*tensor.ne[1]; + tensor.file_offset = 4096 + (size_t) layer*3*1024 + (size_t) projection*1024; + tensor.file_size = tensor.nb[2]*tensor.ne[2]; + return tensor; + }); + REQUIRE(tensors.size() == LLAMA_DSV41_N_LAYER*3); + REQUIRE(names.front() == "blk.0.ffn_gate_exps.weight"); + REQUIRE(names.back() == "blk.39.ffn_down_exps.weight"); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + REQUIRE(tensors[3*il + 0].layer == il); + REQUIRE(tensors[3*il + 1].layer == il); + REQUIRE(tensors[3*il + 2].layer == il); + REQUIRE(tensors[3*il + 0].file_index == (size_t) il % 3); + REQUIRE(tensors[3*il + 0].file_offset == 4096 + (size_t) il*3*1024); + REQUIRE(tensors[3*il + 0].ne[2] == LLAMA_DSV41_N_EXPERT); + REQUIRE(tensors[3*il + 0].nb[2] == 3041280); + REQUIRE(tensors[3*il + 1].nb[2] == 3041280); + REQUIRE(tensors[3*il + 2].nb[2] == 3870720); + } + REQUIRE(std::none_of(names.begin(), names.end(), [](const std::string & name) { + return name.find("shexp") != std::string::npos; + })); + size_t one_slot_bytes = 0; + for (const auto & tensor : tensors) { + one_slot_bytes += tensor.nb[2]; + } + REQUIRE(one_slot_bytes == 398131200); +} + +void test_configuration(const fixture & f) { + llama_dsv41_expert_runtime_params params; + params.direct_io = false; + require_throws([&] { + llama_dsv41_expert_runtime runtime( + f.tensors, params, [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }); + }); + + params.cache_slots = 1; + params.cache_bytes = f.cache_bytes(1) - 1; + require_throws([&] { + llama_dsv41_expert_runtime runtime( + f.tensors, params, [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }); + }); + + auto runtime = f.make_runtime(1); + runtime.acquire_context(); + require_throws([&] { runtime.acquire_context(); }); + runtime.release_context(); + runtime.acquire_context(); + runtime.release_context(); +} + +void test_remap_upload_and_eviction(const fixture & f) { + struct upload_record { + std::string name; + size_t offset; + std::vector bytes; + }; + std::vector uploads; + auto runtime = f.make_runtime(3, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + uploads.push_back({ tensor->name, offset, std::vector( + static_cast(data), static_cast(data) + size) }); + ggml_backend_tensor_set(tensor, data, offset, size); + }); + + const auto remapped = runtime.remap(0, { 3, 3, 1, 2 }); + REQUIRE(remapped == std::vector({ 2, 2, 0, 1 })); + REQUIRE(uploads.size() == 9); + REQUIRE(uploads[0].bytes.back() == 0x11); + REQUIRE(uploads[1].bytes.back() == 0x21); + REQUIRE(uploads[2].bytes.back() == 0x31); + require_throws([&] { runtime.remap(0, { 0 }); }); + runtime.release(0); + + uploads.clear(); + REQUIRE(runtime.remap(0, { 3 }).front() == 2); + REQUIRE(uploads.empty()); + runtime.release(0); + + REQUIRE(runtime.remap(0, { 4 }).front() == 0); + runtime.release(0); +} + +void test_capacity_and_upload_failure(const fixture & f) { + auto runtime = f.make_runtime(2); + require_throws([&] { runtime.remap(0, { 0, 1, 2 }); }); + REQUIRE(runtime.remap(0, { 0, 0, 1 }).size() == 3); + runtime.release(0); + + size_t calls = 0; + auto failing = f.make_runtime(1, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + if (++calls == 5) { + throw std::runtime_error("synthetic upload failure"); + } + ggml_backend_tensor_set(tensor, data, offset, size); + }); + REQUIRE(failing.remap(0, { 0 }).front() == 0); + failing.release(0); + require_throws([&] { failing.remap(0, { 1 }); }); + const size_t after_failure = calls; + REQUIRE(failing.remap(0, { 0 }).front() == 0); + REQUIRE(calls == after_failure + 3); + failing.release(0); +} + +void test_publication_allocation_failure(const fixture & f) { + size_t attempts = 0; + auto runtime = f.make_runtime(1, {}, [&] { + if (attempts++ == 0) { + throw std::bad_alloc(); + } + }); + + require_throws([&] { runtime.remap(0, { 0 }); }); + REQUIRE(runtime.remap(0, { 1 }).front() == 0); + runtime.release(0); +} + +void test_graph_callbacks(const fixture & f) { + auto runtime = f.make_runtime(2); + ggml_init_params params = { + /*.mem_size =*/ 2*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 3, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, fixture::n_embd, 3, 1); + ggml_set_input(selected); + ggml_set_input(input); + ggml_tensor * remapped = llama_dsv41_build_expert_remap(ctx, selected, runtime, 0, sched, backend); + ggml_tensor * expert_values = ggml_mul_mat_id( + ctx, runtime.cache_tensor(0, LLAMA_EXPERT_PROJECTION_GATE), input, remapped); + ggml_tensor * release = llama_dsv41_build_expert_release(ctx, expert_values, runtime, 0, sched, backend); + ggml_set_output(remapped); + ggml_set_output(release); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, release); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + const int32_t original[] = { 3, 1, 1 }; + std::vector input_data(fixture::n_embd*3, 1.0f); + ggml_backend_tensor_set(selected, original, 0, sizeof(original)); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size()*sizeof(float)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + const std::string error = runtime.consume_error(); + if (!error.empty()) { + throw std::runtime_error(error); + } + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, remapped) == backend); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, release) == backend); + REQUIRE(runtime.remap(0, { 0 }).front() >= 0); + runtime.release(0); + runtime.release_all(); + + const int32_t over_capacity[] = { 3, 1, 2 }; + ggml_backend_tensor_set(selected, over_capacity, 0, sizeof(over_capacity)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(!runtime.consume_error().empty()); + int32_t safe_ids[3] = { -1, -1, -1 }; + ggml_backend_tensor_get(remapped, safe_ids, 0, sizeof(safe_ids)); + REQUIRE(safe_ids[0] == 0 && safe_ids[1] == 0 && safe_ids[2] == 0); + REQUIRE(runtime.remap(0, { 0 }).front() >= 0); + runtime.release(0); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +void test_graph_upload_failure_sentinel(const fixture & f) { + size_t calls = 0; + auto runtime = f.make_runtime(1, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + if (++calls == 5) { + throw std::runtime_error("synthetic graph upload failure"); + } + ggml_backend_tensor_set(tensor, data, offset, size); + }); + ggml_init_params params = { + /*.mem_size =*/ 2*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, fixture::n_embd, 1, 1); + ggml_set_input(selected); + ggml_set_input(input); + ggml_tensor * remapped = llama_dsv41_build_expert_remap(ctx, selected, runtime, 0, sched, backend); + ggml_tensor * expert_values = ggml_mul_mat_id( + ctx, runtime.cache_tensor(0, LLAMA_EXPERT_PROJECTION_GATE), input, remapped); + ggml_tensor * release = llama_dsv41_build_expert_release(ctx, expert_values, runtime, 0, sched, backend); + ggml_set_output(remapped); + ggml_set_output(release); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, release); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + std::vector input_data(fixture::n_embd, 1.0f); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size()*sizeof(float)); + const int32_t first[] = { 0 }; + ggml_backend_tensor_set(selected, first, 0, sizeof(first)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(runtime.consume_error().empty()); + + const int32_t failed[] = { 1 }; + ggml_backend_tensor_set(selected, failed, 0, sizeof(failed)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(runtime.consume_error().find("synthetic graph upload failure") != std::string::npos); + int32_t safe_id = -1; + ggml_backend_tensor_get(remapped, &safe_id, 0, sizeof(safe_id)); + REQUIRE(safe_id == 0); + + REQUIRE(runtime.remap(0, { 2 }).front() == 0); + runtime.release(0); + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +void test_grovemoe_lookup_ids() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_tensor * explicit_slots = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_set_input(selected); + ggml_set_input(explicit_slots); + const llm_moe_expert_ids grovemoe = llm_build_moe_expert_ids( + ctx, LLM_ARCH_GROVEMOE, selected, nullptr, 2, 8, 4); + const llm_moe_expert_ids deepseek = llm_build_moe_expert_ids( + ctx, LLM_ARCH_DEEPSEEK41, selected, explicit_slots, 8, 8, 0); + REQUIRE(grovemoe.routing == selected); + REQUIRE(grovemoe.lookup != selected); + REQUIRE(deepseek.routing == selected); + REQUIRE(deepseek.lookup == explicit_slots); + ggml_set_output(grovemoe.routing); + ggml_set_output(grovemoe.lookup); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, grovemoe.routing); + ggml_build_forward_expand(graph, grovemoe.lookup); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + const int32_t original[] = { 7, 1 }; + ggml_backend_tensor_set(selected, original, 0, sizeof(original)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + int32_t routing_ids[2] = {}; + int32_t chunk_ids[2] = {}; + ggml_backend_tensor_get(grovemoe.routing, routing_ids, 0, sizeof(routing_ids)); + ggml_backend_tensor_get(grovemoe.lookup, chunk_ids, 0, sizeof(chunk_ids)); + REQUIRE(routing_ids[0] == 7 && routing_ids[1] == 1); + REQUIRE(chunk_ids[0] == 1 && chunk_ids[1] == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +struct sync_test_context { + llama_dsv41_expert_runtime * runtime = nullptr; + int synchronize_count = 0; + bool saw_pinned = false; +}; + +const char * sync_test_backend_name(ggml_backend_t) { + return "dsv41-sync-test"; +} + +void sync_test_backend_synchronize(ggml_backend_t backend) { + auto * state = static_cast(backend->context); + state->synchronize_count++; + try { + state->runtime->remap(0, { 2 }); + } catch (const std::exception &) { + state->saw_pinned = true; + } +} + +const char * sync_test_device_name(ggml_backend_dev_t) { + return "dsv41-sync-test"; +} + +enum ggml_backend_dev_type sync_test_device_type(ggml_backend_dev_t) { + return GGML_BACKEND_DEVICE_TYPE_CPU; +} + +bool sync_test_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { + return true; +} + +bool sync_test_device_supports_buft(ggml_backend_dev_t, ggml_backend_buffer_type_t buft) { + return buft == ggml_backend_cpu_buffer_type(); +} + +void test_release_after_sync(const fixture & f) { + auto runtime = f.make_runtime(1); + REQUIRE(runtime.remap(0, { 1 }).front() == 0); + + sync_test_context state = { &runtime }; + ggml_backend_device device = {}; + device.iface.get_name = sync_test_device_name; + device.iface.get_type = sync_test_device_type; + device.iface.supports_op = sync_test_device_supports_op; + device.iface.supports_buft = sync_test_device_supports_buft; + ggml_backend backend = {}; + backend.iface.get_name = sync_test_backend_name; + backend.iface.synchronize = sync_test_backend_synchronize; + backend.device = &device; + backend.context = &state; + ggml_backend_t backends[] = { &backend }; + ggml_backend_buffer_type_t bufts[] = { ggml_backend_cpu_buffer_type() }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, 1, 16, false, false); + REQUIRE(sched != nullptr); + + runtime.release_all_after_sync(sched); + REQUIRE(state.synchronize_count == 1); + REQUIRE(state.saw_pinned); + REQUIRE(runtime.remap(0, { 2 }).front() == 0); + runtime.release(0); + ggml_backend_sched_free(sched); +} + +void test_original_and_slot_ids() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * cache = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, 2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 2, 1); + ggml_tensor * slots = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_tensor * probs = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 4, 1); + ggml_tensor * original = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + for (ggml_tensor * tensor : { cache, input, slots, probs, original }) { + ggml_set_input(tensor); + } + + ggml_tensor * expert_values = ggml_mul_mat_id(ctx, cache, input, slots); + ggml_tensor * routing_weights = ggml_get_rows(ctx, probs, original); + ggml_set_output(expert_values); + ggml_set_output(routing_weights); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, expert_values); + ggml_build_forward_expand(graph, routing_weights); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + const float cache_data[] = { 10.0f, 20.0f }; + const float input_data[] = { 1.0f, 1.0f }; + const int32_t slot_data[] = { 1, 0 }; + const float prob_data[] = { 1.0f, 2.0f, 3.0f, 4.0f }; + const int32_t original_data[] = { 3, 1 }; + ggml_backend_tensor_set(cache, cache_data, 0, sizeof(cache_data)); + ggml_backend_tensor_set(input, input_data, 0, sizeof(input_data)); + ggml_backend_tensor_set(slots, slot_data, 0, sizeof(slot_data)); + ggml_backend_tensor_set(probs, prob_data, 0, sizeof(prob_data)); + ggml_backend_tensor_set(original, original_data, 0, sizeof(original_data)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + + float expert_result[2] = {}; + float weight_result[2] = {}; + ggml_backend_tensor_get(expert_values, expert_result, 0, sizeof(expert_result)); + ggml_backend_tensor_get(routing_weights, weight_result, 0, sizeof(weight_result)); + REQUIRE(expert_result[0] == 20.0f && expert_result[1] == 10.0f); + REQUIRE(weight_result[0] == 4.0f && weight_result[1] == 2.0f); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +} + +int main() { + try { + test_registration(); + fixture f; + test_configuration(f); + test_remap_upload_and_eviction(f); + test_capacity_and_upload_failure(f); + test_publication_allocation_failure(f); + test_graph_callbacks(f); + test_graph_upload_failure_sentinel(f); + test_grovemoe_lookup_ids(); + test_release_after_sync(f); + test_original_and_slot_ids(); + } catch (const std::exception & error) { + std::fprintf(stderr, "test-deepseek41-expert: %s\n", error.what()); + return 1; + } + return 0; +} diff --git a/tests/test-deepseek41-memory.cpp b/tests/test-deepseek41-memory.cpp new file mode 100644 index 000000000000..057b9e4848f4 --- /dev/null +++ b/tests/test-deepseek41-memory.cpp @@ -0,0 +1,736 @@ +#include "../src/llama-batch.h" +#include "../src/llama-io.h" +#include "../src/llama-memory-dsv41.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#endif + +static void check(bool condition, const char * message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message); + std::exit(1); + } +} + +static void expect_invalid(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::invalid_argument &) { + return; + } + check(false, message); +} + +static void expect_runtime(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_dsv41_memory_config small_config( + uint32_t n_ctx = 256, + uint32_t n_seq = 3, + uint32_t n_ubatch = 128) { + llama_dsv41_memory_config config; + config.n_ctx = n_ctx; + config.n_seq = n_seq; + config.n_ubatch = n_ubatch; + config.kv_width = 8; + config.index_width = 4; + config.candidate_topk_blocks = 4; + config.candidate_block_size = 8; + config.ratios.resize(LLAMA_DSV41_N_LAYER); + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + config.ratios[il] = llama_dsv41_compress_ratio(il); + } + config.buft_for_layer = [](int32_t) { return ggml_backend_cpu_buffer_type(); }; + return config; +} + +static llama_ubatch make_ubatch(llama_pos start, uint32_t count, llama_seq_id seq_id) { + llama_batch_allocr allocator(1); + llama_ubatch ubatch = allocator.ubatch_reserve(count, 1); + ubatch.data->seq_id_data.resize(count); + ubatch.data->seq_id_unq = { seq_id }; + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.seq_idx[seq_id] = 0; + for (uint32_t i = 0; i < count; ++i) { + ubatch.token[i] = 10 + i; + ubatch.pos[i] = start + i; + ubatch.n_seq_id[i] = 1; + ubatch.data->seq_id_data[i] = seq_id; + ubatch.seq_id[i] = &ubatch.data->seq_id_data[i]; + ubatch.output[i] = 1; + } + return ubatch; +} + +static llama_ubatch make_coupled_ubatch( + llama_pos start, + uint32_t count, + llama_seq_id first, + llama_seq_id second) { + llama_ubatch ubatch = make_ubatch(start, count, first); + ubatch.data->seq_id_data.resize((size_t) count*2); + ubatch.data->seq_id_unq = { first, second }; + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.n_seqs_unq = 2; + ubatch.seq_idx[first] = 0; + ubatch.seq_idx[second] = 1; + for (uint32_t i = 0; i < count; ++i) { + ubatch.n_seq_id[i] = 2; + ubatch.data->seq_id_data[2*i] = first; + ubatch.data->seq_id_data[2*i + 1] = second; + ubatch.seq_id[i] = ubatch.data->seq_id_data.data() + 2*i; + } + return ubatch; +} + +static size_t state_sequence_tensor_bytes(const llama_memory_dsv41 & memory) { + size_t result = 0; + const auto add = [&](const ggml_tensor * tensor) { + result += sizeof(uint64_t); + result += tensor->ne[2] == (int64_t) memory.config().n_seq ? + tensor->nb[2] : tensor->nb[1]; + }; + for (uint32_t il = 0; il < memory.config().n_layer; ++il) { + add(memory.raw_k(il)); + } + for (uint32_t source : memory.config().kv_sources) { + add(memory.compressed_kv(source)); + add(memory.index_keys(source)); + if (memory.config().ratios[source] == 2) { + add(memory.compressor_carry_kv(source)); + add(memory.compressor_carry_score(source)); + } + } + add(memory.committed_candidate_ids()); + add(memory.position_state()); + return result; +} + +class vector_writer : public llama_io_write_i { +public: + void write(const void * src, size_t size) override { + const uint8_t * bytes = static_cast(src); + data.insert(data.end(), bytes, bytes + size); + } + + void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + const size_t old_size = data.size(); + data.resize(old_size + size); + ggml_backend_tensor_get(tensor, data.data() + old_size, offset, size); + } + + size_t n_bytes() override { + return data.size(); + } + + std::vector data; +}; + +class vector_reader : public llama_io_read_i { +public: + explicit vector_reader(const std::vector & data) : data(data) {} + + void read(void * dst, size_t size) override { + if (offset > data.size() || size > data.size() - offset) { + throw std::runtime_error("test state buffer is truncated"); + } + std::memcpy(dst, data.data() + offset, size); + offset += size; + } + + void read_tensor(ggml_tensor * tensor, size_t tensor_offset, size_t size) override { + (void) tensor; + (void) tensor_offset; + (void) size; + throw std::runtime_error("DeepSeek V4.1 state restore must read tensor bytes before publication"); + } + + size_t n_bytes() override { + return offset; + } + +private: + const std::vector & data; + size_t offset = 0; +}; + +class device_writer : public llama_io_write_i { +public: + void write(const void * src, size_t size) override { + const uint8_t * bytes = static_cast(src); + metadata.insert(metadata.end(), bytes, bytes + size); + } + + void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + tensors.emplace_back(size); + ggml_backend_tensor_get(tensor, tensors.back().data(), offset, size); + tensor_bytes += size; + } + + size_t n_bytes() override { + return metadata.size(); + } + + std::vector metadata; + std::vector> tensors; + size_t tensor_bytes = 0; +}; + +class device_reader : public llama_io_read_i { +public: + explicit device_reader(const device_writer & writer) : + metadata(writer.metadata), + tensors(writer.tensors) { + } + + void read(void * dst, size_t size) override { + if (offset > metadata.size() || size > metadata.size() - offset) { + throw std::runtime_error("test on-device metadata is truncated"); + } + std::memcpy(dst, metadata.data() + offset, size); + offset += size; + } + + void read_tensor(ggml_tensor * tensor, size_t tensor_offset, size_t size) override { + if (i_tensor >= tensors.size() || tensors[i_tensor].size() != size) { + throw std::runtime_error("test on-device tensor layout differs"); + } + ggml_backend_tensor_set(tensor, tensors[i_tensor].data(), tensor_offset, size); + ++i_tensor; + } + + size_t n_bytes() override { + return offset; + } + + size_t tensor_reads() const { + return i_tensor; + } + +private: + const std::vector & metadata; + const std::vector> & tensors; + size_t offset = 0; + size_t i_tensor = 0; +}; + +static void test_transaction_commit_rollback() { + llama_memory_dsv41 memory(small_config()); + llama_ubatch ubatch = make_ubatch(0, 1, 0); + + auto full_base = memory.init_full(); + auto * full = dynamic_cast(full_base.get()); + check(full != nullptr && full->get_ubatch().n_tokens == memory.config().n_ubatch, + "full memory context did not expose a bounded reserve plan"); + check(full->plan().sources.size() == memory.config().kv_sources.size(), + "full memory context source plan mismatch"); + llama_ubatch reserve_decode = make_ubatch(0, 1, 0); + check(full->topology(reserve_decode, 1).n_tokens == 1 && + full->plan().positions.size() == 1, + "full memory context did not resize its synthetic decode plan"); + llama_ubatch reserve_prefill = make_ubatch(0, 7, 0); + check(full->graph_plan(reserve_prefill).positions.size() == 7 && + full->topology(reserve_prefill, 3).n_outputs == 3, + "full memory context did not resize its synthetic prefill plan"); + + const size_t row_bytes = memory.raw_k(0)->nb[1]; + std::vector original(row_bytes, 0x31); + std::vector changed(row_bytes, 0x72); + std::vector actual(row_bytes); + ggml_backend_tensor_set(memory.raw_k(0), original.data(), 0, row_bytes); + + llama_memory_dsv41_context rollback_context(&memory, std::vector { ubatch }); + check(rollback_context.apply(), "transaction prepare failed"); + check(memory.seq_pos_max(0) == -1, "prepare published the position early"); + ggml_backend_tensor_set(memory.raw_k(0), changed.data(), 0, row_bytes); + rollback_context.stage_candidate_ids(0, { 7 }); + rollback_context.rollback(); + ggml_backend_tensor_get(memory.raw_k(0), actual.data(), 0, row_bytes); + check(actual == original, "rollback did not restore a graph-written raw row"); + check(memory.seq_pos_max(0) == -1, "rollback changed the committed position"); + check(memory.sequence_candidate_ids(0).empty(), "rollback published candidate IDs"); + + llama_memory_dsv41_context commit_context(&memory, std::vector { ubatch }); + check(commit_context.apply(), "commit prepare failed"); + commit_context.stage_candidate_ids(0, { 3 }); + const auto topology = commit_context.topology(ubatch, 1); + check(topology.n_tokens == 1 && topology.n_seqs == 1 && topology.n_outputs == 1, + "graph topology batch identity mismatch"); + check(topology.start_positions == std::vector({ 0 }), + "graph topology start position mismatch"); + auto reusable_topology = topology; + reusable_topology.seq_ids = { 2 }; + reusable_topology.start_positions = { 17 }; + reusable_topology.transaction_generation++; + check(topology.same_topology(reusable_topology), + "graph reuse rejected refreshed sequence and position inputs"); + reusable_topology.candidate_width++; + check(!topology.same_topology(reusable_topology), + "graph reuse accepted a changed candidate workspace shape"); + commit_context.commit(); + check(memory.seq_pos_max(0) == 0, "commit did not publish the position"); + check(memory.sequence_candidate_ids(0) == std::vector({ 3 }), + "commit did not publish candidate IDs"); + + llama_memory_dsv41_context candidate_rollback( + &memory, std::vector { make_ubatch(1, 1, 0) }); + check(candidate_rollback.apply(), "candidate rollback prepare failed"); + candidate_rollback.stage_candidate_ids(0, { 9 }); + candidate_rollback.rollback(); + check(memory.sequence_candidate_ids(0) == std::vector({ 3 }), + "candidate rollback changed committed IDs"); + + llama_memory_dsv41_context over_capacity( + &memory, + std::vector { + make_ubatch(memory.config().n_ctx, 1, 1), + }); + check(!over_capacity.apply(), "over-capacity transaction was accepted"); + check(memory.seq_pos_max(1) == -1, "over-capacity prepare mutated sequence state"); + + llama_memory_dsv41_context coupled( + &memory, + std::vector { + make_coupled_ubatch(0, 1, 1, 2), + }); + check(!coupled.apply(), "coupled sequence evaluation was accepted"); + check(memory.seq_pos_max(1) == -1 && memory.seq_pos_max(2) == -1, + "rejected coupled sequence evaluation mutated state"); +} + +static void test_window_compression_and_sequences() { + llama_memory_dsv41 memory(small_config()); + llama_ubatch prefill = make_ubatch(0, 128, 0); + llama_memory_dsv41_context prefill_context( + &memory, std::vector { prefill }); + check(prefill_context.apply(), "128-token prefill prepare failed"); + const auto & prefill_plan = prefill_context.plan(); + check(prefill_plan.raw.n_visible.back() == 128, "raw window width at 127 mismatch"); + check(prefill_plan.raw.read_idxs[127*128] == 0 && + prefill_plan.raw.read_idxs[127*128 + 127] == 127, + "raw window order at 127 mismatch"); + check(prefill_plan.sources[0].ratio == 2 && + prefill_plan.sources[0].compression.n_visible.back() == 64, + "ratio-2 visibility at 127 mismatch"); + prefill_context.stage_candidate_ids(127, { 4, 3, 2, 1 }); + prefill_context.commit(); + check(memory.seq_pos_max(0) == 127, "prefill position mismatch"); + + const size_t raw_row_bytes = memory.raw_k(0)->nb[1]; + const size_t carry_plane_bytes = memory.compressor_carry_kv(2)->nb[2]; + std::vector raw_before(raw_row_bytes, 0x21); + std::vector raw_after(raw_row_bytes, 0x72); + std::vector carry_before(carry_plane_bytes, 0x32); + std::vector carry_after(carry_plane_bytes, 0x83); + std::vector actual; + ggml_backend_tensor_set(memory.raw_k(0), raw_before.data(), 0, raw_row_bytes); + ggml_backend_tensor_set(memory.compressor_carry_kv(2), carry_before.data(), 0, carry_plane_bytes); + + llama_ubatch decode = make_ubatch(128, 1, 0); + llama_memory_dsv41_context decode_context( + &memory, std::vector { decode }); + check(decode_context.apply(), "boundary decode prepare failed"); + const auto & decode_plan = decode_context.plan(); + check(decode_plan.raw.write_idxs == std::vector({ 0 }), + "raw ring write did not wrap at 128"); + check(decode_plan.raw.read_idxs.front() == 1 && + decode_plan.raw.read_idxs[127] == 0, + "raw ring order at 128 mismatch"); + check(decode_plan.sources[0].compression.n_visible == std::vector({ 64 }), + "ratio-2 carry visibility at 128 mismatch"); + check(decode_plan.sources[0].compression.state_persist_dst_idxs == std::vector({ 0 }), + "ratio-2 carry publication row mismatch"); + ggml_backend_tensor_set(memory.raw_k(0), raw_after.data(), 0, raw_row_bytes); + ggml_backend_tensor_set(memory.compressor_carry_kv(2), carry_after.data(), 0, carry_plane_bytes); + decode_context.stage_candidate_ids(0, { 8, 7, 6, 5 }); + decode_context.commit(); + check(memory.seq_pos_max(0) == 128, "decode position mismatch"); + check(memory.seq_rm(0, 128, -1), "immediate rollback failed"); + check(memory.seq_pos_max(0) == 127, "immediate rollback position mismatch"); + check(memory.sequence_candidate_ids(0) == std::vector({ 4, 3, 2, 1 }), + "immediate rollback did not restore candidate state"); + actual.resize(raw_row_bytes); + ggml_backend_tensor_get(memory.raw_k(0), actual.data(), 0, raw_row_bytes); + check(actual == raw_before, "immediate rollback did not restore the raw ring row"); + actual.resize(carry_plane_bytes); + ggml_backend_tensor_get(memory.compressor_carry_kv(2), actual.data(), 0, carry_plane_bytes); + check(actual == carry_before, "immediate rollback did not restore compressor carry"); + + llama_memory_dsv41_context recommit_context( + &memory, std::vector { decode }); + check(recommit_context.apply(), "decode after rollback prepare failed"); + recommit_context.stage_candidate_ids(0, { 8, 7, 6, 5 }); + recommit_context.commit(); + memory.seq_cp(0, 1, -1, -1); + check(memory.seq_pos_max(1) == 128, "full sequence copy lost the position"); + check(memory.sequence_candidate_ids(1) == std::vector({ 8, 7, 6, 5 }), + "full sequence copy lost candidate state"); + memory.seq_keep(1); + check(memory.seq_pos_max(0) == -1 && memory.seq_pos_max(1) == 128, + "sequence keep retained another sequence"); + check(!memory.seq_rm(1, 128, -1), "copied sequence rollback was accepted without a retained snapshot"); + check(memory.seq_pos_max(1) == 128, "rejected copied-sequence rollback mutated the position"); + check(!memory.seq_rm(1, 64, -1), "non-immediate suffix rollback was accepted"); + check(memory.seq_pos_max(1) == 128, "rejected suffix rollback mutated the position"); + check(!memory.seq_rm(1, 64, 96), "interior range removal was accepted"); + check(!memory.seq_rm(-1, 0, -1), "partial wildcard removal was accepted"); + check(memory.seq_pos_max(1) == 128, "rejected wildcard removal mutated the position"); + expect_invalid( + [&] { memory.seq_cp(1, 2, 64, -1); }, + "partial sequence copy was accepted"); + check(memory.seq_rm(-1, -1, -1), "wildcard full-memory removal failed"); + check(memory.seq_pos_max(0) == -1 && + memory.seq_pos_max(1) == -1 && + memory.seq_pos_max(2) == -1, + "wildcard full-memory removal retained sequence state"); + llama_memory_dsv41_context negative_wildcard_context( + &memory, std::vector { make_ubatch(0, 1, 2) }); + check(negative_wildcard_context.apply(), "negative wildcard setup failed"); + negative_wildcard_context.commit(); + check(memory.seq_rm(-2, -1, -1), "negative wildcard full-memory removal failed"); + check(memory.seq_pos_max(2) == -1, "negative wildcard removal retained sequence state"); +} + +static void test_long_prefill_raw_publication() { + llama_memory_dsv41 memory(small_config(256, 1, 256)); + llama_memory_dsv41_context context( + &memory, + std::vector { make_ubatch(0, 256, 0) }); + check(context.apply(), "long prefill prepare failed"); + const auto & raw = context.plan().raw; + check(raw.persist_src_idxs.size() == LLAMA_DSV41_N_SWA, + "long prefill retained duplicate raw ring writes"); + check(raw.persist_src_idxs.front() == 128 && + raw.persist_src_idxs.back() == 255, + "long prefill did not retain the final raw window"); + std::vector unique = raw.write_idxs; + std::sort(unique.begin(), unique.end()); + check(std::adjacent_find(unique.begin(), unique.end()) == unique.end(), + "long prefill raw ring write indexes are not unique"); + context.rollback(); +} + +static void test_state_save_load() { + llama_memory_dsv41 memory(small_config()); + llama_ubatch ubatch = make_ubatch(0, 3, 0); + llama_memory_dsv41_context context(&memory, std::vector { ubatch }); + check(context.apply(), "state test prepare failed"); + context.stage_candidate_ids(2, { 11 }); + context.commit(); + + const size_t row_bytes = memory.raw_k(0)->nb[1]; + std::vector expected(row_bytes, 0x5a); + std::vector actual(row_bytes); + ggml_backend_tensor_set(memory.raw_k(0), expected.data(), 2*row_bytes, row_bytes); + + vector_writer writer; + memory.state_write(writer, 0); + expect_runtime( + [&] { + vector_reader nonempty_reader(writer.data); + memory.state_read(nonempty_reader, 0); + }, + "state restore overwrote a non-empty destination"); + memory.clear(true); + vector_reader reader(writer.data); + memory.state_read(reader, 2); + check(memory.seq_pos_max(2) == 2, "state restore lost the position"); + check(memory.sequence_candidate_ids(2) == std::vector({ 11 }), + "state restore lost candidate IDs"); + ggml_backend_tensor_get( + memory.raw_k(0), + actual.data(), + ((size_t) 2*memory.config().raw_window + 2)*row_bytes, + row_bytes); + check(actual == expected, "state restore lost raw cache data"); + + std::vector truncated = writer.data; + truncated.pop_back(); + expect_runtime( + [&] { + vector_reader truncated_reader(truncated); + memory.state_read(truncated_reader, 1); + }, + "truncated state restore was accepted"); + check(memory.seq_pos_max(1) == -1 && memory.seq_pos_max(2) == 2, + "failed state restore changed committed sequence state"); + + vector_writer full_writer; + memory.state_write(full_writer, -1); + memory.clear(true); + std::vector incomplete = full_writer.data; + const uint32_t incomplete_count = memory.config().n_seq - 1; + std::memcpy( + incomplete.data() + sizeof(uint64_t) + sizeof(uint32_t), + &incomplete_count, + sizeof(incomplete_count)); + expect_runtime( + [&] { + vector_reader incomplete_reader(incomplete); + memory.state_read(incomplete_reader, -1); + }, + "incomplete full state restore was accepted"); + for (uint32_t seq = 0; seq < memory.config().n_seq; ++seq) { + check(memory.seq_pos_max(seq) == -1, "incomplete full restore changed memory"); + } + + std::vector duplicate = full_writer.data; + const size_t header_size = + sizeof(uint64_t) + 3*sizeof(uint32_t); + const size_t first_record_size = + sizeof(llama_seq_id) + sizeof(llama_pos) + sizeof(uint32_t) + + sizeof(uint8_t) + state_sequence_tensor_bytes(memory); + const llama_seq_id duplicate_id = 0; + std::memcpy( + duplicate.data() + header_size + first_record_size, + &duplicate_id, + sizeof(duplicate_id)); + expect_runtime( + [&] { + vector_reader duplicate_reader(duplicate); + memory.state_read(duplicate_reader, -1); + }, + "duplicate full state sequence was accepted"); + for (uint32_t seq = 0; seq < memory.config().n_seq; ++seq) { + check(memory.seq_pos_max(seq) == -1, "duplicate full restore changed memory"); + } + + vector_reader full_reader(full_writer.data); + memory.state_read(full_reader, -1); + check(memory.seq_pos_max(2) == 2, "full state restore lost sequence state"); + + device_writer on_device_writer; + memory.state_write(on_device_writer, 2, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + vector_writer host_writer; + memory.state_write(host_writer, 2); + check(on_device_writer.metadata.size() + on_device_writer.tensor_bytes == host_writer.data.size(), + "on-device state metadata includes tensor payload bytes"); + memory.clear(true); + device_reader on_device_reader(on_device_writer); + memory.state_read(on_device_reader, 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + check(on_device_reader.tensor_reads() == on_device_writer.tensors.size(), + "on-device state restore did not consume every tensor"); + check(memory.seq_pos_max(1) == 2 && + memory.sequence_candidate_ids(1) == std::vector({ 11 }), + "on-device state restore lost sequence metadata"); + ggml_backend_tensor_get( + memory.raw_k(0), + actual.data(), + ((size_t) memory.config().raw_window + 2)*row_bytes, + row_bytes); + check(actual == expected, "on-device state restore lost raw cache data"); +} + +static void test_accounting() { + const auto config = small_config(); + llama_memory_dsv41 memory(small_config()); + memory.set_graph_workspace_size(1234); + const auto bytes = memory.accounting(); + const uint64_t raw_expected = + (uint64_t) config.n_layer*config.kv_width*config.raw_window*config.n_seq*sizeof(uint16_t); + const uint64_t compressed_rows = + (uint64_t) 3*((config.n_ctx + 1)/2) + config.n_ctx; + const uint64_t compressed_expected = + compressed_rows*config.kv_width*config.n_seq*sizeof(uint16_t); + const uint64_t index_expected = + compressed_rows*config.index_width*config.n_seq*sizeof(uint16_t); + const uint64_t carry_expected = + (uint64_t) 3*2*config.kv_width*config.n_seq*sizeof(uint16_t)*2; + const uint64_t candidate_scores_expected = + (uint64_t) (config.n_ctx/config.candidate_block_size)*config.n_ubatch*sizeof(float); + const uint64_t candidate_ids_expected = + (uint64_t) config.candidate_topk_blocks*(config.n_ubatch + config.n_seq)*sizeof(int32_t); + const uint64_t position_rows = + 1 + config.raw_window + 3*((config.n_ctx + 1)/2 + 2) + config.n_ctx; + const uint64_t position_expected = position_rows*config.n_seq*sizeof(int32_t); + + check(bytes.raw_kv == raw_expected, "raw cache accounting mismatch"); + check(bytes.compressed_kv == compressed_expected, "compressed cache accounting mismatch"); + check(bytes.index_keys == index_expected, "index key accounting mismatch"); + check(bytes.compressor_carry == carry_expected, "compressor carry accounting mismatch"); + check(bytes.candidate_scores == candidate_scores_expected, "candidate score accounting mismatch"); + check(bytes.candidate_ids == candidate_ids_expected, "candidate ID accounting mismatch"); + check(bytes.position_state == position_expected, "position accounting mismatch"); + check(bytes.graph_workspace == 1234, "graph workspace accounting mismatch"); + const auto breakdown = memory.memory_breakdown(); + check(breakdown.size() == 1 && + breakdown.begin()->second >= bytes.total() - bytes.graph_workspace, + "memory breakdown does not include all allocated state"); + + auto no_alloc_config = small_config(); + no_alloc_config.no_alloc = true; + llama_memory_dsv41 no_alloc_memory(std::move(no_alloc_config)); + check(no_alloc_memory.raw_k(0)->buffer == nullptr, + "no-allocation memory probe allocated cache storage"); + const auto no_alloc_bytes = no_alloc_memory.accounting(); + check(no_alloc_bytes.total() == bytes.total() - bytes.graph_workspace, + "no-allocation memory accounting differs from allocated state"); +} + +#if !defined(_WIN32) +static llama_engram_layout make_engram_layout() { + llama_engram_layout layout; + layout.encoding = LLAMA_DSV41_ENGRAM_ENCODING; + layout.layer_ids = { 1, 14 }; + layout.token_map.resize(32); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = i; + } + layout.compressed_vocab_size = 32; + layout.pad_id = 2; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + layout.multipliers[layer][i] = 101 + 8*layer + 2*i; + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 2; + layout.rows[layer] += 2; + } + } + return layout; +} + +struct engram_test_file { + std::string path = "test-dsv41-memory-engram.bin"; + int fd = -1; + std::array extents; + + explicit engram_test_file(const llama_engram_layout & layout) { + fd = open(path.c_str(), O_CREAT | O_TRUNC | O_RDWR, 0600); + check(fd >= 0, "failed to create Engram memory test file"); + uint64_t offset = 4096; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + extents[layer] = { + path, + offset, + layout.rows[layer], + LLAMA_ENGRAM_ROW_BYTES, + layout.rows[layer], + GGML_TYPE_I8, + }; + for (uint32_t row = 0; row < layout.rows[layer]; ++row) { + uint8_t data[LLAMA_ENGRAM_ROW_BYTES]; + std::fill(data, data + LLAMA_ENGRAM_ROW_BYTES, (uint8_t) (8 + row)); + check( + pwrite(fd, data, sizeof(data), offset + (uint64_t) row*sizeof(data)) == + (ssize_t) sizeof(data), + "failed to write Engram memory test row"); + } + offset += (uint64_t) layout.rows[layer]*LLAMA_ENGRAM_ROW_BYTES + 4096; + } + } + + ~engram_test_file() { + if (fd >= 0) { + close(fd); + } + unlink(path.c_str()); + } +}; + +static void test_engram_transaction() { + const llama_engram_layout layout = make_engram_layout(); + engram_test_file file(layout); + auto config = small_config(2048, 1, 4); + auto runtime = std::make_unique( + layout, file.extents, config.n_ubatch); + llama_dsv41_engram_runtime * runtime_ptr = runtime.get(); + config.engram = std::move(runtime); + llama_memory_dsv41 memory(std::move(config)); + + const size_t row_bytes = memory.raw_k(0)->nb[1]; + std::vector changed(row_bytes, 0x6b); + std::vector actual(row_bytes); + llama_memory_dsv41_context stale( + &memory, std::vector { make_ubatch(0, 1, 0) }); + check(stale.apply(), "stale Engram transaction prepare failed"); + ggml_backend_tensor_set(memory.raw_k(0), changed.data(), 0, row_bytes); + stale.stage_candidate_ids(0, { 9 }); + runtime_ptr->seq_remove(0); + expect_runtime([&] { stale.commit(); }, "stale Engram transaction commit was accepted"); + stale.rollback(); + ggml_backend_tensor_get(memory.raw_k(0), actual.data(), 0, row_bytes); + check(std::all_of(actual.begin(), actual.end(), [](uint8_t value) { return value == 0; }), + "failed commit did not restore graph-written state"); + check(memory.seq_pos_max(0) == -1 && memory.sequence_candidate_ids(0).empty(), + "failed commit published sequence state"); + + llama_memory_dsv41_context first( + &memory, std::vector { make_ubatch(0, 2, 0) }); + check(first.apply() && first.engram_transaction() != nullptr, + "Engram transaction was not exposed"); + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + const auto packed = first.engram_row_ids(layer); + const uint32_t * strided = first.engram_transaction()->row_ids(layer); + check(packed.size() == 2*LLAMA_ENGRAM_COLS, "packed Engram row ID shape mismatch"); + for (uint32_t token = 0; token < 2; ++token) { + check(std::memcmp( + packed.data() + token*LLAMA_ENGRAM_COLS, + strided + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(int32_t)) == 0, + "packed Engram row IDs changed the token stride"); + } + } + first.stage_candidate_ids(1, { 1 }); + first.commit(); + + llama_memory_dsv41_context rolled_back( + &memory, std::vector { make_ubatch(2, 1, 0) }); + check(rolled_back.apply(), "Engram rollback prepare failed"); + rolled_back.rollback(); + + llama_memory_dsv41_context second( + &memory, std::vector { make_ubatch(2, 1, 0) }); + check(second.apply(), "Engram state advanced during rollback"); + second.stage_candidate_ids(0, { 2 }); + second.commit(); + check(memory.seq_pos_max(0) == 2, "Engram commit position mismatch"); + check(memory.seq_rm(0, 2, -1), "Engram suffix rollback at transaction boundary failed"); + check(memory.seq_pos_max(0) == 1, "Engram suffix rollback did not restore position"); + + for (llama_pos pos = 2; pos < 1024; ++pos) { + llama_memory_dsv41_context decode( + &memory, std::vector { make_ubatch(pos, 1, 0) }); + check(decode.apply(), "long Engram decode prepare failed"); + decode.commit(); + check(memory.retained_rollback_count() == 1, + "Engram decode retained more than the immediate rollback state"); + } + check(memory.seq_rm(0, 1023, -1), "long Engram decode immediate rollback failed"); + check(memory.seq_pos_max(0) == 1022, "long Engram decode rollback restored the wrong position"); + check(memory.retained_rollback_count() == 0, + "Engram rollback retained an obsolete rollback state"); +} +#endif + +int main() { + test_transaction_commit_rollback(); + test_window_compression_and_sequences(); + test_long_prefill_raw_publication(); + test_state_save_load(); + test_accounting(); +#if !defined(_WIN32) + test_engram_transaction(); +#endif + return 0; +} diff --git a/tests/test-deepseek41-runtime.cpp b/tests/test-deepseek41-runtime.cpp new file mode 100644 index 000000000000..2b2579a3eaaa --- /dev/null +++ b/tests/test-deepseek41-runtime.cpp @@ -0,0 +1,703 @@ +#include "../src/llama-dsv41.h" +#include "../src/llama-arch.h" +#include "../tools/deepseek-v41-trace/trace-components.h" + +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +std::string llama_dsv41_graph_trace_name(const char * trace, uint32_t layer); +ggml_tensor * llama_dsv41_graph_append_zero_row(ggml_context * ctx, ggml_tensor * tensor); +ggml_tensor * llama_dsv41_graph_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type); + +static void check(bool condition, const std::string & message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message.c_str()); + std::exit(1); + } +} + +static void expect_throw(const std::function & fn, const std::string & message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_dsv41_config valid_config() { + llama_dsv41_config config = {}; + config.n_ctx_train = LLAMA_DSV41_N_CTX; + config.n_embd = LLAMA_DSV41_N_EMBD; + config.n_layer = LLAMA_DSV41_N_LAYER; + config.n_vocab = LLAMA_DSV41_N_VOCAB; + config.n_head = LLAMA_DSV41_N_HEAD; + config.n_head_kv = LLAMA_DSV41_N_HEAD_KV; + config.n_head_dim = LLAMA_DSV41_N_HEAD_DIM; + config.n_rot = LLAMA_DSV41_N_ROT; + config.n_lora_q = LLAMA_DSV41_N_LORA_Q; + config.n_lora_o = LLAMA_DSV41_N_LORA_O; + config.n_o_group = LLAMA_DSV41_N_O_GROUP; + config.n_ff_dense = LLAMA_DSV41_N_FF_DENSE; + config.n_ff_expert = LLAMA_DSV41_N_FF_EXP; + config.n_expert = LLAMA_DSV41_N_EXPERT; + config.n_expert_used = LLAMA_DSV41_N_EXPERT_USED; + config.n_expert_shared = LLAMA_DSV41_N_EXPERT_SHARED; + config.indexer_n_head = LLAMA_DSV41_N_INDEX_HEAD; + config.indexer_head_size = LLAMA_DSV41_N_INDEX_HEAD_DIM; + config.indexer_top_k = LLAMA_DSV41_N_INDEX_TOP_K; + config.hc_count = LLAMA_DSV41_HC_MULT; + config.hc_sinkhorn_iters = LLAMA_DSV41_HC_SINKHORN_ITERS; + config.raw_window = LLAMA_DSV41_N_SWA; + config.candidate_source_layer = LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + config.candidate_topk_blocks = LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS; + config.candidate_block_size = LLAMA_DSV41_CANDIDATE_BLOCK_SIZE; + config.f_norm_rms_eps = 1.0e-20f; + config.hc_eps = 1.0e-6f; + config.swiglu_clamp = 10.0f; + config.routed_scale = 1.5f; + config.rope_theta = 10000.0f; + config.compress_rope_theta = 160000.0f; + config.yarn_factor = 16.0f; + config.yarn_beta_fast = 32.0f; + config.yarn_beta_slow = 1.0f; + config.yarn_original_context = 65536; + config.expert_weights_norm = true; + config.hidden_act = "silu"; + config.scoring_func = "sqrtsoftplus"; + config.topk_method = "noaux_tc"; + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + config.compress_ratios.push_back(llama_dsv41_compress_ratio(il)); + } + config.kv_sources = { 2, 8, 14, 20 }; + config.index_sources = { 2, 8, 14, 20, 24, 28, 32, 36 }; + config.engram_layers = { 1, 14 }; + config.engram_rows = { 384006168, 384016682 }; + config.engram_encoding = LLAMA_DSV41_ENGRAM_ENCODING; + config.engram_compressed_vocab_size = LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB; + config.engram_pad_id = LLAMA_DSV41_ENGRAM_PAD_ID; + config.engram_token_map_size = LLAMA_DSV41_N_VOCAB; + config.engram_primes_size = LLAMA_DSV41_ENGRAM_PRIMES_COUNT; + config.engram_multipliers_size = LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT; + return config; +} + +static void test_hparams() { + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + + llama_dsv41_validate_config(valid_config()); + check(valid_config().engram_rows == std::vector({ 384006168, 384016682 }), "published Engram rows mismatch"); + + llama_dsv41_config config = valid_config(); + config.compress_ratios.insert(config.compress_ratios.end(), 3, 0); + expect_throw([&]() { llama_dsv41_validate_config(config); }, "43-entry source-config compression layout was accepted"); + + config = valid_config(); + config.compress_ratios[20] = 2; + expect_throw([&]() { llama_dsv41_validate_config(config); }, "invalid ratio-1 boundary was accepted"); + + config = valid_config(); + config.kv_sources = { 2, 8, 20 }; + expect_throw([&]() { llama_dsv41_validate_config(config); }, "invalid KV source map was accepted"); + + config = valid_config(); + config.engram_primes_size = 24; + expect_throw([&]() { llama_dsv41_validate_config(config); }, "truncated Engram prime table was accepted"); + + check(llm_arch_is_hybrid(LLM_ARCH_DEEPSEEK41), "DeepSeek V4.1 must use hybrid context handling"); + check(!llm_arch_supports_rs_rollback(LLM_ARCH_DEEPSEEK41), + "DeepSeek V4.1 must not advertise partial rollback support"); +} + +static void test_source_maps() { + check(llama_dsv41_compress_ratio(0) == 0, "layer 0 ratio mismatch"); + check(llama_dsv41_compress_ratio(1) == 0, "layer 1 ratio mismatch"); + check(llama_dsv41_compress_ratio(2) == 2, "layer 2 ratio mismatch"); + check(llama_dsv41_compress_ratio(19) == 2, "layer 19 ratio mismatch"); + check(llama_dsv41_compress_ratio(20) == 1, "layer 20 ratio mismatch"); + + check(llama_dsv41_kv_source_layer(0) == -1, "layer 0 unexpectedly has a KV source"); + check(llama_dsv41_kv_source_layer(2) == 2, "layer 2 KV source mismatch"); + check(llama_dsv41_kv_source_layer(7) == 2, "layer 7 KV source mismatch"); + check(llama_dsv41_kv_source_layer(8) == 8, "layer 8 KV source mismatch"); + check(llama_dsv41_kv_source_layer(19) == 14, "layer 19 KV source mismatch"); + check(llama_dsv41_kv_source_layer(39) == 20, "layer 39 KV source mismatch"); + + check(llama_dsv41_index_source_layer(19) == 14, "layer 19 index source mismatch"); + check(llama_dsv41_index_source_layer(20) == 20, "layer 20 index source mismatch"); + check(llama_dsv41_index_source_layer(23) == 20, "layer 23 index source mismatch"); + check(llama_dsv41_index_source_layer(24) == 24, "layer 24 index source mismatch"); + check(llama_dsv41_index_source_layer(39) == 36, "layer 39 index source mismatch"); +} + +static void test_compression() { + const auto ratio_2 = llama_dsv41_build_compression_plan({ 0, 1, 2 }, 2, 1024); + check(ratio_2.n_visible == std::vector({ 0, 1, 1 }), "ratio-2 visible counts mismatch"); + check(ratio_2.write_idxs == std::vector({ 0 }), "ratio-2 write index mismatch"); + check(ratio_2.write_pos == std::vector({ 0 }), "ratio-2 compressed position mismatch"); + check(ratio_2.state_persist_dst_idxs == std::vector({ 0, 1 }), "ratio-2 state rows mismatch"); + + const auto ratio_1 = llama_dsv41_build_compression_plan({ 19, 20 }, 1, 1024); + check(ratio_1.n_visible == std::vector({ 20, 21 }), "ratio-1 visible counts mismatch"); + check(ratio_1.write_idxs == std::vector({ 19, 20 }), "ratio-1 write indexes mismatch"); + check(ratio_1.write_pos == std::vector({ 19, 20 }), "ratio-1 compressed positions mismatch"); + + const auto layer_0 = llama_dsv41_build_layer_plan(0, { 0 }, 1024); + check(layer_0.ratio == 0 && layer_0.compression.write_idxs.empty(), "layer 0 must use raw attention only"); + const auto layer_2 = llama_dsv41_build_layer_plan(2, { 0, 1 }, 1024); + check(layer_2.ratio == 2 && layer_2.owns_kv_source, "layer 2 compression ownership mismatch"); + check(layer_2.compression.write_idxs == std::vector({ 0 }), "layer 2 graph compression mismatch"); + const auto layer_20 = llama_dsv41_build_layer_plan(20, { 20 }, 1024); + check(layer_20.ratio == 1 && layer_20.owns_kv_source, "layer 20 compression ownership mismatch"); + check(layer_20.builds_candidates && !layer_20.uses_candidates, "layer 20 candidate propagation mismatch"); + const auto layer_21 = llama_dsv41_build_layer_plan(21, { 21 }, 1024); + check(!layer_21.uses_candidates && layer_21.reuses_index_selection, "layer 21 index reuse mismatch"); + const auto layer_24 = llama_dsv41_build_layer_plan(24, { 24 }, 1024); + check(!layer_24.owns_kv_source && layer_24.owns_index_source, "layer 24 source ownership mismatch"); + check(layer_24.uses_candidates, "layer 24 must consume layer-20 candidates"); + check(llama_dsv41_build_layer_plan(39, { 39 }, 1024).collapses_output, "final layer output collapse missing"); + expect_throw([&]() { llama_dsv41_build_layer_plan(20, { 20, 22 }, 1024); }, "non-contiguous graph plan was accepted"); +} + +static void test_state() { + llama_dsv41_cache_state state(1024); + for (llama_pos pos = 0; pos <= 129; ++pos) { + state.append(pos); + } + check(state.position() == 129, "cache position mismatch"); + check(state.raw_slots()[0] == 128 && state.raw_slots()[1] == 129, "raw ring state mismatch"); + check(state.compressed_slots(2)[0] == 0, "ratio-2 first compressed row mismatch"); + check(state.compressed_slots(2)[64] == 128, "ratio-2 boundary row mismatch"); + check(state.pending_slots(2) == std::vector({ 128, 129 }), "ratio-2 pending rows mismatch"); + check(state.compressed_slots(20)[129] == 129, "ratio-1 direct row mismatch"); + state.set_candidate_blocks({ 4, 1 }); + check(state.candidate_blocks() == std::vector({ 4, 1 }), "candidate state mismatch"); + expect_throw([&]() { state.append(131); }, "non-contiguous cache append was accepted"); + state.clear(); + check(state.position() == -1 && state.raw_slots()[0] == -1, "cache clear mismatch"); + + llama_dsv41_cache_state small(1); + small.append(0); + expect_throw([&]() { small.append(1); }, "compressed cache overflow was accepted"); + check(small.position() == 0 && small.raw_slots()[1] == -1, "failed cache append mutated state"); + + const auto bytes = llama_dsv41_account_memory(32768, 1, 8192, 2, 2, 1234); + check(bytes.raw_kv > 0 && bytes.compressed_kv > 0 && bytes.index_keys > 0, "cache memory accounting is incomplete"); + check(bytes.compressor_carry > 0 && bytes.candidate_scores > 0 && bytes.candidate_ids > 0, "state memory accounting is incomplete"); + check(bytes.total() == bytes.raw_kv + bytes.compressed_kv + bytes.index_keys + bytes.compressor_carry + + bytes.candidate_scores + bytes.candidate_ids + bytes.position_state + bytes.graph_workspace, + "memory accounting total mismatch"); +} + +static void test_raw_ring() { + const auto at_127 = llama_dsv41_raw_ring_order(127, 128); + check(at_127.size() == 128 && at_127.front() == 0 && at_127.back() == 127, "raw ring at 127 mismatch"); + const auto at_128 = llama_dsv41_raw_ring_order(128, 128); + check(at_128.front() == 1 && at_128.back() == 0, "raw ring at 128 mismatch"); + const auto at_129 = llama_dsv41_raw_ring_order(129, 128); + check(at_129.front() == 2 && at_129.back() == 1, "raw ring at 129 mismatch"); +} + +static void test_candidates() { + for (uint32_t n_visible : { 1u, 7u, 8u, 9u, 127u, 16385u, 16392u, 17017u }) { + std::vector scores(n_visible); + for (uint32_t i = 0; i < n_visible; ++i) { + scores[i] = -(float) i; + } + const auto blocks = llama_dsv41_select_candidate_blocks(scores, n_visible, 8, 2048); + const int32_t final_block = (int32_t) ((n_visible - 1)/8); + check(std::find(blocks.begin(), blocks.end(), final_block) != blocks.end(), "final visible candidate block was dropped"); + check(blocks.size() == std::min(2048, (n_visible + 7)/8), "candidate block count mismatch"); + const auto rows = llama_dsv41_candidate_rows(blocks, n_visible, 8); + check(std::find(rows.begin(), rows.end(), (int32_t) n_visible - 1) != rows.end(), "final visible row was filtered"); + check(std::all_of(rows.begin(), rows.end(), [&](int32_t row) { return row >= 0 && (uint32_t) row < n_visible; }), "candidate rows crossed causal visibility"); + } + + const auto tie = llama_dsv41_select_candidate_blocks(std::vector(24, 1.0f), 24, 8, 2); + check(tie == std::vector({ 2, 0 }), "candidate tie-break or final block retention mismatch"); + + std::vector partial_scores(9, -100.0f); + partial_scores[0] = 100.0f; + const auto partial = llama_dsv41_select_candidate_blocks(partial_scores, 9, 8, 1); + check(partial == std::vector({ 1 }), "final partial candidate block was not forced"); + + std::vector full_scores(16392, -100.0f); + full_scores[0] = 100.0f; + const auto full = llama_dsv41_select_candidate_blocks(full_scores, 16392, 8, 2048); + check(std::find(full.begin(), full.end(), 2048) != full.end(), "final full candidate block was not forced"); + + const auto inf_tie = llama_dsv41_select_candidate_blocks( + std::vector(24, std::numeric_limits::infinity()), 24, 8, 2); + check(std::find(inf_tie.begin(), inf_tie.end(), 2) != inf_tie.end(), "final candidate block lost an infinity tie"); + + const auto inf_boundary = llama_dsv41_select_candidate_blocks( + std::vector(16392, std::numeric_limits::infinity()), 16392, 8, 2048); + check(std::find(inf_boundary.begin(), inf_boundary.end(), 2048) != inf_boundary.end(), + "final full candidate block lost an infinity tie"); + std::vector unique_blocks = inf_boundary; + std::sort(unique_blocks.begin(), unique_blocks.end()); + check(std::adjacent_find(unique_blocks.begin(), unique_blocks.end()) == unique_blocks.end(), + "candidate selection contains duplicate blocks"); + check(std::all_of(unique_blocks.begin(), unique_blocks.end(), [](int32_t block) { + return block >= 0 && block <= 2048; + }), "candidate selection contains an out-of-range block"); +} + +static void test_output_collapse() { + const std::vector residual = { + 1.0f, 2.0f, + 3.0f, 4.0f, + 5.0f, 6.0f, + 7.0f, 8.0f, + }; + const auto result = llama_dsv41_output_collapse(residual, { 0.1f, 0.2f, 0.3f, 0.4f }, 2, 4); + check(result.size() == 2, "output collapse width mismatch"); + check(std::abs(result[0] - 5.0f) < 1.0e-6f, "output collapse first value mismatch"); + check(std::abs(result[1] - 6.0f) < 1.0e-6f, "output collapse second value mismatch"); +} + +static void test_candidate_graph_selection(uint32_t n_visible, uint32_t n_candidate) { + ggml_init_params params = { + /*.mem_size =*/ 4*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create candidate graph context"); + + ggml_tensor * scores = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_visible, 1); + std::fill_n( + static_cast(scores->data), + ggml_nelements(scores), + std::numeric_limits::infinity()); + ggml_tensor * block_scores = ggml_pool_1d( + ctx, scores, GGML_OP_POOL_MAX, 8, 8, 0); + ggml_tensor * final_blocks = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 1, 1); + const int32_t final_block = (int32_t) block_scores->ne[0] - 1; + static_cast(final_blocks->data)[0] = final_block; + ggml_tensor * candidate_blocks = llama_dsv41_build_candidate_blocks( + ctx, block_scores, final_blocks, n_candidate); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, candidate_blocks); + check( + ggml_graph_compute_with_ctx(ctx, gf, 1) == GGML_STATUS_SUCCESS, + "candidate selection graph execution failed"); + + const int32_t * ids = static_cast(candidate_blocks->data); + std::vector selected(ids, ids + n_candidate); + check(selected.front() == final_block, "final candidate block is not first"); + check(std::count(selected.begin(), selected.end(), final_block) == 1, + "final candidate block was not retained exactly once"); + std::vector unique = selected; + std::sort(unique.begin(), unique.end()); + check(std::adjacent_find(unique.begin(), unique.end()) == unique.end(), + "candidate graph selected duplicate blocks"); + check(std::all_of(unique.begin(), unique.end(), [&](int32_t block) { + return block >= 0 && block <= final_block; + }), "candidate graph selected an out-of-range block"); + + ggml_free(ctx); +} + +static void test_graph_contract() { + enum stage { + STAGE_ENGRAM, + STAGE_ATTN_HC, + STAGE_CARRIED_PRE, + STAGE_ATTN, + STAGE_ATTN_POST, + STAGE_FFN_HC, + STAGE_ATTN_PRE, + STAGE_ROUTED_EXPERTS, + STAGE_SHARED_EXPERT, + STAGE_FFN_POST, + STAGE_CARRY_FFN_PRE, + }; + + const std::vector common = { + STAGE_ATTN_HC, + STAGE_CARRIED_PRE, + STAGE_ATTN, + STAGE_ATTN_POST, + STAGE_FFN_HC, + STAGE_ATTN_PRE, + STAGE_ROUTED_EXPERTS, + STAGE_SHARED_EXPERT, + STAGE_FFN_POST, + STAGE_CARRY_FFN_PRE, + }; + uint32_t ratio_count[3] = {}; + uint32_t kv_sources = 0; + uint32_t index_sources = 0; + uint32_t candidate_sources = 0; + std::vector candidate_trace_layers; + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + std::vector stages = common; + if (il == 1 || il == 14) { + stages.insert(stages.begin(), STAGE_ENGRAM); + check(stages[0] == STAGE_ENGRAM && stages[1] == STAGE_ATTN_HC, + "Engram must precede attention HC"); + } + check(stages[stages.size() - 1] == STAGE_CARRY_FFN_PRE, + "FFN pre must be carried to the next layer"); + check(stages[1 + (il == 1 || il == 14)] == STAGE_CARRIED_PRE, + "attention must collapse with carried pre"); + check(stages[5 + (il == 1 || il == 14)] == STAGE_ATTN_PRE, + "FFN must collapse with current attention pre"); + check(stages[7 + (il == 1 || il == 14)] == STAGE_SHARED_EXPERT, + "shared expert must be added after routed experts"); + + const uint32_t ratio = llama_dsv41_compress_ratio(il); + ratio_count[ratio]++; + kv_sources += llama_dsv41_kv_source_layer(il) == (int32_t) il; + index_sources += llama_dsv41_index_source_layer(il) == (int32_t) il; + candidate_sources += il == LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + + check( + llama_dsv41_graph_trace_name("expert.ids", il) == + "dsv41.trace.expert.ids.l" + std::to_string(il), + "expert ID trace name is unstable"); + const auto expert_ids = dsv41_trace_parse_name(llama_dsv41_graph_trace_name("expert.ids", il)); + check(expert_ids && expert_ids->component == "expert.ids" && + expert_ids->layer == (int) il && + std::string(expert_ids->semantic_id_space) == "original", + "exporter does not recognize original expert ID trace"); + check( + llama_dsv41_graph_trace_name("expert.weights", il) == + "dsv41.trace.expert.weights.l" + std::to_string(il), + "expert weight trace name is unstable"); + const auto expert_weights = dsv41_trace_parse_name(llama_dsv41_graph_trace_name("expert.weights", il)); + check(expert_weights && expert_weights->component == "expert.weights" && + expert_weights->layer == (int) il, + "exporter does not recognize expert weight trace"); + check( + llama_dsv41_graph_trace_name("attn.source", il) == + "dsv41.trace.attn.source.l" + std::to_string(il), + "attention source trace name is unstable"); + const auto attention_source = dsv41_trace_parse_name(llama_dsv41_graph_trace_name("attn.source", il)); + check(attention_source && attention_source->component == "attn.source" && + attention_source->layer == (int) il, + "exporter does not recognize attention source trace"); + if (il > LLAMA_DSV41_CANDIDATE_SOURCE_LAYER && + llama_dsv41_index_source_layer(il) == (int32_t) il) { + candidate_trace_layers.push_back(il); + check( + llama_dsv41_graph_trace_name("attn.candidates", il) == + "dsv41.trace.attn.candidates.l" + std::to_string(il), + "attention candidate trace name is unstable"); + const auto candidates = dsv41_trace_parse_name(llama_dsv41_graph_trace_name("attn.candidates", il)); + check(candidates && candidates->component == "attn.candidates" && + candidates->layer == (int) il, + "exporter does not recognize propagated candidate trace"); + } + } + check(ratio_count[0] == 2 && ratio_count[1] == 20 && ratio_count[2] == 18, + "ratio 0/1/2 layer counts mismatch"); + check(kv_sources == 4, "KV source ownership count mismatch"); + check(index_sources == 8, "index source ownership count mismatch"); + check(candidate_sources == 1, "candidate source ownership count mismatch"); + check(candidate_trace_layers == std::vector({ 24, 28, 32, 36 }), + "attention candidate trace layer coverage mismatch"); + check( + llama_dsv41_graph_trace_name("attn.candidate_blocks", 20) == + "dsv41.trace.attn.candidate_blocks.l20", + "candidate block trace name is unstable"); + const auto candidate_blocks = + dsv41_trace_parse_name(llama_dsv41_graph_trace_name("attn.candidate_blocks", 20)); + check(candidate_blocks && candidate_blocks->component == "attn.candidate_blocks" && + candidate_blocks->layer == 20, + "exporter does not recognize candidate block trace"); + check( + llama_dsv41_graph_trace_name("engram.row_ids", 1) == + "dsv41.trace.engram.row_ids.l1" && + llama_dsv41_graph_trace_name("engram.row_ids", 14) == + "dsv41.trace.engram.row_ids.l14", + "Engram row trace names are unstable"); + for (uint32_t layer : { 1u, 14u }) { + const auto engram = dsv41_trace_parse_name(llama_dsv41_graph_trace_name("engram.row_ids", layer)); + check(engram && engram->component == "engram.row_ids" && + engram->layer == (int) layer, + "exporter does not recognize Engram row trace"); + } + expect_throw( + [] { + dsv41_trace_select_name("dsv41.trace.attn.candidates.l20"); + }, + "exporter accepted an unexpected candidate trace layer"); + expect_throw( + [] { + dsv41_trace_select_name("dsv41.trace.attn.candidates.layer24"); + }, + "exporter accepted a malformed trace layer suffix"); + expect_throw( + [] { + dsv41_trace_select_name("dsv41.trace.unknown.l24"); + }, + "exporter accepted an unknown reserved trace tensor name"); + check(!dsv41_trace_select_name("dsv41_attn_candidates_l24"), + "exporter treated an ordinary graph tensor as reserved"); + check(llama_dsv41_build_layer_plan(39, { 39 }, 1024).collapses_output, + "final layer must preserve streams for carried-pre output collapse"); +} + +static void test_graph_construction() { + ggml_init_params params = { + /*.mem_size =*/ 4*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create graph test context"); + + ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 2, 3); + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 2, 3); + ggml_tensor * pooled = llama_dsv41_build_ratio_pool(ctx, kv, gate, 2); + check(pooled->ne[0] == 512 && pooled->ne[1] == 3, "ratio-2 graph output shape mismatch"); + + ggml_tensor * direct = llama_dsv41_build_ratio_pool( + ctx, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 1, 3), nullptr, 1); + check(direct->ne[0] == 512 && direct->ne[1] == 3, "ratio-1 graph output shape mismatch"); + + ggml_tensor * compressed_scores = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 64, 2); + ggml_tensor * raw_scores = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 128, 64, 2); + ggml_tensor * probs = llama_dsv41_build_shared_softmax(ctx, raw_scores, compressed_scores); + check(probs->ne[0] == 160 && probs->ne[1] == 64 && probs->ne[2] == 2, "shared-softmax graph shape mismatch"); + + ggml_tensor * raw_order = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); + ggml_tensor * compressed_order = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + const float raw_values[] = { 1.0f, 2.0f }; + const float compressed_value = 3.0f; + std::memcpy(raw_order->data, raw_values, sizeof(raw_values)); + std::memcpy(compressed_order->data, &compressed_value, sizeof(compressed_value)); + ggml_tensor * ordered_probs = llama_dsv41_build_shared_softmax(ctx, raw_order, compressed_order); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, ordered_probs); + check(ggml_graph_compute_with_ctx(ctx, gf, 1) == GGML_STATUS_SUCCESS, "shared-softmax graph execution failed"); + const float * ordered = static_cast(ordered_probs->data); + check(ordered[0] < ordered[1] && ordered[1] < ordered[2], "shared-softmax segment order mismatch"); + + ggml_tensor * f16_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 32, 1); + std::vector cache_values(32, 1.0f); + ggml_fp32_to_fp16_row( + cache_values.data(), + static_cast(f16_cache->data), + cache_values.size()); + ggml_tensor * cache_with_sentinel = + llama_dsv41_graph_append_zero_row(ctx, f16_cache); + ggml_tensor * sentinel_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2); + static_cast(sentinel_ids->data)[0] = 1; + static_cast(sentinel_ids->data)[1] = 0; + ggml_tensor * sentinel_rows = ggml_get_rows(ctx, cache_with_sentinel, sentinel_ids); + ggml_tensor * completion_zero = + llama_dsv41_graph_completion_zero(ctx, sentinel_rows, GGML_TYPE_F16); + ggml_cgraph * support_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(support_gf, sentinel_rows); + ggml_build_forward_expand(support_gf, completion_zero); + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create graph support backend"); + for (int i = 0; i < ggml_graph_n_nodes(support_gf); ++i) { + ggml_tensor * node = ggml_graph_node(support_gf, i); + if (node->op == GGML_OP_SCALE) { + check(node->src[0]->type == GGML_TYPE_F32, + "DeepSeek V4.1 graph contains a non-F32 SCALE input"); + } + check(ggml_backend_supports_op(backend, node), + "CPU backend does not support a DeepSeek V4.1 dependency node"); + } + ggml_backend_free(backend); + check(cache_with_sentinel->ne[1] == 2, + "compressed cache sentinel row was not allocated"); + check(ggml_graph_compute_with_ctx(ctx, support_gf, 1) == GGML_STATUS_SUCCESS, + "compressed sentinel graph execution failed"); + for (uint32_t i = 0; i < 32; ++i) { + check(ggml_get_f32_1d(sentinel_rows, i) == 0.0f, + "compressed sentinel row is not zero"); + check(ggml_get_f32_1d(sentinel_rows, 32 + i) == 1.0f, + "compressed real row changed"); + } + + ggml_tensor * prior_ring = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, 2); + ggml_tensor * current_k = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, 1); + ggml_tensor * read_idx = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, 1); + ggml_tensor * write_idx = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, 1); + static_cast(prior_ring->data)[0] = 10.0f; + static_cast(prior_ring->data)[1] = 20.0f; + static_cast(current_k->data)[0] = 30.0f; + static_cast(read_idx->data)[0] = 0; + static_cast(write_idx->data)[0] = 0; + ggml_tensor * prior_read = ggml_get_rows( + ctx, prior_ring, read_idx); + ggml_tensor * attention = ggml_add( + ctx, prior_read, current_k); + ggml_tensor * completion = ggml_argsort_top_k( + ctx, ggml_view_1d(ctx, attention, 1, 0), 1); + completion = ggml_scale( + ctx, ggml_cast(ctx, completion, GGML_TYPE_F32), 0.0f); + ggml_tensor * delayed_k = ggml_add( + ctx, current_k, completion); + ggml_tensor * ring_update = ggml_set_rows( + ctx, prior_ring, delayed_k, write_idx); + ggml_cgraph * ring_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(ring_gf, ring_update); + check( + ggml_graph_compute_with_ctx(ctx, ring_gf, 1) == + GGML_STATUS_SUCCESS, + "ordered raw-ring update graph execution failed"); + check( + static_cast(attention->data)[0] == 40.0f && + static_cast(prior_ring->data)[0] == 30.0f, + "raw-ring write did not wait for the prior-ring read"); + + ggml_tensor * selected_ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 3, 2); + const int32_t selected_values[] = { 5, 1, 3, 4, 0, 2 }; + std::memcpy(selected_ids->data, selected_values, sizeof(selected_values)); + ggml_tensor * selected_order = ggml_argsort( + ctx, ggml_cast(ctx, selected_ids, GGML_TYPE_F32), + GGML_SORT_ORDER_ASC); + ggml_tensor * selected_sorted = ggml_get_rows( + ctx, ggml_reshape_3d(ctx, selected_ids, 1, 3, 2), + selected_order); + selected_sorted = ggml_cont( + ctx, ggml_reshape_2d(ctx, selected_sorted, 3, 2)); + ggml_tensor * routing_probs = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 6, 2); + const float routing_values[] = { + 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, + 20.0f, 21.0f, 22.0f, 23.0f, 24.0f, 25.0f, + }; + std::memcpy( + routing_probs->data, routing_values, + sizeof(routing_values)); + ggml_tensor * selected_weights = ggml_get_rows( + ctx, + ggml_reshape_3d(ctx, routing_probs, 1, 6, 2), + selected_sorted); + selected_weights = ggml_cont( + ctx, ggml_reshape_2d(ctx, selected_weights, 3, 2)); + ggml_cgraph * selected_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(selected_gf, selected_sorted); + ggml_build_forward_expand(selected_gf, selected_weights); + check( + ggml_graph_compute_with_ctx(ctx, selected_gf, 1) == + GGML_STATUS_SUCCESS, + "selected ID ordering graph execution failed"); + const int32_t selected_expected[] = { 1, 3, 5, 0, 2, 4 }; + check( + std::memcmp( + selected_sorted->data, selected_expected, + sizeof(selected_expected)) == 0, + "selected IDs are not accumulated in original ID order"); + const float selected_weight_expected[] = { + 11.0f, 13.0f, 15.0f, 20.0f, 22.0f, 24.0f, + }; + check( + std::memcmp( + selected_weights->data, selected_weight_expected, + sizeof(selected_weight_expected)) == 0, + "routing weights are not paired with sorted original IDs"); + + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 4, 2); + ggml_tensor * pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 2); + ggml_tensor * collapsed = llama_dsv41_build_output_collapse(ctx, residual, pre, 32, 4, 2); + check(collapsed->type == GGML_TYPE_BF16, "output collapse BF16 boundary is missing"); + check(collapsed->ne[0] == 32 && collapsed->ne[1] == 2, "output collapse graph shape mismatch"); + + ggml_tensor * output_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 32); + ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 64); + ggml_tensor * logits = llama_dsv41_build_output(ctx, residual, pre, output_norm, output, 1.0e-20f, 4); + check(logits->ne[0] == 64 && logits->ne[1] == 2, "final output graph shape mismatch"); + + ggml_tensor * original_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 6, 2); + ggml_tensor * slot_ids = ggml_cont(ctx, original_ids); + check(ggml_is_contiguous(original_ids), "original expert IDs must be contiguous"); + check(ggml_is_contiguous(slot_ids) && slot_ids != original_ids, + "slot IDs must be a distinct contiguous remap"); + + ggml_tensor * routed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + ggml_tensor * shared = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + ggml_tensor * combined = ggml_add(ctx, routed, shared); + check(combined->src[0] == routed && combined->src[1] == shared, + "shared expert output is not added to routed output"); + + ggml_tensor * exec_residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 4, 1); + ggml_tensor * exec_pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 1); + ggml_tensor * exec_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 32); + ggml_tensor * exec_output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + std::fill_n(static_cast(exec_residual->data), ggml_nelements(exec_residual), 1.0f); + std::fill_n(static_cast(exec_pre->data), ggml_nelements(exec_pre), 0.25f); + std::fill_n(static_cast(exec_norm->data), ggml_nelements(exec_norm), 1.0f); + std::fill_n(static_cast(exec_output->data), ggml_nelements(exec_output), 1.0f); + ggml_tensor * exec_collapse = llama_dsv41_build_output_collapse( + ctx, exec_residual, exec_pre, 32, 4, 1); + ggml_tensor * exec_norm_input = llama_dsv41_build_output_norm_input( + ctx, exec_collapse); + check(exec_collapse->type == GGML_TYPE_BF16, + "production output collapse is not BF16"); + check(exec_norm_input->type == GGML_TYPE_F32, + "production output RMSNorm input is not F32"); + ggml_tensor * exec_normalized = ggml_rms_norm( + ctx, exec_norm_input, 1.0e-20f); + check(exec_normalized->src[0] == exec_norm_input, + "production RMSNorm does not consume the F32 collapse"); + exec_normalized = ggml_mul(ctx, exec_normalized, exec_norm); + ggml_tensor * exec_logits = ggml_mul_mat( + ctx, exec_output, exec_normalized); + ggml_cgraph * exec_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(exec_gf, exec_logits); + ggml_backend_t exec_backend = ggml_backend_cpu_init(); + check(exec_backend != nullptr, "failed to create output graph CPU backend"); + for (int i = 0; i < ggml_graph_n_nodes(exec_gf); ++i) { + check(ggml_backend_supports_op(exec_backend, ggml_graph_node(exec_gf, i)), + "CPU backend does not support the production output graph"); + } + ggml_backend_free(exec_backend); + check(ggml_graph_compute_with_ctx(ctx, exec_gf, 1) == GGML_STATUS_SUCCESS, "final output graph execution failed"); + const float * exec_values = static_cast(exec_logits->data); + check(std::isfinite(exec_values[0]) && std::isfinite(exec_values[1]), "final output graph produced non-finite logits"); + check(std::abs(exec_values[0] - 32.0f) < 1.0e-4f && std::abs(exec_values[1] - 32.0f) < 1.0e-4f, + "final output graph numeric mismatch"); + + ggml_free(ctx); +} + +int main() { + test_hparams(); + test_source_maps(); + test_compression(); + test_state(); + test_raw_ring(); + test_candidates(); + test_output_collapse(); + test_candidate_graph_selection(24, 1); + test_candidate_graph_selection(24, 2); + test_candidate_graph_selection(16392, 2048); + test_graph_contract(); + test_graph_construction(); + return 0; +} diff --git a/tests/test-deepseek41-schema.cpp b/tests/test-deepseek41-schema.cpp new file mode 100644 index 000000000000..4ae136d90c1a --- /dev/null +++ b/tests/test-deepseek41-schema.cpp @@ -0,0 +1,86 @@ +#include "../src/llama-arch.h" + +#include +#include +#include +#include +#include + +static void check(bool condition, const std::string & message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message.c_str()); + std::exit(1); + } +} + +int main() { + check(llm_arch_from_string("deepseek41") == LLM_ARCH_DEEPSEEK41, "deepseek41 architecture lookup failed"); + check(std::string(llm_arch_name(LLM_ARCH_DEEPSEEK41)) == "deepseek41", "deepseek41 architecture name failed"); + + const LLM_KV kv(LLM_ARCH_DEEPSEEK41); + const std::vector> keys = { + { LLM_KV_DSV41_CONFIG, "deepseek41.config" }, + { LLM_KV_DSV41_VOCAB_SIZE, "deepseek41.vocab_size" }, + { LLM_KV_DSV41_HIDDEN_SIZE, "deepseek41.hidden_size" }, + { LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, "deepseek41.moe_intermediate_size" }, + { LLM_KV_DSV41_NUM_HIDDEN_LAYERS, "deepseek41.num_hidden_layers" }, + { LLM_KV_DSV41_NUM_ATTENTION_HEADS, "deepseek41.num_attention_heads" }, + { LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, "deepseek41.num_key_value_heads" }, + { LLM_KV_DSV41_HEAD_DIM, "deepseek41.head_dim" }, + { LLM_KV_DSV41_QK_ROPE_HEAD_DIM, "deepseek41.qk_rope_head_dim" }, + { LLM_KV_DSV41_Q_LORA_RANK, "deepseek41.q_lora_rank" }, + { LLM_KV_DSV41_O_LORA_RANK, "deepseek41.o_lora_rank" }, + { LLM_KV_DSV41_O_GROUPS, "deepseek41.o_groups" }, + { LLM_KV_DSV41_N_ROUTED_EXPERTS, "deepseek41.n_routed_experts" }, + { LLM_KV_DSV41_N_SHARED_EXPERTS, "deepseek41.n_shared_experts" }, + { LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, "deepseek41.num_experts_per_tok" }, + { LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, "deepseek41.max_position_embeddings" }, + { LLM_KV_DSV41_SLIDING_WINDOW, "deepseek41.sliding_window" }, + { LLM_KV_DSV41_INDEX_N_HEADS, "deepseek41.index_n_heads" }, + { LLM_KV_DSV41_INDEX_HEAD_DIM, "deepseek41.index_head_dim" }, + { LLM_KV_DSV41_INDEX_TOPK, "deepseek41.index_topk" }, + { LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, "deepseek41.candidate_source_layer_id" }, + { LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, "deepseek41.candidate_topk_blocks" }, + { LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, "deepseek41.candidate_block_size" }, + { LLM_KV_DSV41_HC_MULT, "deepseek41.hc_mult" }, + { LLM_KV_DSV41_HC_SINKHORN_ITERS, "deepseek41.hc_sinkhorn_iters" }, + { LLM_KV_DSV41_ROPE_THETA, "deepseek41.rope_theta" }, + { LLM_KV_DSV41_COMPRESS_ROPE_THETA, "deepseek41.compress_rope_theta" }, + { LLM_KV_DSV41_RMS_NORM_EPS, "deepseek41.rms_norm_eps" }, + { LLM_KV_DSV41_HC_EPS, "deepseek41.hc_eps" }, + { LLM_KV_DSV41_SWIGLU_LIMIT, "deepseek41.swiglu_limit" }, + { LLM_KV_DSV41_ROUTED_SCALING_FACTOR, "deepseek41.routed_scaling_factor" }, + { LLM_KV_DSV41_SCORING_FUNC, "deepseek41.scoring_func" }, + { LLM_KV_DSV41_HIDDEN_ACT, "deepseek41.hidden_act" }, + { LLM_KV_DSV41_TOPK_METHOD, "deepseek41.topk_method" }, + { LLM_KV_DSV41_NORM_TOPK_PROB, "deepseek41.norm_topk_prob" }, + { LLM_KV_DSV41_COMPRESS_RATIOS, "deepseek41.compress_ratios" }, + { LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, "deepseek41.kv_source_layer_ids" }, + { LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, "deepseek41.index_source_layer_ids" }, + { LLM_KV_DSV41_ROPE_SCALING_FACTOR, "deepseek41.rope_scaling.factor" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, "deepseek41.rope_scaling.beta_fast" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, "deepseek41.rope_scaling.beta_slow" }, + { LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, "deepseek41.rope_scaling.original_max_position_embeddings" }, + { LLM_KV_DSV41_ENGRAM_ENCODING, "deepseek41.engram.encoding" }, + { LLM_KV_DSV41_ENGRAM_LAYER_IDS, "deepseek41.engram.layer_ids" }, + { LLM_KV_DSV41_ENGRAM_ROWS, "deepseek41.engram.rows" }, + { LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, "deepseek41.engram.compressed_vocab_size" }, + { LLM_KV_DSV41_ENGRAM_PAD_ID, "deepseek41.engram.pad_id" }, + { LLM_KV_DSV41_ENGRAM_TOKEN_MAP, "deepseek41.engram.token_map" }, + { LLM_KV_DSV41_ENGRAM_PRIMES, "deepseek41.engram.primes" }, + { LLM_KV_DSV41_ENGRAM_MULTIPLIERS, "deepseek41.engram.multipliers" }, + }; + + for (const auto & item : keys) { + check(kv(item.first) == item.second, std::string("metadata key mismatch: ") + item.second); + } + + const LLM_TN tn(LLM_ARCH_DEEPSEEK41); + check(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", 2).str() == "blk.2.attn_kv_a_norm.weight", "V4.1 KV A norm tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_EMBD, "weight", 1).str() == "blk.1.engram_embd.weight", "Engram embedding tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_Q_NORM, "weight", 1).str() == "blk.1.engram_q_norm.weight", "Engram query norm tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_K_NORM, "weight", 1).str() == "blk.1.engram_k_norm.weight", "Engram key norm tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_KV, "weight", 1).str() == "blk.1.engram_kv.weight", "Engram projection tensor name failed"); + + return 0; +} diff --git a/tests/test-engram.cpp b/tests/test-engram.cpp new file mode 100644 index 000000000000..e741f231262b --- /dev/null +++ b/tests/test-engram.cpp @@ -0,0 +1,444 @@ +#include "../src/llama-bounded-file.h" +#include "../src/llama-engram.h" +#include "../src/llama-ple-disk.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#endif + +static void check(bool condition, const char * message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message); + std::exit(1); + } +} + +static void expect_invalid(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::invalid_argument &) { + return; + } + check(false, message); +} + +static void expect_domain(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::domain_error &) { + return; + } + check(false, message); +} + +static void expect_runtime(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_engram_layout make_layout() { + llama_engram_layout layout; + layout.encoding = "e4m3_e8m0_32_row264"; + layout.layer_ids = { 1, 14 }; + layout.token_map.resize(256); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = (uint32_t) i / 2; + } + layout.compressed_vocab_size = 128; + layout.pad_id = 1; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t depth = 0; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + layout.multipliers[layer][depth] = 35184372088831ull - 2 * (depth + 4 * layer); + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 16000057; + layout.rows[layer] += layout.primes[layer][col]; + } + } + return layout; +} + +static void reference_hash( + const llama_engram_layout & layout, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * output) { + for (size_t pos = 0; pos < count; ++pos) { + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + uint64_t offset = 0; + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + uint64_t hash = 0; + bool blocked = false; + for (size_t shift = 0; shift < col / LLAMA_ENGRAM_HEADS + 2; ++shift) { + const bool before_start = shift > pos; + const size_t source = before_start ? 0 : pos - shift; + blocked = blocked || before_start || (mask != nullptr && mask[source] == 0); + const uint32_t id = blocked ? layout.pad_id : layout.token_map[tokens[source]]; + hash ^= (uint64_t) id * layout.multipliers[layer][shift]; + } + *output++ = (uint32_t) (hash % layout.primes[layer][col] + offset); + offset += layout.primes[layer][col]; + } + } + } +} + +static void test_layout_validation() { + llama_engram_layout layout = make_layout(); + llama_engram_hasher valid(layout); + check(valid.layout().rows == layout.rows, "valid Engram layout changed"); + + llama_engram_layout bad = layout; + bad.encoding = "e4m3"; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad encoding was accepted"); + bad = layout; + bad.layer_ids[1] = bad.layer_ids[0]; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "duplicate layer IDs were accepted"); + bad = layout; + bad.token_map.clear(); + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "empty token map was accepted"); + bad = layout; + bad.compressed_vocab_size = 0; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "empty compressed vocabulary was accepted"); + bad = layout; + bad.pad_id = bad.compressed_vocab_size; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad pad ID was accepted"); + bad = layout; + bad.token_map[10] = bad.compressed_vocab_size; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad token map entry was accepted"); + bad = layout; + bad.multipliers[0][0]--; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "even multiplier was accepted"); + bad = layout; + bad.multipliers[0][0] = UINT64_MAX; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "overflowing multiplier was accepted"); + bad = layout; + bad.primes[0][0] = 1; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad prime was accepted"); + bad = layout; + bad.rows[0]--; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad row extent was accepted"); + bad = layout; + bad.primes[0].fill(UINT32_MAX); + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "overflowing row extent was accepted"); +} + +static void test_hash() { + const llama_engram_layout layout = make_layout(); + const llama_engram_hasher hasher(layout); + constexpr size_t count = 513; + constexpr size_t width = LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS; + std::vector tokens(count); + std::vector mask(count); + std::vector expected(count * width); + std::vector actual(count * width); + for (size_t i = 0; i < count; ++i) { + tokens[i] = (int32_t) ((i * 97 + i / 3) % layout.token_map.size()); + mask[i] = i % 17 != 0 && (i < 125 || i > 131); + } + + for (int masked = 0; masked < 2; ++masked) { + const uint8_t * current_mask = masked != 0 ? mask.data() : nullptr; + reference_hash(layout, tokens.data(), current_mask, count, expected.data()); + llama_engram_history expected_history; + expected_history.reset(); + hasher.hash(expected_history, tokens.data(), current_mask, count, actual.data()); + check(actual == expected, "single-chunk Engram hash differs from full-history reference"); + for (size_t chunk = 1; chunk <= count; ++chunk) { + llama_engram_history history; + history.reset(); + for (size_t pos = 0; pos < count; pos += chunk) { + const size_t size = std::min(chunk, count - pos); + hasher.hash(history, tokens.data() + pos, + current_mask != nullptr ? current_mask + pos : nullptr, + size, actual.data() + pos * width); + } + check(actual == expected, "rolling Engram hash differs from full-history reference"); + check(history.tail == expected_history.tail, "chunking changed final Engram history"); + } + } + + llama_engram_history history; + history.reset(); + const llama_engram_history before = history; + uint32_t output[2 * width]; + std::fill(output, output + 2 * width, UINT32_MAX); + const int32_t invalid_high[] = { 0, (int32_t) layout.token_map.size() }; + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, 2, output); }, + "high invalid token was accepted"); + check(history.tail == before.tail, "invalid token mutated Engram history"); + check(output[0] == UINT32_MAX, "invalid token changed Engram output"); + const int32_t invalid_low[] = { 0, -1 }; + expect_invalid([&] { hasher.hash(history, invalid_low, nullptr, 2, output); }, + "negative token was accepted"); + check(history.tail == before.tail, "negative token mutated Engram history"); + hasher.hash(history, nullptr, nullptr, 0, nullptr); + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, SIZE_MAX, output); }, + "overflowing hash count was accepted"); + history.tail[0] = (int32_t) layout.compressed_vocab_size; + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, 1, output); }, + "invalid history was accepted"); +} + +static float decode_reference(uint8_t code, uint8_t scale) { + const int exponent = (code >> 3) & 15; + double value = exponent != 0 ? + (1.0 + (code & 7) / 8.0) * std::pow(2.0, exponent - 7) : + (code & 7) / 512.0; + if ((code & 128) != 0) { + value = -value; + } + float result = (float) (value * std::pow(2.0, (int) scale - 127)); + uint32_t bits; + memcpy(&bits, &result, sizeof(bits)); + bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u; + memcpy(&result, &bits, sizeof(result)); + return result; +} + +static void test_decode() { + uint8_t row[LLAMA_ENGRAM_ROW_BYTES]; + float output[LLAMA_ENGRAM_DIM]; + for (uint32_t code = 0; code < 256; ++code) { + memset(row, (int) code, LLAMA_ENGRAM_DIM); + for (uint32_t scale = 0; scale < 256; ++scale) { + memset(row + LLAMA_ENGRAM_DIM, (int) scale, LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM); + const float expected = decode_reference((uint8_t) code, (uint8_t) scale); + const bool valid = (code & 127) != 127 && scale != 255 && std::isfinite(expected); + if (valid) { + llama_engram_decode_row(row, output); + for (float value : output) { + check(memcmp(&value, &expected, sizeof(value)) == 0, "Engram decode differs from BF16 reference"); + } + } else { + output[0] = 123456.0f; + expect_domain([&] { llama_engram_decode_row(row, output); }, "invalid Engram value was accepted"); + check(output[0] == 123456.0f, "invalid Engram value changed output"); + } + } + } + + memset(row, 0, sizeof(row)); + memset(row + LLAMA_ENGRAM_DIM, 127, LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM); + row[0] = 128; + llama_engram_decode_row(row, output); + check(output[0] == 0.0f && std::signbit(output[0]), "negative zero was not preserved"); + check(output[1] == 0.0f && !std::signbit(output[1]), "positive zero was not preserved"); +} + +#if !defined(_WIN32) +static void write_full(int fd, const void * data, size_t size, uint64_t offset) { + const ssize_t written = pwrite(fd, data, size, (off_t) offset); + check(written == (ssize_t) size, "failed to write Engram test row"); +} + +static void test_disk_rows() { + char path[] = "/tmp/llama-engram-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create Engram test file"); + + const uint64_t offset = (1ull << 33) + 32; + uint8_t raw[3][LLAMA_ENGRAM_ROW_BYTES]; + for (size_t row = 0; row < 3; ++row) { + for (size_t i = 0; i < LLAMA_ENGRAM_DIM; ++i) { + raw[row][i] = (uint8_t) i; + } + raw[row][127] = 0; + raw[row][255] = 128; + for (size_t i = 0; i < LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM; ++i) { + raw[row][LLAMA_ENGRAM_DIM + i] = (uint8_t) (126 + row); + } + } + write_full(fd, raw, sizeof(raw), offset); + + llama_engram_table table(path, offset, 3); + check(table.n_rows() == 3, "Engram table row count changed"); + const uint32_t rows[] = { 2, 0, 2, 1 }; + float output[4 * LLAMA_ENGRAM_DIM]; + table.read(rows, 4, output); + for (size_t row = 0; row < 4; ++row) { + float expected[LLAMA_ENGRAM_DIM]; + llama_engram_decode_row(raw[rows[row]], expected); + check(memcmp(output + row * LLAMA_ENGRAM_DIM, expected, sizeof(expected)) == 0, + "Engram row order was not preserved"); + } + + constexpr size_t tokens = 2051; + constexpr size_t stride = LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS; + constexpr size_t width = LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM; + std::vector batch_rows(tokens * stride, UINT32_MAX); + std::vector batch(tokens * width + 1); + for (size_t token = 0; token < tokens; ++token) { + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + batch_rows[token * stride + col] = (uint32_t) ((token * 7 + col * 11) % 3); + } + } + const size_t sizes[] = { 1, 2, 31, 65, 257, 2047, 2048, 2049, tokens }; + float decoded[3][LLAMA_ENGRAM_DIM]; + for (size_t row = 0; row < 3; ++row) { + llama_engram_decode_row(raw[row], decoded[row]); + } + for (size_t count : sizes) { + batch[count * width] = 123456.0f; + table.read_batch(batch_rows.data(), count, stride, batch.data()); + for (size_t token = 0; token < count; ++token) { + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + const uint32_t row = batch_rows[token * stride + col]; + check(memcmp(batch.data() + token * width + col * LLAMA_ENGRAM_DIM, + decoded[row], sizeof(decoded[row])) == 0, + "batched Engram row differs from direct decode"); + } + } + check(batch[count * width] == 123456.0f, "batched Engram read overflowed output"); + } + + table.read(nullptr, 0, nullptr); + table.read_batch(nullptr, 0, 0, nullptr); + expect_invalid([&] { table.read_batch(batch_rows.data(), 1, LLAMA_ENGRAM_COLS - 1, batch.data()); }, + "short Engram batch stride was accepted"); + expect_invalid([&] { table.read_batch(batch_rows.data(), 2, SIZE_MAX, batch.data()); }, + "overflowing Engram batch stride was accepted"); + expect_invalid([&] { table.read_batch(batch_rows.data(), SIZE_MAX, stride, batch.data()); }, + "overflowing Engram batch count was accepted"); + + const uint32_t bad_row = 3; + output[0] = 123456.0f; + expect_invalid([&] { table.read(&bad_row, 1, output); }, "invalid Engram row was accepted"); + check(output[0] == 123456.0f, "invalid Engram row changed output"); + batch_rows[(tokens - 1) * stride] = bad_row; + batch[0] = 123456.0f; + expect_invalid([&] { table.read_batch(batch_rows.data(), tokens, stride, batch.data()); }, + "invalid batched Engram row was accepted"); + check(batch[0] == 123456.0f, "invalid batched Engram row changed output"); + batch_rows[(tokens - 1) * stride] = 0; + + uint8_t invalid = 127; + write_full(fd, &invalid, 1, offset); + const uint32_t first_row = 0; + expect_domain([&] { table.read(&first_row, 1, output); }, "E4M3 NaN row was accepted"); + write_full(fd, raw, sizeof(raw), offset); + invalid = 255; + write_full(fd, &invalid, 1, offset + LLAMA_ENGRAM_DIM); + expect_domain([&] { table.read(&first_row, 1, output); }, "E8M0 scale 255 row was accepted"); + write_full(fd, raw, sizeof(raw), offset); + + check(ftruncate(fd, (off_t) (offset + 260)) == 0, "failed to truncate Engram test file"); + expect_runtime([&] { table.read(&first_row, 1, output); }, "short Engram read was accepted"); + close(fd); + + expect_invalid([&] { llama_engram_table invalid_table(path, offset, 1); }, + "truncated Engram extent was accepted"); + expect_invalid([&] { llama_engram_table invalid_table(path, UINT64_MAX - 1, 3); }, + "overflowing Engram extent was accepted"); + expect_invalid([&] { llama_engram_table invalid_table(path, offset, 0); }, + "empty Engram table was accepted"); + check(unlink(path) == 0, "failed to remove Engram test file"); +} + +#if defined(__linux__) +static void test_direct_tail_read() { + char path[] = "/tmp/llama-direct-tail-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create direct tail test file"); + + constexpr uint64_t offset = 4096; + const uint8_t expected[] = { 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2 }; + write_full(fd, expected, sizeof(expected), offset); + close(fd); + + llama_bounded_file::params params; + params.direct_io = true; + params.direct_io_required = true; + llama_bounded_file file(path, params); + check(file.direct_io(), "direct tail test did not use O_DIRECT"); + llama_bounded_file::buffer scratch = file.make_buffer(sizeof(expected)); + uint8_t actual[sizeof(expected)] = {}; + file.read(offset, actual, sizeof(actual), scratch); + check(memcmp(actual, expected, sizeof(actual)) == 0, "valid direct tail read failed"); + check(unlink(path) == 0, "failed to remove direct tail test file"); +} +#endif + +static void test_ple_disk_reader() { + char path[] = "/tmp/llama-ple-reader-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create PLE test file"); + + constexpr uint64_t offset = 32; + constexpr size_t columns = 4; + constexpr size_t rows = 3; + const float table[rows][columns] = { + { 1.0f, 2.0f, 3.0f, 4.0f }, + { -1.0f, -2.0f, -3.0f, -4.0f }, + { 0.5f, 0.25f, 0.125f, 0.0625f }, + }; + write_full(fd, table, sizeof(table), offset); + close(fd); + + for (int direct = 0; direct < 2; ++direct) { + llama_ple_disk::params params; + params.n_threads = 2; + params.cache_bytes = sizeof(table); + params.direct_io = direct != 0; + llama_ple_disk disk(path, offset, GGML_TYPE_F32, columns, rows, params); + const int32_t ids[] = { 2, 0, 2, 1 }; + float output[4][columns]; + disk.gather(ids, 4, output[0]); + for (size_t i = 0; i < 4; ++i) { + check(memcmp(output[i], table[ids[i]], sizeof(output[i])) == 0, + "shared bounded reader changed PLE row output"); + } + } + + llama_ple_disk::params params; + params.n_threads = 4; + params.cache_bytes = 0; + params.direct_io = false; + llama_ple_disk disk(path, offset, GGML_TYPE_F32, columns, rows, params); + check(truncate(path, (off_t) (offset + sizeof(float))) == 0, "failed to truncate PLE test file"); + const int32_t ids[] = { 0, 1, 2 }; + float output[3][columns]; + expect_runtime([&] { disk.gather(ids, 3, output[0]); }, "threaded PLE read failure did not propagate"); + + const int repair_fd = open(path, O_WRONLY); + check(repair_fd >= 0, "failed to reopen PLE test file"); + write_full(repair_fd, table, sizeof(table), offset); + close(repair_fd); + disk.gather(ids, 3, output[0]); + check(memcmp(output, table, sizeof(table)) == 0, "PLE reader did not recover after worker failure"); + + check(unlink(path) == 0, "failed to remove PLE test file"); +} +#endif + +int main() { + test_layout_validation(); + test_hash(); + test_decode(); +#if !defined(_WIN32) + test_disk_rows(); +#if defined(__linux__) + test_direct_tail_read(); +#endif + test_ple_disk_reader(); +#endif + std::puts("Engram layout, hash and bounded disk rows: PASS"); + return 0; +} diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp new file mode 100644 index 000000000000..fbf5f0b10954 --- /dev/null +++ b/tests/test-expert-store.cpp @@ -0,0 +1,651 @@ +#include "../src/llama-expert-store.h" +#include "../src/llama-model-loader.h" + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(cond) do { if (!(cond)) { throw std::runtime_error("requirement failed: " #cond); } } while (0) + +namespace { + +struct temp_file { + std::filesystem::path path; + + explicit temp_file(const char * suffix) { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-expert-store-" + std::to_string(stamp) + "-" + std::to_string(++sequence) + suffix); + } + + ~temp_file() { + std::error_code ec; + std::filesystem::remove(path, ec); + } +}; + +struct fixture { + static constexpr int64_t n_embd = 512; + static constexpr int64_t n_ff = 256; + static constexpr int64_t n_expert = 4; + + temp_file file { ".gguf" }; + std::vector tensors; + + fixture() { + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff) * n_embd; + const size_t data_size = 2 * gate_plane * n_expert + down_plane * n_expert; + + ggml_init_params ggml_params = { + /*.mem_size =*/ data_size + 8 * ggml_tensor_overhead() + 4096, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(ggml_params); + REQUIRE(ctx != nullptr); + + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * down = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_ff, n_embd, n_expert); + ggml_set_name(gate, "blk.0.ffn_gate_exps.weight"); + ggml_set_name(up, "blk.0.ffn_up_exps.weight"); + ggml_set_name(down, "blk.0.ffn_down_exps.weight"); + + fill_tensor(gate, 0x10); + fill_tensor(up, 0x20); + fill_tensor(down, 0x30); + + gguf_context * gguf = gguf_init_empty(); + REQUIRE(gguf != nullptr); + gguf_set_val_str(gguf, "general.architecture", "deepseek41"); + gguf_add_tensor(gguf, gate); + gguf_add_tensor(gguf, up); + gguf_add_tensor(gguf, down); + REQUIRE(gguf_write_to_file(gguf, file.path.string().c_str(), false)); + gguf_free(gguf); + ggml_free(ctx); + + std::vector splits; + llama_model_loader loader( + nullptr, nullptr, nullptr, file.path.string(), splits, nullptr, + LLAMA_LOAD_MODE_MMAP, false, false, false, nullptr, nullptr); + + REQUIRE(loader.ctx_map.empty()); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_gate_exps.weight", 0, LLAMA_EXPERT_PROJECTION_GATE, { n_embd, n_ff, n_expert })); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_up_exps.weight", 0, LLAMA_EXPERT_PROJECTION_UP, { n_embd, n_ff, n_expert })); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_down_exps.weight", 0, LLAMA_EXPERT_PROJECTION_DOWN, { n_ff, n_embd, n_expert })); + loader.done_getting_tensors(); + + REQUIRE(loader.external.any()); + REQUIRE(loader.ctx_map.empty()); + for (const auto & tensor : tensors) { + REQUIRE(tensor.file_index == 0); + REQUIRE(loader.external.has(loader.require_tensor_meta(tensor.name))); + } + loader.init_mappings(true); + REQUIRE(loader.mappings.size() == 1); + REQUIRE(loader.ctx_map.empty()); + const auto & ranges = loader.external.for_file(0); + REQUIRE(ranges.size() == 3); + for (size_t i = 0; i < ranges.size(); ++i) { + REQUIRE(ranges[i].first == tensors[i].file_offset); + REQUIRE(ranges[i].second == tensors[i].file_offset + tensors[i].nb[2] * tensors[i].ne[2]); + } + } + + static void fill_tensor(ggml_tensor * tensor, uint8_t tag) { + memset(tensor->data, 0, ggml_nbytes(tensor)); + for (int64_t expert = 0; expert < tensor->ne[2]; ++expert) { + uint8_t * plane = static_cast(tensor->data) + expert * tensor->nb[2]; + plane[tensor->nb[2] - 1] = tag + expert; + } + } + + llama_expert_store make_store(size_t slots, size_t bytes) const { + llama_expert_store_params params; + params.cache_slots = slots; + params.cache_bytes = bytes; + params.direct_io = false; + return llama_expert_store(tensors, params); + } + + size_t max_plane_size() const { + size_t result = 0; + for (const auto & tensor : tensors) { + result = std::max(result, tensor.nb[2]); + } + return result; + } + + size_t all_projection_bytes() const { + size_t result = 0; + for (const auto & tensor : tensors) { + result += tensor.nb[2]; + } + return result; + } +}; + +template +void require_throws(F && fn) { + bool threw = false; + try { + fn(); + } catch (const std::exception &) { + threw = true; + } + REQUIRE(threw); +} + +const llama_expert_store::payload & find_payload( + const std::vector & payloads, + llama_expert_projection projection, + int32_t expert_id) { + for (const auto & payload : payloads) { + if (payload.projection == projection && payload.expert_id == expert_id) { + return payload; + } + } + throw std::runtime_error("payload not found"); +} + +void test_layout_and_offsets(const fixture & f) { + const auto & gate = f.tensors[0]; + const auto & up = f.tensors[1]; + const auto & down = f.tensors[2]; + + REQUIRE(gate.type == GGML_TYPE_IQ2_XXS); + REQUIRE(up.type == GGML_TYPE_IQ2_XXS); + REQUIRE(down.type == GGML_TYPE_Q2_K); + REQUIRE(gate.nb[2] == ggml_row_size(GGML_TYPE_IQ2_XXS, fixture::n_embd) * fixture::n_ff); + REQUIRE(down.nb[2] == ggml_row_size(GGML_TYPE_Q2_K, fixture::n_ff) * fixture::n_embd); + REQUIRE(gate.file_offset + 3 * gate.nb[2] > gate.file_offset); + + llama_expert_store store = f.make_store(3, f.all_projection_bytes()); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.resident_bytes() == 0); + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 2 } }, + { 0, LLAMA_EXPERT_PROJECTION_DOWN, { 2 } }, + }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 3); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_GATE, 2).data[gate.nb[2] - 1] == 0x12); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_UP, 2).data[up.nb[2] - 1] == 0x22); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_DOWN, 2).data[down.nb[2] - 1] == 0x32); +} + +void test_alignment_and_large_offsets() { + const auto aligned = llama_expert_store_align_read(4097, 5000, 4096, 20000); + REQUIRE(aligned.offset == 4096); + REQUIRE(aligned.prefix == 1); + REQUIRE(aligned.size == 8192); + + const uint64_t large_offset = (uint64_t(1) << 32) + 123; + const auto large = llama_expert_store_align_read(large_offset, 777, 4096, large_offset + 777); + REQUIRE(large.offset > std::numeric_limits::max()); + REQUIRE(large.prefix == 123); + REQUIRE(large.size == 4096); + + require_throws([] { + llama_expert_store_align_read(0, 1, 3000, 1); + }); + require_throws([] { + llama_expert_store_align_read(UINT64_MAX - 4, 8, 4096, UINT64_MAX); + }); + + size_t granularity_first = 1; + size_t granularity_last = 1; + llama_mlock::align_range(&granularity_first, &granularity_last); + const size_t lock_granularity = granularity_last; + REQUIRE(granularity_first == 0); + REQUIRE(lock_granularity > 1); + + size_t lock_first = lock_granularity + 1; + size_t lock_last = 2 * lock_granularity; + llama_mlock::align_range(&lock_first, &lock_last); + REQUIRE(lock_first == lock_granularity); + REQUIRE(lock_last == 2 * lock_granularity); +} + +void test_external_mapping_access_policy() { + const size_t page = 4096; + const llama_mmap::ranges external = { + { 0, page }, + { 2 * page + 1, 4 * page - 1 }, + { 5 * page, 6 * page }, + { 7 * page, 8 * page }, + { 9 * page, 10 * page }, + }; + + REQUIRE(llama_mmap::use_sequential_file_advice(false)); + REQUIRE(!llama_mmap::use_sequential_file_advice(true)); + REQUIRE(llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, true).empty()); + + const auto lazy_ranges = llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, false); + REQUIRE(lazy_ranges.size() == 4); + REQUIRE(lazy_ranges.front() == std::make_pair(page, 2 * page + 1)); + REQUIRE(lazy_ranges.back() == std::make_pair(8 * page, 9 * page)); +} + +#if defined(__linux__) +int file_advice_calls = 0; + +int fail_file_advice(int, int) { + ++file_advice_calls; + return EIO; +} + +void test_external_mapping_advice_failure() { + temp_file file { ".bin" }; + { + std::ofstream out(file.path, std::ios::binary); + REQUIRE(out.good()); + out.seekp(8191); + out.put('\0'); + } + + llama_file input(file.path.string(), "rb"); + file_advice_calls = 0; + bool continued_after_advice = false; + require_throws([&] { + llama_mmap mapping(&input, 0, false, { { 4096, 8192 } }, true, fail_file_advice); + continued_after_advice = true; + }); + REQUIRE(file_advice_calls == 1); + REQUIRE(!continued_after_advice); + + llama_mmap legacy_mapping(&input, 0, false, {}, false, fail_file_advice); + REQUIRE(file_advice_calls == 2); + REQUIRE(legacy_mapping.addr() != nullptr); +} +#endif + +void test_published_layout_accounting() { + const int64_t n_embd = 7680; + const int64_t n_ff = 1536; + const int64_t n_expert = 384; + const int64_t n_layer = 40; + + const uint64_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const uint64_t up_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const uint64_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff) * n_embd; + const uint64_t slot_bytes = (gate_plane + up_plane + down_plane) * n_layer; + const uint64_t routed_bytes = slot_bytes * n_expert; + const uint64_t dense_bytes = 10067427328; + const uint64_t engram_bytes = 202758045696; + + REQUIRE(gate_plane == 3041280); + REQUIRE(up_plane == 3041280); + REQUIRE(down_plane == 3870720); + REQUIRE(gate_plane + up_plane + down_plane == 9953280); + REQUIRE(slot_bytes == 398131200); + REQUIRE(routed_bytes == 152882380800); + REQUIRE(224 * slot_bytes == 89181388800); + REQUIRE(256 * slot_bytes == 101921587200); + REQUIRE(dense_bytes + 224 * slot_bytes == 99248816128); + REQUIRE(dense_bytes + 256 * slot_bytes == 111989014528); + REQUIRE(engram_bytes > routed_bytes); +} + +void test_large_offset_read() { + temp_file sparse { ".bin" }; + const int64_t ne0 = 256; + const int64_t ne1 = 256; + const int64_t n_expert = 1; + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, ne0) * ne1; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, ne0) * ne1; + const uint64_t base = (uint64_t(1) << 32) + 4096; + const uint64_t file_size = base + 2 * gate_plane + down_plane; + + { + std::ofstream out(sparse.path, std::ios::binary | std::ios::trunc); + REQUIRE(out.good()); + out.seekp(static_cast(file_size - 1)); + out.put('\0'); + const std::vector> markers = { + { base + gate_plane - 1, 0x41 }, + { base + 2 * gate_plane - 1, 0x42 }, + { file_size - 1, 0x43 }, + }; + for (const auto & marker : markers) { + out.seekp(static_cast(marker.first)); + out.put(static_cast(marker.second)); + } + } + + auto make_tensor = [&](const char * name, llama_expert_projection projection, ggml_type type, uint64_t offset) { + llama_expert_store_tensor tensor; + tensor.name = name; + tensor.fname = sparse.path.string(); + tensor.layer = 0; + tensor.projection = projection; + tensor.type = type; + tensor.ne[0] = ne0; + tensor.ne[1] = ne1; + tensor.ne[2] = n_expert; + tensor.nb[0] = ggml_type_size(type); + tensor.nb[1] = ggml_row_size(type, ne0); + tensor.nb[2] = tensor.nb[1] * ne1; + tensor.file_offset = offset; + tensor.file_size = file_size; + return tensor; + }; + + std::vector tensors; + tensors.push_back(make_tensor("gate", LLAMA_EXPERT_PROJECTION_GATE, GGML_TYPE_IQ2_XXS, base)); + tensors.push_back(make_tensor("up", LLAMA_EXPERT_PROJECTION_UP, GGML_TYPE_IQ2_XXS, base + gate_plane)); + tensors.push_back(make_tensor("down", LLAMA_EXPERT_PROJECTION_DOWN, GGML_TYPE_Q2_K, base + 2 * gate_plane)); + + llama_expert_store_params params { 2 * gate_plane + down_plane, 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_DOWN, { 0 } }, + }); + const auto payloads = lease.payloads(); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_GATE, 0).data[gate_plane - 1] == 0x41); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_UP, 0).data[gate_plane - 1] == 0x42); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_DOWN, 0).data[down_plane - 1] == 0x43); +} + +void test_cache_and_remapping(const fixture & f) { + const size_t gate_plane = f.tensors[0].nb[2]; + llama_expert_store store = f.make_store(2, 2 * f.max_plane_size()); + + { + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0, 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1, 0 } }, + }); + REQUIRE(lease.slot_ids().size() == 2); + REQUIRE(lease.slot_ids()[0][0] == lease.slot_ids()[0][1]); + REQUIRE(lease.slot_ids()[1][1] == lease.slot_ids()[0][0]); + REQUIRE(lease.slot_ids()[1][0] != lease.slot_ids()[1][1]); + REQUIRE(store.resident_entries() == 2); + REQUIRE(store.resident_bytes() == 2 * gate_plane); + } + + { + auto hit = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(hit.slot_ids()[0][0] == 0); + } + { + auto miss = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } } }); + REQUIRE(miss.slot_ids()[0][0] == 1); + } + + const llama_expert_store_stats stats = store.stats(); + REQUIRE(stats.hits == 1); + REQUIRE(stats.misses == 3); + REQUIRE(stats.evictions == 1); + REQUIRE(stats.bytes_read == 3 * gate_plane); + + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { -1 } } }); + }); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { fixture::n_expert } } }); + }); +} + +void test_direct_io(const fixture & f) { + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.io_alignment = 4096; + params.direct_io = true; + params.allow_buffered_io = true; + + llama_expert_store store(f.tensors, params); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 3 } } }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 1); + REQUIRE(payloads[0].data[payloads[0].size - 1] == 0x13); + REQUIRE(store.stats().bytes_read >= payloads[0].size); + REQUIRE(store.stats().bytes_read <= payloads[0].size + 2 * params.io_alignment); +} + +#if defined(__linux__) +void test_direct_io_file_tail(const fixture & f) { + const auto & down = f.tensors[2]; + REQUIRE(down.file_offset + down.nb[2] * down.ne[2] == down.file_size); + + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.io_alignment = 4096; + params.direct_io = true; + params.allow_buffered_io = false; + + llama_expert_store store(f.tensors, params); + REQUIRE(store.direct_io_active()); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_DOWN, { fixture::n_expert - 1 } } }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 1); + REQUIRE(payloads[0].data[payloads[0].size - 1] == 0x33); + REQUIRE(store.direct_io_active()); +} +#endif + +#if defined(_WIN32) +void test_windows_direct_io_policy(const fixture & f) { + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.direct_io = true; + params.allow_buffered_io = false; + + require_throws([&] { + llama_expert_store store(f.tensors, params); + }); + + params.allow_buffered_io = true; + llama_expert_store store(f.tensors, params); + REQUIRE(!store.direct_io_active()); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(lease.payloads().size() == 1); +} +#endif + +void test_pins_and_atomic_failure(const fixture & f) { + llama_expert_store store = f.make_store(1, f.max_plane_size()); + auto pinned = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + const llama_expert_store_stats before = store.stats(); + + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + }); + REQUIRE(store.resident_entries() == 1); + REQUIRE(store.stats().hits == before.hits); + REQUIRE(store.stats().misses == before.misses); + REQUIRE(pinned.payloads()[0].expert_id == 0); + + pinned = {}; + auto replacement = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + REQUIRE(replacement.payloads()[0].expert_id == 1); + REQUIRE(store.stats().evictions == 1); + + llama_expert_store::lease surviving; + { + llama_expert_store short_lived = f.make_store(1, f.max_plane_size()); + surviving = short_lived.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } } }); + } + REQUIRE(surviving.payloads()[0].expert_id == 2); +} + +void test_limits_and_validation(const fixture & f) { + require_throws([&] { + f.make_store(0, f.max_plane_size()); + }); + require_throws([&] { + f.make_store(1, f.tensors[2].nb[2] - 1); + }); + { + llama_expert_store_params params { f.max_plane_size(), 1, 1, false }; + llama_expert_store store(f.tensors, params); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(lease.payloads().size() == 1); + } + { + llama_expert_store store = f.make_store(3, f.max_plane_size()); + require_throws([&] { + store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 0 } }, + }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); + } + { + llama_expert_store store = f.make_store(1, 2 * f.tensors[0].nb[2]); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0, 1 } } }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); + } + + auto bad_type = f.tensors; + bad_type[0].type = GGML_TYPE_Q2_K; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_type), params); + }); + + auto bad_stride = f.tensors; + bad_stride[1].nb[2]++; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_stride), params); + }); + + auto bad_bounds = f.tensors; + bad_bounds[2].file_size = bad_bounds[2].file_offset + bad_bounds[2].nb[2] - 1; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_bounds), params); + }); +} + +void test_payload_validation(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + { + std::fstream io(copy.path, std::ios::binary | std::ios::in | std::ios::out); + REQUIRE(io.good()); + io.seekp(static_cast(tensors[0].file_offset)); + const uint8_t invalid_scale[2] = { 0x00, 0x7c }; + io.write(reinterpret_cast(invalid_scale), sizeof(invalid_scale)); + } + + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + }); + REQUIRE(store.resident_entries() == 0); +} + +void test_truncated_file(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + std::filesystem::resize_file(copy.path, f.tensors[0].file_offset + f.tensors[0].nb[2] - 1); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); +} + +void test_failed_replacement_keeps_resident_entry(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + llama_expert_store_params params { f.max_plane_size(), 1, 4096, false }; + llama_expert_store store(std::move(tensors), params); + { + auto resident = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(resident.payloads()[0].expert_id == 0); + } + const llama_expert_store_stats before = store.stats(); + std::filesystem::resize_file(copy.path, f.tensors[0].file_offset + 2*f.tensors[0].nb[2] - 1); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + }); + REQUIRE(store.resident_entries() == 1); + REQUIRE(store.stats().evictions == before.evictions); + auto hit = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(hit.payloads()[0].expert_id == 0); +} + +} + +int main() { + try { + fixture f; + test_layout_and_offsets(f); + test_alignment_and_large_offsets(); + test_external_mapping_access_policy(); +#if defined(__linux__) + test_external_mapping_advice_failure(); +#endif + test_published_layout_accounting(); + test_large_offset_read(); + test_cache_and_remapping(f); + test_direct_io(f); +#if defined(__linux__) + test_direct_io_file_tail(f); +#endif +#if defined(_WIN32) + test_windows_direct_io_policy(f); +#endif + test_pins_and_atomic_failure(f); + test_limits_and_validation(f); + test_payload_validation(f); + test_truncated_file(f); + test_failed_replacement_keeps_resident_entry(f); + } catch (const std::exception & e) { + fprintf(stderr, "test-expert-store: %s\n", e.what()); + return 1; + } + return 0; +} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 3496f72e4949..c68a936cfea0 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -9,6 +9,8 @@ // TODO: replace with #include "llama-ext.h" in the future #include "../src/llama-arch.h" +#include "../src/llama-dsv41.h" +#include "../src/llama-dsv41-expert.h" #include "../src/llama-model-saver.h" #include @@ -16,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -323,6 +326,107 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } + if (arch == LLM_ARCH_DEEPSEEK41) { + std::vector compress_ratios; + compress_ratios.reserve(LLAMA_DSV41_N_LAYER); + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + compress_ratios.push_back(llama_dsv41_compress_ratio(il)); + } + std::vector token_map(LLAMA_DSV41_N_VOCAB, 0); + std::vector primes(LLAMA_DSV41_ENGRAM_PRIMES_COUNT, 2); + primes[0] = 384006168 - 2*(LLAMA_ENGRAM_COLS - 1); + primes[LLAMA_ENGRAM_COLS] = 384016682 - 2*(LLAMA_ENGRAM_COLS - 1); + const std::vector multipliers = { + 101, 103, 105, 107, + 109, 111, 113, 115, + }; + + ms.add_kv(LLM_KV_DSV41_CONFIG, "{}"); + ms.add_kv(LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, LLAMA_DSV41_N_CTX); + ms.add_kv(LLM_KV_DSV41_HIDDEN_SIZE, LLAMA_DSV41_N_EMBD); + ms.add_kv(LLM_KV_DSV41_NUM_HIDDEN_LAYERS, LLAMA_DSV41_N_LAYER); + ms.add_kv(LLM_KV_DSV41_VOCAB_SIZE, LLAMA_DSV41_N_VOCAB); + ms.add_kv(LLM_KV_DSV41_NUM_ATTENTION_HEADS, LLAMA_DSV41_N_HEAD); + ms.add_kv(LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, LLAMA_DSV41_N_HEAD_KV); + ms.add_kv(LLM_KV_DSV41_HEAD_DIM, LLAMA_DSV41_N_HEAD_DIM); + ms.add_kv(LLM_KV_DSV41_QK_ROPE_HEAD_DIM, LLAMA_DSV41_N_ROT); + ms.add_kv(LLM_KV_DSV41_Q_LORA_RANK, LLAMA_DSV41_N_LORA_Q); + ms.add_kv(LLM_KV_DSV41_O_LORA_RANK, LLAMA_DSV41_N_LORA_O); + ms.add_kv(LLM_KV_DSV41_O_GROUPS, LLAMA_DSV41_N_O_GROUP); + ms.add_kv(LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, LLAMA_DSV41_N_FF_EXP); + ms.add_kv(LLM_KV_DSV41_N_ROUTED_EXPERTS, LLAMA_DSV41_N_EXPERT); + ms.add_kv(LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, LLAMA_DSV41_N_EXPERT_USED); + ms.add_kv(LLM_KV_DSV41_N_SHARED_EXPERTS, LLAMA_DSV41_N_EXPERT_SHARED); + ms.add_kv(LLM_KV_DSV41_INDEX_N_HEADS, LLAMA_DSV41_N_INDEX_HEAD); + ms.add_kv(LLM_KV_DSV41_INDEX_HEAD_DIM, LLAMA_DSV41_N_INDEX_HEAD_DIM); + ms.add_kv(LLM_KV_DSV41_INDEX_TOPK, LLAMA_DSV41_N_INDEX_TOP_K); + ms.add_kv(LLM_KV_DSV41_HC_MULT, LLAMA_DSV41_HC_MULT); + ms.add_kv(LLM_KV_DSV41_HC_SINKHORN_ITERS, LLAMA_DSV41_HC_SINKHORN_ITERS); + ms.add_kv(LLM_KV_DSV41_SLIDING_WINDOW, LLAMA_DSV41_N_SWA); + ms.add_kv(LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, LLAMA_DSV41_CANDIDATE_SOURCE_LAYER); + ms.add_kv(LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS); + ms.add_kv(LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, LLAMA_DSV41_CANDIDATE_BLOCK_SIZE); + ms.add_kv(LLM_KV_DSV41_RMS_NORM_EPS, 1.0e-20f); + ms.add_kv(LLM_KV_DSV41_HC_EPS, 1.0e-6f); + ms.add_kv(LLM_KV_DSV41_SWIGLU_LIMIT, 10.0f); + ms.add_kv(LLM_KV_DSV41_ROUTED_SCALING_FACTOR, 1.5f); + ms.add_kv(LLM_KV_DSV41_ROPE_THETA, uint32_t(10000)); + ms.add_kv(LLM_KV_DSV41_COMPRESS_ROPE_THETA, uint32_t(160000)); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_FACTOR, 16.0f); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, 32.0f); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, 1.0f); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, 65536.0f); + ms.add_kv(LLM_KV_DSV41_NORM_TOPK_PROB, true); + ms.add_kv(LLM_KV_DSV41_HIDDEN_ACT, "silu"); + ms.add_kv(LLM_KV_DSV41_SCORING_FUNC, "sqrtsoftplus"); + ms.add_kv(LLM_KV_DSV41_TOPK_METHOD, "noaux_tc"); + ms.add_kv(LLM_KV_DSV41_COMPRESS_RATIOS, compress_ratios); + ms.add_kv(LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, std::vector({ 2, 8, 14, 20 })); + ms.add_kv(LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, std::vector({ 2, 8, 14, 20, 24, 28, 32, 36 })); + ms.add_kv(LLM_KV_DSV41_ENGRAM_ENCODING, LLAMA_DSV41_ENGRAM_ENCODING); + ms.add_kv(LLM_KV_DSV41_ENGRAM_LAYER_IDS, std::vector({ 1, 14 })); + ms.add_kv(LLM_KV_DSV41_ENGRAM_ROWS, std::vector({ 384006168, 384016682 })); + ms.add_kv(LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB); + ms.add_kv(LLM_KV_DSV41_ENGRAM_PAD_ID, LLAMA_DSV41_ENGRAM_PAD_ID); + ms.add_kv(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, token_map); + ms.add_kv(LLM_KV_DSV41_ENGRAM_PRIMES, primes); + ms.add_kv(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, multipliers); + + auto add_tensor = [&](const std::string & name, ggml_type type, const std::initializer_list & ne) { + ggml_tensor tensor = {}; + tensor.type = type; + tensor.ne[0] = tensor.ne[1] = tensor.ne[2] = tensor.ne[3] = 1; + size_t i = 0; + for (int64_t dim : ne) { + tensor.ne[i++] = dim; + } + tensor.nb[0] = ggml_type_size(type); + tensor.nb[1] = ggml_row_size(type, tensor.ne[0]); + tensor.nb[2] = tensor.nb[1]*tensor.ne[1]; + tensor.nb[3] = tensor.nb[2]*tensor.ne[2]; + snprintf(tensor.name, sizeof(tensor.name), "%s", name.c_str()); + gguf_add_tensor(ms.gguf_ctx, &tensor); + }; + llama_dsv41_register_expert_tensors( + [&](const std::string & name, + int32_t, + llama_expert_projection projection, + const std::initializer_list & ne) { + add_tensor( + name, + projection == LLAMA_EXPERT_PROJECTION_DOWN ? + GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS, + ne); + return llama_expert_store_tensor(); + }); + const LLM_TN tn(LLM_ARCH_DEEPSEEK41); + for (int32_t il : { 1, 14 }) { + add_tensor( + tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(), + GGML_TYPE_I8, + { LLAMA_ENGRAM_ROW_BYTES, il == 1 ? 384006168 : 384016682 }); + } + } ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab"); // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd); // ms.add_kv(LLM_KV_DENSE_3_FEAT_IN, n_embd); @@ -368,18 +472,20 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f); ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); - for (uint32_t il = 0; il < n_layer; il++) { - ggml_tensor t; - memset(&t, 0, sizeof(ggml_tensor)); - t.type = GGML_TYPE_F16; - ggml_format_name(&t, "conv%" PRIu32 "d.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); - ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); - ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); - ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); + if (arch != LLM_ARCH_DEEPSEEK41) { + for (uint32_t il = 0; il < n_layer; il++) { + ggml_tensor t; + memset(&t, 0, sizeof(ggml_tensor)); + t.type = GGML_TYPE_F16; + ggml_format_name(&t, "conv%" PRIu32 "d.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + } } return ret; } @@ -390,7 +496,7 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, bool no_alloc = false) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -398,9 +504,15 @@ static std::pair get_model_and_ctx( devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); model_params.split_mode = split_mode; + model_params.no_alloc = no_alloc; + if (no_alloc) { + model_params.n_gpu_layers = 0; + model_params.expert_cache_bytes = std::numeric_limits::max(); + model_params.expert_cache_slots = LLAMA_DSV41_N_EXPERT_USED; + } llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = 0; + ctx_params.n_ctx = no_alloc ? 256 : 0; ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; if (!encode) { @@ -421,6 +533,53 @@ static std::pair get_model_and_ctx( return std::make_pair(std::move(model), std::move(lctx)); } +static int test_deepseek41_no_alloc(const size_t seed) { + gguf_context_ptr malformed = get_gguf_ctx(LLM_ARCH_DEEPSEEK41, true); + ggml_tensor malformed_token_embd = {}; + malformed_token_embd.type = GGML_TYPE_F32; + malformed_token_embd.ne[0] = malformed_token_embd.ne[1] = + malformed_token_embd.ne[2] = malformed_token_embd.ne[3] = 1; + malformed_token_embd.nb[0] = ggml_type_size(malformed_token_embd.type); + malformed_token_embd.nb[1] = ggml_row_size( + malformed_token_embd.type, malformed_token_embd.ne[0]); + malformed_token_embd.nb[2] = + malformed_token_embd.nb[1]*malformed_token_embd.ne[1]; + malformed_token_embd.nb[3] = + malformed_token_embd.nb[2]*malformed_token_embd.ne[2]; + ggml_set_name( + &malformed_token_embd, + LLM_TN(LLM_ARCH_DEEPSEEK41)( + LLM_TENSOR_TOKEN_EMBD, "weight").str().c_str()); + gguf_add_tensor(malformed.get(), &malformed_token_embd); + bool rejected = false; + try { + auto unused = get_model_and_ctx( + malformed.get(), nullptr, seed, {}, + LLAMA_SPLIT_MODE_LAYER, false, true); + GGML_UNUSED(unused); + } catch (const std::runtime_error &) { + rejected = true; + } + if (!rejected) { + throw std::runtime_error( + "DeepSeek V4.1 no-allocation loader accepted a malformed tensor"); + } + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_DEEPSEEK41, true); + if (gguf_find_tensor(gguf_ctx.get(), "blk.0.ffn_gate_exps.weight") < 0 || + gguf_find_tensor(gguf_ctx.get(), "blk.1.engram_embd.weight") < 0) { + throw std::runtime_error("DeepSeek V4.1 no-allocation fixture is incomplete"); + } + auto model_and_ctx = get_model_and_ctx( + gguf_ctx.get(), nullptr, seed, {}, + LLAMA_SPLIT_MODE_LAYER, false, true); + if (llama_n_ctx(model_and_ctx.second.get()) != 256) { + throw std::runtime_error("DeepSeek V4.1 no-allocation context size mismatch"); + } + printf("deepseek41: no-allocation architecture graph PASS\n"); + return 0; +} + static std::vector get_logits( llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); @@ -475,6 +634,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DOTS3NOTE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: @@ -561,6 +721,9 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } + if (arch == LLM_ARCH_DEEPSEEK41) { + return false; // The fixed published geometry is too large for the compact generated-model fixture. + } // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_QWEN4EXP) { @@ -870,6 +1033,9 @@ int main(int argc, char ** argv) { if (!out.empty()) { return save_models(arch, seed, verbosity, out); } + if (arch == LLM_ARCH_DEEPSEEK41) { + return test_deepseek41_no_alloc(seed); + } return test_backends(arch, seed, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what()); diff --git a/tests/test_strix_memory_watchdog.py b/tests/test_strix_memory_watchdog.py new file mode 100644 index 000000000000..f9cba175a8d1 --- /dev/null +++ b/tests/test_strix_memory_watchdog.py @@ -0,0 +1,2415 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import fcntl +import hashlib +import io +import json +import os +import signal +import subprocess +import sys +import tempfile +import time +import unittest +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[1] + / "scripts" + / "strix_memory_watchdog.py" +) +SPEC = importlib.util.spec_from_file_location( + "strix_memory_watchdog", SCRIPT_PATH +) +assert SPEC is not None +assert SPEC.loader is not None +watchdog = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = watchdog +SPEC.loader.exec_module(watchdog) + + +def snapshot( + used_bytes: int, + *, + total_bytes: int = 200, + active_swaps: tuple[str, ...] = (), +) -> Any: + return watchdog.HostSnapshot( + total_bytes=total_bytes, + available_bytes=total_bytes - used_bytes, + active_swaps=active_swaps, + ) + + +class SequenceReader: + def __init__(self, values: list[Any]): + self.values = values + self.index = 0 + + def read_snapshot(self) -> Any: + index = min(self.index, len(self.values) - 1) + self.index += 1 + value = self.values[index] + if isinstance(value, Exception): + raise value + return value + + +class FakeClock: + def __init__(self): + self.value = 0.0 + + def monotonic(self) -> float: + return self.value + + def sleep(self, seconds: float) -> None: + self.value += seconds + + +class FakeProcess: + def __init__(self, returncode: int | None = None): + self.pid = 4321 + self.returncode = returncode + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + if self.returncode is None: + raise subprocess.TimeoutExpired("fake", timeout or 0.0) + return self.returncode + + +class Harness: + def __init__( + self, + values: list[Any], + process: FakeProcess, + signal_handler: Any | None = None, + ): + self.reader = SequenceReader(values) + self.process = process + self.signal_handler = signal_handler + self.clock = FakeClock() + self.stream = io.StringIO() + self.launched = False + self.signals: list[int] = [] + fixed_time = datetime(2026, 1, 1, tzinfo=timezone.utc) + self.audit = watchdog.AuditLogger( + self.stream, wall_clock=lambda: fixed_time + ) + + def launcher(self, command: tuple[str, ...], **kwargs: Any) -> FakeProcess: + self.launched = True + self.command = command + self.launch_kwargs = kwargs + return self.process + + def signal_group(self, process_group_id: int, signal_number: int) -> str: + self.signals.append(signal_number) + if self.signal_handler is not None: + self.signal_handler(self.process, signal_number) + return f"{signal.Signals(signal_number).name.lower()}_sent" + + def group_alive(self, process_group_id: int) -> bool: + return self.process.returncode is None + + def run(self, **overrides: Any) -> int: + config = watchdog.WatchdogConfig( + command=("fake-command",), + soft_bytes=100, + emergency_bytes=150, + grace_seconds=2, + sample_interval_seconds=1, + **overrides, + ) + return watchdog.run_watchdog( + config, + reader=self.reader, + audit=self.audit, + launcher=self.launcher, + signal_group=self.signal_group, + group_alive=self.group_alive, + monotonic=self.clock.monotonic, + sleeper=self.clock.sleep, + ) + + def records(self) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in self.stream.getvalue().splitlines() + ] + + +class TestProcfsParsing(unittest.TestCase): + def test_parses_meminfo_as_integer_bytes_and_allows_zero_swap(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + (root / "meminfo").write_text( + "MemTotal: 131072 kB\n" + "MemFree: 4096 kB\n" + "MemAvailable: 32768 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + + result = watchdog.ProcfsReader(root).read_snapshot() + + self.assertEqual(result.total_bytes, 131072 * 1024) + self.assertEqual(result.available_bytes, 32768 * 1024) + self.assertEqual(result.used_bytes, 98304 * 1024) + self.assertEqual(result.active_swaps, ()) + + def test_rejects_active_swap_entry(self) -> None: + content = ( + "Filename Type Size Used Priority\n" + "/swapfile file 1048572 0 -2\n" + ) + self.assertEqual( + watchdog.ProcfsReader._parse_swaps(content), + ("/swapfile",), + ) + + def test_rejects_malformed_or_missing_procfs_data(self) -> None: + with self.assertRaisesRegex( + watchdog.ProcfsError, "malformed MemAvailable" + ): + watchdog.ProcfsReader._parse_meminfo( + "MemTotal: 10 kB\nMemAvailable: unknown\n" + ) + with self.assertRaisesRegex( + watchdog.ProcfsError, "missing MemAvailable" + ): + watchdog.ProcfsReader._parse_meminfo("MemTotal: 10 kB\n") + with self.assertRaisesRegex( + watchdog.ProcfsError, "malformed swaps header" + ): + watchdog.ProcfsReader._parse_swaps("") + with tempfile.TemporaryDirectory() as temp_dir: + with self.assertRaisesRegex( + watchdog.ProcfsError, "cannot read" + ): + watchdog.ProcfsReader( + Path(temp_dir) + ).read_snapshot() + + +class TestWatchdogBehavior(unittest.TestCase): + @staticmethod + def _process_is_running(process_id: int) -> bool: + result = subprocess.run( + ["ps", "-o", "stat=", "-p", str(process_id)], + capture_output=True, + check=False, + text=True, + ) + return result.returncode == 0 and not result.stdout.lstrip().startswith( + "Z" + ) + + @staticmethod + def _write_procfs_fixture(root: Path) -> None: + (root / "meminfo").write_text( + "MemTotal: 131072 kB\nMemAvailable: 65536 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + + @staticmethod + def _lease_arguments(root: Path) -> list[str]: + return [ + "--lease-path", + str(root / "lease.json"), + "--heartbeat-path", + str(root / "heartbeat.json"), + "--audit-path", + str(root / "persistent-audit.jsonl"), + ] + + @staticmethod + def _proc_stat( + process_id: int, + parent_id: int, + process_group_id: int, + start_time_ticks: int, + ) -> str: + fields = [ + "S", + str(parent_id), + str(process_group_id), + *(["0"] * 16), + str(start_time_ticks), + ] + return f"{process_id} (python) {' '.join(fields)}\n" + + def test_parent_signals_leave_no_child_or_grandchild(self) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGHUP,signal.SIG_IGN);" + "signal.signal(signal.SIGINT,signal.SIG_IGN);" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n');" + " time.sleep(30)\n" + ) + for signal_number in ( + signal.SIGHUP, + signal.SIGINT, + signal.SIGTERM, + ): + with self.subTest(signal=signal.Signals(signal_number).name): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + self._write_procfs_fixture(root) + stderr_path = root / "stderr.jsonl" + with stderr_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail( + "child process group did not start" + ) + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + time.sleep(0.05) + wrapper.send_signal(signal_number) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + self.assertEqual( + wrapper.returncode, 128 + signal_number + ) + records = [ + json.loads(line) + for line in stderr_path.read_text( + encoding="utf-8" + ).splitlines() + ] + self.assertEqual( + records[-1]["classification"], "parent_signal" + ) + signal_records = [ + record + for record in records + if record["event"] == "process_group_signal" + ] + forwarded = [ + record["signal"] for record in signal_records + ] + self.assertEqual( + forwarded[0], + signal.Signals(signal_number).name, + ) + self.assertEqual(forwarded[-1], "SIGKILL") + self.assertEqual( + signal_records[0]["child_status"], "running" + ) + self.assertIsNone( + signal_records[0]["child_returncode"] + ) + self.assertLess( + signal_records[0]["timestamp"], + signal_records[-1]["timestamp"], + ) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse( + self._process_is_running(process_id) + ) + lease = json.loads( + (root / "lease.json").read_text( + encoding="utf-8" + ) + ) + heartbeat = json.loads( + (root / "heartbeat.json").read_text( + encoding="utf-8" + ) + ) + persistent_records = [ + json.loads(line) + for line in ( + root / "persistent-audit.jsonl" + ).read_text(encoding="utf-8").splitlines() + ] + self.assertEqual(lease["state"], "final") + self.assertEqual( + lease["final"]["classification"], + "parent_signal", + ) + self.assertEqual(heartbeat["state"], "final") + self.assertEqual( + persistent_records[-1]["classification"], + "parent_signal", + ) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_parent_signal_grace_outlives_guardian_pulse_timeout( + self, + ) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "open(sys.argv[1],'w').write(str(os.getpid()));" + "time.sleep(30)" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pid" + stderr_path = root / "stderr.jsonl" + self._write_procfs_fixture(root) + with stderr_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.4", + "--sample-interval-seconds", + "0.05", + "--heartbeat-max-age-seconds", + "0.1", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + child_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail("child process did not become ready") + time.sleep(0.01) + child_pid = int( + pid_file.read_text(encoding="utf-8") + ) + started = time.monotonic() + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + elapsed = time.monotonic() - started + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in stderr_path.read_text( + encoding="utf-8" + ).splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertGreaterEqual(elapsed, 0.35) + self.assertEqual( + wrapper.returncode, + 128 + signal.SIGTERM, + records, + ) + self.assertEqual(signals, ["SIGTERM", "SIGKILL"]) + self.assertEqual( + records[-1]["classification"], "parent_signal" + ) + self.assertFalse(self._process_is_running(child_pid)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_parent_signal_allows_exit_after_pulse_deadline(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " time.sleep(0.25)\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "open(sys.argv[1],'w').write(str(os.getpid()))\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pid" + stderr_path = root / "stderr.jsonl" + self._write_procfs_fixture(root) + with stderr_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.4", + "--sample-interval-seconds", + "0.05", + "--heartbeat-max-age-seconds", + "0.1", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + try: + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + self.fail("child process did not become ready") + time.sleep(0.01) + started = time.monotonic() + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + elapsed = time.monotonic() - started + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + + records = [ + json.loads(line) + for line in stderr_path.read_text( + encoding="utf-8" + ).splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertGreaterEqual(elapsed, 0.2) + self.assertLess(elapsed, 0.4) + self.assertEqual(wrapper.returncode, 128 + signal.SIGTERM) + self.assertEqual(signals, ["SIGTERM"]) + self.assertEqual(records[-1]["child_returncode"], 0) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_control_failure_still_kills_and_reaps_group( + self, + ) -> None: + class FailingFinalAudit: + def __init__(self, stream: Any, fail_at: int): + self.stream = stream + self.write_count = 0 + self.fail_at = fail_at + + def write(self, value: str) -> int: + self.write_count += 1 + if self.write_count == self.fail_at: + raise OSError("audit write failed") + return self.stream.write(value) + + def flush(self) -> None: + self.stream.flush() + + def fileno(self) -> int: + return self.stream.fileno() + + def close(self) -> None: + self.stream.close() + + class FailingFinalLease: + def finalize(self, record: dict[str, Any]) -> None: + raise watchdog.ArtifactError( + "lease", "final lease write failed" + ) + + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n');" + " time.sleep(30)\n" + ) + for mode in ("closed", "blocked"): + for artifact_failure in ( + "term_audit", + "kill_audit", + "final_audit", + "lease", + ): + with self.subTest( + mode=mode, + artifact_failure=artifact_failure, + ): + self._assert_guardian_control_failure_cleanup( + mode, + artifact_failure, + child_code, + FailingFinalAudit, + FailingFinalLease, + ) + + def _assert_guardian_control_failure_cleanup( + self, + mode: str, + artifact_failure: str, + child_code: str, + failing_final_audit: type, + failing_final_lease: type, + ) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + pid_path = Path(temp_dir) / "pids" + process = subprocess.Popen( + [ + sys.executable, + "-c", + child_code, + str(pid_path), + ], + start_new_session=True, + ) + read_fd, write_fd = os.pipe() + os.set_blocking(write_fd, False) + if mode == "closed": + os.close(write_fd) + write_fd = -1 + else: + try: + while True: + os.write(write_fd, b"x" * 65536) + except BlockingIOError: + pass + guardian = watchdog.GuardianProcess( + process, + process.pid, + write_fd, + ) + stream = io.StringIO() + audit = watchdog.AuditLogger(stream) + if artifact_failure.endswith("_audit"): + persistent_path = Path(temp_dir) / "persistent.jsonl" + persistent_stream = persistent_path.open( + "w", encoding="utf-8" + ) + audit.persistent_stream = failing_final_audit( + persistent_stream, + { + "term_audit": 1, + "kill_audit": 2, + "final_audit": 3, + }[artifact_failure], + ) + else: + audit.lease_manager = failing_final_lease() + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if time.monotonic() >= deadline: + self.fail("child process group did not start") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + result = watchdog._graceful_cleanup( + audit, + guardian, + snapshot(50), + 50, + "parent_signal", + 128 + signal.SIGTERM, + "wrapper received SIGTERM", + signal.SIGTERM, + 0.4, + watchdog._signal_process_group, + watchdog._process_group_alive, + time.monotonic, + time.sleep, + ) + finally: + if process.poll() is None: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5) + guardian.close() + os.close(read_fd) + + records = [ + json.loads(line) + for line in stream.getvalue().splitlines() + ] + signals = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(result, watchdog.EXIT_SIGNAL_ERROR) + self.assertEqual(signals, ["SIGTERM", "SIGKILL"]) + self.assertEqual( + records[-1]["classification"], "signal_error" + ) + self.assertEqual(records[-1]["exit_code"], 7) + self.assertEqual( + records[-1]["threshold_reason"], + "guardian control failed during graceful cleanup", + ) + self.assertEqual( + records[-1]["secondary_errors"][0]["component"], + ( + "audit" + if artifact_failure.endswith("_audit") + else "lease" + ), + ) + self.assertIn( + ( + "audit write failed" + if artifact_failure.endswith("_audit") + else "final lease write failed" + ), + records[-1]["secondary_errors"][0]["detail"], + ) + self.assertEqual( + records[-1]["child_returncode"], + -signal.SIGKILL, + ) + assert child_pid is not None + assert grandchild_pid is not None + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_pipe_close_kills_group_without_fd_leak(self) -> None: + child_code = ( + "import json,os,subprocess,sys,time\n" + "targets=[]\n" + "for name in os.listdir('/proc/self/fd'):\n" + " try: targets.append(os.readlink('/proc/self/fd/'+name))\n" + " except OSError: pass\n" + "grandchild=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'])\n" + "open(sys.argv[1],'w').write(json.dumps({" + "'child':os.getpid(),'grandchild':grandchild.pid," + "'fds':targets}))\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + state_path = root / "state.json" + guardian = watchdog._launch_guardian( + ( + sys.executable, + "-c", + child_code, + str(state_path), + ), + os.environ.copy(), + 0.5, + 1.0, + signal.pthread_sigmask(signal.SIG_BLOCK, ()), + ) + control_target = os.readlink( + f"/proc/self/fd/{guardian.pulse_fd}" + ) + deadline = time.monotonic() + 5 + while not state_path.exists(): + if time.monotonic() >= deadline: + self.fail("guardian payload did not become ready") + time.sleep(0.01) + state = json.loads(state_path.read_text(encoding="utf-8")) + os.close(guardian.pulse_fd) + guardian.wait(timeout=5) + + self.assertNotIn(control_target, state["fds"]) + for process_id in (state["child"], state["grandchild"]): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_documents_setsid_escape_limit(self) -> None: + child_code = ( + "import os,subprocess,sys,time\n" + "escaped=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'],start_new_session=True)\n" + "open(sys.argv[1],'w').write(str(escaped.pid))\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + state_path = Path(temp_dir) / "escaped-pid" + guardian = watchdog._launch_guardian( + ( + sys.executable, + "-c", + child_code, + str(state_path), + ), + os.environ.copy(), + 0.5, + 1.0, + signal.pthread_sigmask(signal.SIG_BLOCK, ()), + ) + deadline = time.monotonic() + 5 + while not state_path.exists(): + if time.monotonic() >= deadline: + self.fail("escaped payload did not become ready") + time.sleep(0.01) + escaped_pid = int( + state_path.read_text(encoding="utf-8") + ) + os.close(guardian.pulse_fd) + guardian.wait(timeout=5) + self.assertTrue(self._process_is_running(escaped_pid)) + os.kill(escaped_pid, signal.SIGKILL) + deadline = time.monotonic() + 2 + while ( + self._process_is_running(escaped_pid) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(escaped_pid)) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guard_kills_group_after_watchdog_loss_or_stall(self) -> None: + child_code = ( + "import importlib.util,os,pathlib,subprocess,sys,time\n" + "script=pathlib.Path(sys.argv[1])\n" + "spec=importlib.util.spec_from_file_location('guard_watchdog',script)\n" + "module=importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name]=module\n" + "spec.loader.exec_module(module)\n" + "module.start_process_group_lease_guard(" + "script,expected_procfs_root=pathlib.Path(sys.argv[2]))\n" + "grandchild=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'])\n" + "open(sys.argv[3],'w').write(" + "f'{os.getpid()} {grandchild.pid}\\n')\n" + "time.sleep(30)\n" + ) + for mode in ("sigkill", "sigstop"): + with self.subTest(mode=mode): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_path = root / "pids" + self._write_procfs_fixture(root) + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--heartbeat-max-age-seconds", + "0.3", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(SCRIPT_PATH), + str(root), + str(pid_path), + ], + stderr=subprocess.PIPE, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if wrapper.poll() is not None: + assert wrapper.stderr is not None + self.fail(wrapper.stderr.read()) + if time.monotonic() >= deadline: + self.fail( + "guarded payload did not become ready" + ) + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + if mode == "sigkill": + wrapper.kill() + else: + os.kill(wrapper.pid, signal.SIGSTOP) + heartbeat_path = root / "heartbeat.json" + heartbeat = json.loads( + heartbeat_path.read_text(encoding="utf-8") + ) + heartbeat["updated_monotonic_ns"] = ( + time.monotonic_ns() + ) + watchdog._write_json_atomic( + heartbeat_path, heartbeat + ) + time.sleep(0.7) + os.kill(wrapper.pid, signal.SIGCONT) + wrapper.wait(timeout=5) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse( + self._process_is_running(process_id) + ) + if wrapper.stderr is not None: + wrapper.stderr.close() + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_payload_guard_fails_closed_on_artifact_error(self) -> None: + child_code = ( + "import importlib.util,os,pathlib,subprocess,sys,time\n" + "script=pathlib.Path(sys.argv[1])\n" + "spec=importlib.util.spec_from_file_location('guard_watchdog',script)\n" + "module=importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name]=module\n" + "spec.loader.exec_module(module)\n" + "module.start_process_group_lease_guard(" + "script,expected_procfs_root=pathlib.Path(sys.argv[2]))\n" + "def fail(*_args,**_kwargs):\n" + " raise module.ArtifactError('script','unreadable')\n" + "module.validate_active_lease=fail\n" + "grandchild=subprocess.Popen([sys.executable,'-c'," + "'import time;time.sleep(30)'])\n" + "open(sys.argv[3],'w').write(" + "f'{os.getpid()} {grandchild.pid}\\n')\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_path = root / "pids" + self._write_procfs_fixture(root) + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--heartbeat-max-age-seconds", + "0.3", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(SCRIPT_PATH), + str(root), + str(pid_path), + ], + stderr=subprocess.PIPE, + text=True, + ) + child_pid = None + grandchild_pid = None + try: + deadline = time.monotonic() + 5 + while not pid_path.exists(): + if wrapper.poll() is not None: + assert wrapper.stderr is not None + self.fail(wrapper.stderr.read()) + if time.monotonic() >= deadline: + self.fail("guarded payload did not become ready") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_path.read_text( + encoding="utf-8" + ).split() + ) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + if wrapper.stderr is not None: + wrapper.stderr.close() + if child_pid is not None: + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + assert child_pid is not None + assert grandchild_pid is not None + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + + def test_child_sigterm_handler_exits_without_escalation(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " open(sys.argv[2],'w').write('handled\\n')\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "open(sys.argv[1],'w').write(f'{os.getpid()}\\n')\n" + "time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + ready_path = root / "ready" + handled_path = root / "handled" + audit_path = root / "audit.jsonl" + self._write_procfs_fixture(root) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.5", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(ready_path), + str(handled_path), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not ready_path.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("SIGTERM child did not become ready") + time.sleep(0.01) + child_pid = int( + ready_path.read_text(encoding="utf-8").strip() + ) + try: + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(wrapper.returncode, 128 + signal.SIGTERM) + self.assertTrue(handled_path.exists()) + self.assertEqual(forwarded, ["SIGTERM"]) + self.assertEqual(records[-1]["child_returncode"], 0) + + def test_leader_exit_cleans_up_surviving_grandchild(self) -> None: + child_code = ( + "import os,signal,sys,time;" + "signal.signal(signal.SIGTERM,signal.SIG_IGN);" + "grandchild=os.fork();" + "\nif grandchild == 0:\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n')\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + audit_path = root / "audit.jsonl" + self._write_procfs_fixture(root) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("leader process did not write child PIDs") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + try: + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + self.assertEqual(wrapper.returncode, 0) + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual(records[-1]["child_returncode"], 0) + self.assertEqual(forwarded, ["SIGTERM", "SIGKILL"]) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + + def test_soft_limit_descendant_escalation_is_grace_timeout(self) -> None: + child_code = ( + "import os,signal,sys,time\n" + "def stop(_signal,_frame):\n" + " raise SystemExit(0)\n" + "signal.signal(signal.SIGTERM,stop)\n" + "grandchild=os.fork()\n" + "if grandchild == 0:\n" + " signal.signal(signal.SIGTERM,signal.SIG_IGN)\n" + " time.sleep(30)\n" + "else:\n" + " open(sys.argv[1],'w').write(" + "f'{os.getpid()} {grandchild}\\n')\n" + " time.sleep(30)\n" + ) + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + pid_file = root / "pids" + audit_path = root / "audit.jsonl" + (root / "meminfo").write_text( + "MemTotal: 3145728 kB\n" + "MemAvailable: 2621440 kB\n", + encoding="utf-8", + ) + (root / "swaps").write_text( + "Filename Type Size Used Priority\n", + encoding="utf-8", + ) + with audit_path.open("w", encoding="utf-8") as audit: + wrapper = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + "--soft-gib", + "1", + "--emergency-gib", + "2", + "--grace-seconds", + "0.2", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + child_code, + str(pid_file), + ], + stderr=audit, + text=True, + ) + deadline = time.monotonic() + 5 + while not pid_file.exists(): + if time.monotonic() >= deadline: + wrapper.kill() + wrapper.wait(timeout=5) + self.fail("soft-limit process group did not start") + time.sleep(0.01) + child_pid, grandchild_pid = ( + int(value) + for value in pid_file.read_text( + encoding="utf-8" + ).split() + ) + next_meminfo = root / "meminfo.next" + next_meminfo.write_text( + "MemTotal: 3145728 kB\n" + "MemAvailable: 1572864 kB\n", + encoding="utf-8", + ) + next_meminfo.replace(root / "meminfo") + try: + wrapper.wait(timeout=5) + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=5) + try: + os.killpg(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + records = [ + json.loads(line) + for line in audit_path.read_text( + encoding="utf-8" + ).splitlines() + ] + forwarded = [ + record["signal"] + for record in records + if record["event"] == "process_group_signal" + ] + final = records[-1] + self.assertEqual(wrapper.returncode, watchdog.EXIT_GRACE_TIMEOUT) + self.assertEqual(final["classification"], "grace_timeout") + self.assertEqual(final["child_returncode"], 0) + self.assertEqual(forwarded, ["SIGTERM", "SIGKILL"]) + for process_id in (child_pid, grandchild_pid): + deadline = time.monotonic() + 2 + while ( + self._process_is_running(process_id) + and time.monotonic() < deadline + ): + time.sleep(0.01) + self.assertFalse(self._process_is_running(process_id)) + + def test_configuration_rejects_non_finite_timing(self) -> None: + config = watchdog.WatchdogConfig( + command=("fake-command",), + grace_seconds=float("nan"), + ) + with self.assertRaisesRegex(ValueError, "grace period"): + config.validate() + + def test_configuration_rejects_weakened_liveness_timing(self) -> None: + cases = ( + ( + {"grace_seconds": 31.0}, + "grace period", + ), + ( + {"sample_interval_seconds": 1.1}, + "sample interval", + ), + ( + { + "sample_interval_seconds": 1.0, + "heartbeat_max_age_seconds": 5.1, + }, + "heartbeat max age", + ), + ) + for overrides, message in cases: + with self.subTest(overrides=overrides): + config = watchdog.WatchdogConfig( + command=("fake-command",), + **overrides, + ) + with self.assertRaisesRegex(ValueError, message): + config.validate() + + def test_stderr_failure_does_not_bypass_cleanup(self) -> None: + class FailingStderr(io.StringIO): + def write(self, value: str) -> int: + raise OSError("stderr closed") + + process = FakeProcess() + + def exit_on_kill( + target: FakeProcess, signal_number: int + ) -> None: + if signal_number == signal.SIGKILL: + target.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50)], + process, + signal_handler=exit_on_kill, + ) + harness.audit = watchdog.AuditLogger(FailingStderr()) + with tempfile.TemporaryDirectory() as temp_dir: + persistent_path = Path(temp_dir) / "audit.jsonl" + harness.audit.open_persistent(persistent_path) + result = watchdog._graceful_cleanup( + harness.audit, + process, + snapshot(50), + 50, + "internal_error", + watchdog.EXIT_INTERNAL_ERROR, + "test cleanup", + signal.SIGTERM, + 0.1, + harness.signal_group, + harness.group_alive, + harness.clock.monotonic, + harness.clock.sleep, + ) + harness.audit.close() + records = [ + json.loads(line) + for line in persistent_path.read_text( + encoding="utf-8" + ).splitlines() + ] + + self.assertEqual( + harness.signals, [signal.SIGTERM, signal.SIGKILL] + ) + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(records[-1]["classification"], "lease_error") + + def test_audit_write_and_close_failures_do_not_bypass_cleanup( + self, + ) -> None: + class FailingPersistent(io.StringIO): + def __init__(self) -> None: + super().__init__() + self.close_called = False + + def write(self, value: str) -> int: + raise OSError("persistent write failed") + + def close(self) -> None: + if self.close_called: + super().close() + return + self.close_called = True + raise OSError("persistent close failed") + + process = FakeProcess() + + def exit_on_kill( + target: FakeProcess, signal_number: int + ) -> None: + if signal_number == signal.SIGKILL: + target.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50)], + process, + signal_handler=exit_on_kill, + ) + persistent = FailingPersistent() + harness.audit.persistent_stream = persistent + result = watchdog._graceful_cleanup( + harness.audit, + process, + snapshot(50), + 50, + "internal_error", + watchdog.EXIT_INTERNAL_ERROR, + "test cleanup", + signal.SIGTERM, + 0.1, + harness.signal_group, + harness.group_alive, + harness.clock.monotonic, + harness.clock.sleep, + ) + + self.assertTrue(persistent.close_called) + self.assertEqual( + harness.signals, [signal.SIGTERM, signal.SIGKILL] + ) + self.assertEqual(process.returncode, -signal.SIGKILL) + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual( + harness.records()[-1]["classification"], "lease_error" + ) + + def test_final_record_survives_artifact_failures(self) -> None: + class FailingLease: + def finalize(self, record: dict[str, Any]) -> None: + raise watchdog.ArtifactError("lease", "write failed") + + class FailingAudit(io.StringIO): + def write(self, value: str) -> int: + raise OSError("write failed") + + for component in ("lease", "audit"): + with self.subTest(component=component): + stream = io.StringIO() + audit = watchdog.AuditLogger(stream) + if component == "lease": + setattr(audit, "lease_manager", FailingLease()) + else: + audit.persistent_stream = FailingAudit() + result = watchdog._emit_final( + audit, + "child_exit", + 0, + "child exited", + snapshot(50), + 50, + ) + records = [ + json.loads(line) + for line in stream.getvalue().splitlines() + ] + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual( + records[-1]["classification"], "lease_error" + ) + self.assertTrue(audit.finalized) + self.assertEqual( + audit.final_exit_code, watchdog.EXIT_LEASE_ERROR + ) + + def test_emergency_signal_precedes_artifact_write(self) -> None: + events: list[str] = [] + + class BlockingAudit(watchdog.AuditLogger): + def __init__(self) -> None: + super().__init__(io.StringIO()) + self.calls = 0 + + def emit( + self, event: str, **fields: object + ) -> dict[str, object]: + self.calls += 1 + events.append(f"audit:{event}") + if self.calls == 3: + raise watchdog.ArtifactError( + "audit", "simulated blocked fsync" + ) + return super().emit(event, **fields) + + def exit_on_kill( + process: FakeProcess, signal_number: int + ) -> None: + events.append(f"signal:{signal_number}") + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(160)], + FakeProcess(), + signal_handler=exit_on_kill, + ) + harness.audit = BlockingAudit() + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(events[2], f"signal:{signal.SIGKILL}") + self.assertEqual(events[3], "audit:process_group_signal") + self.assertEqual(harness.process.returncode, -signal.SIGKILL) + + def test_cleanup_reaps_after_persistent_audit_failure(self) -> None: + class FailingSignalAudit(watchdog.AuditLogger): + def emit( + self, event: str, **fields: object + ) -> dict[str, object]: + if event == "process_group_signal": + raise watchdog.ArtifactError( + "audit", "simulated persistent write failure" + ) + return super().emit(event, **fields) + + process = FakeProcess() + signals: list[int] = [] + clock = FakeClock() + + def signal_group( + process_group_id: int, signal_number: int + ) -> str: + signals.append(signal_number) + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + return f"{signal.Signals(signal_number).name.lower()}_sent" + + result = watchdog._graceful_cleanup( + FailingSignalAudit(io.StringIO()), + process, + snapshot(50), + 50, + "parent_signal", + 128 + signal.SIGTERM, + "wrapper received SIGTERM", + signal.SIGTERM, + 0.1, + signal_group, + lambda _process_group_id: process.returncode is None, + clock.monotonic, + clock.sleep, + ) + + self.assertEqual(result, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(signals, [signal.SIGTERM, signal.SIGKILL]) + self.assertEqual(process.returncode, -signal.SIGKILL) + + def test_invalid_artifact_path_emits_configuration_final(self) -> None: + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--lease-path", + "~strix-watchdog-user-does-not-exist/lease.json", + "--heartbeat-path", + "/tmp/heartbeat.json", + "--audit-path", + "/tmp/audit.jsonl", + "--", + sys.executable, + "-c", + "pass", + ], + capture_output=True, + check=False, + text=True, + ) + + self.assertEqual(result.returncode, watchdog.EXIT_PROCFS_ERROR) + final = json.loads(result.stderr.splitlines()[-1]) + self.assertEqual(final["classification"], "configuration_error") + + def test_cli_fixture_launches_command_and_propagates_exit(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + child_result_path = root / "child-result.json" + self._write_procfs_fixture(root) + child_code = ( + "import json,os,sys,time\n" + "keys=('STRIX_MEMORY_WATCHDOG_LEASE_PATH'," + "'STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH'," + "'STRIX_MEMORY_WATCHDOG_AUDIT_PATH')\n" + "deadline=time.monotonic()+5\n" + "while True:\n" + " try:\n" + " with open(os.environ[keys[0]],encoding='utf-8') as stream:\n" + " lease=json.load(stream)\n" + " break\n" + " except (OSError,json.JSONDecodeError):\n" + " if time.monotonic()>=deadline: raise\n" + " time.sleep(0.01)\n" + "assert os.getpgrp()==lease['child_process_group_id']\n" + "open(sys.argv[1],'w').write(json.dumps({" + "key:os.environ[key] for key in keys}))\n" + "raise SystemExit(23)\n" + ) + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--sample-interval-seconds", + "0.01", + "--", + sys.executable, + "-c", + child_code, + str(child_result_path), + ], + capture_output=True, + check=False, + text=True, + ) + lease = json.loads( + (root / "lease.json").read_text(encoding="utf-8") + ) + persistent_records = [ + json.loads(line) + for line in ( + root / "persistent-audit.jsonl" + ).read_text(encoding="utf-8").splitlines() + ] + child_result = json.loads( + child_result_path.read_text(encoding="utf-8") + ) + + self.assertEqual(result.returncode, 23) + records = [ + json.loads(line) for line in result.stderr.splitlines() + ] + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual(records[-1]["child_returncode"], 23) + self.assertEqual(lease["format"], watchdog.LEASE_FORMAT) + self.assertEqual(lease["version"], watchdog.LEASE_VERSION) + self.assertEqual(lease["state"], "final") + self.assertEqual(lease["soft_bytes"], 116 * 1024**3) + self.assertEqual( + lease["emergency_bytes"], 118 * 1024**3 + ) + self.assertEqual( + lease["child_command_sha256"], + watchdog._command_sha256( + ( + sys.executable, + "-c", + child_code, + str(child_result_path), + ) + ), + ) + self.assertEqual( + lease["final"]["classification"], "child_exit" + ) + self.assertEqual( + persistent_records[-1]["classification"], "child_exit" + ) + self.assertEqual( + child_result["STRIX_MEMORY_WATCHDOG_LEASE_PATH"], + str((root / "lease.json").resolve()), + ) + self.assertEqual( + child_result["STRIX_MEMORY_WATCHDOG_HEARTBEAT_PATH"], + str((root / "heartbeat.json").resolve()), + ) + self.assertEqual( + child_result["STRIX_MEMORY_WATCHDOG_AUDIT_PATH"], + str((root / "persistent-audit.jsonl").resolve()), + ) + + @unittest.skipUnless( + sys.platform.startswith("linux"), + "Linux guardian lifecycle", + ) + def test_guardian_preserves_payload_signal_status(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self._write_procfs_fixture(root) + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--sample-interval-seconds", + "0.01", + "--", + sys.executable, + "-c", + ( + "import os,signal,time;" + "time.sleep(0.1);" + "os.kill(os.getpid(),signal.SIGTERM)" + ), + ], + capture_output=True, + check=False, + text=True, + timeout=5, + ) + records = [ + json.loads(line) + for line in result.stderr.splitlines() + ] + + self.assertEqual(result.returncode, 128 + signal.SIGTERM) + self.assertEqual(records[-1]["classification"], "child_exit") + self.assertEqual( + records[-1]["child_returncode"], -signal.SIGTERM + ) + self.assertEqual(records[-1]["child_status"], "signaled") + + def test_existing_lease_fails_closed_and_stops_child(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self._write_procfs_fixture(root) + lease_path = root / "lease.json" + lease_path.write_text("untrusted\n", encoding="utf-8") + result = subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + "--procfs-root", + str(root), + *self._lease_arguments(root), + "--grace-seconds", + "0.1", + "--sample-interval-seconds", + "0.05", + "--", + sys.executable, + "-c", + "import time;time.sleep(30)", + ], + capture_output=True, + check=False, + text=True, + timeout=5, + ) + records = [ + json.loads(line) for line in result.stderr.splitlines() + ] + child_pid = records[-1]["child_pid"] + self.assertEqual( + lease_path.read_text(encoding="utf-8"), "untrusted\n" + ) + + self.assertEqual(result.returncode, watchdog.EXIT_LEASE_ERROR) + self.assertEqual(records[-1]["classification"], "lease_error") + self.assertFalse(self._process_is_running(child_pid)) + + def test_active_lease_validation_rejects_tamper_and_stale(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + process_root = root / "proc" + watchdog_pid = 1200 + guardian_pid = 1250 + child_pid = 1300 + current_pid = 1400 + watchdog_start_ticks = 456789 + script_path = root / "watchdog.py" + script_path.write_text("print('watchdog')\n", encoding="utf-8") + lease_path = root / "lease.json" + heartbeat_path = root / "heartbeat.json" + audit_path = root / "audit.jsonl" + command = [sys.executable, "run_matrix.py"] + argv = [ + sys.executable, + "watchdog.py", + "--procfs-root", + "/proc", + "--lease-path", + str(lease_path), + "--heartbeat-path", + str(heartbeat_path), + "--audit-path", + str(audit_path), + "--", + *command, + ] + cmdline = b"\0".join(os.fsencode(value) for value in argv) + for process_id, parent_id, group_id, start_ticks in ( + (watchdog_pid, 1, watchdog_pid, watchdog_start_ticks), + (guardian_pid, watchdog_pid, guardian_pid, 456790), + (child_pid, guardian_pid, guardian_pid, 456791), + (current_pid, child_pid, guardian_pid, 456792), + ): + process_dir = process_root / str(process_id) + process_dir.mkdir(parents=True) + (process_dir / "stat").write_text( + self._proc_stat( + process_id, + parent_id, + group_id, + start_ticks, + ), + encoding="utf-8", + ) + (process_root / str(watchdog_pid) / "cwd").symlink_to( + root, target_is_directory=True + ) + (process_root / str(watchdog_pid) / "exe").symlink_to( + Path(sys.executable).resolve() + ) + (process_root / str(watchdog_pid) / "cmdline").write_bytes( + cmdline + ) + + audit_line = ( + '{"event":"child_started",' + '"timestamp":"2026-01-01T00:00:00Z"}\n' + ) + audit_descriptor = os.open( + audit_path, + os.O_CREAT | os.O_EXCL | os.O_RDWR, + 0o600, + ) + os.write(audit_descriptor, audit_line.encode("utf-8")) + os.fsync(audit_descriptor) + fcntl.flock( + audit_descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB + ) + audit_status = os.fstat(audit_descriptor) + fd_root = process_root / str(watchdog_pid) / "fd" + fd_root.mkdir() + (fd_root / "9").symlink_to(audit_path) + lease = { + "format": watchdog.LEASE_FORMAT, + "version": watchdog.LEASE_VERSION, + "lease_id": "test-lease", + "state": "active", + "watchdog_pid": watchdog_pid, + "watchdog_start_time_utc": "2026-01-01T00:00:00.000Z", + "watchdog_start_time_ticks": watchdog_start_ticks, + "watchdog_command_sha256": hashlib.sha256( + cmdline + ).hexdigest(), + "watchdog_executable_path": str( + Path(sys.executable).resolve() + ), + "watchdog_script_path": str(script_path), + "watchdog_script_sha256": hashlib.sha256( + script_path.read_bytes() + ).hexdigest(), + "soft_bytes": watchdog.DEFAULT_SOFT_BYTES, + "emergency_bytes": watchdog.DEFAULT_EMERGENCY_BYTES, + "strict_ceiling_bytes": watchdog.STRICT_CEILING_BYTES, + "grace_seconds": watchdog.DEFAULT_GRACE_SECONDS, + "sample_interval_seconds": ( + watchdog.DEFAULT_SAMPLE_INTERVAL_SECONDS + ), + "guardian_pid": guardian_pid, + "child_pid": child_pid, + "child_process_group_id": guardian_pid, + "command": command, + "child_command_sha256": watchdog._command_sha256( + command + ), + "heartbeat_path": str(heartbeat_path), + "max_heartbeat_age_seconds": 5.0, + "audit_path": str(audit_path), + "audit_device": audit_status.st_dev, + "audit_inode": audit_status.st_ino, + "audit_uid": audit_status.st_uid, + "audit_mode": 0o600, + "audit_fd": 9, + "procfs_root": "/proc", + } + heartbeat = { + "format": watchdog.HEARTBEAT_FORMAT, + "version": watchdog.HEARTBEAT_VERSION, + "lease_id": "test-lease", + "sequence": 4, + "state": "active", + "updated_at": "2026-01-01T00:00:01.000Z", + "updated_monotonic_ns": 9_000_000_000, + "watchdog_pid": watchdog_pid, + "watchdog_start_time_ticks": ( + watchdog_start_ticks + ), + "child_pid": child_pid, + "child_process_group_id": guardian_pid, + "sample": { + "audit_record_sha256": hashlib.sha256( + audit_line.encode("utf-8") + ).hexdigest() + }, + } + watchdog._write_json_atomic( + lease_path, lease, create=True + ) + watchdog._write_json_atomic( + heartbeat_path, heartbeat, create=True + ) + + try: + validation_args = { + "expected_script_path": script_path, + "expected_executable_path": Path(sys.executable), + "expected_command": command, + "expected_heartbeat_path": heartbeat_path, + "expected_audit_path": audit_path, + "expected_max_heartbeat_age_seconds": 5.0, + "current_process_id": current_pid, + "process_procfs_root": process_root, + "monotonic_ns": lambda: 10_000_000_000, + "pidfd_open": lambda _pid: os.open( + os.devnull, os.O_RDONLY + ), + } + validated = watchdog.validate_active_lease( + lease_path, **validation_args + ) + self.assertEqual(validated["lease_id"], "test-lease") + + def publish_lease( + value: dict[str, Any], + process_argv: list[str] = argv, + ) -> None: + process_cmdline = b"\0".join( + os.fsencode(argument) + for argument in process_argv + ) + (process_root / str(watchdog_pid) / "cmdline").write_bytes( + process_cmdline + ) + value["watchdog_command_sha256"] = hashlib.sha256( + process_cmdline + ).hexdigest() + watchdog._write_json_atomic(lease_path, value) + + for name, bad_argv in ( + ( + "helper inert argument", + [ + sys.executable, + "helper.py", + str(script_path), + *argv[2:], + ], + ), + ( + "python command string", + [ + sys.executable, + "-c", + "pass", + str(script_path), + *argv[2:], + ], + ), + ( + "python module", + [ + sys.executable, + "-m", + "helper", + str(script_path), + *argv[2:], + ], + ), + ( + "interpreter option before script", + [ + sys.executable, + "-O", + str(script_path), + *argv[2:], + ], + ), + ): + with self.subTest(name): + publish_lease(dict(lease), bad_argv) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "executable argv position", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("wrong command-line policy"): + bad_argv = list(argv) + procfs_index = bad_argv.index("--procfs-root") + 1 + bad_argv[procfs_index] = "/tmp/not-proc" + publish_lease(dict(lease), bad_argv) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "command-line policy", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("wrong lease timing policy"): + bad_lease = dict(lease) + bad_lease["grace_seconds"] = 29.0 + publish_lease(bad_lease) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "lease timing policy", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("wrong monitored command"): + bad_argv = [*argv[:-1], "other_matrix.py"] + bad_lease = dict(lease) + bad_lease["command"] = [ + sys.executable, + "other_matrix.py", + ] + bad_lease["child_command_sha256"] = ( + watchdog._command_sha256( + bad_lease["command"] + ) + ) + publish_lease(bad_lease, bad_argv) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "monitored command", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("tampered script SHA"): + tampered = dict(lease) + tampered["watchdog_script_sha256"] = "0" * 64 + publish_lease(tampered) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, "script SHA" + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("stale heartbeat"): + publish_lease(dict(lease)) + heartbeat["updated_monotonic_ns"] = 1 + watchdog._write_json_atomic( + heartbeat_path, heartbeat + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "heartbeat is stale", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("arbitrary heartbeat"): + heartbeat["updated_monotonic_ns"] = 9_000_000_000 + heartbeat["lease_id"] = "helper-lease" + watchdog._write_json_atomic( + heartbeat_path, heartbeat + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "heartbeat identity", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("outside process group"): + heartbeat["lease_id"] = "test-lease" + watchdog._write_json_atomic( + heartbeat_path, heartbeat + ) + (process_root / str(current_pid) / "stat").write_text( + self._proc_stat( + current_pid, + child_pid, + 9999, + 456792, + ), + encoding="utf-8", + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "outside the monitored process group", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + ( + process_root / str(current_pid) / "stat" + ).write_text( + self._proc_stat( + current_pid, + child_pid, + guardian_pid, + 456792, + ), + encoding="utf-8", + ) + + with self.subTest("environment path mismatch"): + bad_validation_args = { + **validation_args, + "expected_heartbeat_path": root / "other.json", + } + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "artifact paths|heartbeat path", + ): + watchdog.validate_active_lease( + lease_path, **bad_validation_args + ) + + with self.subTest("lease inode mismatch"): + publish_lease(dict(lease)) + lease_record = json.loads( + lease_path.read_text(encoding="utf-8") + ) + lease_record["file_inode"] = 0 + lease_path.write_text( + json.dumps(lease_record), encoding="utf-8" + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "identity does not match", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + + with self.subTest("watchdog start tick mismatch"): + publish_lease(dict(lease)) + (process_root / str(watchdog_pid) / "stat").write_text( + self._proc_stat( + watchdog_pid, + 1, + watchdog_pid, + watchdog_start_ticks + 1, + ), + encoding="utf-8", + ) + with self.assertRaisesRegex( + watchdog.LeaseValidationError, + "start time", + ): + watchdog.validate_active_lease( + lease_path, **validation_args + ) + finally: + os.close(audit_descriptor) + + def test_zero_swap_gate_launches_and_propagates_child_exit(self) -> None: + harness = Harness([snapshot(50)], FakeProcess(returncode=37)) + + result = harness.run() + + self.assertEqual(result, 37) + self.assertTrue(harness.launched) + self.assertTrue(harness.launch_kwargs["start_new_session"]) + final = harness.records()[-1] + self.assertEqual(final["classification"], "child_exit") + self.assertEqual(final["total_bytes"], 200) + self.assertEqual(final["available_bytes"], 150) + self.assertEqual(final["used_bytes"], 50) + self.assertEqual(final["peak_used_bytes"], 50) + self.assertEqual(final["child_status"], "exited") + self.assertEqual(final["process_group_status"], "leader_exited") + + def test_signaled_child_exit_uses_shell_exit_convention(self) -> None: + harness = Harness( + [snapshot(50)], + FakeProcess(returncode=-signal.SIGTERM), + ) + + result = harness.run() + + self.assertEqual(result, 128 + signal.SIGTERM) + + def test_active_swap_rejects_startup_without_launch(self) -> None: + harness = Harness( + [snapshot(50, active_swaps=("/swapfile",))], + FakeProcess(), + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SWAP_ACTIVE) + self.assertFalse(harness.launched) + self.assertEqual( + harness.records()[-1]["classification"], + "startup_swap_active", + ) + + def test_soft_limit_sends_sigterm(self) -> None: + def exit_on_term(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGTERM: + process.returncode = -signal.SIGTERM + + harness = Harness( + [snapshot(50), snapshot(110)], + FakeProcess(), + signal_handler=exit_on_term, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + result = harness.run( + lease_path=root / "lease.json", + heartbeat_path=root / "heartbeat.json", + audit_path=root / "audit.jsonl", + ) + harness.audit.close() + heartbeat = json.loads( + (root / "heartbeat.json").read_text(encoding="utf-8") + ) + lease = json.loads( + (root / "lease.json").read_text(encoding="utf-8") + ) + + self.assertEqual(result, watchdog.EXIT_SOFT_LIMIT) + self.assertEqual(harness.signals, [signal.SIGTERM]) + self.assertEqual( + harness.records()[-1]["classification"], "soft_limit" + ) + self.assertEqual(heartbeat["state"], "final") + self.assertEqual(heartbeat["sequence"], 3) + self.assertEqual(lease["final"]["classification"], "soft_limit") + + def test_emergency_limit_sends_sigkill(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(160)], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_EMERGENCY_LIMIT) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "emergency_limit" + ) + + def test_grace_timeout_escalates_to_sigkill(self) -> None: + def ignore_term(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), snapshot(110)], + FakeProcess(), + signal_handler=ignore_term, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_GRACE_TIMEOUT) + self.assertEqual( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + self.assertEqual(harness.clock.value, 2.0) + self.assertEqual( + harness.records()[-1]["classification"], "grace_timeout" + ) + + def test_swap_appearing_during_execution_kills_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [ + snapshot(50), + snapshot(60, active_swaps=("/swapfile",)), + ], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_SWAP_ACTIVE) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "swap_appeared" + ) + + def test_runtime_procfs_error_kills_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), watchdog.ProcfsError("missing meminfo")], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_PROCFS_ERROR) + self.assertEqual(harness.signals, [signal.SIGKILL]) + self.assertEqual( + harness.records()[-1]["classification"], "procfs_error" + ) + + def test_unexpected_monitor_error_cleans_up_process_group(self) -> None: + def exit_on_kill(process: FakeProcess, signal_number: int) -> None: + if signal_number == signal.SIGKILL: + process.returncode = -signal.SIGKILL + + harness = Harness( + [snapshot(50), RuntimeError("unexpected")], + FakeProcess(), + signal_handler=exit_on_kill, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + result = harness.run( + lease_path=root / "lease.json", + heartbeat_path=root / "heartbeat.json", + audit_path=root / "audit.jsonl", + ) + harness.audit.close() + lease = json.loads( + (root / "lease.json").read_text(encoding="utf-8") + ) + persistent_records = [ + json.loads(line) + for line in (root / "audit.jsonl").read_text( + encoding="utf-8" + ).splitlines() + ] + + self.assertEqual(result, watchdog.EXIT_INTERNAL_ERROR) + self.assertEqual( + harness.signals, + [signal.SIGTERM, signal.SIGKILL], + ) + final = harness.records()[-1] + self.assertEqual(final["classification"], "internal_error") + self.assertIn("RuntimeError: unexpected", final["error"]) + self.assertEqual(lease["final"]["classification"], "internal_error") + self.assertEqual( + persistent_records[-1]["classification"], "internal_error" + ) + + def test_launch_failure_is_explicit(self) -> None: + harness = Harness([snapshot(50)], FakeProcess()) + + def fail_launch( + command: tuple[str, ...], **kwargs: Any + ) -> FakeProcess: + raise FileNotFoundError(2, "No such file or directory") + + setattr(harness, "launcher", fail_launch) + + result = harness.run() + + self.assertEqual(result, watchdog.EXIT_LAUNCH_ERROR) + self.assertEqual( + harness.records()[-1]["classification"], "launch_error" + ) + + +if __name__ == "__main__": + unittest.main()