ik_llama.cpp/docs/development/on-demand-tensor-reload.md

# On-Demand Tensor Reload

## Overview

This patch introduces **selective tensor hot-swapping** for `ik_llama.cpp` models, now with full support for `graph`/`layer` split mode.
It allows individual tensors (or groups of tensors) to be reloaded from their original on-disk GGUF files **without tearing down the process, the `llama_model`, or the `llama_context`**. Tensors may reside on any backend—GPU, CPU, or split across multiple GPUs—and the reload logic preserves that placement.

This is primarily intended for:

* Iterative experimentation and LoRA-like surgical updates.
* Dynamic MoE (Mixture-of-Experts) expert swapping.
* **Mixed-quantization perplexity benchmarks**, where the bulk of a model lives in one quant (e.g., Q4_X) on GPU while individual experts are hot-swapped one-by-one into a different quant (e.g., IQ1_KT) to measure isolated quality impact.

---

## Motivation

Standard `ik_llama.cpp` workflows require restarting the entire executable to pick up new weights. For large models distributed across multiple GPUs—or models that spill into CPU memory—this incurs significant downtime. This patch solves that by:

1. **Tracking provenance**: At load time, every tensor is mapped back to its source GGUF shard, byte offset, and modification time.
2. **Detecting changes**: At runtime, it cheaply `stat()`s the source files to see if a tensor’s backing data has changed.
3. **Surgical replacement**: Only the changed tensors are re-mapped/re-allocated. The rest of the model stays resident in GPU/CPU memory.
4. **Graph safety**: Cached CUDA graphs are invalidated and the context’s cached compute graphs (`ctx->prev` / `ctx->prev_mtp`) are reset so that the next evaluation rebuilds the graph with the new buffer pointers, sizes, or types.

---

## High-Level Architecture

The patch adds a `reload_info` registry to `llama_model` (defined in `src/llama-reload-info.h`). The lifecycle has five phases:

### 1. Registration Phase (`llama_model_load`)
During model loading, every weight that is successfully mapped gets an entry in `model.reload->tensor_reload_sources` **only when the environment variable `LLAMA_HOTSWAP_ENABLED` is set**:

```cpp
struct tensor_reload_source {
    std::string   path;          // Absolute path to the GGUF shard
    size_t        data_offset;   // Byte offset of the tensor data in the file
    size_t        nbytes;        // Current byte size
    int64_t       last_mtime;    // Last modification time (seconds)
    int64_t       last_mtime_ns; // Nanosecond precision on Linux

    // Snapshots of the *original* loaded state so we can reattach later
    ggml_backend_buffer_t original_buffer;
    void                * original_data;
    ggml_type             original_type;
    int64_t               original_ne[GGML_MAX_DIMS];
    size_t                original_nb[GGML_MAX_DIMS];
    ggml_split_tensor_t * original_extra;
    std::vector<split_info> original_splits;
    std::vector<std::string> sibling_names; // MoE siblings
    reload_state          state;
};
```

### 2. Snapshot Phase (`snapshot_all_reload_tensors`)
The first time a reload is requested, an **eager snapshot** is taken of every registered tensor and its MoE siblings. This captures the original buffer handles, split descriptors, and strides. This snapshot is essential for:

* **Reattachment**: If a tensor was detached to a private buffer because it grew, but later shrinks back to its original size/type, it can be reattached to the original shared buffer, avoiding memory fragmentation.
* **MoE consistency**: MoE layers often have three sibling tensors (`ffn_down_exps`, `ffn_up_exps`, `ffn_gate_exps`) that must share the same split topology across GPUs.

### 3. Detection Phase (`reload_changed_tensors`)
When the user (or the server health-check loop) calls `llama_reload_changed_tensors()`:

1. It iterates over the registry and `stat()`s each source file.
2. If `mtime` (or `mtime_ns`) differs, it re-parses the GGUF header (`gguf_find_tensor_meta`) to get the new `offset`, `nbytes`, `ggml_type`, and on-disk shape (`ne`).
3. **Shape verification**: If the on-disk dimensions differ from the model tensor (`file_ne[i] != tensor->ne[i]`), the tensor is skipped entirely; the reload logic refuses to change logical shapes.
4. It builds a **sorted job list**: tensors that are **returning to their original snapshot** are processed first. This maximizes the chance of freeing private buffers before allocating new ones, reducing memory pressure.

