UPSTREAM PR #18747: ggml, llama : add CPU paged attention for memory-efficient KV cache

BRANCH_RULES.md

@@ -0,0 +1,93 @@
+# Branch Management Rules for llama.cpp Development
+
+## Branch Hierarchy
+
+```
+origin/master (upstream)
+    │
+    └── production-consolidated (PROTECTED - production/benchmarks)
+            │
+            └── research/* (experimental work)
+```
+
+## Branch Purposes
+
+| Branch | Purpose | Build From | Modify? |
+|--------|---------|------------|---------|
+| `production-consolidated` | Production benchmarks, stable features | This branch | NO - cherry-pick only |
+| `research/*` | Experimental features, testing | `production-consolidated` | YES |
+| `feature/*` | Isolated feature development | `origin/master` | YES |
+
+## Rules
+
+### Rule 1: Never Modify production-consolidated Directly
+- All changes go through cherry-pick from tested feature branches
+- Create `research/your-experiment` for testing
+- Only cherry-pick to production-consolidated after validation
+
+### Rule 2: Always Rebuild After Branch Switch
+```bash
+# ALWAYS do this after switching branches:
+rm -rf build
+cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_NATIVE=ON -DGGML_OPENMP=ON -DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS -DLLAMA_CURL=OFF
+cmake --build build -j$(nproc)
+```
+
+### Rule 3: Verify Build Before Benchmarking
+```bash
+# Check for undefined symbols (should return nothing):
+nm build/bin/llama-cli | grep "U.*llama_"
+
+# Verify required flags exist:
+./build/bin/llama-cli --help | grep moe-n-expert
+```
+
+### Rule 4: Research Branch Workflow
+```bash
+# Start new research from production-consolidated:
+git checkout production-consolidated
+git checkout -b research/my-experiment
+
+# Work on experiment...
+# When done and validated, cherry-pick to production-consolidated:
+git checkout production-consolidated
+git cherry-pick <commit-hash>
+git push fork production-consolidated
+```
+
+### Rule 5: Tag Working States
+```bash
+# After successful benchmark session:
+git tag benchmark-$(date +%Y-%m-%d) -m "Working benchmark state"
+git push fork --tags
+```
+
+### Rule 6: Document Feature Dependencies
+Before cherry-picking, verify the commit doesn't depend on uncommitted infrastructure.
+Check with:
+```bash
+# See what files the commit touches:
+git show <commit> --stat
+
+# Check if any new types/functions are referenced but not defined:
+git show <commit> | grep -E "^[\+].*\(" | head -20
+```
+
+## Current Branch Status (2026-01-10)
+
+| Branch | Status | Features |
+|--------|--------|----------|
+| `production-consolidated` | STABLE | MoE hard mask, layer skip, lookahead fixes, parallel repack |
+| `feature/moe-hard-mask` | MERGED | → production-consolidated |
+| `feature/eagle-penultimate-layer` | DO NOT USE | Has undefined symbol issues |
+| `mtp-branch` | INCOMPATIBLE | Requires infrastructure not in origin/master |
+
+## Recovering from Build Issues
+
+If you encounter `undefined symbol` errors or SIGSEGV:
+
+1. **Check current branch**: `git branch --show-current`
+2. **Check build date**: `ls -la build/bin/llama-cli`
+3. **Clean rebuild**: `rm -rf build && cmake -B build ... && cmake --build build -j$(nproc)`
+4. **Verify symbols**: `nm build/bin/llama-cli | grep "U.*llama"`
+5. **If still broken**: `git checkout production-consolidated` and rebuild

common/arg.cpp

@@ -1244,6 +1244,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.n_keep = value;
         }
     ));
+    add_opt(common_arg(
+        {"--n-layer-exit"}, "N",
+        string_format("exit after N layers (default: %d, 0 = compute all layers)\n"
+            "for layer skip / early exit speculation (CAS-Spec, CLaSp)", params.n_layer_exit),
+        [](common_params & params, int value) {
+            params.n_layer_exit = value;
+        }
+    ).set_env("LLAMA_ARG_N_LAYER_EXIT"));
     add_opt(common_arg(
         {"--swa-full"},
         string_format("use full-size SWA cache (default: %s)\n"
@@ -1305,6 +1313,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
                                    string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str()));
                            }
                        }).set_env("LLAMA_ARG_FLASH_ATTN"));