### 4. Reload Phase (`reload_tensor`)
For each changed tensor, the patch performs a careful in-place update.

#### 0. Shape Verification
Before any metadata or buffer changes, the code verifies that the on-disk `ne[0..3]` exactly match the current model tensor. If any dimension differs, the reload is aborted with a log message and the tensor is left untouched.

#### A. Returning Check
The first decision is whether the tensor's new on-disk type matches its **original** snapshot type (`curr_type == src.original_type`).

* **Returning to original**: The tensor is reattached to its original shared buffer and original split descriptors. Any private buffer allocated during a previous reload is freed (only if the tensor's state is `DETACHED` or `FALLBACK_CPU`). State becomes `ON_ORIGINAL`.
* **Changed**: Proceed to metadata update and buffer reallocation.

#### B. Metadata Update & Block-Size Alignment
If the tensor’s `ggml_type` changed (e.g., Q4_X → IQ1_KT), the main tensor descriptor and all its split descriptors are updated with new `type` and `nb` values. The logical shape (`ne`) is guaranteed unchanged by the preceding shape verification. However, for fused/multi-GPU splits the per-device boundaries must be recalculated.

**Critical constraint for fused/multi-GPU splits:**  
Different quants use different block sizes:
* **Q4_X / Q4_0**: block size **32**
* **IQ1_KT**: block size **256**

When a tensor changes between these types, `apply_tensor_type_change()` re-rounds every GPU slice’s `ne[0]` to the nearest multiple of the new block size. If this redistribution is not propagated to all siblings in the same MoE layer, the CUDA split backend dispatches rows to the wrong devices and **matmul fails**.

#### C. Buffer Lifecycle
The patch tracks each tensor with a `reload_state` enum (`UNINITIALIZED`, `ON_ORIGINAL`, `DETACHED`, `FALLBACK_CPU`). Buffers are only freed if the state is not `ON_ORIGINAL`, ensuring shared original buffers are never corrupted.

| Scenario | Action |
|----------|--------|
| Returning to original snapshot | **Reattach** to `original_buffer`, restore original splits, free old private buffer if any. |
| Changed type/size while previously on original | **Detach** from the shared buffer to a newly allocated private buffer so the shared region isn’t corrupted for other tensors. |
| Changed type/size while already detached | Free old private buffer, allocate new one. |
| Allocation fails on target backend | **CPU fallback**: allocate on `ggml_backend_cpu_buffer_type()` and clear split metadata. State becomes `FALLBACK_CPU`. |

#### D. Split Tensor (Multi-GPU) Handling
For split tensors, the patch:
- Recomputes per-device bounds using the new block-size alignment.
- Reallocates per-device split buffers if necessary.
- **Resyncs MoE siblings**: If `ffn_down_exps` changes its split topology, `ffn_up_exps` and `ffn_gate_exps` in the same layer are forced to adopt identical per-device `ne[0]` distributions and strides. This is required by the CUDA split-backend contract.

#### E. Data Copy
Finally, the tensor bytes are read from the updated file and copied into the (possibly new) backend buffer via `ggml_backend_tensor_set`.

---

## Hybrid CPU/GPU Inference

When running with `--split-mode layer --fit --gpu-layers 99` (or any configuration where the model does not fully fit in VRAM), some tensors naturally land in CPU memory. The hot-swap system fully supports this:

* **CPU tensors are reloadable**: The reload logic reads the new data from disk and copies it into the CPU backend buffer exactly as it would for CUDA buffers.
* **Fallback allocator**: If a GPU buffer allocation fails during a reload (e.g., because an IQ1_KT expert is larger than the original Q4_X expert), the system automatically falls back to a CPU buffer for that tensor.

This allows you to keep, for example, 90 % of an MoE model on 13 GPUs while a few large expert tensors cycle through CPU RAM, or to benchmark quants that vary in size per-expert without worrying about exact VRAM fitting.

---

## API & Environment Variables

### Public C API
```cpp
// include/llama.h
LLAMA_API bool llama_reload_changed_tensors(struct llama_context * ctx);
```

Returns `true` if at least one tensor was reloaded. When this happens, the function also resets the context’s cached compute graphs (`ctx->prev` and `ctx->prev_mtp`) so that the next evaluation performs a full graph rebuild with the new tensor pointers.

### Environment Variables

| Variable | Purpose |
|----------|---------|
| `LLAMA_HOTSWAP_ENABLED` | Enables the hot-swap loop in `perplexity` and the health-check hook in `server`. |
| `LLAMA_PERPLEXITY_PRE_RELOAD_SCRIPT` | Path to an executable script run between perplexity iterations (e.g., to regenerate/re-quantize a tensor file). |

---

## Integration Points

### `examples/perplexity/perplexity.cpp`
When `LLAMA_HOTSWAP_ENABLED` is set, the tool runs in a loop:

1. Perform an initial `llama_reload_changed_tensors()` to apply any pending changes before the first evaluation.
2. Compute perplexity (or Hellaswag, etc.).
3. Print timings and write logs.
4. Execute the optional pre-reload script.
5. Call `llama_reload_changed_tensors(ctx)`. If no tensors changed, exit; otherwise repeat from step 2.

### `examples/server/server.cpp`
On every health-check (`/health`) request, if `LLAMA_HOTSWAP_ENABLED` is set, the server calls `llama_reload_changed_tensors()`. This provides a convenient, external trigger: simply `touch` or overwrite a tensor’s source GGUF file and poll `/health` to apply the change.

---

## MoE Sibling Resync

MoE weights are often stored as three separate tensors that must be split identically across GPUs. The patch automatically detects these families by suffix:

- `.ffn_down_exps.weight`
- `.ffn_up_exps.weight`
- `.ffn_gate_exps.weight`

When one member of the family is reloaded and its per-device split dimensions change—especially when crossing quant types with different block sizes (Q4_X=32 vs IQ1_KT=256)—`resync_moe_sibling_splits()` is invoked. The logic follows these steps:

1. **Fast path**: If the reference tensor is returning to its original snapshot, the siblings are also reattached to their original snapshots via `reattach_split_tensor_to_shared()`—no data movement is required.
2. **Phase A – Detach**: Siblings are detached from shared buffers (freeing only non-original buffers) and new main handles are allocated. Split tensors receive a dummy `data` pointer because the split backend uses `extra->splits`.
3. **Phase B – Propagate dimensions**: The reference tensor’s per-device `ne[0]` distribution is copied to the siblings, and strides (`nb[]`) are recomputed using a temporary `ggml_context`. This step is mandatory because the valid split boundaries depend on the quantization block size.
4. **Phase C – Allocate GPU splits**: New per-device GPU buffers are allocated for each sibling split.
5. **Phase D – CPU fallback (if needed)**: If any GPU allocation fails, the **entire** sibling group is moved to CPU buffers to maintain consistency.
6. **Phase E – Write back**: The original sibling data (which has not changed, only the layout) is written back into the new buffers via `ggml_backend_tensor_set`.

---

## Buffer Lifecycle Details

### Reattachment to Shared Buffers
If a tensor was originally loaded in a large shared GGUF buffer alongside other tensors, and it was previously detached because it grew, the patch attempts to **reattach** it when it returns to its original size and type. This is done by restoring:

- `tensor->buffer = original_buffer`
- `tensor->data   = original_data`
- `tensor->extra  = original_extra` (restoring all split descriptors)

This prevents unbounded memory growth during iterative experiments where tensors oscillate between two states.

### State Machine
Because `ggml` does not provide native reference counting on buffers, the patch uses a per-tensor state machine to avoid corrupting shared allocations:

* `ON_ORIGINAL`: The tensor still lives in its initial shared buffer. This buffer is **never** freed during reload.
* `DETACHED`: The tensor was moved to a privately allocated buffer. This buffer **is** freed before the next reload.
* `FALLBACK_CPU`: The tensor was moved to CPU memory after a GPU allocation failure.

Only buffers belonging to tensors in the `DETACHED` or `FALLBACK_CPU` states are released, ensuring that shared original buffers remain valid for all other tensors that still reference them.

---

## Limitations & Safety Notes

1. **File path stability**: The source file must remain at the same path. Renaming or removing shards will cause `stat()` or `open()` to fail.
2. **No locking**: There is no file-locking protocol. The user must ensure the GGUF file is not being written to while `ik_llama.cpp` is reading it.
3. **Graph rebuild cost**: While cheaper than a full process restart, rebuilding the CUDA graph (or CPU graph) incurs a one-time latency spike after a reload.
4. **Platform specifics**: Nanosecond mtime checks use `st_mtim.tv_nsec` and are guarded by `#ifdef __linux__`.
5. **Thread safety**: `llama_reload_changed_tensors` is **not** thread-safe with active inference. Ensure the context is idle before calling (the perplexity example naturally guarantees this; the server example only invokes it during the synchronous `/health` handler).

---

## Usage Example: Per-Expert Quantization Sweep (Q4_X ↔ IQ1_KT)

This example benchmarks a massive MoE model where the base weights are **Q4_X**. The tool iteratively replaces individual `ffn_down_exps.weight` tensors with **IQ1_KT** equivalents to measure the isolated perplexity impact of each expert's quantization level.

A sanity check is embedded in the source directory: one of the "IQ1_KT" shard files is actually the original **Q4_X** tensor. When the rotation reaches that slot, the reloaded tensor is byte-for-byte identical to the baseline, so the PPL must match exactly—confirming that the hot-swap machinery introduces no loss.

### 1. Helper script (`tensor-swap.sh`)
Place the rotation script in your model directory (e.g., `/opt/THIREUS/Kimi-K2.6/Q4_X/`). It maintains `.bak` files so that each iteration restores the previous tensor before installing the next candidate.

```bash
#!/bin/bash
set -euo pipefail

TARGET_GLOB="*Q4_X*gguf"
SOURCE_DIR="../smol-IQ1-KT-mist.bin"
TENSOR_NAME_PATTERN="blk\.[0-9]+\.ffn_down_exps\.weight"

# ... (see full script in patch) ...
```

The script scans for target files matching `*Q4_X*gguf` containing `blk.[N].ffn_down_exps.weight`, then pulls replacements from `../smol-IQ1-KT-mist.bin/` by matching the `SPECIAL_TENSOR-NNNN-of-XXXX.gguf` shard number.

### 2. Launch perplexity with hot-swap enabled

```bash
ulimit -n 9999
ulimit -l unlimited

export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7,8,9,10,11,12"
export LLAMA_HOTSWAP_ENABLED=1
export LLAMA_PERPLEXITY_PRE_RELOAD_SCRIPT=./tensor-swap.sh
export LLAMA_DEBUG=1

# --offload-policy -1,off \

GGML_CUDA_NO_PINNED=1 \
/opt/ik_llama.cpp/ik_llama.cpp/build/bin/llama-perplexity \
    --chunks 8 \
    -f /opt/ik_llama.cpp/wiki.test.raw \
    --model /opt/THIREUS/Kimi-K2.6/Q4_X/Kimi-K2.6-THIREUS-Q4_X-SPECIAL_TENSOR-00001-of-01097.gguf \
    --alias THIREUS/Kimi-K2.6-Q4_X.bin \
    -b 512 -ub 512 \
    --ctx-size 512 \
    --fit \
    --fit-margin 4200 \
    --gpu-fit-margin 0,4400,12,4400 \
    --temp 0.0 --top-k 0 --top-p 1.0 \
    -ctk f16 \
    -ctv q8_0 \
    -amb 128 \
    -mea 128 \
    -wgt 1 \
    --mlock \
    --split-mode layer \
    --graph-reduce-type f16 \
    --threads $(grep ^cpu\\scores /proc/cpuinfo | uniq | awk '{print $4}' | xargs -I{} echo "{}-0" | bc) \
    -sas \
    --gpu-layers 99 \
    --no-offload-only-active-experts \
    --host 0.0.0.0 \
    --port 8080 \
    --log-enable \
    --logdir /var/log/ \
    --jinja \
    --special \
    --prompt-cache "$HOME/.cache/ik_llama.cpp/prompt-cache.bin" --prompt-cache-all \
    --slot-save-path "$HOME/.cache/ik_llama.cpp/slot.bin" \
    --lookup-cache-dynamic "$HOME/.cache/ik_llama.cpp/slot.bin" \
    --keep -1 \
    --slot-prompt-similarity 0.35 \
    --metrics \
    -cuda fusion=1
```

### 3. What happens

1. The model loads with **Q4_X** weights distributed across 13 GPUs using layer splitting.
2. The first pass computes the baseline perplexity over 8 chunks.
3. `tensor-swap.sh` runs between iterations:
   * Restores the previously swapped tensor from `.bak` to its original Q4_X state.
   * Copies the next IQ1_KT expert shard into place.
4. `llama_reload_changed_tensors()` detects the `mtime` changes, re-parses the GGUF headers, and reloads the affected `ffn_down_exps.weight` tensor(s).
   * The restored tensor **returns to its original Q4_X snapshot** and reattaches to its shared buffer.
   * The newly swapped tensor is loaded into a private buffer with the new IQ1_KT data.
   * Because Q4_X and IQ1_KT have different block sizes (32 vs 256), the split backend redistributes per-device boundaries and resyncs the MoE siblings (`ffn_up_exps` and `ffn_gate_exps`) to the same layout.
5. The CUDA graphs are invalidated and the next perplexity iteration begins.
6. When the rotation hits the sanity-check slot (where the source file is actually the original Q4_X tensor), the perplexity returns to the exact baseline value, confirming the reload is lossless.

### 4. Expected behavior

```text
snapshot_all_reload_tensors: eager snapshot of all reload tensors + siblings
perplexity: calculating perplexity over 8 chunks, n_ctx=512, batch_size=512, n_seq=1
[1]1.0622,[2]1.2068,[3]1.2327,[4]1.1873,[5]1.1487,[6]1.1283,[7]1.1214,[8]1.1109,
Final estimate: PPL = 1.1109

main: executing pre-reload script: ./tensor-swap.sh
main: [pre-reload] Swapped index 0 (tensor #00918)
reloaded tensor 'blk.1.ffn_down_exps.weight'

perplexity: calculating perplexity over 8 chunks ...
Final estimate: PPL = 1.1105

main: executing pre-reload script: ./tensor-swap.sh
main: [pre-reload] Restored index 0. Advancing to index 1.
main: [pre-reload] Swapped index 1 (tensor #00921)
reloaded tensor 'blk.1.ffn_down_exps.weight'
reloaded tensor 'blk.2.ffn_down_exps.weight'

perplexity: calculating perplexity over 8 chunks ...
Final estimate: PPL = 1.1080
```

Notice that when the script restores a tensor to its original Q4_X shard, the reload reattaches it to the shared buffer with zero copy. When the sanity-check slot is reached, the PPL returns to the exact baseline, proving the mechanism is sound.

---

## Summary of Changed Files

| File | Change |
|------|--------|
| `examples/perplexity/perplexity.cpp` | Hot-swap loop + pre-reload script execution. |
| `examples/server/server.cpp` | Trigger reload on `/health` when env var is set. |
| `ggml/include/ggml-cuda.h` | Add `ggml_backend_cuda_invalidate_graphs()`. |
| `ggml/include/ggml.h` | Conditional `GGML_MAX_SRC` override. |
| `ggml/src/CMakeLists.txt` | Propagate `GGML_MAX_SRC` compile definition. |
| `ggml/src/ggml-cuda.cu` | Implement graph invalidation; debug prints for split tensors. |
| `ggml/src/ggml.c` | Debug print in `ggml_mul_mat_id` for shape mismatches. |
| `include/llama.h` | Declare `llama_reload_changed_tensors()`. |
| `src/llama-mmap.cpp/h` | Expose `llama_file::get_path()` so reload registry knows the source file path. |
| `src/llama-model.h` | Add `std::unique_ptr<reload_info> reload` to `llama_model`. |
| `src/llama-reload-info.h` | **New.** Defines `tensor_reload_source`, `reload_state`, and `reload_info` registry. |
| `src/llama-reload.cpp` | **New.** Core implementation: GGUF header parser, snapshot, reload, MoE resync, buffer management, CPU fallback, shape verification. |
| `src/llama.cpp` | Wire reload registry into `llama_model_load`; reset cached compute graphs (`ctx->prev` / `ctx->prev_mtp`) on reload; export C API. |
| `src/CMakeLists.txt` | Propagate `GGML_MAX_SRC` compile definition. |