+    add_opt(common_arg(
+        {"--paged-attn"}, "N",
+        string_format("enable paged attention with block size N tokens (default: %d, 0 = disabled)", params.paged_attn_block_size),
+        [](common_params & params, int value) {
+            params.paged_attn_block_size = value;
+        }
+    ).set_env("LLAMA_PAGED_ATTN"));
+    add_opt(common_arg(
+        {"--paged-attn-max-blocks"}, "N",
+        string_format("max blocks for paged attention memory reduction (default: %d, 0 = unlimited)", params.paged_attn_max_blocks),
+        [](common_params & params, int value) {
+            params.paged_attn_max_blocks = value;
+        }
+    ).set_env("LLAMA_PAGED_ATTN_MAX_BLOCKS"));
     add_opt(common_arg(
         {"-p", "--prompt"}, "PROMPT",
         "prompt to start generation with; for system message, use -sys",
@@ -1901,6 +1923,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.yarn_beta_fast = std::stof(value);
         }
     ).set_env("LLAMA_ARG_YARN_BETA_FAST"));
+    add_opt(common_arg(
+        {"--moe-n-expert"}, "N",
+        string_format("MoE: override number of active experts (default: %d = model default)\n"
+                      "for MoE self-draft speculation, use 1 for draft context", params.moe_n_expert_override),
+        [](common_params & params, int value) {
+            params.moe_n_expert_override = value;
+        }
+    ).set_env("LLAMA_ARG_MOE_N_EXPERT"));
     add_opt(common_arg(
         {"-gan", "--grp-attn-n"}, "N",
         string_format("group-attention factor (default: %d)", params.grp_attn_n),

common/common.cpp

@@ -17,6 +17,7 @@
 #include <codecvt>
 #include <chrono>
 #include <cstdarg>
+#include <cstdlib>
 #include <cstring>
 #include <ctime>
 #include <filesystem>
@@ -1411,6 +1412,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.yarn_beta_fast    = params.yarn_beta_fast;
     cparams.yarn_beta_slow    = params.yarn_beta_slow;
     cparams.yarn_orig_ctx     = params.yarn_orig_ctx;
+    cparams.moe_n_expert_override = params.moe_n_expert_override;
+    cparams.n_layer_exit      = params.n_layer_exit;
     cparams.pooling_type      = params.pooling_type;
     cparams.attention_type    = params.attention_type;
     cparams.flash_attn_type   = params.flash_attn_type;
@@ -1425,6 +1428,15 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;
 
+    // Set paged attention environment variables if CLI flags are used
+    // This allows CLI flags to override any existing env vars
+    if (params.paged_attn_block_size > 0) {
+        setenv("LLAMA_PAGED_ATTN", std::to_string(params.paged_attn_block_size).c_str(), 1);
+    }
+    if (params.paged_attn_max_blocks > 0) {
+        setenv("LLAMA_PAGED_ATTN_MAX_BLOCKS", std::to_string(params.paged_attn_max_blocks).c_str(), 1);
+    }
+
     return cparams;
 }
 

common/common.h

@@ -328,6 +328,8 @@ struct common_params {
     float   yarn_beta_fast        = -1.0f; // YaRN low correction dim
     float   yarn_beta_slow        = -1.0f; // YaRN high correction dim
     int32_t yarn_orig_ctx         =     0; // YaRN original context length
+    int32_t moe_n_expert_override =     0; // MoE self-draft: override n_expert_used (0 = use model default)
+    int32_t n_layer_exit          =     0; // exit after this many layers, 0 = all (for layer skip speculation)
 
     // offload params
     std::vector<ggml_backend_dev_t> devices; // devices to use for offloading
@@ -356,6 +358,10 @@ struct common_params {
     enum llama_attention_type    attention_type    = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings
     enum llama_flash_attn_type   flash_attn_type   = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention
 
+    // Paged attention parameters
+    uint32_t paged_attn_block_size = 0;   // block size in tokens (0 = disabled)
+    uint32_t paged_attn_max_blocks = 0;   // max blocks for memory reduction (0 = unlimited)
+
     struct common_params_sampling    sampling;
     struct common_params_speculative speculative;
     struct common_params_vocoder     vocoder;

examples/lookahead/lookahead.cpp

@@ -50,6 +50,12 @@ int main(int argc, char ** argv) {
     const int N = 5;  // n-gram size
     const int G = 15; // max verification n-grams
 
+    // lookahead requires W + G + 1 sequences for parallel Jacobi decoding
+    params.n_parallel = W + G + 1;
+
+    // unified KV cache is required for coupled sequences in batch splitting
+    params.kv_unified = true;
+
     // init llama.cpp
     llama_backend_init();
     llama_numa_init(params.numa);
@@ -115,7 +121,7 @@ int main(int argc, char ** argv) {
     // seq_id == 0           : the current input token
     // seq_id [1, W]         : tokens from the past N - 1 Jacobi iterations
     // seq_id [W + 1, W + G] : verification n-grams
-    llama_batch batch = llama_batch_init(params.n_ctx, 0, W + G + 1);
+    llama_batch batch = llama_batch_init(llama_n_ctx(ctx), 0, W + G + 1);
 
     // target model sampling context
     struct common_sampler * smpl = common_sampler_init(model, params.sampling);

examples/lookup/lookup.cpp

@@ -106,7 +106,7 @@ int main(int argc, char ** argv){
 
     std::vector<llama_token> draft;
 
-    llama_batch batch_tgt = llama_batch_init(params.n_ctx, 0, 1);
+    llama_batch batch_tgt = llama_batch_init(llama_n_ctx(ctx), 0, 1);
 
     const auto t_dec_start = ggml_time_us();
 

ggml/include/ggml.h

@@ -545,6 +545,7 @@ extern "C" {
         GGML_OP_FILL,
 
         GGML_OP_FLASH_ATTN_EXT,
+        GGML_OP_FLASH_ATTN_EXT_PAGED,  // Paged attention with block table indirection
         GGML_OP_FLASH_ATTN_BACK,
         GGML_OP_SSM_CONV,
         GGML_OP_SSM_SCAN,
@@ -2342,6 +2343,21 @@ extern "C" {
             struct ggml_tensor * a,
             struct ggml_tensor * sinks);
 
+    // Paged flash attention with indirect KV access via block table
+    // block_table: I32 tensor [max_blocks, n_seqs] mapping logical → physical blocks
+    // op_params[4] encodes block_size
+    GGML_API struct ggml_tensor * ggml_flash_attn_ext_paged(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * block_table,  // I32 [max_blocks, n_seqs] logical→physical mapping
+            float                 scale,
+            float                 max_bias,
+            float                 logit_softcap,
+            int32_t               block_size);
+
     // TODO: needs to be adapted to ggml_flash_attn_ext
     GGML_API struct ggml_tensor * ggml_flash_attn_back(
            struct ggml_context * ctx,

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1963,6 +1963,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_flash_attn_ext(params, tensor);
             } break;
+        case GGML_OP_FLASH_ATTN_EXT_PAGED:
+            {
+                ggml_compute_forward_flash_attn_ext_paged(params, tensor);
+            } break;
         case GGML_OP_FLASH_ATTN_BACK:
             {
                 int32_t t = ggml_get_op_params_i32(tensor, 0);
@@ -2333,6 +2337,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_ARGSORT:
         case GGML_OP_TOP_K:
         case GGML_OP_FLASH_ATTN_EXT:
+        case GGML_OP_FLASH_ATTN_EXT_PAGED:
         case GGML_OP_FLASH_ATTN_BACK:
         case GGML_OP_SSM_CONV:
         case GGML_OP_SSM_SCAN:
@@ -2865,6 +2870,7 @@ struct ggml_cplan ggml_graph_plan(
                         cur += sizeof(int32_t)*node->src[0]->ne[0]*n_tasks;
                     } break;
                 case GGML_OP_FLASH_ATTN_EXT:
+                case GGML_OP_FLASH_ATTN_EXT_PAGED:
                     {
                         const int64_t ne10 = node->src[1]->ne[0]; // DK
                         const int64_t ne20 = node->src[2]->ne[0]; // DV