diff --git a/.devops/cuda.Dockerfile b/.devops/cuda.Dockerfile index c9a498d538b3..3dc51d6abe1b 100644 --- a/.devops/cuda.Dockerfile +++ b/.devops/cuda.Dockerfile @@ -13,18 +13,6 @@ ARG APP_REVISION=N/A ARG NODE_VERSION=24 -FROM docker.io/node:$NODE_VERSION AS web - -ARG APP_VERSION - -WORKDIR /app/tools/ui - -COPY tools/ui/package.json tools/ui/package-lock.json ./ -RUN npm ci - -COPY tools/ui/ ./ -RUN LLAMA_BUILD_NUMBER="$APP_VERSION" npm run build - FROM ${BASE_CUDA_DEV_CONTAINER} AS build ARG GCC_VERSION @@ -40,12 +28,10 @@ WORKDIR /app COPY . . -COPY --from=web /app/tools/ui/dist tools/ui/dist - RUN if [ "${CUDA_DOCKER_ARCH}" != "default" ]; then \ export CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=${CUDA_DOCKER_ARCH}"; \ fi && \ - cmake -B build -DGGML_NATIVE=OFF -DGGML_CUDA=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DLLAMA_BUILD_TESTS=OFF ${CMAKE_ARGS} -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined . && \ + cmake -B build -DGGML_NATIVE=OFF -DGGML_CUDA=ON -DGGML_CUDA_NO_VMM=ON -DGGML_CUDA_FA_ALL_QUANTS=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_UI=OFF ${CMAKE_ARGS} -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined . && \ cmake --build build --config Release -j$(nproc) RUN mkdir -p /app/lib && \ diff --git a/.gitignore b/.gitignore index 9b589615a402..9ad4c9796fc7 100644 --- a/.gitignore +++ b/.gitignore @@ -130,6 +130,7 @@ flake.lock /hellaswag_val_full.txt /winogrande-debiased-eval.csv /wikitext-2-raw/ +/benchmarks/results/ # Test models for lora adapters diff --git a/README.md b/README.md index aae3bcd35ad9..bb5c40d7ab29 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,93 @@ +# Adaptive KV Streaming for llama.cpp + +This branch adds an experimental, block-granular KV cache streaming path to the CUDA `llama-server`. It is intended for running long contexts when model weights leave too little VRAM for the complete KV cache. + +With `--kv-stream-stage-mib N`, the authoritative KV tensors are stored in pinned host memory while a bounded CUDA pool is shared by resident KV pages and a transfer ring. The runtime adapts that split as the context grows: it keeps as many pages resident as the budget allows, reclaims resident space for staging when more streaming is required, and prefetches later layers while the current layer computes. This avoids relying on uncontrolled Unified Memory page thrashing and preserves exact attention over the full context. + +Detailed project story, design, implementation, and benchmark results are in +[Running Qwen 27B on 16G VRAM with Full Context Length: Building Adaptive KV Cache Streaming for llama.cpp](https://medium.com/@raymond860909/running-qwen-27b-on-16g-vram-with-full-context-length-building-adaptive-kv-cache-streaming-for-bf1e819116e9). + +> [!WARNING] +> This is research code optimized and production-validated primarily for an RTX 5070 Ti with 16 GB VRAM, `unsloth/Qwen3.8-27B-GGUF` `UD-Q3_K_XL`, a 262144-token context, Flash Attention, a Q8_0 K cache, a Q4_0 V cache, and one server slot. +> CUDA correctness tests cover every KV type currently accepted by the CLI, including native and F16-conversion fallback paths. Production performance for other models, KV combinations, parallel slots, and non-CUDA backends is not yet broadly characterized. + +## Build the modified server + +Install a C++ compiler, CMake, and the CUDA toolkit, then run this command from the repository root: + +```bash +cmake -S . -B build -DGGML_CUDA=ON -DGGML_CUDA_FA_ALL_QUANTS=ON -DCMAKE_BUILD_TYPE=Release && cmake --build build --config Release --target llama-server -j +``` + +The executable is created at `build/bin/llama-server`. + +Example using the tested cache configuration: + +```bash +./build/bin/llama-server \ + --model /path/to/model.gguf \ + --ctx-size 262144 \ + -fa on \ + -ctk q8_0 \ + -ctv q4_0 \ + -ngl all \ + -b 512 \ + -ub 512 \ + -np 1 \ + --kv-stream-stage-mib 2304 +``` + +The best value for `--kv-stream-stage-mib` depends on the model, context capacity, GPU, and other VRAM consumers. Start conservatively and increase it while checking startup and peak VRAM use. + +### Batch and micro-batch sizes + +`-b` sets the logical prompt batch size and `-ub` sets the largest physical batch submitted to one graph. This branch no longer requires `256/256`; `-ub` may be any positive value no larger than `-b`. + +The Qwen3.8 MMA prefill path processes the full physical batch and allocates only the partial workspace that kernel actually emits. Generic vector and F16-conversion fallback paths use a bounded 256-query workspace: each staged KV span is consumed by all query tiles before its ring slot is released, so wider micro-batches do not multiply KV host-to-device transfers. + +The Q8_0/Q4_0 Qwen configuration has been exercised with `b/ub` values `256/256`, `512/512`, `768/512`, and `1024/1024`, including non-divisible final micro-batches. A 122880-token production-shaped run at `512/512` completed with adaptive streaming active. Wider values can require more graph and accumulator memory, so validate them on the target GPU. + +### Optional Unified Memory for model weights + +Adaptive KV streaming works with or without Unified Memory. Leave `GGML_CUDA_ENABLE_UNIFIED_MEMORY` unset for ordinary CUDA device allocations. To make GPU-offloaded model buffers CUDA managed allocations, launch the same server with the environment variable enabled: + +```bash +GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 \ +./build/bin/llama-server \ + --model /path/to/model.gguf \ + --ctx-size 262144 \ + -fa on \ + -ctk q8_0 \ + -ctv q4_0 \ + -ngl all \ + -b 512 \ + -ub 512 \ + -np 1 \ + --kv-stream-stage-mib 2304 +``` + +With this flag, CUDA-backed model buffers, including GPU-offloaded weights, are allocated with `cudaMallocManaged` and their pages can migrate between VRAM and host memory. The adaptive resident-page and transfer-ring pool is intentionally different: it is still allocated with `cudaMalloc`, so that fixed-size pool remains physically allocated in VRAM instead of becoming managed memory. UVM is therefore optional for this branch and does not change the KV streaming pool into pageable storage. + +## Recreate the benchmark graph + +The benchmark driver automatically selects the largest practical adaptive KV pool for each configured context capacity, sweeps from 8K through the requested maximum, and generates the CSV, PNG, and SVG results: + +```bash +python3 -m pip install matplotlib + +python3 benchmarks/benchmark_kv_stream.py \ + --model /path/to/model.gguf \ + --max-context 192K \ + --batch-size 512 \ + --ubatch-size 512 +``` + +The only required arguments are the model GGUF and maximum context. See [benchmarks/README.md](benchmarks/README.md) for the pool-probing algorithm, generated files, optional settings, and resumable output directories. + +--- + +## Upstream llama.cpp README + # llama.cpp ![llama](https://raw.githubusercontent.com/ggml-org/llama.brand/refs/heads/master/cover/llama-cpp/cover-llama-cpp-dark.svg) diff --git a/arm117-artifacts/Aoff-np1/gpu-mem.txt b/arm117-artifacts/Aoff-np1/gpu-mem.txt new file mode 100644 index 000000000000..e724d264810f --- /dev/null +++ b/arm117-artifacts/Aoff-np1/gpu-mem.txt @@ -0,0 +1,2 @@ +0, 14149 MiB, 16311 MiB +1, 9080 MiB, 12288 MiB diff --git a/arm117-artifacts/Aoff-np1/gpu-procs.txt b/arm117-artifacts/Aoff-np1/gpu-procs.txt new file mode 100644 index 000000000000..599b85ef5509 --- /dev/null +++ b/arm117-artifacts/Aoff-np1/gpu-procs.txt @@ -0,0 +1,3 @@ +1311054 /tmp/opencode/arm117/port/build/bin/llama-server -m /mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf --rpc 127.0.0.1:50052 -ts 27,38 -ngl 99 --rope-scaling yarn --rope-scale 5 --yarn-orig-ctx 32768 -fa on -ctk q8_0 -ctv q5_1 -ctkd q8_0 -ctvd q5_1 --no-kv-unified --cache-prompt --cache-reuse 64 --cache-idle-slots --cache-ram 1024 --ubatch-size 512 --cont-batching -np 1 -c 65536 --parallel-ctx-threshold 100000 --spec-type draft-mtp --prio-batch 1 --kv-stream-stage-mib 0 --jinja --host 0.0.0.0 --port 8080 --metrics --slots --log-verbosity 4 +1311032 /tmp/opencode/arm117/port/build/bin/ggml-rpc-server --host 127.0.0.1 --port 50052 -d CUDA1 +1311054 /tmp/opencode/arm117/port/build/bin/llama-server -m /mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf --rpc 127.0.0.1:50052 -ts 27,38 -ngl 99 --rope-scaling yarn --rope-scale 5 --yarn-orig-ctx 32768 -fa on -ctk q8_0 -ctv q5_1 -ctkd q8_0 -ctvd q5_1 --no-kv-unified --cache-prompt --cache-reuse 64 --cache-idle-slots --cache-ram 1024 --ubatch-size 512 --cont-batching -np 1 -c 65536 --parallel-ctx-threshold 100000 --spec-type draft-mtp --prio-batch 1 --kv-stream-stage-mib 0 --jinja --host 0.0.0.0 --port 8080 --metrics --slots --log-verbosity 4 diff --git a/arm117-artifacts/Aoff-np1/health.json b/arm117-artifacts/Aoff-np1/health.json new file mode 100644 index 000000000000..1a36cf5fc27a --- /dev/null +++ b/arm117-artifacts/Aoff-np1/health.json @@ -0,0 +1 @@ +{"status":"ok"} \ No newline at end of file diff --git a/arm117-artifacts/Aoff-np1/log-stream-lines.txt b/arm117-artifacts/Aoff-np1/log-stream-lines.txt new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/arm117-artifacts/Aoff-np1/slots.json b/arm117-artifacts/Aoff-np1/slots.json new file mode 100644 index 000000000000..1fde192dbff9 --- /dev/null +++ b/arm117-artifacts/Aoff-np1/slots.json @@ -0,0 +1 @@ +[{"id":0,"n_ctx":65536,"speculative":true,"is_processing":false,"id_task":33,"n_prompt_tokens":2260,"n_prompt_tokens_processed":0,"n_prompt_tokens_cache":0,"params":{"seed":4294967295,"temperature":0.0,"dynatemp_range":0.0,"dynatemp_exponent":1.0,"top_k":20,"top_p":0.949999988079071,"min_p":0.05000000074505806,"top_n_sigma":-1.0,"xtc_probability":0.0,"xtc_threshold":0.10000000149011612,"typical_p":1.0,"repeat_last_n":64,"repeat_penalty":1.0,"presence_penalty":0.0,"frequency_penalty":0.0,"dry_multiplier":0.0,"dry_base":1.75,"dry_allowed_length":2,"dry_penalty_last_n":64,"mirostat":0,"mirostat_tau":5.0,"mirostat_eta":0.10000000149011612,"adaptive_target":-1.0,"adaptive_decay":0.8999999761581421,"max_tokens":96,"n_predict":96,"n_keep":0,"n_discard":0,"ignore_eos":false,"stream":false,"n_probs":0,"min_keep":0,"chat_format":"peg-native","reasoning_format":"deepseek","reasoning_in_content":false,"generation_prompt":"<|im_start|>assistant\n\n","samplers":["penalties","dry","top_n_sigma","top_k","typ_p","top_p","min_p","xtc","temperature"],"speculative.types":"none,draft-mtp","timings_per_token":false,"post_sampling_probs":false,"backend_sampling":false,"lora":[]},"next_token":[{"has_next_token":false,"has_new_line":false,"n_remain":-1,"n_decoded":0}]}] \ No newline at end of file diff --git a/arm117-artifacts/Aon-np1/gpu-mem.txt b/arm117-artifacts/Aon-np1/gpu-mem.txt new file mode 100644 index 000000000000..d92658c7ca1d --- /dev/null +++ b/arm117-artifacts/Aon-np1/gpu-mem.txt @@ -0,0 +1,2 @@ +0, 15184 MiB, 16311 MiB +1, 9142 MiB, 12288 MiB diff --git a/arm117-artifacts/Aon-np1/gpu-procs.txt b/arm117-artifacts/Aon-np1/gpu-procs.txt new file mode 100644 index 000000000000..1a269272e480 --- /dev/null +++ b/arm117-artifacts/Aon-np1/gpu-procs.txt @@ -0,0 +1,3 @@ +1312993 /tmp/opencode/arm117/port/build/bin/llama-server -m /mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf --rpc 127.0.0.1:50052 -ts 27,38 -ngl 99 --rope-scaling yarn --rope-scale 5 --yarn-orig-ctx 32768 -fa on -ctk q8_0 -ctv q5_1 -ctkd q8_0 -ctvd q5_1 --no-kv-unified --cache-prompt --cache-reuse 64 --cache-idle-slots --cache-ram 1024 --ubatch-size 512 --cont-batching -np 1 -c 65536 --parallel-ctx-threshold 100000 --spec-type draft-mtp --prio-batch 1 --kv-stream-stage-mib 2048 --jinja --host 0.0.0.0 --port 8080 --metrics --slots --log-verbosity 4 +1312976 /tmp/opencode/arm117/port/build/bin/ggml-rpc-server --host 127.0.0.1 --port 50052 -d CUDA1 +1312993 /tmp/opencode/arm117/port/build/bin/llama-server -m /mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf --rpc 127.0.0.1:50052 -ts 27,38 -ngl 99 --rope-scaling yarn --rope-scale 5 --yarn-orig-ctx 32768 -fa on -ctk q8_0 -ctv q5_1 -ctkd q8_0 -ctvd q5_1 --no-kv-unified --cache-prompt --cache-reuse 64 --cache-idle-slots --cache-ram 1024 --ubatch-size 512 --cont-batching -np 1 -c 65536 --parallel-ctx-threshold 100000 --spec-type draft-mtp --prio-batch 1 --kv-stream-stage-mib 2048 --jinja --host 0.0.0.0 --port 8080 --metrics --slots --log-verbosity 4 diff --git a/arm117-artifacts/Aon-np1/health.json b/arm117-artifacts/Aon-np1/health.json new file mode 100644 index 000000000000..1a36cf5fc27a --- /dev/null +++ b/arm117-artifacts/Aon-np1/health.json @@ -0,0 +1 @@ +{"status":"ok"} \ No newline at end of file diff --git a/arm117-artifacts/Aon-np1/log-stream-lines.txt b/arm117-artifacts/Aon-np1/log-stream-lines.txt new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/arm117-artifacts/Aon-np1/slots.json b/arm117-artifacts/Aon-np1/slots.json new file mode 100644 index 000000000000..977716ea743b --- /dev/null +++ b/arm117-artifacts/Aon-np1/slots.json @@ -0,0 +1 @@ +[{"id":0,"n_ctx":65536,"speculative":true,"is_processing":false,"id_task":20,"n_prompt_tokens":2260,"n_prompt_tokens_processed":0,"n_prompt_tokens_cache":0,"params":{"seed":4294967295,"temperature":0.0,"dynatemp_range":0.0,"dynatemp_exponent":1.0,"top_k":20,"top_p":0.949999988079071,"min_p":0.05000000074505806,"top_n_sigma":-1.0,"xtc_probability":0.0,"xtc_threshold":0.10000000149011612,"typical_p":1.0,"repeat_last_n":64,"repeat_penalty":1.0,"presence_penalty":0.0,"frequency_penalty":0.0,"dry_multiplier":0.0,"dry_base":1.75,"dry_allowed_length":2,"dry_penalty_last_n":64,"mirostat":0,"mirostat_tau":5.0,"mirostat_eta":0.10000000149011612,"adaptive_target":-1.0,"adaptive_decay":0.8999999761581421,"max_tokens":96,"n_predict":96,"n_keep":0,"n_discard":0,"ignore_eos":false,"stream":false,"n_probs":0,"min_keep":0,"chat_format":"peg-native","reasoning_format":"deepseek","reasoning_in_content":false,"generation_prompt":"<|im_start|>assistant\n\n","samplers":["penalties","dry","top_n_sigma","top_k","typ_p","top_p","min_p","xtc","temperature"],"speculative.types":"none,draft-mtp","timings_per_token":false,"post_sampling_probs":false,"backend_sampling":false,"lora":[]},"next_token":[{"has_next_token":false,"has_new_line":false,"n_remain":-1,"n_decoded":0}]}] \ No newline at end of file diff --git a/arm117-artifacts/cell-a/req-p1-Aoff-np1-rep2.json b/arm117-artifacts/cell-a/req-p1-Aoff-np1-rep2.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aoff-np1-rep2.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aoff-np1.json b/arm117-artifacts/cell-a/req-p1-Aoff-np1.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aoff-np1.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aoff-np2-rep2.json b/arm117-artifacts/cell-a/req-p1-Aoff-np2-rep2.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aoff-np2-rep2.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aoff-np2-rep3.json b/arm117-artifacts/cell-a/req-p1-Aoff-np2-rep3.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aoff-np2-rep3.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aoff-np2.json b/arm117-artifacts/cell-a/req-p1-Aoff-np2.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aoff-np2.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aon-np1-rep2.json b/arm117-artifacts/cell-a/req-p1-Aon-np1-rep2.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aon-np1-rep2.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aon-np1.json b/arm117-artifacts/cell-a/req-p1-Aon-np1.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aon-np1.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aon-np2-rep2.json b/arm117-artifacts/cell-a/req-p1-Aon-np2-rep2.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aon-np2-rep2.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/req-p1-Aon-np2.json b/arm117-artifacts/cell-a/req-p1-Aon-np2.json new file mode 100644 index 000000000000..50a2d15609c6 --- /dev/null +++ b/arm117-artifacts/cell-a/req-p1-Aon-np2.json @@ -0,0 +1 @@ +{"messages": [{"role": "user", "content": "Filler sentence number 0 for the parity kernel prompt body. Filler sentence number 1 for the parity kernel prompt body. Filler sentence number 2 for the parity kernel prompt body. Filler sentence number 3 for the parity kernel prompt body. Filler sentence number 4 for the parity kernel prompt body. Filler sentence number 5 for the parity kernel prompt body. Filler sentence number 6 for the parity kernel prompt body. Filler sentence number 7 for the parity kernel prompt body. Filler sentence number 8 for the parity kernel prompt body. Filler sentence number 9 for the parity kernel prompt body. Filler sentence number 10 for the parity kernel prompt body. Filler sentence number 11 for the parity kernel prompt body. Filler sentence number 12 for the parity kernel prompt body. Filler sentence number 13 for the parity kernel prompt body. Filler sentence number 14 for the parity kernel prompt body. Filler sentence number 15 for the parity kernel prompt body. Filler sentence number 16 for the parity kernel prompt body. Filler sentence number 17 for the parity kernel prompt body. Filler sentence number 18 for the parity kernel prompt body. Filler sentence number 19 for the parity kernel prompt body. Filler sentence number 20 for the parity kernel prompt body. Filler sentence number 21 for the parity kernel prompt body. Filler sentence number 22 for the parity kernel prompt body. Filler sentence number 23 for the parity kernel prompt body. Filler sentence number 24 for the parity kernel prompt body. Filler sentence number 25 for the parity kernel prompt body. Filler sentence number 26 for the parity kernel prompt body. Filler sentence number 27 for the parity kernel prompt body. Filler sentence number 28 for the parity kernel prompt body. Filler sentence number 29 for the parity kernel prompt body. Filler sentence number 30 for the parity kernel prompt body. Filler sentence number 31 for the parity kernel prompt body. Filler sentence number 32 for the parity kernel prompt body. Filler sentence number 33 for the parity kernel prompt body. Filler sentence number 34 for the parity kernel prompt body. Filler sentence number 35 for the parity kernel prompt body. Filler sentence number 36 for the parity kernel prompt body. Filler sentence number 37 for the parity kernel prompt body. Filler sentence number 38 for the parity kernel prompt body. Filler sentence number 39 for the parity kernel prompt body. Filler sentence number 40 for the parity kernel prompt body. Filler sentence number 41 for the parity kernel prompt body. Filler sentence number 42 for the parity kernel prompt body. Filler sentence number 43 for the parity kernel prompt body. Filler sentence number 44 for the parity kernel prompt body. Filler sentence number 45 for the parity kernel prompt body. Filler sentence number 46 for the parity kernel prompt body. Filler sentence number 47 for the parity kernel prompt body. Filler sentence number 48 for the parity kernel prompt body. Filler sentence number 49 for the parity kernel prompt body. Filler sentence number 50 for the parity kernel prompt body. Filler sentence number 51 for the parity kernel prompt body. Filler sentence number 52 for the parity kernel prompt body. Filler sentence number 53 for the parity kernel prompt body. Filler sentence number 54 for the parity kernel prompt body. Filler sentence number 55 for the parity kernel prompt body. Filler sentence number 56 for the parity kernel prompt body. Filler sentence number 57 for the parity kernel prompt body. Filler sentence number 58 for the parity kernel prompt body. Filler sentence number 59 for the parity kernel prompt body. Filler sentence number 60 for the parity kernel prompt body. Filler sentence number 61 for the parity kernel prompt body. Filler sentence number 62 for the parity kernel prompt body. Filler sentence number 63 for the parity kernel prompt body. Filler sentence number 64 for the parity kernel prompt body. Filler sentence number 65 for the parity kernel prompt body. Filler sentence number 66 for the parity kernel prompt body. Filler sentence number 67 for the parity kernel prompt body. Filler sentence number 68 for the parity kernel prompt body. Filler sentence number 69 for the parity kernel prompt body. Filler sentence number 70 for the parity kernel prompt body. Filler sentence number 71 for the parity kernel prompt body. Filler sentence number 72 for the parity kernel prompt body. Filler sentence number 73 for the parity kernel prompt body. Filler sentence number 74 for the parity kernel prompt body. Filler sentence number 75 for the parity kernel prompt body. Filler sentence number 76 for the parity kernel prompt body. Filler sentence number 77 for the parity kernel prompt body. Filler sentence number 78 for the parity kernel prompt body. Filler sentence number 79 for the parity kernel prompt body. Filler sentence number 80 for the parity kernel prompt body. Filler sentence number 81 for the parity kernel prompt body. Filler sentence number 82 for the parity kernel prompt body. Filler sentence number 83 for the parity kernel prompt body. Filler sentence number 84 for the parity kernel prompt body. Filler sentence number 85 for the parity kernel prompt body. Filler sentence number 86 for the parity kernel prompt body. Filler sentence number 87 for the parity kernel prompt body. Filler sentence number 88 for the parity kernel prompt body. Filler sentence number 89 for the parity kernel prompt body. Filler sentence number 90 for the parity kernel prompt body. Filler sentence number 91 for the parity kernel prompt body. Filler sentence number 92 for the parity kernel prompt body. Filler sentence number 93 for the parity kernel prompt body. Filler sentence number 94 for the parity kernel prompt body. Filler sentence number 95 for the parity kernel prompt body. Filler sentence number 96 for the parity kernel prompt body. Filler sentence number 97 for the parity kernel prompt body. Filler sentence number 98 for the parity kernel prompt body. Filler sentence number 99 for the parity kernel prompt body. Filler sentence number 100 for the parity kernel prompt body. Filler sentence number 101 for the parity kernel prompt body. Filler sentence number 102 for the parity kernel prompt body. Filler sentence number 103 for the parity kernel prompt body. Filler sentence number 104 for the parity kernel prompt body. Filler sentence number 105 for the parity kernel prompt body. Filler sentence number 106 for the parity kernel prompt body. Filler sentence number 107 for the parity kernel prompt body. Filler sentence number 108 for the parity kernel prompt body. Filler sentence number 109 for the parity kernel prompt body. Filler sentence number 110 for the parity kernel prompt body. Filler sentence number 111 for the parity kernel prompt body. Filler sentence number 112 for the parity kernel prompt body. Filler sentence number 113 for the parity kernel prompt body. Filler sentence number 114 for the parity kernel prompt body. Filler sentence number 115 for the parity kernel prompt body. Filler sentence number 116 for the parity kernel prompt body. Filler sentence number 117 for the parity kernel prompt body. Filler sentence number 118 for the parity kernel prompt body. Filler sentence number 119 for the parity kernel prompt body. Filler sentence number 120 for the parity kernel prompt body. Filler sentence number 121 for the parity kernel prompt body. Filler sentence number 122 for the parity kernel prompt body. Filler sentence number 123 for the parity kernel prompt body. Filler sentence number 124 for the parity kernel prompt body. Filler sentence number 125 for the parity kernel prompt body. Filler sentence number 126 for the parity kernel prompt body. Filler sentence number 127 for the parity kernel prompt body. Filler sentence number 128 for the parity kernel prompt body. Filler sentence number 129 for the parity kernel prompt body. Filler sentence number 130 for the parity kernel prompt body. Filler sentence number 131 for the parity kernel prompt body. Filler sentence number 132 for the parity kernel prompt body. Filler sentence number 133 for the parity kernel prompt body. Filler sentence number 134 for the parity kernel prompt body. Filler sentence number 135 for the parity kernel prompt body. Filler sentence number 136 for the parity kernel prompt body. Filler sentence number 137 for the parity kernel prompt body. Filler sentence number 138 for the parity kernel prompt body. Filler sentence number 139 for the parity kernel prompt body. Filler sentence number 140 for the parity kernel prompt body. Filler sentence number 141 for the parity kernel prompt body. Filler sentence number 142 for the parity kernel prompt body. Filler sentence number 143 for the parity kernel prompt body. Filler sentence number 144 for the parity kernel prompt body. Filler sentence number 145 for the parity kernel prompt body. Filler sentence number 146 for the parity kernel prompt body. Filler sentence number 147 for the parity kernel prompt body. Filler sentence number 148 for the parity kernel prompt body. Filler sentence number 149 for the parity kernel prompt body.\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A"}], "max_tokens": 96, "temperature": 0, "cache_prompt": true} diff --git a/arm117-artifacts/cell-a/resp-p1-Aoff-np1-rep2.json b/arm117-artifacts/cell-a/resp-p1-Aoff-np1-rep2.json new file mode 100644 index 000000000000..d20df7bc2564 --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aoff-np1-rep2.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789111099,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":0}},"id":"chatcmpl-N83SDRz0nYa0hwsNloXZzoNxgDZHknDD","timings":{"cache_n":0,"prompt_n":2208,"prompt_ms":3057.821,"prompt_per_token_ms":1.3848826992753622,"prompt_per_second":722.0828164892582,"predicted_n":53,"predicted_ms":1190.644,"predicted_per_token_ms":22.897,"predicted_per_second":43.673843734987116,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aoff-np1.json b/arm117-artifacts/cell-a/resp-p1-Aoff-np1.json new file mode 100644 index 000000000000..0bfe83d8cb6d --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aoff-np1.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789110988,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":2204}},"id":"chatcmpl-0NE7AZ6rnUsi3NgNnRcFa2K7SUoDyJNi","timings":{"cache_n":2204,"prompt_n":4,"prompt_ms":180.348,"prompt_per_token_ms":45.087,"prompt_per_second":22.17934216071151,"predicted_n":53,"predicted_ms":1198.634,"predicted_per_token_ms":23.050653846153846,"predicted_per_second":43.38271732655673,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aoff-np2-rep2.json b/arm117-artifacts/cell-a/resp-p1-Aoff-np2-rep2.json new file mode 100644 index 000000000000..c22d8edfd064 --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aoff-np2-rep2.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789111126,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":2204}},"id":"chatcmpl-p4X7ms4fDePivAORVOEhqz3s8vq1bPBZ","timings":{"cache_n":2204,"prompt_n":4,"prompt_ms":176.953,"prompt_per_token_ms":44.23825,"prompt_per_second":22.604872480263122,"predicted_n":53,"predicted_ms":1179.323,"predicted_per_token_ms":22.679288461538462,"predicted_per_second":44.09309408872718,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aoff-np2-rep3.json b/arm117-artifacts/cell-a/resp-p1-Aoff-np2-rep3.json new file mode 100644 index 000000000000..929d65d10eea --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aoff-np2-rep3.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789111183,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":0}},"id":"chatcmpl-kIQaYfLUvQnZsF1PP19ygcfGNsVLLWR0","timings":{"cache_n":0,"prompt_n":2208,"prompt_ms":3049.499,"prompt_per_token_ms":1.3811136775362318,"prompt_per_second":724.0533608963309,"predicted_n":53,"predicted_ms":1203.121,"predicted_per_token_ms":23.13694230769231,"predicted_per_second":43.220922916315146,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aoff-np2.json b/arm117-artifacts/cell-a/resp-p1-Aoff-np2.json new file mode 100644 index 000000000000..9b7da8b19932 --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aoff-np2.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789111124,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":0}},"id":"chatcmpl-l9f7Sme4TJCGaKjV13vpIaVnBbAy3Any","timings":{"cache_n":0,"prompt_n":2208,"prompt_ms":3048.721,"prompt_per_token_ms":1.380761322463768,"prompt_per_second":724.2381313344187,"predicted_n":53,"predicted_ms":1201.784,"predicted_per_token_ms":23.111230769230772,"predicted_per_second":43.269006743308275,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aon-np1-rep2.json b/arm117-artifacts/cell-a/resp-p1-Aon-np1-rep2.json new file mode 100644 index 000000000000..301bcc0949a1 --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aon-np1-rep2.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789111071,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":2204}},"id":"chatcmpl-bjqTbPkhtHCGtgD629eQSvcMzsMMQkPU","timings":{"cache_n":2204,"prompt_n":4,"prompt_ms":182.089,"prompt_per_token_ms":45.52225,"prompt_per_second":21.96727973683199,"predicted_n":53,"predicted_ms":1195.775,"predicted_per_token_ms":22.99567307692308,"predicted_per_second":43.486441847337495,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aon-np1.json b/arm117-artifacts/cell-a/resp-p1-Aon-np1.json new file mode 100644 index 000000000000..e7acd8b8c32f --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aon-np1.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"We need answer user's request. User provided filler sentences and then: \"Now reply with exactly this text and nothing else: PARITY-KERNEL-A\". Need final exactly PARITY-KERNEL-A. Ensure no extra.\n"}}],"created":1789111063,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":53,"prompt_tokens":2208,"total_tokens":2261,"prompt_tokens_details":{"cached_tokens":0}},"id":"chatcmpl-fy47muKLt832J7rivGnCkQWkwj4qIsCn","timings":{"cache_n":0,"prompt_n":2208,"prompt_ms":3086.847,"prompt_per_token_ms":1.3980285326086956,"prompt_per_second":715.2929834228906,"predicted_n":53,"predicted_ms":1198.814,"predicted_per_token_ms":23.054115384615386,"predicted_per_second":43.376203481107154,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aon-np2-rep2.json b/arm117-artifacts/cell-a/resp-p1-Aon-np2-rep2.json new file mode 100644 index 000000000000..c5202cba8202 --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aon-np2-rep2.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"The user is asking me to reply with exactly the text \"PARITY-KERNEL-A\" and nothing else. The instruction is clear and explicit: reply with exactly \"PARITY-KERNEL-A\" and nothing else.\n"}}],"created":1789111155,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":52,"prompt_tokens":2208,"total_tokens":2260,"prompt_tokens_details":{"cached_tokens":2204}},"id":"chatcmpl-MMj6ZrGElRvUe8t6DNW3KdYAnBH80T8c","timings":{"cache_n":2204,"prompt_n":4,"prompt_ms":178.986,"prompt_per_token_ms":44.7465,"prompt_per_second":22.348116612472484,"predicted_n":52,"predicted_ms":1177.532,"predicted_per_token_ms":23.088862745098037,"predicted_per_second":43.310924883570046,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/resp-p1-Aon-np2.json b/arm117-artifacts/cell-a/resp-p1-Aon-np2.json new file mode 100644 index 000000000000..3cf7d313a458 --- /dev/null +++ b/arm117-artifacts/cell-a/resp-p1-Aon-np2.json @@ -0,0 +1 @@ +{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"PARITY-KERNEL-A","reasoning_content":"The user is asking me to reply with exactly the text \"PARITY-KERNEL-A\" and nothing else. The instruction is clear and explicit: reply with exactly \"PARITY-KERNEL-A\" and nothing else.\n"}}],"created":1789111153,"model":"/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf","system_fingerprint":"b10875-017a512d","object":"chat.completion","usage":{"completion_tokens":52,"prompt_tokens":2208,"total_tokens":2260,"prompt_tokens_details":{"cached_tokens":0}},"id":"chatcmpl-KQLi4rdrZrSFJIzuxRNEKh1CF7At2THM","timings":{"cache_n":0,"prompt_n":2208,"prompt_ms":3064.175,"prompt_per_token_ms":1.3877604166666668,"prompt_per_second":720.5854756990053,"predicted_n":52,"predicted_ms":1196.311,"predicted_per_token_ms":23.457078431372548,"predicted_per_second":42.63105496814792,"draft_n":45,"draft_n_accepted":37}} \ No newline at end of file diff --git a/arm117-artifacts/cell-a/text-p1-Aoff-np1-rep2.txt b/arm117-artifacts/cell-a/text-p1-Aoff-np1-rep2.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aoff-np1-rep2.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aoff-np1.txt b/arm117-artifacts/cell-a/text-p1-Aoff-np1.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aoff-np1.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aoff-np2-rep2.txt b/arm117-artifacts/cell-a/text-p1-Aoff-np2-rep2.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aoff-np2-rep2.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aoff-np2-rep3.txt b/arm117-artifacts/cell-a/text-p1-Aoff-np2-rep3.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aoff-np2-rep3.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aoff-np2.txt b/arm117-artifacts/cell-a/text-p1-Aoff-np2.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aoff-np2.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aon-np1-rep2.txt b/arm117-artifacts/cell-a/text-p1-Aon-np1-rep2.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aon-np1-rep2.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aon-np1.txt b/arm117-artifacts/cell-a/text-p1-Aon-np1.txt new file mode 100644 index 000000000000..1cd2b4706ad7 --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aon-np1.txt @@ -0,0 +1,2 @@ +We need answer user's request. User provided filler sentences and then: "Now reply with exactly this text and nothing else: PARITY-KERNEL-A". Need final exactly PARITY-KERNEL-A. Ensure no extra. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aon-np2-rep2.txt b/arm117-artifacts/cell-a/text-p1-Aon-np2-rep2.txt new file mode 100644 index 000000000000..441dc2d221cd --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aon-np2-rep2.txt @@ -0,0 +1,2 @@ +The user is asking me to reply with exactly the text "PARITY-KERNEL-A" and nothing else. The instruction is clear and explicit: reply with exactly "PARITY-KERNEL-A" and nothing else. +PARITY-KERNEL-A diff --git a/arm117-artifacts/cell-a/text-p1-Aon-np2.txt b/arm117-artifacts/cell-a/text-p1-Aon-np2.txt new file mode 100644 index 000000000000..441dc2d221cd --- /dev/null +++ b/arm117-artifacts/cell-a/text-p1-Aon-np2.txt @@ -0,0 +1,2 @@ +The user is asking me to reply with exactly the text "PARITY-KERNEL-A" and nothing else. The instruction is clear and explicit: reply with exactly "PARITY-KERNEL-A" and nothing else. +PARITY-KERNEL-A diff --git a/arm117-artifacts/perf-base.txt b/arm117-artifacts/perf-base.txt new file mode 100644 index 000000000000..8d6fd8db84a9 --- /dev/null +++ b/arm117-artifacts/perf-base.txt @@ -0,0 +1,33 @@ +=== multiturn-growth-test === + port: 8080 + sessions: 1 + turns: 12 + new tok/turn: ~8000 (~5333 words) + output/turn: 750 tokens + target depth: turn 12 ≈ 103250 resident tokens + +Launched 1 sessions, waiting for completion... + + summary: turn1=31.54 turn9=21.89 mean=24.26 tok/s prompt_tok: 6682 -> 59714 (8.9x) + +=== all sessions complete, results === + +session 1 | 12 turns | 5333 words/turn new | 750 output tokens +---------------------------------------------------------------------- + turn 1/12 wall= 23.78s prompt_tok= 6682 comp_tok= 750 tok/s= 31.54 msgs=2 + turn 2/12 wall= 28.35s prompt_tok= 13308 comp_tok= 750 tok/s= 26.46 msgs=4 + turn 3/12 wall= 30.28s prompt_tok= 19927 comp_tok= 750 tok/s= 24.77 msgs=6 + turn 4/12 wall= 35.27s prompt_tok= 26550 comp_tok= 750 tok/s= 21.27 msgs=8 + turn 5/12 wall= 30.88s prompt_tok= 33175 comp_tok= 750 tok/s= 24.29 msgs=10 + turn 6/12 wall= 31.14s prompt_tok= 39799 comp_tok= 750 tok/s= 24.08 msgs=12 + turn 7/12 wall= 32.17s prompt_tok= 46424 comp_tok= 750 tok/s= 23.31 msgs=14 + turn 8/12 wall= 36.14s prompt_tok= 53044 comp_tok= 750 tok/s= 20.75 msgs=16 + turn 9/12 wall= 34.26s prompt_tok= 59714 comp_tok= 750 tok/s= 21.89 msgs=18 + turn 10/12 FAILED: HTTP Error 400: Bad Request + turn 11/12 FAILED: HTTP Error 400: Bad Request + turn 12/12 FAILED: HTTP Error 400: Bad Request + + summary: turn1=31.54 turn9=21.89 mean=24.26 tok/s prompt_tok: 6682 -> 59714 (8.9x) + +=== overlap check === +concurrency check: N/A (single session) diff --git a/arm117-artifacts/perf-baseline-binary.txt b/arm117-artifacts/perf-baseline-binary.txt new file mode 100644 index 000000000000..8d6fd8db84a9 --- /dev/null +++ b/arm117-artifacts/perf-baseline-binary.txt @@ -0,0 +1,33 @@ +=== multiturn-growth-test === + port: 8080 + sessions: 1 + turns: 12 + new tok/turn: ~8000 (~5333 words) + output/turn: 750 tokens + target depth: turn 12 ≈ 103250 resident tokens + +Launched 1 sessions, waiting for completion... + + summary: turn1=31.54 turn9=21.89 mean=24.26 tok/s prompt_tok: 6682 -> 59714 (8.9x) + +=== all sessions complete, results === + +session 1 | 12 turns | 5333 words/turn new | 750 output tokens +---------------------------------------------------------------------- + turn 1/12 wall= 23.78s prompt_tok= 6682 comp_tok= 750 tok/s= 31.54 msgs=2 + turn 2/12 wall= 28.35s prompt_tok= 13308 comp_tok= 750 tok/s= 26.46 msgs=4 + turn 3/12 wall= 30.28s prompt_tok= 19927 comp_tok= 750 tok/s= 24.77 msgs=6 + turn 4/12 wall= 35.27s prompt_tok= 26550 comp_tok= 750 tok/s= 21.27 msgs=8 + turn 5/12 wall= 30.88s prompt_tok= 33175 comp_tok= 750 tok/s= 24.29 msgs=10 + turn 6/12 wall= 31.14s prompt_tok= 39799 comp_tok= 750 tok/s= 24.08 msgs=12 + turn 7/12 wall= 32.17s prompt_tok= 46424 comp_tok= 750 tok/s= 23.31 msgs=14 + turn 8/12 wall= 36.14s prompt_tok= 53044 comp_tok= 750 tok/s= 20.75 msgs=16 + turn 9/12 wall= 34.26s prompt_tok= 59714 comp_tok= 750 tok/s= 21.89 msgs=18 + turn 10/12 FAILED: HTTP Error 400: Bad Request + turn 11/12 FAILED: HTTP Error 400: Bad Request + turn 12/12 FAILED: HTTP Error 400: Bad Request + + summary: turn1=31.54 turn9=21.89 mean=24.26 tok/s prompt_tok: 6682 -> 59714 (8.9x) + +=== overlap check === +concurrency check: N/A (single session) diff --git a/arm117-artifacts/perf-on.txt b/arm117-artifacts/perf-on.txt new file mode 100644 index 000000000000..ada4447f2467 --- /dev/null +++ b/arm117-artifacts/perf-on.txt @@ -0,0 +1,33 @@ +=== multiturn-growth-test === + port: 8080 + sessions: 1 + turns: 12 + new tok/turn: ~8000 (~5333 words) + output/turn: 750 tokens + target depth: turn 12 ≈ 103250 resident tokens + +Launched 1 sessions, waiting for completion... + + summary: turn1=31.52 turn9=21.41 mean=17.25 tok/s prompt_tok: 6682 -> 60506 (9.1x) + +=== all sessions complete, results === + +session 1 | 12 turns | 5333 words/turn new | 750 output tokens +---------------------------------------------------------------------- + turn 1/12 wall= 23.79s prompt_tok= 6682 comp_tok= 750 tok/s= 31.52 msgs=2 + turn 2/12 wall= 31.59s prompt_tok= 13308 comp_tok= 688 tok/s= 21.78 msgs=4 + turn 3/12 wall= 31.25s prompt_tok= 19960 comp_tok= 750 tok/s= 24.00 msgs=6 + turn 4/12 wall= 22.01s prompt_tok= 27219 comp_tok= 359 tok/s= 16.31 msgs=8 + turn 5/12 wall= 15.64s prompt_tok= 33883 comp_tok= 181 tok/s= 11.57 msgs=10 + turn 6/12 wall= 15.81s prompt_tok= 40538 comp_tok= 139 tok/s= 8.79 msgs=12 + turn 7/12 wall= 16.65s prompt_tok= 47195 comp_tok= 151 tok/s= 9.07 msgs=14 + turn 8/12 wall= 19.10s prompt_tok= 53852 comp_tok= 206 tok/s= 10.78 msgs=16 + turn 9/12 wall= 35.03s prompt_tok= 60506 comp_tok= 750 tok/s= 21.41 msgs=18 + turn 10/12 FAILED: HTTP Error 400: Bad Request + turn 11/12 FAILED: HTTP Error 400: Bad Request + turn 12/12 FAILED: HTTP Error 400: Bad Request + + summary: turn1=31.52 turn9=21.41 mean=17.25 tok/s prompt_tok: 6682 -> 60506 (9.1x) + +=== overlap check === +concurrency check: N/A (single session) diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000000..a6d722f28b6b --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,91 @@ +# Adaptive KV streaming benchmark + +`benchmark_kv_stream.py` performs the complete context-matched benchmark and +creates its graph. The only required inputs are the GGUF model and the largest +context capacity to test: + +```bash +python3 benchmarks/benchmark_kv_stream.py \ + --model /path/to/model.gguf \ + --max-context 192K +``` + +The default server is `build/bin/llama-server`. Use `--server` when the +binary is elsewhere. Matplotlib is the only Python dependency: + +```bash +python3 -m pip install matplotlib +``` + +## What the script does + +For every configured context capacity from 8K through `--max-context`, in 8K +steps, the script: + +1. Starts a fresh adaptive KV streaming server with a small probe pool. +2. Measures free VRAM after model initialization and a warm-up request. +3. Assigns all measured free VRAM to the pool, rounded down to 32 MiB. +4. Runs a full prompt and 256-token decode to validate that pool. +5. If the candidate fails, reduces it by 64 MiB and retries. +6. Records prefill speed, decode speed, selected pool size, and VRAM telemetry. +7. Updates the CSV and Matplotlib graph after every successful point. + +There is no fixed VRAM safety reserve. Actual server execution is the +validation: allocation failures are handled by automatic pool backoff. + +The prompt length at each point is the configured context capacity minus the +256 decode tokens. For example, the 192K point starts the server with +`--ctx-size 196608`, prefills 196352 tokens, and then decodes 256 tokens. +If the maximum is not a multiple of 8K, the exact maximum is appended as the +last point. + +The driver uses the configuration currently supported and validated by this +branch: + +- Flash Attention enabled +- K cache `q8_0` +- V cache `q4_0` +- all model layers on the GPU +- one server slot +- 256-token batch and micro-batch by default +- ordinary CUDA allocation, without UVM + +Use `--batch-size` and `--ubatch-size` to benchmark other logical and physical batch sizes. The micro-batch must not exceed the logical batch. Both values are included in result metadata and the resume signature. + +```bash +python3 benchmarks/benchmark_kv_stream.py \ + --model /path/to/model.gguf \ + --max-context 192K \ + --batch-size 512 \ + --ubatch-size 512 +``` + +## Results and resuming + +By default, a timestamped directory is created under +`benchmarks/results/adaptive-kv-sweep-*`. It contains: + +- `results.jsonl`: metadata, pool probes, retries, and measurements +- `results.csv`: one successful measurement per context capacity +- `kv-stream-sweep.png` and `kv-stream-sweep.svg`: decode, prefill, and pool + size plots +- `logs/`: one server log per probe and benchmark attempt + +Use an explicit output directory to resume an interrupted sweep: + +```bash +python3 benchmarks/benchmark_kv_stream.py \ + --model /path/to/model.gguf \ + --max-context 192K \ + --output-dir benchmarks/results/my-sweep +``` + +Re-run the same command after an interruption. Completed contexts are skipped. +The script rejects a resume if the model or benchmark settings differ, avoiding +mixed data in one result set. + +Run `python3 benchmarks/benchmark_kv_stream.py --help` for optional GPU, +timeout, pool-step, output, and server arguments. + +Do not run another GPU workload during the sweep. Its allocations would change +the automatically selected pool and invalidate comparisons between points. diff --git a/benchmarks/benchmark_kv_stream.py b/benchmarks/benchmark_kv_stream.py new file mode 100755 index 000000000000..f98f1a990a45 --- /dev/null +++ b/benchmarks/benchmark_kv_stream.py @@ -0,0 +1,985 @@ +#!/usr/bin/env python3 +"""Run and plot an automatic context-matched adaptive KV streaming sweep.""" + +from __future__ import annotations + +import argparse +import csv +from dataclasses import dataclass +import datetime as dt +import json +import os +from pathlib import Path +import re +import signal +import subprocess +import sys +import time +import urllib.error +import urllib.request + + +ROOT = Path(__file__).resolve().parents[1] +SERVER_BIN = "Release/llama-server.exe" if os.name == "nt" else "llama-server" +CONTEXT_STEP = 8192 +UVM_ENV_NAMES = ( + "GGML_CUDA_ENABLE_UNIFIED_MEMORY", + "GGML_CUDA_PREFER_MODEL_WEIGHTS", + "GGML_CUDA_PREFER_KV_HOST", + "GGML_CUDA_KV_ACCESSED_BY_GPU", +) + +KV_STREAM_TRACE_RE = re.compile( + r"kv_stream_adapt: active (\d+), resident (\d+), ring (\d+), " + r"samples (\d+), misses (\d+), copy busy ([0-9.]+)%, peak (\d+)" +) + + + +def parse_token_count(value: str) -> int: + text = value.strip().lower() + multiplier = 1 + for suffix, factor in ( + ("kib", 1024), + ("ki", 1024), + ("k", 1024), + ("mib", 1024 * 1024), + ("mi", 1024 * 1024), + ("m", 1024 * 1024), + ): + if text.endswith(suffix): + text = text[:-len(suffix)] + multiplier = factor + break + try: + result = int(text) * multiplier + except ValueError as exc: + raise argparse.ArgumentTypeError(f"invalid token count: {value}") from exc + if result <= 0: + raise argparse.ArgumentTypeError("token count must be positive") + return result + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--model", type=Path, required=True, help="Qwen3.8-27B GGUF model") + parser.add_argument( + "--max-context", + type=parse_token_count, + required=True, + help="largest configured context capacity, for example 192K or 196608", + ) + parser.add_argument( + "--min-context", + type=parse_token_count, + default=CONTEXT_STEP, + help="smallest configured context capacity", + ) + parser.add_argument( + "--context-step", + type=parse_token_count, + default=CONTEXT_STEP, + help="context-capacity increment", + ) + parser.add_argument( + "--server", + type=Path, + default=ROOT / "build/bin" / SERVER_BIN, + help="adaptive KV streaming llama-server binary", + ) + parser.add_argument( + "--output-dir", + type=Path, + help="result directory; reuse it to resume an interrupted sweep", + ) + parser.add_argument("--decode-tokens", type=int, default=256) + parser.add_argument( + "--batch-size", type=int, default=256, + help="logical maximum prompt batch size", + ) + parser.add_argument( + "--ubatch-size", type=int, default=256, + help="physical maximum batch size; must not exceed --batch-size", + ) + parser.add_argument("--cache-type-k", default="q8_0") + parser.add_argument("--cache-type-v", default="q4_0") + parser.add_argument("--probe-pool-mib", type=int, default=64) + parser.add_argument("--pool-step-mib", type=int, default=32) + parser.add_argument("--pool-backoff-mib", type=int, default=64) + parser.add_argument("--pool-retries", type=int, default=8) + parser.add_argument("--max-pool-mib", type=int) + parser.add_argument( + "--fixed-pool-mib", + type=int, + help="use exactly this KV pool size and skip per-context probing", + ) + parser.add_argument( + "--trace-kv-stream", + action="store_true", + help="enable and parse adaptive KV residency trace logging", + ) + parser.add_argument("--gpu-index", type=int, default=0) + parser.add_argument( + "--cuda-visible-devices", + help="CUDA_VISIBLE_DEVICES value for the server; defaults to the inherited value", + ) + parser.add_argument("--nvidia-smi", default="nvidia-smi") + parser.add_argument("--port", type=int, default=12355) + parser.add_argument("--startup-timeout", type=int, default=240) + parser.add_argument("--request-timeout", type=int, default=1800) + parser.add_argument("--release-timeout", type=int, default=90) + parser.add_argument("--release-slack-mib", type=int, default=64) + parser.add_argument("--prompt-suffix", default="The capital of France is") + parser.add_argument("--fill-token-id", type=int) + parser.add_argument( + "--extra-server-arg", + action="append", + default=[], + metavar="ARG", + help="append a server argument (repeat; use --extra-server-arg=--flag)", + ) + args = parser.parse_args(argv) + args.model = args.model.resolve() + args.server = args.server.resolve() + if not args.server.is_file(): + exe_server = args.server.with_suffix(".exe").resolve() + if exe_server.is_file(): + args.server = exe_server + return args + + +def context_capacities( + max_context: int, + min_context: int = CONTEXT_STEP, + step: int = CONTEXT_STEP, +) -> list[int]: + capacities = list(range(min_context, max_context + 1, step)) + if not capacities or capacities[-1] != max_context: + capacities.append(max_context) + return capacities + + +@dataclass(frozen=True) +class GpuMemory: + total_mib: int + used_mib: int + free_mib: int + + +def query_gpu_memory(nvidia_smi: str, gpu_index: int) -> GpuMemory: + try: + result = subprocess.run( + [ + nvidia_smi, + "-i", + str(gpu_index), + "--query-gpu=memory.total,memory.used,memory.free", + "--format=csv,noheader,nounits", + ], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + except OSError as exc: + raise RuntimeError(f"cannot run nvidia-smi executable: {nvidia_smi}") from exc + if result.returncode != 0: + raise RuntimeError(f"nvidia-smi failed: {result.stderr.strip()}") + try: + total, used, free = ( + int(part.strip()) for part in result.stdout.splitlines()[0].split(",") + ) + except (IndexError, ValueError) as exc: + raise RuntimeError(f"unexpected nvidia-smi output: {result.stdout!r}") from exc + return GpuMemory(total, used, free) + + +def estimate_pool_mib( + probe_pool_mib: int, + free_mib: int, + step_mib: int, + max_pool_mib: int | None = None, +) -> int: + candidate = probe_pool_mib + free_mib + candidate = candidate // step_mib * step_mib + if max_pool_mib is not None: + candidate = min(candidate, max_pool_mib // step_mib * step_mib) + if candidate < probe_pool_mib: + raise RuntimeError(f"only {free_mib} MiB is free after the probe") + return candidate + + +def http_json(url: str, payload: dict | None, timeout: int) -> dict: + request = urllib.request.Request( + url, + data=None if payload is None else json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.load(response) + except urllib.error.HTTPError as exc: + body = exc.read().decode(errors="replace") + raise RuntimeError(f"HTTP {exc.code}: {body[:1000]}") from exc + + +def clean_server_env( + cuda_visible_devices: str | None, + trace_kv_stream: bool = False, +) -> dict[str, str]: + env = os.environ.copy() + for name in UVM_ENV_NAMES: + env.pop(name, None) + env.pop("GGML_CUDA_KV_STREAM_FIXED_RING_SLOTS", None) + env.pop("LLAMA_KV_STREAM_TRACE", None) + if trace_kv_stream: + env["LLAMA_KV_STREAM_TRACE"] = "1" + if cuda_visible_devices is not None: + env["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + return env + + +def server_command( + args: argparse.Namespace, + context_capacity: int, + pool_mib: int, +) -> list[str]: + return [ + str(args.server), + "-m", + str(args.model), + "--alias", + "adaptive-kv-stream", + "--host", + "127.0.0.1", + "--port", + str(args.port), + "--ctx-size", + str(context_capacity), + "-fa", + "on", + "-ctk", + args.cache_type_k, + "-ctv", + args.cache_type_v, + "-ngl", + "all", + "-b", + str(args.batch_size), + "-ub", + str(args.ubatch_size), + "-np", + "1", + "--no-mmproj", + "--no-warmup", + "--reasoning-format", + "none", + "--kv-stream-stage-mib", + str(pool_mib), + *args.extra_server_arg, + ] + + +class Server: + def __init__( + self, + args: argparse.Namespace, + context_capacity: int, + pool_mib: int, + log_path: Path, + ) -> None: + self.port = args.port + self.log_path = log_path + self.log_file = log_path.open("wb") + self.process: subprocess.Popen | None = None + try: + self.process = subprocess.Popen( + server_command(args, context_capacity, pool_mib), + cwd=args.server.parent, + env=clean_server_env( + args.cuda_visible_devices, + args.trace_kv_stream, + ), + stdout=self.log_file, + stderr=subprocess.STDOUT, + ) + except Exception: + self.log_file.close() + raise + + deadline = time.monotonic() + args.startup_timeout + last_error = "server did not become ready" + while time.monotonic() < deadline: + status = self.process.poll() + if status is not None: + self.log_file.flush() + log_tail = self.log_tail() + self.stop() + raise RuntimeError(f"server exited with status {status}: {log_tail}") + try: + health = http_json(self.url("/health"), None, 2) + if health.get("status") == "ok": + return + except Exception as exc: + last_error = str(exc) + time.sleep(0.25) + self.stop() + raise RuntimeError(f"{last_error}: {self.log_tail()}") + + def url(self, path: str) -> str: + return f"http://127.0.0.1:{self.port}{path}" + + def log_tail(self, lines: int = 30) -> str: + try: + return "\n".join( + self.log_path.read_text(errors="replace").splitlines()[-lines:] + ) + except OSError: + return "" + + def stop(self) -> None: + if self.process is not None and self.process.poll() is None: + terminate_signal = ( + signal.CTRL_C_EVENT if os.name == "nt" else signal.SIGINT + ) + # On Windows, CTRL_C_EVENT reaches the shared console and would + # interrupt this benchmark process itself; mask SIGINT while we + # send it so the child stops but we keep running. + previous_handler = None + if os.name == "nt": + previous_handler = signal.signal(signal.SIGINT, lambda *_: None) + try: + self.process.send_signal(terminate_signal) + try: + self.process.wait(timeout=8) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait(timeout=10) + finally: + if previous_handler is not None: + signal.signal(signal.SIGINT, previous_handler) + if not self.log_file.closed: + self.log_file.close() + + +def wait_for_release( + args: argparse.Namespace, + baseline_used_mib: int, +) -> None: + deadline = time.monotonic() + args.release_timeout + while time.monotonic() < deadline: + memory = query_gpu_memory(args.nvidia_smi, args.gpu_index) + if memory.used_mib <= baseline_used_mib + args.release_slack_mib: + return + time.sleep(1) + memory = query_gpu_memory(args.nvidia_smi, args.gpu_index) + raise RuntimeError( + f"GPU memory did not return to baseline: " + f"{memory.used_mib} MiB used, baseline {baseline_used_mib} MiB" + ) + + +def prepare_server( + args: argparse.Namespace, + server: Server, +) -> tuple[list[int], int]: + tokenized = http_json( + server.url("/tokenize"), + {"content": args.prompt_suffix, "add_special": False}, + 30, + ) + suffix = tokenized.get("tokens") + if not suffix or not all(isinstance(token, int) for token in suffix): + raise RuntimeError(f"unexpected tokenize response: {tokenized}") + fill_token_id = ( + args.fill_token_id if args.fill_token_id is not None else suffix[0] + ) + http_json( + server.url("/completion"), + { + "prompt": [fill_token_id] * 16, + "n_predict": 4, + "ignore_eos": True, + "cache_prompt": False, + "temperature": 0, + "reasoning_format": "none", + "response_fields": ["timings"], + }, + args.request_timeout, + ) + return suffix, fill_token_id + + +def append_jsonl(path: Path, row: dict) -> None: + with path.open("a") as stream: + stream.write(json.dumps(row, sort_keys=True) + "\n") + print(json.dumps(row, sort_keys=True), flush=True) + + +def load_measurements(path: Path) -> dict[int, dict]: + rows: dict[int, dict] = {} + if not path.is_file(): + return rows + with path.open() as stream: + for line in stream: + row = json.loads(line) + if row.get("type") == "measurement" and row.get("status") == "ok": + rows[int(row["context_capacity"])] = row + return rows + + +def load_metadata(path: Path) -> dict | None: + if not path.is_file(): + return None + with path.open() as stream: + for line in stream: + row = json.loads(line) + if row.get("type") == "metadata": + return row + return None + + +def resume_signature(args: argparse.Namespace, capacities: list[int]) -> dict: + return { + "model": str(args.model.resolve()), + "server": str(args.server.resolve()), + "max_context": args.max_context, + "min_context": args.min_context, + "context_step": args.context_step, + "cache_type_k": args.cache_type_k, + "cache_type_v": args.cache_type_v, + "fixed_pool_mib": args.fixed_pool_mib, + "trace_kv_stream": args.trace_kv_stream, + "capacities": capacities, + "decode_tokens": args.decode_tokens, + "batch_size": args.batch_size, + "ubatch_size": args.ubatch_size, + "probe_pool_mib": args.probe_pool_mib, + "pool_step_mib": args.pool_step_mib, + "pool_backoff_mib": args.pool_backoff_mib, + "pool_retries": args.pool_retries, + "max_pool_mib": args.max_pool_mib, + "gpu_index": args.gpu_index, + "cuda_visible_devices": args.cuda_visible_devices, + "fill_token_id": args.fill_token_id, + "prompt_suffix": args.prompt_suffix, + "extra_server_args": args.extra_server_arg, + } + + +def validate_resume(metadata: dict | None, signature: dict, path: Path) -> None: + if metadata is None: + raise SystemExit(f"existing result file has no metadata record: {path}") + mismatches = [ + key for key, value in signature.items() if metadata.get(key) != value + ] + if mismatches: + details = ", ".join( + f"{key}={metadata.get(key)!r} (requested {signature[key]!r})" + for key in mismatches + ) + raise SystemExit( + f"cannot resume {path} with different settings: {details}" + ) + + +def git_revision() -> str: + result = subprocess.run( + ["git", "rev-parse", "--short=12", "HEAD"], + cwd=ROOT, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + check=False, + ) + return result.stdout.strip() or "unknown" + + +def probe_pool( + args: argparse.Namespace, + context_capacity: int, + baseline_used_mib: int, + logs_dir: Path, + results_path: Path, +) -> int: + print( + f"[{context_capacity}] probing VRAM with " + f"{args.probe_pool_mib} MiB pool", + flush=True, + ) + server: Server | None = None + try: + server = Server( + args, + context_capacity, + args.probe_pool_mib, + logs_dir / f"context-{context_capacity}-probe.log", + ) + prepare_server(args, server) + memory = query_gpu_memory(args.nvidia_smi, args.gpu_index) + selected_pool = estimate_pool_mib( + args.probe_pool_mib, + memory.free_mib, + args.pool_step_mib, + args.max_pool_mib, + ) + append_jsonl( + results_path, + { + "type": "pool_probe", + "context_capacity": context_capacity, + "probe_pool_mib": args.probe_pool_mib, + "selected_pool_mib": selected_pool, + "vram_total_mib": memory.total_mib, + "vram_used_mib": memory.used_mib, + "vram_free_mib": memory.free_mib, + }, + ) + return selected_pool + finally: + if server is not None: + server.stop() + wait_for_release(args, baseline_used_mib) +def parse_kv_stream_trace(log_path: Path) -> dict: + text = log_path.read_text(errors="replace") + samples = [] + for match in KV_STREAM_TRACE_RE.finditer(text): + active_tokens = int(match.group(1)) + resident_pages = int(match.group(2)) + samples.append( + { + "active_tokens": active_tokens, + "active_pages": (active_tokens + 255) // 256, + "resident_pages": resident_pages, + "ring_slots": int(match.group(3)), + } + ) + streamed = [ + sample + for sample in samples + if sample["active_pages"] > sample["resident_pages"] + ] + return { + "streaming_active": bool(streamed), + "stream_first_active_tokens": ( + min(sample["active_tokens"] for sample in streamed) + if streamed + else None + ), + "stream_trace_samples": len(samples), + "stream_max_active_pages": max( + (sample["active_pages"] for sample in samples), default=None + ), + "stream_min_resident_pages": min( + (sample["resident_pages"] for sample in samples), default=None + ), + "stream_max_ring_slots": max( + (sample["ring_slots"] for sample in samples), default=None + ), + "stream_repartitions": text.count("adaptive KV partition:"), + } + + + + +def run_measurement( + args: argparse.Namespace, + context_capacity: int, + pool_mib: int, + baseline_used_mib: int, + logs_dir: Path, +) -> dict: + prompt_tokens = context_capacity - args.decode_tokens + server: Server | None = None + try: + server = Server( + args, + context_capacity, + pool_mib, + logs_dir / f"context-{context_capacity}-pool-{pool_mib}.log", + ) + suffix, fill_token_id = prepare_server(args, server) + prefix_count = prompt_tokens - len(suffix) + if prefix_count < 0: + raise RuntimeError("context capacity is too small for the prompt") + prompt = [fill_token_id] * prefix_count + suffix + before = query_gpu_memory(args.nvidia_smi, args.gpu_index) + started = time.monotonic() + response = http_json( + server.url("/completion"), + { + "prompt": prompt, + "n_predict": args.decode_tokens, + "ignore_eos": True, + "cache_prompt": False, + "temperature": 0, + "seed": 1, + "reasoning_format": "none", + "response_fields": ["timings"], + }, + args.request_timeout, + ) + timings = response.get("timings") or {} + if timings.get("predicted_n") != args.decode_tokens: + raise RuntimeError( + f"incomplete decode: expected {args.decode_tokens}, " + f"received {timings.get('predicted_n')}" + ) + after = query_gpu_memory(args.nvidia_smi, args.gpu_index) + measurement = { + "type": "measurement", + "status": "ok", + "context_capacity": context_capacity, + "prompt_tokens": prompt_tokens, + "decode_tokens": args.decode_tokens, + "pool_mib": pool_mib, + "fill_token_id": fill_token_id, + "prompt_ms": timings.get("prompt_ms"), + "prefill_tps": timings.get("prompt_per_second"), + "predicted_ms": timings.get("predicted_ms"), + "decode_tps": timings.get("predicted_per_second"), + "wall_seconds": time.monotonic() - started, + "vram_before_mib": before.used_mib, + "vram_after_mib": after.used_mib, + "vram_free_after_mib": after.free_mib, + } + if args.trace_kv_stream: + measurement.update(parse_kv_stream_trace(server.log_path)) + return measurement + finally: + if server is not None: + server.stop() + wait_for_release(args, baseline_used_mib) + + +def benchmark_with_backoff( + args: argparse.Namespace, + context_capacity: int, + selected_pool_mib: int, + baseline_used_mib: int, + logs_dir: Path, + results_path: Path, +) -> dict: + pool_mib = selected_pool_mib + last_error = "benchmark did not run" + for attempt in range(args.pool_retries + 1): + print( + f"[{context_capacity}] benchmark attempt {attempt + 1}: " + f"pool={pool_mib} MiB", + flush=True, + ) + try: + return run_measurement( + args, + context_capacity, + pool_mib, + baseline_used_mib, + logs_dir, + ) + except Exception as exc: + last_error = str(exc) + append_jsonl( + results_path, + { + "type": "measurement_attempt", + "status": "failed", + "context_capacity": context_capacity, + "pool_mib": pool_mib, + "attempt": attempt + 1, + "error": last_error[-4000:], + }, + ) + next_pool = ( + (pool_mib - args.pool_backoff_mib) + // args.pool_step_mib + * args.pool_step_mib + ) + if next_pool < args.probe_pool_mib: + break + pool_mib = next_pool + raise RuntimeError( + f"context {context_capacity} failed after pool backoff: {last_error}" + ) + + +def write_csv(path: Path, rows: dict[int, dict]) -> None: + fields = [ + "context_capacity", + "prompt_tokens", + "decode_tokens", + "pool_mib", + "prefill_tps", + "decode_tps", + "prompt_ms", + "predicted_ms", + "wall_seconds", + "vram_before_mib", + "vram_after_mib", + "vram_free_after_mib", + "streaming_active", + "stream_first_active_tokens", + "stream_trace_samples", + "stream_max_active_pages", + "stream_min_resident_pages", + "stream_max_ring_slots", + "stream_repartitions", + ] + with path.open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + for context in sorted(rows): + writer.writerow({field: rows[context].get(field) for field in fields}) + + +def require_matplotlib(): + try: + import matplotlib.pyplot as plt + except ModuleNotFoundError as exc: + raise SystemExit( + "Matplotlib is required. Install it with: " + "python3 -m pip install matplotlib" + ) from exc + return plt + + +def plot_results(output_dir: Path, rows: dict[int, dict], plt) -> None: + if not rows: + return + contexts = sorted(rows) + x = [context / 1024 for context in contexts] + + fig, (decode_ax, pool_ax) = plt.subplots( + 2, + 1, + figsize=(12.5, 9), + height_ratios=(3, 1), + sharex=True, + constrained_layout=True, + ) + prefill_ax = decode_ax.twinx() + decode_ax.plot( + x, + [rows[context]["decode_tps"] for context in contexts], + color="#1f77b4", + marker="o", + linewidth=2.4, + label="Decode speed", + ) + prefill_ax.plot( + x, + [rows[context]["prefill_tps"] for context in contexts], + color="#ff7f0e", + marker="o", + linestyle="--", + linewidth=2.2, + label="Prefill speed", + ) + pool_ax.plot( + x, + [rows[context]["pool_mib"] for context in contexts], + color="#666666", + marker="o", + linewidth=1.8, + label="Selected KV pool", + ) + + decode_ax.set_title("Adaptive KV streaming context sweep") + decode_ax.set_ylabel("Decode speed (tokens/s)") + prefill_ax.set_ylabel("Prefill speed (tokens/s)") + pool_ax.set_xlabel("Configured context capacity (Ki tokens)") + pool_ax.set_ylabel("Pool (MiB)") + decode_ax.set_ylim(bottom=0) + prefill_ax.set_ylim(bottom=0) + pool_ax.set_ylim(bottom=0) + decode_ax.grid(True, alpha=0.25) + pool_ax.grid(True, alpha=0.25) + handles_a, labels_a = decode_ax.get_legend_handles_labels() + handles_b, labels_b = prefill_ax.get_legend_handles_labels() + decode_ax.legend(handles_a + handles_b, labels_a + labels_b, loc="best") + pool_ax.legend(loc="best") + + png_path = output_dir / "kv-stream-sweep.png" + fig.savefig(png_path, dpi=180) + fig.savefig(png_path.with_suffix(".svg")) + plt.close(fig) + + +def validate_args(args: argparse.Namespace) -> None: + if not args.model.is_file(): + raise SystemExit(f"model not found: {args.model}") + server_ok = args.server.is_file() and ( + os.name == "nt" or os.access(args.server, os.X_OK) + ) + if not server_ok: + raise SystemExit(f"server is not executable: {args.server}") + numeric_positive = ( + args.decode_tokens, + args.batch_size, + args.ubatch_size, + args.probe_pool_mib, + args.pool_step_mib, + args.pool_backoff_mib, + args.startup_timeout, + args.request_timeout, + args.release_timeout, + ) + if any(value <= 0 for value in numeric_positive): + raise SystemExit("decode, batch, pool, and timeout settings must be positive") + if args.min_context <= 0 or args.context_step <= 0: + raise SystemExit("minimum context and context step must be positive") + if args.min_context > args.max_context: + raise SystemExit("minimum context must not exceed maximum context") + if args.ubatch_size > args.batch_size: + raise SystemExit("ubatch size must not exceed batch size") + if args.fixed_pool_mib is not None and args.fixed_pool_mib <= 0: + raise SystemExit("fixed pool must be positive") + if ( + args.pool_retries < 0 + or args.release_slack_mib < 0 + ): + raise SystemExit("retry count and release slack must not be negative") + if args.max_pool_mib is not None and args.max_pool_mib < args.probe_pool_mib: + raise SystemExit("maximum pool must be at least the probe pool") + if args.gpu_index < 0 or not 1 <= args.port <= 65535: + raise SystemExit("GPU index or port is invalid") + if args.max_context < CONTEXT_STEP: + raise SystemExit("maximum context must be at least 8192") + if args.decode_tokens >= CONTEXT_STEP: + raise SystemExit("decode token count must be smaller than 8192") + if not args.prompt_suffix: + raise SystemExit("prompt suffix must not be empty") + if args.fill_token_id is not None and args.fill_token_id < 0: + raise SystemExit("fill token ID must not be negative") + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + validate_args(args) + plt = require_matplotlib() + + stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S") + if args.output_dir is None: + args.output_dir = ( + ROOT / "benchmarks/results" / f"adaptive-kv-sweep-{stamp}" + ) + args.output_dir.mkdir(parents=True, exist_ok=True) + logs_dir = args.output_dir / "logs" + logs_dir.mkdir(parents=True, exist_ok=True) + results_path = args.output_dir / "results.jsonl" + csv_path = args.output_dir / "results.csv" + + baseline = query_gpu_memory(args.nvidia_smi, args.gpu_index) + capacities = context_capacities( + args.max_context, + args.min_context, + args.context_step, + ) + signature = resume_signature(args, capacities) + if results_path.exists(): + validate_resume(load_metadata(results_path), signature, results_path) + rows = load_measurements(results_path) + + if not results_path.exists(): + append_jsonl( + results_path, + { + "type": "metadata", + "timestamp": dt.datetime.now(dt.timezone.utc).isoformat(), + "revision": git_revision(), + **signature, + "cache_type_k": args.cache_type_k, + "cache_type_v": args.cache_type_v, + "flash_attention": True, + "parallel": 1, + "baseline_vram_used_mib": baseline.used_mib, + "baseline_vram_total_mib": baseline.total_mib, + "uvm": False, + }, + ) + + print( + f"GPU baseline: {baseline.used_mib}/{baseline.total_mib} MiB used", + flush=True, + ) + print( + f"Sweep: {len(capacities)} points, {args.min_context} through " + f"{args.max_context} tokens, K={args.cache_type_k}, V={args.cache_type_v}", + flush=True, + ) + + interrupted = False + failed = False + try: + for index, context_capacity in enumerate(capacities, start=1): + if context_capacity in rows: + print( + f"[{index}/{len(capacities)}] {context_capacity}: " + "already complete", + flush=True, + ) + continue + print( + f"[{index}/{len(capacities)}] context capacity " + f"{context_capacity}", + flush=True, + ) + try: + if args.fixed_pool_mib is None: + selected_pool = probe_pool( + args, context_capacity, baseline.used_mib, + logs_dir, results_path, + ) + measurement = benchmark_with_backoff( + args, context_capacity, selected_pool, + baseline.used_mib, logs_dir, results_path, + ) + else: + selected_pool = args.fixed_pool_mib + measurement = run_measurement( + args, context_capacity, selected_pool, + baseline.used_mib, logs_dir, + ) + append_jsonl(results_path, measurement) + rows[context_capacity] = measurement + write_csv(csv_path, rows) + plot_results(args.output_dir, rows, plt) + except Exception as exc: + failed = True + append_jsonl( + results_path, + { + "type": "point_failure", + "status": "failed", + "context_capacity": context_capacity, + "error": str(exc)[-4000:], + }, + ) + print( + "Stopping at the first failed context. Re-run with the " + "same --output-dir to resume after correcting the issue.", + file=sys.stderr, + flush=True, + ) + break + except KeyboardInterrupt: + interrupted = True + print("Sweep interrupted; completed points remain resumable.", flush=True) + finally: + if rows: + write_csv(csv_path, rows) + plot_results(args.output_dir, rows, plt) + + print(f"JSONL: {results_path}", flush=True) + if rows: + print(f"CSV: {csv_path}", flush=True) + print(f"Plot: {args.output_dir / 'kv-stream-sweep.png'}", flush=True) + if interrupted: + return 130 + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_kv_type_matrix.py b/benchmarks/benchmark_kv_type_matrix.py new file mode 100755 index 000000000000..c96a088112c7 --- /dev/null +++ b/benchmarks/benchmark_kv_type_matrix.py @@ -0,0 +1,387 @@ +#!/usr/bin/env python3 +"""Benchmark and plot a matrix of adaptive-streaming K/V cache types.""" + +from __future__ import annotations + +import argparse +import csv +import datetime as dt +from pathlib import Path +import subprocess +import sys + + +ROOT = Path(__file__).resolve().parents[1] +RUNNER = Path(__file__).with_name("benchmark_kv_stream.py") +DEFAULT_TYPES = ("bf16", "q8_0", "q4_0") + + +def parse_token_count(value: str) -> int: + text = value.strip().lower() + for suffix, multiplier in ( + ("kib", 1024), + ("ki", 1024), + ("k", 1024), + ): + if text.endswith(suffix): + text = text[:-len(suffix)] + break + else: + multiplier = 1 + try: + count = int(text) * multiplier + except ValueError as exc: + raise argparse.ArgumentTypeError(f"invalid token count: {value}") from exc + if count <= 0: + raise argparse.ArgumentTypeError("token count must be positive") + return count + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument( + "--server", + type=Path, + default=ROOT / "build-kv-cuda/bin/llama-server", + ) + parser.add_argument("--min-context", type=parse_token_count, default=40 * 1024) + parser.add_argument("--max-context", type=parse_token_count, default=160 * 1024) + parser.add_argument("--context-step", type=parse_token_count, default=8 * 1024) + parser.add_argument("--decode-tokens", type=int, default=256) + parser.add_argument( + "--pool-mib", + type=int, + help="fixed pool size; omit to probe the maximum usable pool at every point", + ) + parser.add_argument( + "--types", + default=",".join(DEFAULT_TYPES), + help="comma-separated cache types used for both K and V", + ) + parser.add_argument("--port", type=int, default=12355) + parser.add_argument("--startup-timeout", type=int, default=240) + parser.add_argument("--request-timeout", type=int, default=1800) + parser.add_argument( + "--output-dir", + type=Path, + help="result directory; reuse it to resume an interrupted matrix", + ) + parser.add_argument( + "--plot-only", + action="store_true", + help="skip measurements and regenerate outputs from completed pair CSVs", + ) + return parser.parse_args(argv) + + +def cache_types(value: str) -> list[str]: + result = [item.strip().lower() for item in value.split(",") if item.strip()] + if not result: + raise SystemExit("at least one cache type is required") + if len(set(result)) != len(result): + raise SystemExit("cache types must not contain duplicates") + return result + + +def pair_name(type_k: str, type_v: str) -> str: + return f"k-{type_k}_v-{type_v}" + + +def runner_command( + args: argparse.Namespace, + type_k: str, + type_v: str, + output_dir: Path, +) -> list[str]: + command = [ + sys.executable, + str(RUNNER), + "--model", + str(args.model), + "--server", + str(args.server), + "--min-context", + str(args.min_context), + "--max-context", + str(args.max_context), + "--context-step", + str(args.context_step), + "--decode-tokens", + str(args.decode_tokens), + "--cache-type-k", + type_k, + "--cache-type-v", + type_v, + "--trace-kv-stream", + "--port", + str(args.port), + "--startup-timeout", + str(args.startup_timeout), + "--request-timeout", + str(args.request_timeout), + "--output-dir", + str(output_dir), + ] + if args.pool_mib is not None: + command.extend(("--fixed-pool-mib", str(args.pool_mib))) + return command + + +def parse_bool(value: str | bool | None) -> bool: + if isinstance(value, bool): + return value + return str(value).strip().lower() in ("1", "true", "yes") + + +def collect_rows(output_dir: Path, types: list[str]) -> list[dict]: + rows = [] + for type_k in types: + for type_v in types: + path = output_dir / pair_name(type_k, type_v) / "results.csv" + if not path.is_file(): + continue + with path.open(newline="") as stream: + for row in csv.DictReader(stream): + row["cache_type_k"] = type_k + row["cache_type_v"] = type_v + row["pair"] = f"K={type_k.upper()}, V={type_v.upper()}" + row["streaming_active"] = parse_bool(row.get("streaming_active")) + rows.append(row) + return rows + + +def write_combined_csv(path: Path, rows: list[dict]) -> None: + fields = [ + "cache_type_k", + "cache_type_v", + "pair", + "context_capacity", + "prompt_tokens", + "decode_tokens", + "pool_mib", + "prefill_tps", + "decode_tps", + "prompt_ms", + "predicted_ms", + "wall_seconds", + "vram_before_mib", + "vram_after_mib", + "vram_free_after_mib", + "streaming_active", + "stream_first_active_tokens", + "stream_trace_samples", + "stream_max_active_pages", + "stream_min_resident_pages", + "stream_max_ring_slots", + "stream_repartitions", + ] + with path.open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + for row in sorted( + rows, + key=lambda item: ( + item["cache_type_k"], + item["cache_type_v"], + int(item["context_capacity"]), + ), + ): + writer.writerow({field: row.get(field) for field in fields}) + + +def write_onsets(path: Path, rows: list[dict], types: list[str]) -> None: + with path.open("w", newline="") as stream: + writer = csv.DictWriter( + stream, + fieldnames=( + "cache_type_k", + "cache_type_v", + "first_streaming_context", + "first_streaming_context_kib", + "first_stream_active_tokens", + ), + ) + writer.writeheader() + for type_k in types: + for type_v in types: + pair_rows = sorted( + ( + row + for row in rows + if row["cache_type_k"] == type_k + and row["cache_type_v"] == type_v + and row["streaming_active"] + ), + key=lambda row: int(row["context_capacity"]), + ) + first = pair_rows[0] if pair_rows else None + context = int(first["context_capacity"]) if first else None + writer.writerow( + { + "cache_type_k": type_k, + "cache_type_v": type_v, + "first_streaming_context": context, + "first_streaming_context_kib": ( + context / 1024 if context is not None else None + ), + "first_stream_active_tokens": ( + first.get("stream_first_active_tokens") if first else None + ), + } + ) + + +def plot_results(path: Path, rows: list[dict], types: list[str]) -> None: + import matplotlib.pyplot as plt + from matplotlib.lines import Line2D + + pairs = [(type_k, type_v) for type_k in types for type_v in types] + cmap = plt.get_cmap("tab10") + fig, (prefill_ax, decode_ax) = plt.subplots( + 2, + 1, + figsize=(15, 11), + sharex=True, + constrained_layout=True, + ) + pair_handles = [] + for index, (type_k, type_v) in enumerate(pairs): + pair_rows = sorted( + ( + row for row in rows + if row["cache_type_k"] == type_k and row["cache_type_v"] == type_v + ), + key=lambda row: int(row["context_capacity"]), + ) + if not pair_rows: + continue + color = cmap(index % 10) + label = f"K={type_k.upper()}, V={type_v.upper()}" + pair_handles.append(Line2D([0], [0], color=color, lw=2.2, label=label)) + x = [int(row["context_capacity"]) / 1024 for row in pair_rows] + for axis, metric, style in ( + (prefill_ax, "prefill_tps", "--"), + (decode_ax, "decode_tps", "-"), + ): + y = [float(row[metric]) for row in pair_rows] + axis.plot(x, y, color=color, linestyle=style, linewidth=2.0) + inactive = [ + (x_value, y_value) + for x_value, y_value, row in zip(x, y, pair_rows) + if not row["streaming_active"] + ] + active = [ + (x_value, y_value) + for x_value, y_value, row in zip(x, y, pair_rows) + if row["streaming_active"] + ] + if inactive: + axis.scatter( + *zip(*inactive), s=30, facecolors="white", + edgecolors=[color], linewidths=1.3, zorder=3, + ) + if active: + axis.scatter( + *zip(*active), s=34, facecolors=[color], + edgecolors="black", linewidths=0.55, zorder=4, + ) + onset = active[0] + axis.scatter( + [onset[0]], [onset[1]], marker="*", s=190, + facecolors=[color], edgecolors="black", + linewidths=0.8, zorder=5, + ) + + selected_types = " / ".join(cache_type.upper() for cache_type in types) + prefill_ax.set_title( + f"Adaptive KV streaming: {selected_types} cache-type matrix") + prefill_ax.set_ylabel("Prefill speed (tokens/s)") + decode_ax.set_ylabel("Decode speed (tokens/s)") + decode_ax.set_xlabel("Configured context capacity (Ki tokens)") + for axis in (prefill_ax, decode_ax): + axis.grid(True, alpha=0.25) + axis.set_ylim(bottom=0) + + state_handles = [ + Line2D( + [0], [0], marker="o", linestyle="none", markerfacecolor="white", + markeredgecolor="black", label="Fully resident", + ), + Line2D( + [0], [0], marker="o", linestyle="none", markerfacecolor="#777777", + markeredgecolor="black", label="Streaming active", + ), + Line2D( + [0], [0], marker="*", linestyle="none", markersize=13, + markerfacecolor="#777777", markeredgecolor="black", + label="First streaming point", + ), + ] + prefill_ax.legend(handles=state_handles, loc="lower left") + fig.legend( + handles=pair_handles, + loc="outside upper center", + ncol=3, + title="KV cache types", + ) + fig.savefig(path, dpi=190) + fig.savefig(path.with_suffix(".svg")) + plt.close(fig) + + +def validate_args(args: argparse.Namespace) -> None: + if not args.model.is_file(): + raise SystemExit(f"model not found: {args.model}") + if not args.server.is_file(): + raise SystemExit(f"server not found: {args.server}") + if args.min_context > args.max_context: + raise SystemExit("minimum context must not exceed maximum context") + if min(args.context_step, args.decode_tokens) <= 0: + raise SystemExit("step and decode tokens must be positive") + if args.pool_mib is not None and args.pool_mib <= 0: + raise SystemExit("fixed pool must be positive") + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + validate_args(args) + types = cache_types(args.types) + if args.output_dir is None: + stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S") + args.output_dir = ROOT / "benchmarks/results" / f"kv-type-matrix-{stamp}" + args.output_dir.mkdir(parents=True, exist_ok=True) + + pairs = [(type_k, type_v) for type_k in types for type_v in types] + if not args.plot_only: + for index, (type_k, type_v) in enumerate(pairs, start=1): + pair_dir = args.output_dir / pair_name(type_k, type_v) + print( + f"\n=== [{index}/{len(pairs)}] K={type_k}, V={type_v} ===", + flush=True, + ) + subprocess.run( + runner_command(args, type_k, type_v, pair_dir), + cwd=ROOT, + check=True, + ) + rows = collect_rows(args.output_dir, types) + write_combined_csv(args.output_dir / "kv-type-matrix.csv", rows) + write_onsets(args.output_dir / "streaming-onsets.csv", rows, types) + plot_results(args.output_dir / "kv-type-matrix.png", rows, types) + + rows = collect_rows(args.output_dir, types) + if not rows: + raise SystemExit("no completed measurements found") + write_combined_csv(args.output_dir / "kv-type-matrix.csv", rows) + write_onsets(args.output_dir / "streaming-onsets.csv", rows, types) + plot_results(args.output_dir / "kv-type-matrix.png", rows, types) + print(f"Results: {args.output_dir}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/test_benchmark_kv_stream.py b/benchmarks/test_benchmark_kv_stream.py new file mode 100644 index 000000000000..e68c73ee4b2b --- /dev/null +++ b/benchmarks/test_benchmark_kv_stream.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +"""Unit tests for the automatic adaptive KV benchmark driver.""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +from pathlib import Path +import sys +import tempfile +import unittest +from unittest import mock + + +SCRIPT = Path(__file__).with_name("benchmark_kv_stream.py") +SPEC = importlib.util.spec_from_file_location("benchmark_kv_stream", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +BENCHMARK = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = BENCHMARK +SPEC.loader.exec_module(BENCHMARK) + + +class BenchmarkKvStreamTest(unittest.TestCase): + def test_parse_token_count(self) -> None: + self.assertEqual(BENCHMARK.parse_token_count("192K"), 192 * 1024) + self.assertEqual(BENCHMARK.parse_token_count("262144"), 262144) + with self.assertRaises(argparse.ArgumentTypeError): + BENCHMARK.parse_token_count("bad") + + def test_parse_args_resolves_launched_paths(self) -> None: + model = Path("models/model.gguf") + server = Path("build/bin/llama-server") + args = BENCHMARK.parse_args( + [ + "--model", str(model), + "--server", str(server), + "--max-context", "8K", + "--batch-size", "768", + "--ubatch-size", "512", + ] + ) + self.assertEqual(args.model, model.resolve()) + self.assertEqual(args.server, server.resolve()) + self.assertEqual(args.batch_size, 768) + self.assertEqual(args.ubatch_size, 512) + + def test_validate_args_rejects_ubatch_larger_than_batch(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + model = root / "model.gguf" + server = root / "llama-server" + model.touch() + server.touch(mode=0o755) + args = BENCHMARK.parse_args( + [ + "--model", str(model), + "--server", str(server), + "--max-context", "8K", + "--batch-size", "256", + "--ubatch-size", "512", + ] + ) + with self.assertRaisesRegex(SystemExit, "must not exceed"): + BENCHMARK.validate_args(args) + + def test_context_capacities_include_non_aligned_maximum(self) -> None: + self.assertEqual( + BENCHMARK.context_capacities(20000), + [8192, 16384, 20000], + ) + + def test_context_capacities_honor_custom_start_and_step(self) -> None: + self.assertEqual( + BENCHMARK.context_capacities(163840, 40960, 8192), + list(range(40960, 163841, 8192)), + ) + + + def test_pool_estimate_uses_all_free_memory_and_rounds_down(self) -> None: + self.assertEqual(BENCHMARK.estimate_pool_mib(64, 3500, 32), 3552) + self.assertEqual( + BENCHMARK.estimate_pool_mib(64, 3500, 32, max_pool_mib=2048), + 2048, + ) + + def test_clean_server_env_removes_memory_policy_overrides(self) -> None: + inherited = { + "GGML_CUDA_ENABLE_UNIFIED_MEMORY": "1", + "GGML_CUDA_PREFER_MODEL_WEIGHTS": "1", + "GGML_CUDA_KV_STREAM_FIXED_RING_SLOTS": "8", + "KEEP_ME": "yes", + } + with mock.patch.dict(os.environ, inherited, clear=True): + env = BENCHMARK.clean_server_env("2") + self.assertNotIn("GGML_CUDA_ENABLE_UNIFIED_MEMORY", env) + self.assertNotIn("GGML_CUDA_PREFER_MODEL_WEIGHTS", env) + self.assertNotIn("GGML_CUDA_KV_STREAM_FIXED_RING_SLOTS", env) + self.assertEqual(env["KEEP_ME"], "yes") + self.assertEqual(env["CUDA_VISIBLE_DEVICES"], "2") + traced = BENCHMARK.clean_server_env(None, trace_kv_stream=True) + self.assertEqual(traced["LLAMA_KV_STREAM_TRACE"], "1") + + def test_server_command_uses_tested_configuration(self) -> None: + args = argparse.Namespace( + server=Path("/tmp/llama-server"), + model=Path("/tmp/model.gguf"), + port=12355, + extra_server_arg=["--verbosity", "3"], + cache_type_k="bf16", + cache_type_v="q8_0", + batch_size=768, + ubatch_size=512, + ) + command = BENCHMARK.server_command(args, 131072, 2304) + self.assertEqual(command[0], "/tmp/llama-server") + self.assertIn("131072", command) + self.assertIn("2304", command) + self.assertEqual(command[command.index("-ctk") + 1], "bf16") + self.assertEqual(command[command.index("-ctv") + 1], "q8_0") + self.assertEqual(command[command.index("-b") + 1], "768") + self.assertEqual(command[command.index("-ub") + 1], "512") + self.assertEqual(command[command.index("-np") + 1], "1") + self.assertEqual(command[-2:], ["--verbosity", "3"]) + + def test_trace_parser_marks_only_pages_beyond_resident_partition(self) -> None: + log = ( + "I kv_stream_adapt: active 65536, resident 256, ring 32, " + "samples 1, misses 0, copy busy 0.0%, peak 1\n" + "W kv_stream_adapt: adaptive KV partition: resident pages/layer " + "256 -> 248, ring slots 32 -> 160, miss 50.0%, copy busy 25.0%\n" + "I kv_stream_adapt: active 65792, resident 248, ring 160, " + "samples 2, misses 1, copy busy 25.0%, peak 10\n" + ) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "server.log" + path.write_text(log) + parsed = BENCHMARK.parse_kv_stream_trace(path) + self.assertTrue(parsed["streaming_active"]) + self.assertEqual(parsed["stream_first_active_tokens"], 65792) + self.assertEqual(parsed["stream_trace_samples"], 2) + self.assertEqual(parsed["stream_repartitions"], 1) + self.assertEqual(parsed["stream_max_ring_slots"], 160) + + + def test_resume_rejects_changed_settings(self) -> None: + + signature = {"model": "/tmp/model.gguf", "max_context": 16384} + BENCHMARK.validate_resume(signature.copy(), signature, Path("results.jsonl")) + with self.assertRaisesRegex(SystemExit, "different settings"): + BENCHMARK.validate_resume( + {"model": "/tmp/other.gguf", "max_context": 16384}, + signature, + Path("results.jsonl"), + ) + + def test_csv_and_plot_accept_partial_sweep(self) -> None: + rows = { + 8192: { + "context_capacity": 8192, + "prompt_tokens": 7936, + "decode_tokens": 256, + "pool_mib": 3552, + "prefill_tps": 1400.0, + "decode_tps": 50.0, + }, + 16384: { + "context_capacity": 16384, + "prompt_tokens": 16128, + "decode_tokens": 256, + "pool_mib": 3520, + "prefill_tps": 1300.0, + "decode_tps": 45.0, + }, + } + try: + plt = BENCHMARK.require_matplotlib() + except SystemExit: + self.skipTest("Matplotlib is not installed") + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) + BENCHMARK.write_csv(output / "results.csv", rows) + BENCHMARK.plot_results(output, rows, plt) + self.assertTrue((output / "results.csv").is_file()) + self.assertTrue((output / "kv-stream-sweep.png").is_file()) + self.assertTrue((output / "kv-stream-sweep.svg").is_file()) + + +if __name__ == "__main__": + unittest.main() diff --git a/benchmarks/test_kv_stream_serial_server.py b/benchmarks/test_kv_stream_serial_server.py new file mode 100644 index 000000000000..d317535369c3 --- /dev/null +++ b/benchmarks/test_kv_stream_serial_server.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 + +import argparse +import json +from pathlib import Path +import signal +import subprocess +import time +import urllib.request + + +ROOT = Path(__file__).resolve().parents[1] +DEFAULT_MODEL = Path( + "/home/raymond/LLM/llama-cache/models--unsloth--Qwen3.8-27B-GGUF/" + "blobs/8c2a45ff85e7674ca185ec8eb6cdeab0e617ed9d8018caed0b64380eb2a67a5e" +) + + +def request(port: int, path: str, payload: dict | None, timeout: int = 600) -> dict: + data = None if payload is None else json.dumps(payload).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}{path}", data=data, + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as response: + return json.load(response) + + +class Server: + def __init__(self, model: Path, port: int, cache_mib: int, log: Path): + command = [ + str(ROOT / "build-kv-cuda/bin/llama-server"), + "-m", str(model), "--host", "127.0.0.1", "--port", str(port), + "--ctx-size", "8448", "-fa", "on", "-ctk", "q8_0", "-ctv", "q4_0", + "-ngl", "all", "-b", "256", "-ub", "256", "-np", "1", + "--no-mmproj", "--no-warmup", "--reasoning-format", "none", + "--kv-stream-stage-mib", "128", "--cache-ram", str(cache_mib), + ] + self.log_file = log.open("wb") + self.process = subprocess.Popen( + command, stdout=self.log_file, stderr=subprocess.STDOUT) + deadline = time.monotonic() + 180 + while time.monotonic() < deadline: + if self.process.poll() is not None: + self.log_file.close() + raise RuntimeError(f"server exited with {self.process.returncode}") + try: + if request(port, "/health", None, 2).get("status") == "ok": + self.port = port + return + except Exception: + time.sleep(0.25) + self.process.send_signal(signal.SIGINT) + self.process.wait(timeout=15) + self.log_file.close() + raise RuntimeError("server did not become ready") + + def stop(self): + if self.process.poll() is None: + self.process.send_signal(signal.SIGINT) + self.process.wait(timeout=15) + self.log_file.close() + + +def completion(server: Server, prompt: list[int], cache_prompt: bool) -> dict: + return request(server.port, "/completion", { + "prompt": prompt, + "n_predict": 16, + "ignore_eos": True, + "cache_prompt": cache_prompt, + "temperature": 0, + "seed": 1, + "reasoning_format": "none", + }) + + +def patterned(size: int, tokens: tuple[int, ...]) -> list[int]: + return [tokens[i % len(tokens)] for i in range(size)] + + +def run_serial(model: Path, port: int, output: Path): + server = Server(model, port, 0, output / "serial.log") + try: + short = patterned(1024, (23066, 1000, 2000)) + medium = patterned(4096, (23066, 3000, 4000, 5000)) + streamed = patterned(6144, (23066, 6000, 7000, 8000, 9000)) + short_first = completion(server, short, False) + medium_first = completion(server, medium, False) + streamed_result = completion(server, streamed, False) + short_second = completion(server, short, False) + medium_second = completion(server, medium, False) + serial_changed = ( + short_second["content"] != short_first["content"] or + medium_second["content"] != medium_first["content"]) + streamed_invalid = not streamed_result["content"].strip("/") + if serial_changed or streamed_invalid: + details = { + "short_first": short_first["content"], + "short_second": short_second["content"], + "medium_first": medium_first["content"], + "streamed": streamed_result["content"], + "medium_second": medium_second["content"], + } + raise RuntimeError(f"serial unrelated-prefill output changed: {json.dumps(details)}") + print("serial unrelated-prefill test: PASS", flush=True) + finally: + server.stop() + + +def run_prompt_cache(model: Path, port: int, output: Path): + server = Server(model, port, 1536, output / "prompt-cache.log") + try: + cached = patterned(4096, (23066, 1100, 2100, 3100)) + unrelated = patterned(2048, (23066, 4100, 5100, 6100)) + expected = completion(server, cached, True)["content"] + completion(server, unrelated, True) + restored = completion(server, cached, True) + if restored["content"] != expected: + raise RuntimeError("prompt-cache restore output changed") + if restored["timings"].get("cache_n", 0) == 0: + raise RuntimeError("prompt-cache restore did not reuse cached tokens") + print(f"prompt-cache restore test: PASS (cache_n={restored['timings']['cache_n']})", flush=True) + finally: + server.stop() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, default=DEFAULT_MODEL) + parser.add_argument("--port", type=int, default=12358) + parser.add_argument("--output", type=Path, default=ROOT / "benchmarks/results/serial-server") + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=True) + run_serial(args.model, args.port, args.output) + run_prompt_cache(args.model, args.port, args.output) + + +if __name__ == "__main__": + main() diff --git a/common/arg.cpp b/common/arg.cpp index 71df1ec5aa89..16c207bfdaae 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2420,6 +2420,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.no_kv_offload = !value; } ).set_env("LLAMA_ARG_KV_OFFLOAD")); + add_opt(common_arg( + {"--kv-stream-stage-mib"}, "N", + string_format("block-streaming KV resident + staging pool in MiB; 0 disables it (default: %u)", params.kv_stream_stage_mib), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("KV stream stage size must be non-negative"); + } + params.kv_stream_stage_mib = value; + } + ).set_env("LLAMA_ARG_KV_STREAM_STAGE_MIB")); add_opt(common_arg( {"--repack"}, {"-nr", "--no-repack"}, diff --git a/common/common.cpp b/common/common.cpp index d162a38800e0..8150b709893b 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1751,6 +1751,7 @@ 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; + cparams.kv_stream_stage_mib = params.kv_stream_stage_mib; return cparams; } diff --git a/common/common.h b/common/common.h index 6ff382534b98..15200bf4ae50 100644 --- a/common/common.h +++ b/common/common.h @@ -587,6 +587,7 @@ struct common_params { ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V + uint32_t kv_stream_stage_mib = 0; // block-streaming staging budget, 0 = disabled [EXPERIMENTAL] common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO; diff --git a/common/speculative.cpp b/common/speculative.cpp index 851a47b9a584..c3d140ed4d2d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2480,6 +2480,9 @@ common_params common_base_params_to_speculative(const common_params & params) { result.cache_type_k = params_spec.cache_type_k; result.cache_type_v = params_spec.cache_type_v; + // The first block-streaming implementation owns only the target cache. + // MTP keeps its ordinary cache until both contexts can share one pool. + result.kv_stream_stage_mib = 0; result.n_outputs_max = params.n_parallel; result.n_outputs_max_per_seq = 1; diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index 1cd81eeaebcd..a7b43c06741d 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -33,6 +33,106 @@ GGML_BACKEND_API bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backen // pinned host buffer for use with the CPU backend for faster copies between CPU and GPU GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void); +// Per-context resources for block-granular KV streaming. The buffer type owns +// authoritative pinned-host tensor storage; the runtime owns CUDA staging. +typedef struct ggml_backend_cuda_kv_stream_runtime * ggml_backend_cuda_kv_stream_runtime_t; + +struct ggml_backend_cuda_kv_stream_params { + int device; + size_t stage_bytes; + uint32_t stage_slots; + size_t pool_bytes; + size_t conversion_bytes; + uint32_t resident_layer_count; + uint32_t page_tokens; + uint32_t decode_span_pages; +}; + +struct ggml_backend_cuda_kv_stream_type_capabilities { + bool classified; + bool storage; + bool online_write; + bool decode_f16; + bool direct_attention; + bool requires_initialization; + bool requires_importance_matrix; + bool auxiliary; +}; + +enum ggml_backend_cuda_kv_stream_attention_mode { + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_UNSUPPORTED = 0, + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT = 1, + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16 = 2, +}; + +GGML_BACKEND_API struct ggml_backend_cuda_kv_stream_type_capabilities +ggml_backend_cuda_kv_stream_get_type_capabilities(enum ggml_type type); +GGML_BACKEND_API enum ggml_backend_cuda_kv_stream_attention_mode +ggml_backend_cuda_kv_stream_get_attention_mode(enum ggml_type type_k, enum ggml_type type_v); + +struct ggml_backend_cuda_kv_stream_stats { + uint64_t resident_hits; + uint64_t resident_misses; + uint64_t streamed_pages; + uint64_t host_to_device_bytes; + uint64_t resident_attention_spans; + uint64_t resident_pages_attended; + uint64_t streamed_attention_spans; + uint64_t streamed_pages_attended; + uint64_t mma_prefill_attention_spans; + uint64_t asynchronous_page_uploads; + uint64_t host_to_device_copy_commands; + uint64_t compute_stream_waits; + uint64_t stage_slot_reuses; + uint64_t cross_layer_prefetches; + uint64_t deadline_samples; + uint64_t deadline_misses; + uint32_t ring_peak_occupancy; + uint64_t staged_set_rows; + uint64_t staged_set_rows_bytes; +}; + +GGML_BACKEND_API ggml_backend_cuda_kv_stream_runtime_t ggml_backend_cuda_kv_stream_runtime_new( + struct ggml_backend_cuda_kv_stream_params params); +GGML_BACKEND_API void ggml_backend_cuda_kv_stream_runtime_free( + ggml_backend_cuda_kv_stream_runtime_t runtime); +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_kv_stream_buffer_type( + ggml_backend_cuda_kv_stream_runtime_t runtime); +GGML_BACKEND_API size_t ggml_backend_cuda_kv_stream_stage_bytes( + ggml_backend_cuda_kv_stream_runtime_t runtime); +GGML_BACKEND_API uint32_t ggml_backend_cuda_kv_stream_stage_slots( + ggml_backend_cuda_kv_stream_runtime_t runtime); +GGML_BACKEND_API uint32_t ggml_backend_cuda_kv_stream_resident_pages_per_layer( + ggml_backend_cuda_kv_stream_runtime_t runtime); +GGML_BACKEND_API bool ggml_backend_cuda_kv_stream_reconfigure( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t active_pages_per_layer, + uint32_t stage_slots); +GGML_BACKEND_API bool ggml_backend_cuda_kv_stream_repartition( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t stage_slots); +GGML_BACKEND_API bool ggml_backend_cuda_kv_stream_set_decode_layout( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t active_pages_per_layer); +GGML_BACKEND_API bool ggml_backend_cuda_kv_stream_mark_dirty_rows( + ggml_backend_cuda_kv_stream_runtime_t runtime, + const int64_t * rows, + size_t count); +GGML_BACKEND_API struct ggml_backend_cuda_kv_stream_stats ggml_backend_cuda_kv_stream_get_stats( + ggml_backend_cuda_kv_stream_runtime_t runtime); +GGML_BACKEND_API bool ggml_backend_cuda_kv_stream_stage_upload( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t slot, + size_t offset, + const void * source, + size_t size); +GGML_BACKEND_API bool ggml_backend_cuda_kv_stream_stage_download( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t slot, + size_t offset, + void * destination, + size_t size); + GGML_BACKEND_API int ggml_backend_cuda_get_device_count(void); GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size); GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 9918c03947c7..5b1c1eac329b 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1252,6 +1252,7 @@ struct ggml_cuda_graph { void * node_src_data_ptrs[GGML_MAX_SRC]; int64_t node_src_ne[GGML_MAX_SRC][GGML_MAX_DIMS]; size_t node_src_nb[GGML_MAX_SRC][GGML_MAX_DIMS]; + uint64_t kv_stream_generation; }; std::vector node_props; diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 7442bc22af20..2221c630d051 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -976,7 +976,8 @@ template void launch_fattn( ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared, const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse, - const int warp_size = WARP_SIZE + const int warp_size = WARP_SIZE, + float * partial_dst = nullptr, float2 * partial_meta = nullptr ) { constexpr int ncols = ncols1 * ncols2; @@ -990,6 +991,9 @@ void launch_fattn( const ggml_tensor * sinks = dst->src[4]; ggml_tensor * KQV = dst; + const bool output_partial = partial_dst != nullptr; + GGML_ASSERT(output_partial == (partial_meta != nullptr)); + GGML_ASSERT(!output_partial || !stream_k); GGML_ASSERT(Q->type == GGML_TYPE_F32); GGML_ASSERT(KQV->type == GGML_TYPE_F32); @@ -1165,14 +1169,15 @@ void launch_fattn( } } else { // parallel_blocks must not be larger than what the tensor size allows: - parallel_blocks = std::min(parallel_blocks, ntiles_KV); + parallel_blocks = output_partial ? 1 : std::min(parallel_blocks, ntiles_KV); // If ntiles_total % blocks_per_wave != 0 then some efficiency is lost due to tail effects. // Test whether parallel_blocks can be set to a higher value for better efficiency. const int blocks_per_wave = nsm * max_blocks_per_sm; int nwaves_best = 0; int efficiency_percent_best = 0; - for (int parallel_blocks_test = parallel_blocks; parallel_blocks_test <= ntiles_KV; ++parallel_blocks_test) { + for (int parallel_blocks_test = parallel_blocks; + !output_partial && parallel_blocks_test <= ntiles_KV; ++parallel_blocks_test) { const int nblocks_total = ntiles_dst * parallel_blocks_test; const int nwaves = (nblocks_total + blocks_per_wave - 1) / blocks_per_wave; const int efficiency_percent = 100 * nblocks_total / (nwaves*blocks_per_wave); @@ -1189,9 +1194,19 @@ void launch_fattn( } } - blocks_num.x = ntiles_x; - blocks_num.y = parallel_blocks; - blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3]; + if (output_partial) { + // MMA kernels flatten Q tiles, GQA groups, KV heads, and sequences + // into blockIdx.x. A multidimensional grid would duplicate every + // tile once per KV head. One block per complete tile also avoids + // fixups while preserving exact partial numerator/meta output. + blocks_num.x = ntiles_dst; + blocks_num.y = 1; + blocks_num.z = 1; + } else { + blocks_num.x = ntiles_x; + blocks_num.y = parallel_blocks; + blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3]; + } if (parallel_blocks > 1) { dst_tmp.alloc(parallel_blocks*ggml_nelements(KQV)); @@ -1230,7 +1245,8 @@ void launch_fattn( mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, KV_max.ptr, - !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, + output_partial ? partial_dst : (!stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data), + output_partial ? partial_meta : dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], K->ne[0], n_kv, K->ne[2], K->ne[3], nb11, nb12, nb13, diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 126a4c4529b9..63843b9d3dc4 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1154,7 +1154,7 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) -template +template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, @@ -1727,10 +1727,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( dstk_val.y += dstk_val_add.y*KQ_crs; } - if (!needs_fixup && !is_fixup) { - const float KQ_rowsum_j = meta_j[1]; - dstk_val.x /= KQ_rowsum_j; - dstk_val.y /= KQ_rowsum_j; + if constexpr (!output_partial) { + if (!needs_fixup && !is_fixup) { + const float KQ_rowsum_j = meta_j[1]; + dstk_val.x /= KQ_rowsum_j; + dstk_val.y /= KQ_rowsum_j; + } + } else if (k00 == 0 && k == 0) { + const int row = (jt*ncols1 + j_dst)*ne02 + c_dst; + dstk_fixup[row] = make_float2(meta_j[0], meta_j[1]); } if (is_fixup) { @@ -1761,7 +1766,7 @@ static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_u (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16); } -template +template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, @@ -1884,6 +1889,8 @@ static __global__ void flash_attn_ext_f16( const half * mask_h = ncols2 == 1 && !mask ? nullptr : (const half *) (mask + nb33*(sequence % ne33)); float2 * dstk = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2); + float2 * dst_meta_tile = output_partial ? + dst_meta + sequence*ne01.z*ne02 + zt_Q : dst_meta; const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; @@ -1897,13 +1904,13 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta_tile, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta_tile, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } @@ -1931,6 +1938,8 @@ static __global__ void flash_attn_ext_f16( const half * mask_h = ncols2 == 1 && !mask ? nullptr : (const half *) (mask + nb33*(sequence % ne33)); float2 * dstk = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2); + float2 * dst_meta_tile = output_partial ? + dst_meta + sequence*ne01.z*ne02 + zt_Q : dst_meta; const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; @@ -1944,8 +1953,8 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta_tile, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, @@ -1963,8 +1972,10 @@ static __global__ void flash_attn_ext_f16( bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -template -void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +template +static void ggml_cuda_flash_attn_ext_mma_f16_case_impl( + ggml_backend_cuda_context & ctx, ggml_tensor * dst, + float * partial_dst, float2 * partial_meta) { const ggml_tensor * KQV = dst; const int id = ggml_cuda_get_device(); const int cc = ggml_cuda_info().devices[id].cc; @@ -2013,10 +2024,10 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) { + if constexpr (!output_partial && ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) { if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) { constexpr bool use_sparse_kernel = true; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; use_sparse = true; static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -2026,7 +2037,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml } } else { constexpr bool use_sparse_kernel = false; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; if (!shared_memory_limit_raised[id]) { @@ -2038,7 +2049,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) { constexpr bool use_sparse_kernel = false; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -2051,7 +2062,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml } else { constexpr bool use_logit_softcap = true; constexpr bool use_sparse_kernel = false; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -2063,7 +2074,23 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml } launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, + !output_partial, use_sparse, warp_size_host, partial_dst, partial_meta); +} + +template +void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_flash_attn_ext_mma_f16_case_impl( + ctx, dst, nullptr, nullptr); +} + +template +void ggml_cuda_flash_attn_ext_mma_f16_partial_case( + ggml_backend_cuda_context & ctx, ggml_tensor * dst, + float * partial_dst, float2 * partial_meta) { + GGML_ASSERT(partial_dst != nullptr && partial_meta != nullptr); + ggml_cuda_flash_attn_ext_mma_f16_case_impl( + ctx, dst, partial_dst, partial_meta); } diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 519b36b9ff49..4a097296fe56 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -530,6 +530,78 @@ static __global__ void flash_attn_ext_vec( #pragma clang diagnostic pop #endif // __clang__ +template +void ggml_cuda_flash_attn_ext_vec_partial_case_impl( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + float * partial_dst, + float2 * partial_meta, + int nparts) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + + constexpr int ncols = 1; + constexpr int nthreads = 128; + constexpr int nwarps = nthreads/WARP_SIZE; + fattn_kernel_t kernel = flash_attn_ext_vec; + + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (logit_softcap != 0.0f) { + scale /= logit_softcap; + } + + const uint32_t n_head = Q->ne[2]; + const uint32_t n_head_log2 = 1u << uint32_t(floorf(log2f(float(n_head)))); + const float m0 = powf(2.0f, -(max_bias )/n_head_log2); + const float m1 = powf(2.0f, -(max_bias/2.0f)/n_head_log2); + const uint3 ne01 = init_fastdiv_values(Q->ne[1]); + + const dim3 blocks(Q->ne[1], nparts, Q->ne[2]*Q->ne[3]); + const dim3 threads(WARP_SIZE, nwarps, 1); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(kernel, launch_params, + (const char *) Q->data, + (const char *) K->data, + (const char *) V->data, + mask ? (const char *) mask->data : nullptr, + nullptr, + nullptr, + partial_dst, + partial_meta, + scale, max_bias, m0, m1, n_head_log2, logit_softcap, + Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], + K->ne[0], K->ne[1], K->ne[2], K->ne[3], K->nb[1], K->nb[2], K->nb[3], + V->nb[1], V->nb[2], V->nb[3], + mask ? mask->ne[1] : 0, mask ? mask->ne[2] : 0, mask ? mask->ne[3] : 0, + mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0); + CUDA_CHECK(cudaGetLastError()); +} + +template +void ggml_cuda_flash_attn_ext_vec_partial_case( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + float * partial_dst, + float2 * partial_meta, + int nparts) { + float logit_softcap = 0.0f; + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (logit_softcap == 0.0f) { + ggml_cuda_flash_attn_ext_vec_partial_case_impl< + D, type_K, type_V, false>(ctx, dst, partial_dst, partial_meta, nparts); + } else { + ggml_cuda_flash_attn_ext_vec_partial_case_impl< + D, type_K, type_V, true>(ctx, dst, partial_dst, partial_meta, nparts); + } +} + template void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; @@ -573,18 +645,28 @@ void ggml_cuda_flash_attn_ext_vec_case(ggml_backend_cuda_context & ctx, ggml_ten } } -#define DECL_FATTN_VEC_CASE(D, type_K, type_V) \ - template void ggml_cuda_flash_attn_ext_vec_case \ - (ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ - -#define EXTERN_DECL_FATTN_VEC_CASES(D, type_K) \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_F16); \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_0); \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_1); \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_0); \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_1); \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q8_0); \ - extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_BF16); \ +#define DECL_FATTN_VEC_CASE(D, type_K, type_V) \ + template void ggml_cuda_flash_attn_ext_vec_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst); \ + template void ggml_cuda_flash_attn_ext_vec_partial_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst, \ + float * partial_dst, float2 * partial_meta, int nparts) + +#define EXTERN_DECL_FATTN_VEC_CASE(D, type_K, type_V) \ + extern template void ggml_cuda_flash_attn_ext_vec_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst); \ + extern template void ggml_cuda_flash_attn_ext_vec_partial_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst, \ + float * partial_dst, float2 * partial_meta, int nparts) + +#define EXTERN_DECL_FATTN_VEC_CASES(D, type_K) \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_F16); \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_0); \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_1); \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_0); \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_1); \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q8_0); \ + EXTERN_DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_BF16); EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_0) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ae217fbd9df1..294ec9e64cac 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -1,9 +1,17 @@ #include "common.cuh" +#include "convert.cuh" #include "fattn-common.cuh" #include "fattn-mma-f16.cuh" #include "fattn-tile.cuh" #include "fattn-vec.cuh" #include "fattn.cuh" +#include "kv-stream-span-tuner.h" + +#include +#include +#include +#include +#include #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) __launch_bounds__(256, 1) @@ -128,6 +136,2139 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context K->ne[1] >= std::max(4096, 2LL*n_kv_max); #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) } +struct ggml_cuda_kv_stream_resident_cache; + +namespace { + +constexpr size_t KV_STREAM_NO_REQUEST = std::numeric_limits::max(); +constexpr uint32_t KV_STREAM_NO_LAYER = std::numeric_limits::max(); +constexpr uint32_t KV_STREAM_COPY_BATCH_PAGES = 32; +constexpr uint32_t KV_STREAM_DECODE_SPAN_PAGES = KV_STREAM_COPY_BATCH_PAGES; + +struct kv_stream_graph_request { + const char * k_data = nullptr; + const char * v_data = nullptr; + size_t k_nb1 = 0; + size_t k_nb2 = 0; + size_t v_nb1 = 0; + size_t v_nb2 = 0; + int64_t n_head_kv = 0; + int64_t token_begin = 0; + int64_t token_count = 0; + size_t k_row_bytes = 0; + size_t v_row_bytes = 0; + size_t k_head_bytes = 0; + size_t k_bytes = 0; + size_t v_offset = 0; + size_t v_head_bytes = 0; + size_t v_bytes = 0; + uint32_t layer = 0; + uint32_t slot = 0; + uint32_t ready_slot = 0; + bool mutable_tail = false; + bool eligible = false; + bool scheduled = false; + bool consumed = false; + bool deadline_sample = false; +}; + +} // namespace + +struct ggml_cuda_kv_stream_transfer_ring { + char * pool_data = nullptr; + size_t page_bytes = 0; + char * conversion_data = nullptr; + size_t conversion_bytes = 0; + uint32_t capacity_slots = 0; + uint32_t active_slots = 0; + uint32_t forced_decode_span_pages = 0; + uint32_t graph_decode_span_pages = UINT32_MAX; + cudaStream_t copy_stream = nullptr; + cudaEvent_t producer_ready = nullptr; + cudaEvent_t eval_start = nullptr; + cudaEvent_t eval_end = nullptr; + cudaEvent_t copy_sample_start = nullptr; + cudaEvent_t copy_sample_end = nullptr; + std::vector ready; + std::vector consumed; + std::vector slot_used; + std::vector slot_request; + uint32_t * ready_flags_host = nullptr; + uint32_t * ready_flags_device = nullptr; + uint64_t * deadline_counters_host = nullptr; + uint64_t * deadline_counters_device = nullptr; + + bool graph_active = false; + bool graph_decode = true; + ggml_cuda_kv_stream_span_tuner span_tuner; + uint32_t graph_layer_count = 0; + uint32_t current_layer = KV_STREAM_NO_LAYER; + size_t next_request = 0; + ggml_cuda_kv_stream_resident_cache * graph_resident_cache = nullptr; + std::vector graph_requests; + std::unordered_map graph_layer_by_k; + std::unordered_map> graph_request_by_k_page; + + uint64_t asynchronous_page_uploads = 0; + uint64_t host_to_device_copy_commands = 0; + uint64_t compute_stream_waits = 0; + uint64_t stage_slot_reuses = 0; + uint64_t cross_layer_prefetches = 0; + uint32_t current_occupancy = 0; + uint32_t ring_peak_occupancy = 0; + uint32_t current_ring_peak_occupancy = 0; + uint32_t last_ring_peak_occupancy = 0; + uint64_t current_epoch_uploads = 0; + double last_copy_engine_busy_ratio = 0.0; + uint32_t copy_sample_uploads = 0; + bool timing_pending = false; + bool timing_current = false; + bool last_graph_decode = false; + bool last_graph_bounded = false; + bool last_graph_streamed = false; + bool copy_sample_recorded = false; +}; + +ggml_cuda_kv_stream_transfer_ring * ggml_cuda_kv_stream_transfer_ring_new( + void * pool_data, size_t page_bytes, uint32_t stage_slots, + void * conversion_data, size_t conversion_bytes, + uint32_t forced_decode_span_pages) { + if (pool_data == nullptr || page_bytes == 0 || stage_slots == 0 || + (conversion_bytes != 0 && conversion_data == nullptr)) { + return nullptr; + } + + auto * ring = new ggml_cuda_kv_stream_transfer_ring; + ring->pool_data = static_cast(pool_data); + ring->page_bytes = page_bytes; + ring->conversion_data = static_cast(conversion_data); + ring->conversion_bytes = conversion_bytes; + ring->capacity_slots = stage_slots; + ring->active_slots = stage_slots; + ring->ready.resize(stage_slots, nullptr); + ring->forced_decode_span_pages = forced_decode_span_pages; + ring->consumed.resize(stage_slots, nullptr); + ring->slot_used.resize(stage_slots, 0); + ring->slot_request.resize(stage_slots, KV_STREAM_NO_REQUEST); + + auto cleanup = [&]() { + for (cudaEvent_t event : ring->ready) { + if (event != nullptr) { + (void) cudaEventDestroy(event); + } + } + for (cudaEvent_t event : ring->consumed) { + if (event != nullptr) { + (void) cudaEventDestroy(event); + } + } + if (ring->producer_ready != nullptr) { + (void) cudaEventDestroy(ring->producer_ready); + } + if (ring->eval_start != nullptr) { + (void) cudaEventDestroy(ring->eval_start); + } + if (ring->eval_end != nullptr) { + (void) cudaEventDestroy(ring->eval_end); + } + if (ring->copy_sample_start != nullptr) { + (void) cudaEventDestroy(ring->copy_sample_start); + } + if (ring->copy_sample_end != nullptr) { + (void) cudaEventDestroy(ring->copy_sample_end); + } + if (ring->copy_stream != nullptr) { + (void) cudaStreamDestroy(ring->copy_stream); + } + if (ring->ready_flags_host != nullptr) { + (void) cudaFreeHost(ring->ready_flags_host); + } + if (ring->deadline_counters_host != nullptr) { + (void) cudaFreeHost(ring->deadline_counters_host); + } + delete ring; + }; + + if (cudaHostAlloc(reinterpret_cast(&ring->ready_flags_host), + stage_slots*sizeof(uint32_t), cudaHostAllocMapped) != cudaSuccess || + cudaHostGetDevicePointer(reinterpret_cast(&ring->ready_flags_device), + ring->ready_flags_host, 0) != cudaSuccess || + cudaHostAlloc(reinterpret_cast(&ring->deadline_counters_host), + 2*sizeof(uint64_t), cudaHostAllocMapped) != cudaSuccess || + cudaHostGetDevicePointer(reinterpret_cast(&ring->deadline_counters_device), + ring->deadline_counters_host, 0) != cudaSuccess) { + (void) cudaGetLastError(); + cleanup(); + return nullptr; + } + std::fill_n(ring->ready_flags_host, stage_slots, 0u); + std::fill_n(ring->deadline_counters_host, 2, uint64_t(0)); + + if (cudaStreamCreateWithFlags(&ring->copy_stream, cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags(&ring->producer_ready, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreate(&ring->eval_start) != cudaSuccess || + cudaEventCreate(&ring->eval_end) != cudaSuccess || + cudaEventCreate(&ring->copy_sample_start) != cudaSuccess || + cudaEventCreate(&ring->copy_sample_end) != cudaSuccess) { + (void) cudaGetLastError(); + cleanup(); + return nullptr; + } + for (uint32_t slot = 0; slot < stage_slots; ++slot) { + if (cudaEventCreateWithFlags(&ring->ready[slot], cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&ring->consumed[slot], cudaEventDisableTiming) != cudaSuccess) { + (void) cudaGetLastError(); + cleanup(); + return nullptr; + } + } + return ring; +} + +void ggml_cuda_kv_stream_transfer_ring_free(ggml_cuda_kv_stream_transfer_ring * ring) { + if (ring == nullptr) { + return; + } + CUDA_CHECK(cudaStreamSynchronize(ring->copy_stream)); + for (cudaEvent_t event : ring->ready) { + CUDA_CHECK(cudaEventDestroy(event)); + } + for (cudaEvent_t event : ring->consumed) { + CUDA_CHECK(cudaEventDestroy(event)); + } + CUDA_CHECK(cudaEventDestroy(ring->producer_ready)); + CUDA_CHECK(cudaEventDestroy(ring->eval_start)); + CUDA_CHECK(cudaEventDestroy(ring->eval_end)); + CUDA_CHECK(cudaEventDestroy(ring->copy_sample_start)); + CUDA_CHECK(cudaEventDestroy(ring->copy_sample_end)); + CUDA_CHECK(cudaStreamDestroy(ring->copy_stream)); + CUDA_CHECK(cudaFreeHost(ring->ready_flags_host)); + CUDA_CHECK(cudaFreeHost(ring->deadline_counters_host)); + delete ring; +} + +bool ggml_cuda_kv_stream_transfer_ring_set_active_slots( + ggml_cuda_kv_stream_transfer_ring * ring, uint32_t stage_slots) { + if (ring == nullptr || stage_slots == 0 || stage_slots > ring->capacity_slots) { + return false; + } + if (ring->active_slots != stage_slots) { + ring->span_tuner.reset(); + ring->timing_pending = false; + ring->timing_current = false; + } + ring->active_slots = stage_slots; + return true; +} + +void ggml_cuda_kv_stream_transfer_ring_reset_span_tuner( + ggml_cuda_kv_stream_transfer_ring * ring) { + if (ring != nullptr) { + ring->span_tuner.reset(); + ring->timing_pending = false; + ring->timing_current = false; + } +} + +bool ggml_cuda_kv_stream_transfer_ring_observe_decode_latency( + ggml_cuda_kv_stream_transfer_ring * ring, double elapsed_ms) { + if (ring == nullptr) { + return false; + } + if (ring->forced_decode_span_pages != 0) { + return true; + } + const bool was_selected = ring->span_tuner.selected(); + ring->span_tuner.observe( + elapsed_ms, ring->last_graph_decode && ring->last_graph_streamed, + ring->last_graph_bounded); + if (!was_selected && ring->span_tuner.selected()) { + GGML_LOG_WARN( + "%s: selected %s decode spans from end-to-end latency " + "(unbounded %.3f ms, %u-page %.3f ms)\n", + __func__, ring->span_tuner.use_bounded() ? "bounded" : "unbounded", + ring->span_tuner.unbounded_average_ms(), KV_STREAM_DECODE_SPAN_PAGES, + ring->span_tuner.bounded_average_ms()); + } + return true; +} + +ggml_cuda_kv_stream_transfer_stats ggml_cuda_kv_stream_transfer_ring_get_stats( + const ggml_cuda_kv_stream_transfer_ring * ring) { + if (ring == nullptr) { + return {}; + } + return { + ring->asynchronous_page_uploads, + ring->host_to_device_copy_commands, + ring->compute_stream_waits, + ring->stage_slot_reuses, + ring->cross_layer_prefetches, + ring->deadline_counters_host[0], + ring->deadline_counters_host[1], + ring->ring_peak_occupancy, + }; +} + +struct ggml_cuda_kv_stream_resident_cache { + char * pool_data = nullptr; + size_t pool_bytes = 0; + size_t scratch_bytes = 0; + size_t page_bytes = 0; + uint32_t layer_count = 0; + uint32_t page_tokens = 0; + uint32_t resident_pages_per_layer = 0; + uint32_t decode_active_pages = 0; + uint32_t next_layer = 0; + std::vector layer_pages; + std::vector layer_offsets; + std::unordered_map layer_by_k; + std::unordered_map layer_by_data; + std::unordered_map mirror_by_data; + std::vector loaded; + std::vector dirty; + std::vector precise_dirty_tracking; + std::vector dirty_rows; + std::vector mutable_pages; + bool all_pages_mutable = false; + ggml_cuda_kv_stream_resident_stats stats; +}; + +static bool kv_stream_resident_cache_layout( + const ggml_cuda_kv_stream_resident_cache * cache, + size_t scratch_bytes, + uint32_t decode_active_pages, + uint32_t & resident_pages_per_layer, + std::vector & layer_pages, + std::vector & layer_offsets) { + if (cache == nullptr || scratch_bytes > cache->pool_bytes || + scratch_bytes%cache->page_bytes != 0) { + return false; + } + const size_t resident_pages_total = + (cache->pool_bytes - scratch_bytes)/cache->page_bytes; + const size_t pages_per_layer = resident_pages_total/cache->layer_count; + if (pages_per_layer > UINT32_MAX || + pages_per_layer > std::numeric_limits::max()/cache->layer_count) { + return false; + } + + resident_pages_per_layer = uint32_t(pages_per_layer); + const size_t controlled_resident_pages = pages_per_layer*cache->layer_count; + layer_pages.assign(cache->layer_count, resident_pages_per_layer); + if (decode_active_pages > resident_pages_per_layer) { + const uint64_t total_active_pages = + uint64_t(decode_active_pages)*cache->layer_count; + if (total_active_pages < controlled_resident_pages) { + return false; + } + const uint64_t streamed_pages = + total_active_pages - controlled_resident_pages; + const uint64_t ring_pages = scratch_bytes/cache->page_bytes; + if (ring_pages == 0 || streamed_pages == 0) { + return false; + } + + // Reduce the number of split-attention layers while each streamed + // working set fits in both the shared ring and the active pages owned + // by one layer. When even one split per model layer exceeds the ring, + // that layer streams in multiple waves. Spread split layers across + // model order so resident layers provide prefetch windows. + const auto ceil_div = [](uint64_t numerator, uint64_t denominator) { + return numerator/denominator + (numerator%denominator != 0); + }; + const uint64_t splits_for_ring = ceil_div(streamed_pages, ring_pages); + const uint64_t splits_for_layer_capacity = + ceil_div(streamed_pages, decode_active_pages); + const uint64_t ring_bounded_splits = + std::min(cache->layer_count, splits_for_ring); + const uint64_t split_layers_wide = + std::max(ring_bounded_splits, splits_for_layer_capacity); + if (split_layers_wide == 0 || split_layers_wide > cache->layer_count) { + return false; + } + const uint32_t split_layers = uint32_t(split_layers_wide); + const uint64_t base_streamed = streamed_pages/split_layers; + const uint64_t remainder = streamed_pages%split_layers; + layer_pages.assign(cache->layer_count, decode_active_pages); + for (uint32_t split = 0; split < split_layers; ++split) { + const uint32_t layer = uint32_t( + uint64_t(split)*cache->layer_count/split_layers); + const uint64_t layer_streamed = base_streamed + (split < remainder ? 1 : 0); + if (layer_streamed > decode_active_pages) { + return false; + } + layer_pages[layer] = decode_active_pages - uint32_t(layer_streamed); + } + } + + layer_offsets.assign(size_t(cache->layer_count) + 1, 0); + for (uint32_t layer = 0; layer < cache->layer_count; ++layer) { + layer_offsets[layer + 1] = layer_offsets[layer] + layer_pages[layer]; + } + return layer_offsets.back() == controlled_resident_pages; +} + +static size_t kv_stream_resident_index( + const ggml_cuda_kv_stream_resident_cache * cache, + uint32_t layer, + uint32_t page) { + GGML_ASSERT(layer < cache->layer_count && page < cache->layer_pages[layer]); + return cache->layer_offsets[layer] + page; +} + +ggml_cuda_kv_stream_resident_cache * ggml_cuda_kv_stream_resident_cache_new( + void * pool_data, size_t pool_bytes, size_t scratch_bytes, size_t page_bytes, + uint32_t layer_count, uint32_t page_tokens) { + if (pool_data == nullptr || scratch_bytes == 0 || scratch_bytes >= pool_bytes || + page_bytes == 0 || scratch_bytes%page_bytes != 0 || + layer_count == 0 || page_tokens != 256) { + return nullptr; + } + + const size_t resident_pages = (pool_bytes - scratch_bytes)/(page_bytes*layer_count); + if (resident_pages == 0 || resident_pages > UINT32_MAX) { + return nullptr; + } + + auto * cache = new ggml_cuda_kv_stream_resident_cache; + cache->pool_data = static_cast(pool_data); + cache->pool_bytes = pool_bytes; + cache->scratch_bytes = scratch_bytes; + cache->page_bytes = page_bytes; + cache->layer_count = layer_count; + cache->page_tokens = page_tokens; + cache->resident_pages_per_layer = resident_pages; + cache->layer_pages.assign(layer_count, uint32_t(resident_pages)); + cache->layer_offsets.resize(size_t(layer_count) + 1); + for (uint32_t layer = 0; layer <= layer_count; ++layer) { + cache->layer_offsets[layer] = size_t(layer)*resident_pages; + } + cache->loaded.resize(cache->layer_offsets.back(), 0); + cache->dirty.resize(cache->layer_offsets.back(), 0); + cache->precise_dirty_tracking.resize(layer_count, 0); + return cache; +} + +void ggml_cuda_kv_stream_resident_cache_free(ggml_cuda_kv_stream_resident_cache * cache) { + delete cache; +} + +void ggml_cuda_kv_stream_resident_cache_reset(ggml_cuda_kv_stream_resident_cache * cache) { + if (cache == nullptr) { + return; + } + std::fill(cache->loaded.begin(), cache->loaded.end(), 0); + std::fill(cache->dirty.begin(), cache->dirty.end(), 0); + std::fill(cache->precise_dirty_tracking.begin(), cache->precise_dirty_tracking.end(), 0); + cache->dirty_rows.clear(); + cache->mutable_pages.clear(); + cache->all_pages_mutable = false; + cache->layer_by_k.clear(); + cache->layer_by_data.clear(); + cache->mirror_by_data.clear(); + cache->next_layer = 0; + cache->stats = {}; +} + +bool ggml_cuda_kv_stream_resident_cache_reconfigure( + ggml_cuda_kv_stream_resident_cache * cache, + size_t scratch_bytes, + uint32_t active_pages_per_layer) { + if (cache == nullptr) { + return false; + } + + uint32_t pages_per_layer = 0; + std::vector layer_pages; + std::vector layer_offsets; + if (!kv_stream_resident_cache_layout( + cache, scratch_bytes, active_pages_per_layer, + pages_per_layer, layer_pages, layer_offsets)) { + return false; + } + + const bool scratch_changed = cache->scratch_bytes != scratch_bytes; + const bool layout_changed = cache->layer_pages != layer_pages; + cache->decode_active_pages = active_pages_per_layer; + cache->resident_pages_per_layer = pages_per_layer; + if (!scratch_changed && !layout_changed) { + return true; + } + + // Publish the scratch boundary and concentrated decode layout together. + // Both changes alter physical K/V addresses, so one invalidation is + // required; applying them separately would discard and reload the same + // resident working set twice. + cache->scratch_bytes = scratch_bytes; + cache->layer_pages = std::move(layer_pages); + cache->layer_offsets = std::move(layer_offsets); + cache->loaded.assign(cache->layer_offsets.back(), 0); + cache->dirty.assign(cache->layer_offsets.back(), 0); + cache->precise_dirty_tracking.assign(cache->layer_count, 0); + cache->dirty_rows.clear(); + cache->mutable_pages.clear(); + cache->all_pages_mutable = false; + cache->layer_by_k.clear(); + cache->layer_by_data.clear(); + cache->mirror_by_data.clear(); + cache->next_layer = 0; + if (scratch_changed) { + cache->stats = {}; + } + return true; +} + + +bool ggml_cuda_kv_stream_resident_cache_repartition( + ggml_cuda_kv_stream_resident_cache * cache, size_t scratch_bytes) { + if (cache == nullptr) { + return false; + } + uint32_t pages_per_layer = 0; + std::vector layer_pages; + std::vector layer_offsets; + if (!kv_stream_resident_cache_layout( + cache, scratch_bytes, cache->decode_active_pages, + pages_per_layer, layer_pages, layer_offsets)) { + return false; + } + if (cache->scratch_bytes == scratch_bytes && cache->layer_pages == layer_pages) { + return true; + } + cache->scratch_bytes = scratch_bytes; + cache->resident_pages_per_layer = pages_per_layer; + cache->layer_pages = std::move(layer_pages); + cache->layer_offsets = std::move(layer_offsets); + cache->loaded.assign(cache->layer_offsets.back(), 0); + cache->dirty.assign(cache->layer_offsets.back(), 0); + cache->precise_dirty_tracking.assign(cache->layer_count, 0); + cache->dirty_rows.clear(); + cache->layer_by_k.clear(); + cache->layer_by_data.clear(); + cache->mirror_by_data.clear(); + cache->next_layer = 0; + cache->stats = {}; + return true; +} + +bool ggml_cuda_kv_stream_resident_cache_set_decode_layout( + ggml_cuda_kv_stream_resident_cache * cache, + uint32_t active_pages_per_layer) { + if (cache == nullptr) { + return false; + } + uint32_t pages_per_layer = 0; + std::vector layer_pages; + std::vector layer_offsets; + if (!kv_stream_resident_cache_layout( + cache, cache->scratch_bytes, active_pages_per_layer, + pages_per_layer, layer_pages, layer_offsets)) { + return false; + } + if (cache->layer_pages == layer_pages) { + cache->decode_active_pages = active_pages_per_layer; + return true; + } + // Resident storage uses separate contiguous K and V planes per layer. + // A layout change moves both the layer base and the V-plane boundary, so + // migrating page-sized byte slots cannot preserve logical K/V pages. + // Reload lazily from the authoritative host cache instead. + cache->layer_pages = std::move(layer_pages); + cache->layer_offsets = std::move(layer_offsets); + cache->loaded.assign(cache->layer_offsets.back(), 0); + cache->dirty.assign(cache->layer_offsets.back(), 0); + cache->precise_dirty_tracking.assign(cache->layer_count, 0); + cache->dirty_rows.clear(); + cache->mutable_pages.clear(); + cache->all_pages_mutable = false; + cache->layer_by_k.clear(); + cache->layer_by_data.clear(); + cache->mirror_by_data.clear(); + cache->next_layer = 0; + cache->decode_active_pages = active_pages_per_layer; + cache->resident_pages_per_layer = pages_per_layer; + return true; +} + +uint32_t ggml_cuda_kv_stream_resident_cache_pages_per_layer( + const ggml_cuda_kv_stream_resident_cache * cache) { + return cache == nullptr ? 0 : cache->resident_pages_per_layer; +} + +uint32_t ggml_cuda_kv_stream_resident_cache_decode_active_pages( + const ggml_cuda_kv_stream_resident_cache * cache) { + return cache == nullptr ? 0 : cache->decode_active_pages; +} + +ggml_cuda_kv_stream_resident_stats ggml_cuda_kv_stream_resident_cache_get_stats( + const ggml_cuda_kv_stream_resident_cache * cache) { + return cache == nullptr ? ggml_cuda_kv_stream_resident_stats{} : cache->stats; +} + +bool ggml_cuda_kv_stream_resident_cache_mark_dirty_rows( + ggml_cuda_kv_stream_resident_cache * cache, + const int64_t * rows, size_t count) { + if (cache == nullptr || (rows == nullptr && count != 0)) { + return false; + } + + cache->mutable_pages.clear(); + cache->all_pages_mutable = false; + if (count == 0) { + cache->dirty_rows.clear(); + } else { + cache->dirty_rows.assign(rows, rows + count); + } + std::fill(cache->dirty.begin(), cache->dirty.end(), 0); + for (uint32_t layer = 0; layer < cache->layer_count; ++layer) { + cache->precise_dirty_tracking[layer] = 1; + } + for (size_t i = 0; i < count; ++i) { + if (rows[i] < 0) { + cache->all_pages_mutable = true; + std::fill(cache->dirty.begin(), cache->dirty.end(), 1); + return true; + } + const uint32_t page = uint32_t(uint64_t(rows[i])/cache->page_tokens); + if (std::find(cache->mutable_pages.begin(), cache->mutable_pages.end(), page) == + cache->mutable_pages.end()) { + cache->mutable_pages.push_back(page); + } + for (uint32_t layer = 0; layer < cache->layer_count; ++layer) { + if (page < cache->layer_pages[layer]) { + cache->dirty[kv_stream_resident_index(cache, layer, page)] = 1; + } + } + } + return true; +} + +bool ggml_cuda_kv_stream_resident_cache_all_layers_fit( + const ggml_cuda_kv_stream_resident_cache * cache, + uint32_t active_pages) { + if (cache == nullptr || active_pages == 0) { + return false; + } + return std::all_of(cache->layer_pages.begin(), cache->layer_pages.end(), + [active_pages](uint32_t pages) { return pages >= active_pages; }); +} + +bool ggml_cuda_kv_stream_resident_cache_get_mirror( + ggml_cuda_kv_stream_resident_cache * cache, + const ggml_tensor * target, + void ** data) { + if (cache == nullptr || target == nullptr || data == nullptr || + cache->dirty_rows.empty()) { + return false; + } + const auto layer_it = cache->layer_by_data.find(target->data); + const auto mirror_it = cache->mirror_by_data.find(target->data); + if (layer_it == cache->layer_by_data.end() || mirror_it == cache->mirror_by_data.end()) { + return false; + } + const uint64_t capacity = uint64_t(cache->layer_pages[layer_it->second])*cache->page_tokens; + for (const int64_t row : cache->dirty_rows) { + if (row < 0 || uint64_t(row) >= capacity) { + return false; + } + } + *data = mirror_it->second; + return true; +} + +void ggml_cuda_kv_stream_resident_cache_mark_mirrored( + ggml_cuda_kv_stream_resident_cache * cache, + const ggml_tensor * target) { + if (cache == nullptr || target == nullptr) { + return; + } + const auto layer_it = cache->layer_by_data.find(target->data); + if (layer_it == cache->layer_by_data.end()) { + return; + } + const uint32_t layer = layer_it->second; + for (const int64_t row : cache->dirty_rows) { + if (row < 0) { + return; + } + const uint32_t page = uint32_t(uint64_t(row)/cache->page_tokens); + if (page < cache->layer_pages[layer]) { + cache->dirty[kv_stream_resident_index(cache, layer, page)] = 0; + } + } +} + +void ggml_cuda_kv_stream_resident_cache_mark_dirty( + ggml_cuda_kv_stream_resident_cache * cache, + const ggml_tensor * target, const ggml_tensor * indices) { + if (cache == nullptr || target == nullptr || indices == nullptr) { + return; + } + const auto layer_it = cache->layer_by_data.find(target->data); + if (layer_it == cache->layer_by_data.end()) { + return; + } + const uint32_t layer = layer_it->second; + GGML_ASSERT(layer < cache->layer_count); + if (indices->buffer == nullptr || !ggml_backend_buffer_is_host(indices->buffer) || + !ggml_is_contiguous(indices) || + (indices->type != GGML_TYPE_I32 && indices->type != GGML_TYPE_I64)) { + return; + } + + cache->precise_dirty_tracking[layer] = 1; + const size_t begin = cache->layer_offsets[layer]; + const size_t end = cache->layer_offsets[layer + 1]; + const size_t count = ggml_nelements(indices); + for (size_t i = 0; i < count; ++i) { + const int64_t row = indices->type == GGML_TYPE_I32 ? + static_cast(indices->data)[i] : + static_cast(indices->data)[i]; + if (row < 0) { + std::fill(cache->dirty.begin() + begin, cache->dirty.begin() + end, 1); + return; + } + const uint64_t page = uint64_t(row)/cache->page_tokens; + if (page < cache->layer_pages[layer]) { + cache->dirty[begin + page] = 1; + } + } +} + +static uint32_t kv_stream_resident_layer( + ggml_cuda_kv_stream_resident_cache * cache, const void * k_key) { + GGML_ASSERT(cache != nullptr); + auto [it, inserted] = cache->layer_by_k.emplace(k_key, cache->next_layer); + if (inserted) { + GGML_ASSERT(cache->next_layer < cache->layer_count); + ++cache->next_layer; + } + return it->second; +} + +static bool kv_stream_page_mutable( + const ggml_cuda_kv_stream_resident_cache * cache, + uint32_t page) { + return cache->all_pages_mutable || + std::find(cache->mutable_pages.begin(), cache->mutable_pages.end(), page) != + cache->mutable_pages.end(); +} + +namespace { + +constexpr int KV_STREAM_HEAD_DIM = 256; +constexpr int KV_STREAM_MAX_PARTS_PER_CHUNK = 16; +constexpr int KV_STREAM_QUERY_WORKSPACE_TOKENS = 256; + +static int kv_stream_parts_per_chunk() { + static const int parts = []() { + const char * value = getenv("GGML_CUDA_KV_STREAM_PARTS"); + const int parsed = value == nullptr ? 16 : atoi(value); + return parsed == 2 || parsed == 4 || parsed == 8 || parsed == 16 ? parsed : 16; + }(); + return parts; +} + +static int64_t kv_stream_block_tokens(const ggml_tensor * dst, size_t stage_bytes) { + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + + const size_t k_row_bytes = ggml_row_size(K->type, K->ne[0]); + const size_t v_row_bytes = ggml_row_size(V->type, V->ne[0]); + const size_t bytes_per_token = k_row_bytes*K->ne[2] + v_row_bytes*V->ne[2]; + if (bytes_per_token == 0) { + return 0; + } + + int64_t tokens = std::min(K->ne[1], stage_bytes/bytes_per_token); + tokens = tokens/FATTN_KQ_STRIDE*FATTN_KQ_STRIDE; + while (tokens > 0) { + const size_t k_bytes = k_row_bytes*tokens*K->ne[2]; + const size_t v_offset = GGML_PAD(k_bytes, 128); + const size_t v_bytes = v_row_bytes*tokens*V->ne[2]; + if (v_offset <= stage_bytes && v_bytes <= stage_bytes - v_offset) { + return tokens; + } + tokens -= FATTN_KQ_STRIDE; + } + return 0; +} + +template +static __global__ void kv_stream_accumulate_chunk_results( + const float * parts, + const float2 * meta, + float * accumulator, + float2 * accumulator_meta, + int nrows, + bool initialize, + int nparts) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= nrows || tid >= D) { + return; + } + ggml_cuda_pdl_sync(); + + __shared__ float old_maximum; + __shared__ float maximum; + __shared__ float old_scale; + __shared__ float denominator; + + const int base = row*nparts; + if (tid == 0) { + old_maximum = initialize ? -FLT_MAX : accumulator_meta[row].x; + maximum = old_maximum; + for (int part = 0; part < nparts; ++part) { + maximum = fmaxf(maximum, meta[base + part].x); + } + + old_scale = initialize ? 0.0f : expf(old_maximum - maximum); + denominator = initialize ? 0.0f : old_scale*accumulator_meta[row].y; + for (int part = 0; part < nparts; ++part) { + const float weight = expf(meta[base + part].x - maximum); + denominator += weight*meta[base + part].y; + } + accumulator_meta[row] = make_float2(maximum, denominator); + } + __syncthreads(); + + float numerator = initialize ? 0.0f : old_scale*accumulator[row*D + tid]; + for (int part = 0; part < nparts; ++part) { + const float weight = expf(meta[base + part].x - maximum); + numerator += weight*parts[(base + part)*D + tid]; + } + accumulator[row*D + tid] = numerator; +} + +template +static __global__ void kv_stream_normalize_chunk_results( + const float * accumulator, + const float2 * accumulator_meta, + float * dst, + int nrows) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= nrows || tid >= D) { + return; + } + ggml_cuda_pdl_sync(); + + dst[row*D + tid] = accumulator[row*D + tid]/accumulator_meta[row].y; +} + +#ifdef GGML_CUDA_FA_ALL_QUANTS +using kv_stream_native_partial_fn = void (*)( + ggml_backend_cuda_context &, ggml_tensor *, float *, float2 *, int); + +template +static kv_stream_native_partial_fn kv_stream_resolve_native_partial_for_v(ggml_type type_v) { + switch (type_v) { + case GGML_TYPE_F16: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_F16>; + case GGML_TYPE_Q4_0: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_Q4_0>; + case GGML_TYPE_Q4_1: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_Q4_1>; + case GGML_TYPE_Q5_0: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_Q5_0>; + case GGML_TYPE_Q5_1: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_Q5_1>; + case GGML_TYPE_Q8_0: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_Q8_0>; + case GGML_TYPE_BF16: + return &ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, type_K, GGML_TYPE_BF16>; + default: + return nullptr; + } +} + +static kv_stream_native_partial_fn kv_stream_resolve_native_partial( + ggml_type type_k, ggml_type type_v) { +#define KV_STREAM_NATIVE_K_CASE(type_K) \ + case type_K: return kv_stream_resolve_native_partial_for_v(type_v) + switch (type_k) { + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_F16); + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_Q4_0); + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_Q4_1); + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_Q5_0); + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_Q5_1); + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_Q8_0); + KV_STREAM_NATIVE_K_CASE(GGML_TYPE_BF16); + default: return nullptr; + } +#undef KV_STREAM_NATIVE_K_CASE +} +#endif // GGML_CUDA_FA_ALL_QUANTS + +} // namespace + +struct ggml_backend_cuda_kv_stream_type_capabilities +ggml_backend_cuda_kv_stream_get_type_capabilities(ggml_type type) { + ggml_backend_cuda_kv_stream_type_capabilities result{}; + + switch (type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q8_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + result.classified = true; + break; + default: + return result; + } + + result.storage = type != GGML_TYPE_Q8_1 && type != GGML_TYPE_Q8_K; + + switch (type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_IQ4_NL: + result.online_write = true; + break; + default: + break; + } + + switch (type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + result.decode_f16 = true; + break; + default: + break; + } + + switch (type) { + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + result.direct_attention = true; + break; + default: + break; + } + + result.requires_initialization = + type == GGML_TYPE_IQ3_XXS || type == GGML_TYPE_IQ3_S || type == GGML_TYPE_IQ2_S; + result.requires_importance_matrix = + type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M; + result.auxiliary = type == GGML_TYPE_Q8_1 || type == GGML_TYPE_Q8_K; + return result; +} + +ggml_backend_cuda_kv_stream_attention_mode +ggml_backend_cuda_kv_stream_get_attention_mode(ggml_type type_k, ggml_type type_v) { + const auto capabilities_k = ggml_backend_cuda_kv_stream_get_type_capabilities(type_k); + const auto capabilities_v = ggml_backend_cuda_kv_stream_get_type_capabilities(type_v); +#ifdef GGML_CUDA_FA_ALL_QUANTS + if (capabilities_k.direct_attention && capabilities_v.direct_attention) { + return GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT; + } +#endif // GGML_CUDA_FA_ALL_QUANTS + if (capabilities_k.storage && capabilities_v.storage && + capabilities_k.online_write && capabilities_v.online_write && + capabilities_k.decode_f16 && capabilities_v.decode_f16 && + !capabilities_k.requires_importance_matrix && !capabilities_v.requires_importance_matrix) { + return GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16; + } + return GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_UNSUPPORTED; +} + +bool ggml_cuda_kv_stream_page_bytes( + ggml_type type_k, ggml_type type_v, + uint32_t head_dim_k, uint32_t head_dim_v, uint32_t head_count, + uint32_t page_tokens, size_t * page_bytes) { + if (head_dim_k == 0 || head_dim_v == 0 || head_count == 0 || page_tokens == 0 || page_bytes == nullptr) { + return false; + } + + const auto capabilities_k = ggml_backend_cuda_kv_stream_get_type_capabilities(type_k); + const auto capabilities_v = ggml_backend_cuda_kv_stream_get_type_capabilities(type_v); + if (!capabilities_k.storage || !capabilities_v.storage) { + return false; + } + + const uint32_t block_size_k = ggml_blck_size(type_k); + const uint32_t block_size_v = ggml_blck_size(type_v); + if (block_size_k == 0 || block_size_v == 0 || + head_dim_k % block_size_k != 0 || head_dim_v % block_size_v != 0) { + return false; + } + + const size_t maximum = std::numeric_limits::max(); + const size_t row_bytes_k = ggml_row_size(type_k, head_dim_k); + const size_t row_bytes_v = ggml_row_size(type_v, head_dim_v); + if (row_bytes_k > maximum/head_count || row_bytes_v > maximum/head_count) { + return false; + } + const size_t token_bytes_k = row_bytes_k*head_count; + const size_t token_bytes_v = row_bytes_v*head_count; + if (token_bytes_k > maximum/page_tokens || token_bytes_v > maximum/page_tokens) { + return false; + } + const size_t page_bytes_k = token_bytes_k*page_tokens; + const size_t page_bytes_v = token_bytes_v*page_tokens; + if (page_bytes_k > maximum - 127) { + return false; + } + const size_t page_offset_v = (page_bytes_k + 127) & ~size_t(127); + if (page_bytes_v > maximum - page_offset_v) { + return false; + } + + *page_bytes = page_offset_v + page_bytes_v; + return true; +} + +bool ggml_cuda_kv_stream_workspace_bytes( + ggml_type type_k, ggml_type type_v, + uint32_t head_dim_k, uint32_t head_dim_v, uint32_t head_count, + uint32_t page_tokens, size_t * workspace_bytes) { + if (workspace_bytes == nullptr) { + return false; + } + const auto mode = + ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v); + if (mode == GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT) { + *workspace_bytes = 0; + return true; + } + if (mode != GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16) { + return false; + } + return ggml_cuda_kv_stream_page_bytes( + GGML_TYPE_F16, GGML_TYPE_F16, + head_dim_k, head_dim_v, head_count, page_tokens, workspace_bytes); +} + +bool ggml_cuda_flash_attn_ext_streamed_supported(const ggml_tensor * dst, size_t stage_bytes) { + if (dst == nullptr || dst->op != GGML_OP_FLASH_ATTN_EXT) { + return false; + } + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + + return Q != nullptr && K != nullptr && V != nullptr && + Q->type == GGML_TYPE_F32 && + ggml_backend_cuda_kv_stream_get_attention_mode(K->type, V->type) != GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_UNSUPPORTED && + Q->ne[0] == KV_STREAM_HEAD_DIM && V->ne[0] == KV_STREAM_HEAD_DIM && + Q->ne[1] >= 1 && Q->ne[3] == 1 && K->ne[3] == 1 && V->ne[3] == 1 && + K->ne[1] == V->ne[1] && K->ne[2] == V->ne[2] && + K->ne[1] % FATTN_KQ_STRIDE == 0 && + K->nb[0] == ggml_element_size(K) && V->nb[0] == ggml_element_size(V) && + K->nb[1] >= ggml_row_size(K->type, K->ne[0]) && + V->nb[1] >= ggml_row_size(V->type, V->ne[0]) && + K->nb[1] == ggml_row_size(K->type, K->ne[0])*K->ne[2] && + K->nb[2] == ggml_row_size(K->type, K->ne[0]) && + V->nb[1] == ggml_row_size(V->type, V->ne[0])*V->ne[2] && + V->nb[2] == ggml_row_size(V->type, V->ne[0]) && + (mask == nullptr || (mask->type == GGML_TYPE_F16 && ggml_is_contiguous(mask))) && + sinks == nullptr && kv_stream_block_tokens(dst, stage_bytes) > 0; +} + +namespace { + +static __global__ void kv_stream_record_deadline( + const uint32_t * ready_flag, + uint64_t * samples, + uint64_t * misses) { + if (threadIdx.x == 0 && blockIdx.x == 0) { + atomicAdd(reinterpret_cast(samples), 1ULL); + if (*ready_flag == 0) { + atomicAdd(reinterpret_cast(misses), 1ULL); + } + } +} + +static bool kv_stream_collect_timing(ggml_cuda_kv_stream_transfer_ring * ring) { + if (!ring->timing_pending) { + return true; + } + const cudaError_t status = cudaEventQuery(ring->eval_end); + if (status == cudaErrorNotReady) { + return false; + } + CUDA_CHECK(status); + + float eval_ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&eval_ms, ring->eval_start, ring->eval_end)); + double busy_ratio = 0.0; + if (ring->copy_sample_recorded && eval_ms > 0.0f && ring->current_epoch_uploads > 0) { + float copy_ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime( + ©_ms, ring->copy_sample_start, ring->copy_sample_end)); + GGML_ASSERT(ring->copy_sample_uploads > 0); + busy_ratio = std::min(1.0, double(copy_ms)*double(ring->current_epoch_uploads)/ + (double(ring->copy_sample_uploads)*double(eval_ms))); + } + ring->last_copy_engine_busy_ratio = busy_ratio; + ring->last_ring_peak_occupancy = ring->current_ring_peak_occupancy; + ring->timing_pending = false; + return true; +} + +static void kv_stream_graph_upload( + ggml_cuda_kv_stream_transfer_ring * ring, + kv_stream_graph_request & request, + uint32_t slot) { + // Preserve the host cache's compact token-major layout. This makes each + // page one contiguous K transfer plus one contiguous V transfer. + char * k_stage = ring->pool_data + size_t(slot)*request.k_bytes; + char * v_stage = ring->pool_data + + size_t(ring->active_slots)*request.k_bytes + size_t(slot)*request.v_bytes; + if (ring->slot_used[slot]) { + CUDA_CHECK(cudaStreamWaitEvent(ring->copy_stream, ring->consumed[slot], 0)); + ++ring->stage_slot_reuses; + } + if (request.deadline_sample) { + CUDA_CHECK(cudaMemsetAsync( + ring->ready_flags_device + slot, 0, sizeof(uint32_t), ring->copy_stream)); + } + if (ring->timing_current && !ring->copy_sample_recorded) { + CUDA_CHECK(cudaEventRecord(ring->copy_sample_start, ring->copy_stream)); + } + GGML_ASSERT(request.k_nb1 == request.k_row_bytes*request.n_head_kv); + GGML_ASSERT(request.k_nb2 == request.k_row_bytes); + GGML_ASSERT(request.v_nb1 == request.v_row_bytes*request.n_head_kv); + GGML_ASSERT(request.v_nb2 == request.v_row_bytes); + CUDA_CHECK(cudaMemcpyAsync( + k_stage, request.k_data + request.token_begin*request.k_nb1, + request.k_bytes, cudaMemcpyHostToDevice, ring->copy_stream)); + CUDA_CHECK(cudaMemcpyAsync( + v_stage, request.v_data + request.token_begin*request.v_nb1, + request.v_bytes, cudaMemcpyHostToDevice, ring->copy_stream)); + ring->host_to_device_copy_commands += 2; + if (request.deadline_sample) { + CUDA_CHECK(cudaMemsetAsync( + ring->ready_flags_device + slot, 1, sizeof(uint32_t), ring->copy_stream)); + } + if (ring->timing_current && !ring->copy_sample_recorded) { + CUDA_CHECK(cudaEventRecord(ring->copy_sample_end, ring->copy_stream)); + ring->copy_sample_recorded = true; + ring->copy_sample_uploads = 1; + } + CUDA_CHECK(cudaEventRecord(ring->ready[slot], ring->copy_stream)); + request.slot = slot; + request.scheduled = true; + ring->slot_used[slot] = 1; + request.ready_slot = slot; + ring->slot_request[slot] = size_t(&request - ring->graph_requests.data()); + ++ring->current_occupancy; + ring->ring_peak_occupancy = std::max( + ring->ring_peak_occupancy, ring->current_occupancy); + ring->current_ring_peak_occupancy = std::max( + ring->current_ring_peak_occupancy, ring->current_occupancy); + ++ring->asynchronous_page_uploads; + ++ring->current_epoch_uploads; + if (ring->graph_resident_cache != nullptr) { + ring->graph_resident_cache->stats.host_to_device_bytes += + request.k_bytes + request.v_bytes; + } + if (ring->current_layer != KV_STREAM_NO_LAYER && request.layer > ring->current_layer) { + ++ring->cross_layer_prefetches; + } +} + +static bool kv_stream_graph_request_follows( + const kv_stream_graph_request & previous, + const kv_stream_graph_request & request) { + return request.eligible && !request.scheduled && !request.consumed && + request.layer == previous.layer && + request.k_data == previous.k_data && request.v_data == previous.v_data && + request.k_nb1 == previous.k_nb1 && request.v_nb1 == previous.v_nb1 && + request.k_bytes == previous.k_bytes && request.v_bytes == previous.v_bytes && + request.token_count == previous.token_count && + request.token_begin == previous.token_begin + previous.token_count; +} + +static void kv_stream_graph_upload_batch( + ggml_cuda_kv_stream_transfer_ring * ring, + size_t first_request, + uint32_t first_slot, + uint32_t batch_pages) { + GGML_ASSERT(batch_pages > 1); + GGML_ASSERT(first_request + batch_pages <= ring->graph_requests.size()); + GGML_ASSERT(first_slot + batch_pages <= ring->active_slots); + auto & first = ring->graph_requests[first_request]; + + for (uint32_t page = 0; page < batch_pages; ++page) { + auto & request = ring->graph_requests[first_request + page]; + const uint32_t slot = first_slot + page; + GGML_ASSERT(request.eligible && !request.scheduled && !request.consumed); + GGML_ASSERT(ring->slot_request[slot] == KV_STREAM_NO_REQUEST); + if (ring->slot_used[slot]) { + CUDA_CHECK(cudaStreamWaitEvent(ring->copy_stream, ring->consumed[slot], 0)); + ++ring->stage_slot_reuses; + } + if (request.deadline_sample) { + CUDA_CHECK(cudaMemsetAsync( + ring->ready_flags_device + slot, 0, sizeof(uint32_t), ring->copy_stream)); + } + } + + if (ring->timing_current && !ring->copy_sample_recorded) { + CUDA_CHECK(cudaEventRecord(ring->copy_sample_start, ring->copy_stream)); + } + GGML_ASSERT(first.k_nb1 == first.k_row_bytes*first.n_head_kv); + GGML_ASSERT(first.k_nb2 == first.k_row_bytes); + GGML_ASSERT(first.v_nb1 == first.v_row_bytes*first.n_head_kv); + GGML_ASSERT(first.v_nb2 == first.v_row_bytes); + char * k_stage = ring->pool_data + size_t(first_slot)*first.k_bytes; + char * v_stage = ring->pool_data + + size_t(ring->active_slots)*first.k_bytes + size_t(first_slot)*first.v_bytes; + CUDA_CHECK(cudaMemcpyAsync( + k_stage, first.k_data + first.token_begin*first.k_nb1, + size_t(batch_pages)*first.k_bytes, cudaMemcpyHostToDevice, ring->copy_stream)); + CUDA_CHECK(cudaMemcpyAsync( + v_stage, first.v_data + first.token_begin*first.v_nb1, + size_t(batch_pages)*first.v_bytes, cudaMemcpyHostToDevice, ring->copy_stream)); + ring->host_to_device_copy_commands += 2; + + for (uint32_t page = 0; page < batch_pages; ++page) { + auto & request = ring->graph_requests[first_request + page]; + if (request.deadline_sample) { + CUDA_CHECK(cudaMemsetAsync( + ring->ready_flags_device + first_slot + page, 1, + sizeof(uint32_t), ring->copy_stream)); + } + } + if (ring->timing_current && !ring->copy_sample_recorded) { + CUDA_CHECK(cudaEventRecord(ring->copy_sample_end, ring->copy_stream)); + ring->copy_sample_recorded = true; + ring->copy_sample_uploads = batch_pages; + } + CUDA_CHECK(cudaEventRecord(ring->ready[first_slot], ring->copy_stream)); + + for (uint32_t page = 0; page < batch_pages; ++page) { + auto & request = ring->graph_requests[first_request + page]; + const uint32_t slot = first_slot + page; + request.slot = slot; + request.ready_slot = first_slot; + request.scheduled = true; + ring->slot_used[slot] = 1; + ring->slot_request[slot] = first_request + page; + ++ring->current_occupancy; + ++ring->asynchronous_page_uploads; + ++ring->current_epoch_uploads; + if (ring->graph_resident_cache != nullptr) { + ring->graph_resident_cache->stats.host_to_device_bytes += + request.k_bytes + request.v_bytes; + } + if (ring->current_layer != KV_STREAM_NO_LAYER && request.layer > ring->current_layer) { + ++ring->cross_layer_prefetches; + } + } + ring->ring_peak_occupancy = std::max( + ring->ring_peak_occupancy, ring->current_occupancy); + ring->current_ring_peak_occupancy = std::max( + ring->current_ring_peak_occupancy, ring->current_occupancy); +} + +static uint32_t kv_stream_graph_batch_pages( + const ggml_cuda_kv_stream_transfer_ring * ring, + uint32_t first_slot) { + if (!ring->graph_active || first_slot >= ring->active_slots || + ring->next_request >= ring->graph_requests.size() || + ring->slot_request[first_slot] != KV_STREAM_NO_REQUEST) { + return 0; + } + const auto & first = ring->graph_requests[ring->next_request]; + if (!first.eligible) { + return 0; + } + GGML_ASSERT(!first.scheduled && !first.consumed); + + const uint32_t maximum = std::min({ + KV_STREAM_COPY_BATCH_PAGES, + ring->active_slots - first_slot, + uint32_t(ring->graph_requests.size() - ring->next_request), + }); + uint32_t pages = 1; + while (pages < maximum) { + if (ring->slot_request[first_slot + pages] != KV_STREAM_NO_REQUEST || + !kv_stream_graph_request_follows( + ring->graph_requests[ring->next_request + pages - 1], + ring->graph_requests[ring->next_request + pages])) { + break; + } + ++pages; + } + return pages; +} + +static uint32_t kv_stream_graph_schedule_batch( + ggml_cuda_kv_stream_transfer_ring * ring, + uint32_t first_slot) { + const uint32_t batch_pages = kv_stream_graph_batch_pages(ring, first_slot); + if (batch_pages == 0) { + return 0; + } + const size_t first_request = ring->next_request; + // Probe every immutable copy batch at its actual compute deadline. Mutable + // tails are produced by this graph and are excluded from prefetch quality + // feedback so they do not force unnecessary resident-page demotions. + ring->graph_requests[first_request].deadline_sample = + !ring->graph_requests[first_request].mutable_tail; + ring->next_request += batch_pages; + if (batch_pages == 1) { + kv_stream_graph_upload( + ring, ring->graph_requests[first_request], first_slot); + } else { + kv_stream_graph_upload_batch( + ring, first_request, first_slot, batch_pages); + } + return batch_pages; +} + +static void kv_stream_graph_fill_free_slots(ggml_cuda_kv_stream_transfer_ring * ring) { + for (uint32_t slot = 0; slot < ring->active_slots; ++slot) { + if (ring->slot_request[slot] != KV_STREAM_NO_REQUEST) { + continue; + } + const uint32_t scheduled = kv_stream_graph_schedule_batch(ring, slot); + if (scheduled == 0) { + break; + } + slot += scheduled - 1; + } +} + +static bool kv_stream_graph_layer_begin( + ggml_cuda_kv_stream_transfer_ring * ring, + const void * k_key, + cudaStream_t compute_stream) { + if (!ring->graph_active) { + return false; + } + const auto layer_it = ring->graph_layer_by_k.find(k_key); + if (layer_it == ring->graph_layer_by_k.end()) { + return false; + } + + ring->current_layer = layer_it->second; + CUDA_CHECK(cudaEventRecord(ring->producer_ready, compute_stream)); + CUDA_CHECK(cudaStreamWaitEvent(ring->copy_stream, ring->producer_ready, 0)); + for (auto & request : ring->graph_requests) { + if (request.layer == ring->current_layer && request.mutable_tail) { + request.eligible = true; + } + } + kv_stream_graph_fill_free_slots(ring); + return true; +} + +static size_t kv_stream_graph_request_index( + const ggml_cuda_kv_stream_transfer_ring * ring, + const void * k_key, + uint32_t page) { + const auto it = ring->graph_request_by_k_page.find(k_key); + if (it == ring->graph_request_by_k_page.end() || page >= it->second.size()) { + return KV_STREAM_NO_REQUEST; + } + return it->second[page]; +} + +static void kv_stream_graph_release( + ggml_cuda_kv_stream_transfer_ring * ring, + size_t request_index, + cudaStream_t compute_stream) { + GGML_ASSERT(request_index < ring->graph_requests.size()); + auto & request = ring->graph_requests[request_index]; + GGML_ASSERT(request.scheduled && !request.consumed); + const uint32_t slot = request.slot; + GGML_ASSERT(ring->slot_request[slot] == request_index); + CUDA_CHECK(cudaEventRecord(ring->consumed[slot], compute_stream)); + request.consumed = true; + ring->slot_request[slot] = KV_STREAM_NO_REQUEST; + GGML_ASSERT(ring->current_occupancy > 0); + --ring->current_occupancy; +} + +} // namespace + +void ggml_cuda_kv_stream_graph_begin(ggml_cuda_kv_stream_transfer_ring * ring) { + GGML_ASSERT(ring != nullptr); + const bool timing_available = kv_stream_collect_timing(ring); + ring->graph_active = true; + ring->graph_decode = true; + ring->graph_decode_span_pages = ring->forced_decode_span_pages != 0 ? + ring->forced_decode_span_pages : + (ring->span_tuner.use_bounded() ? KV_STREAM_DECODE_SPAN_PAGES : UINT32_MAX); + ring->graph_layer_count = 0; + ring->current_layer = KV_STREAM_NO_LAYER; + ring->next_request = 0; + ring->graph_resident_cache = nullptr; + ring->graph_requests.clear(); + ring->graph_layer_by_k.clear(); + ring->graph_request_by_k_page.clear(); + ring->current_occupancy = 0; + if (timing_available) { + ring->current_ring_peak_occupancy = 0; + ring->current_epoch_uploads = 0; + ring->copy_sample_recorded = false; + ring->copy_sample_uploads = 0; + } + ring->timing_current = timing_available; + std::fill(ring->slot_request.begin(), ring->slot_request.end(), KV_STREAM_NO_REQUEST); +} + +bool ggml_cuda_kv_stream_graph_add_attention( + ggml_cuda_kv_stream_transfer_ring * ring, + ggml_cuda_kv_stream_resident_cache * resident_cache, + const ggml_tensor * dst) { + GGML_ASSERT(ring != nullptr); + const ggml_tensor * K = dst == nullptr ? nullptr : dst->src[1]; + const ggml_tensor * V = dst == nullptr ? nullptr : dst->src[2]; + if (resident_cache == nullptr || dst == nullptr || + !ggml_cuda_flash_attn_ext_streamed_supported(dst, ring->page_bytes)) { + return false; + } + if (ring->graph_resident_cache != nullptr && ring->graph_resident_cache != resident_cache) { + return false; + } + ring->graph_decode = ring->graph_decode && dst->src[0]->ne[1] == 1; + if (dst->src[0]->ne[1] != 1) { + // Graphs are rebuilt across warmup, prompt chunks, and slot reuse. + // Relearn pointer-to-layer identity once per prefill graph while the + // resident page contents are refreshed by the local multi-token path. + if (ring->graph_resident_cache == nullptr) { + resident_cache->layer_by_k.clear(); + resident_cache->layer_by_data.clear(); + resident_cache->next_layer = 0; + ring->graph_resident_cache = resident_cache; + } + const uint32_t resident_layer = kv_stream_resident_layer(resident_cache, K->data); + resident_cache->layer_by_data[K->data] = resident_layer; + resident_cache->layer_by_data[V->data] = resident_layer; + char * layer_base = resident_cache->pool_data + resident_cache->scratch_bytes + + resident_cache->layer_offsets[resident_layer]*resident_cache->page_bytes; + resident_cache->mirror_by_data[K->data] = layer_base; + resident_cache->mirror_by_data[V->data] = layer_base + + size_t(resident_cache->layer_pages[resident_layer])*K->nb[1]*resident_cache->page_tokens; + return false; + } + ring->graph_resident_cache = resident_cache; + // Resident placement is stable across evaluations, while deadlines must + // follow this graph's finalized execution order. + const uint32_t resident_layer = kv_stream_resident_layer(resident_cache, K->data); + resident_cache->layer_by_data[K->data] = resident_layer; + resident_cache->layer_by_data[V->data] = resident_layer; + char * layer_base = resident_cache->pool_data + resident_cache->scratch_bytes + + resident_cache->layer_offsets[resident_layer]*resident_cache->page_bytes; + resident_cache->mirror_by_data[K->data] = layer_base; + resident_cache->mirror_by_data[V->data] = layer_base + + size_t(resident_cache->layer_pages[resident_layer])*K->nb[1]*resident_cache->page_tokens; + + const int64_t block_tokens = resident_cache->page_tokens; + const int nchunks = int((K->ne[1] + block_tokens - 1)/block_tokens); + if (uint32_t(nchunks) <= resident_cache->layer_pages[resident_layer]) { + return true; + } + const uint32_t layer = ring->graph_layer_count++; + ring->graph_layer_by_k[K->data] = layer; + auto & page_requests = ring->graph_request_by_k_page[K->data]; + page_requests.assign(nchunks, KV_STREAM_NO_REQUEST); + + for (int chunk = 0; chunk < nchunks; ++chunk) { + const uint32_t page = uint32_t(chunk); + if (page < resident_cache->layer_pages[resident_layer]) { + continue; + } + const int64_t token_begin = chunk*block_tokens; + const int64_t token_count = std::min(block_tokens, K->ne[1] - token_begin); + const size_t k_row_bytes = ggml_row_size(K->type, K->ne[0]); + const size_t v_row_bytes = ggml_row_size(V->type, V->ne[0]); + + kv_stream_graph_request request; + request.k_data = static_cast(K->data); + request.v_data = static_cast(V->data); + request.k_nb1 = K->nb[1]; + request.k_nb2 = K->nb[2]; + request.v_nb1 = V->nb[1]; + request.v_nb2 = V->nb[2]; + request.n_head_kv = K->ne[2]; + request.token_begin = token_begin; + request.token_count = token_count; + request.k_row_bytes = k_row_bytes; + request.v_row_bytes = v_row_bytes; + request.k_head_bytes = k_row_bytes*token_count; + request.k_bytes = request.k_head_bytes*K->ne[2]; + request.v_offset = GGML_PAD(request.k_bytes, 128); + request.v_head_bytes = v_row_bytes*token_count; + request.v_bytes = request.v_head_bytes*V->ne[2]; + request.layer = layer; + request.mutable_tail = chunk == nchunks - 1 || + kv_stream_page_mutable(resident_cache, page); + request.eligible = !request.mutable_tail; + GGML_ASSERT(request.v_offset + request.v_bytes == ring->page_bytes); + + page_requests[page] = ring->graph_requests.size(); + ring->graph_requests.push_back(request); + } + return true; +} + +void ggml_cuda_kv_stream_graph_finalize( + ggml_cuda_kv_stream_transfer_ring * ring, cudaStream_t compute_stream) { + GGML_ASSERT(ring != nullptr); + ring->last_graph_decode = ring->graph_decode; + ring->last_graph_bounded = ring->graph_decode_span_pages != UINT32_MAX; + ring->last_graph_streamed = !ring->graph_requests.empty(); + if (ring->graph_requests.empty()) { + ring->timing_current = false; + return; + } + if (ring->timing_current) { + CUDA_CHECK(cudaEventRecord(ring->eval_start, compute_stream)); + } + kv_stream_graph_fill_free_slots(ring); +} + +void ggml_cuda_kv_stream_graph_end( + ggml_cuda_kv_stream_transfer_ring * ring, cudaStream_t compute_stream) { + GGML_ASSERT(ring != nullptr); + if (ring->timing_current) { + CUDA_CHECK(cudaEventRecord(ring->eval_end, compute_stream)); + ring->timing_pending = true; + ring->timing_current = false; + } +} + +double ggml_cuda_kv_stream_copy_engine_busy_ratio( + ggml_cuda_kv_stream_transfer_ring * ring) { + if (ring == nullptr) { + return 0.0; + } + (void) kv_stream_collect_timing(ring); + return ring->last_copy_engine_busy_ratio; +} + +uint32_t ggml_cuda_kv_stream_last_ring_peak_occupancy( + const ggml_cuda_kv_stream_transfer_ring * ring) { + return ring == nullptr ? 0 : ring->last_ring_peak_occupancy; +} + +void ggml_cuda_flash_attn_ext_streamed( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + ggml_cuda_kv_stream_transfer_ring * transfer_ring, + ggml_cuda_kv_stream_resident_cache * resident_cache) { + GGML_ASSERT(transfer_ring != nullptr); + void * stage_data = transfer_ring->pool_data; + const size_t stage_bytes = transfer_ring->page_bytes; + GGML_ASSERT(ggml_cuda_flash_attn_ext_streamed_supported(dst, stage_bytes)); + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const auto attention_mode = + ggml_backend_cuda_kv_stream_get_attention_mode(K->type, V->type); + const bool convert_to_f16 = + attention_mode == GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16; +#ifdef GGML_CUDA_FA_ALL_QUANTS + const kv_stream_native_partial_fn native_partial = + attention_mode == GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT ? + kv_stream_resolve_native_partial(K->type, V->type) : nullptr; + GGML_ASSERT(attention_mode != GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT || + native_partial != nullptr); +#endif // GGML_CUDA_FA_ALL_QUANTS + const to_fp16_cuda_t converter_k = convert_to_f16 && K->type != GGML_TYPE_F16 ? + ggml_get_to_fp16_cuda(K->type) : nullptr; + const to_fp16_cuda_t converter_v = convert_to_f16 && V->type != GGML_TYPE_F16 ? + ggml_get_to_fp16_cuda(V->type) : nullptr; + GGML_ASSERT(!convert_to_f16 || K->type == GGML_TYPE_F16 || converter_k != nullptr); + GGML_ASSERT(!convert_to_f16 || V->type == GGML_TYPE_F16 || converter_v != nullptr); + const int64_t block_tokens = resident_cache != nullptr ? + resident_cache->page_tokens : kv_stream_block_tokens(dst, stage_bytes); + const int nchunks = (K->ne[1] + block_tokens - 1)/block_tokens; + const int nrows = ggml_nrows(dst); + const uint32_t maximum_streamed_span_pages = Q->ne[1] == 1 ? + transfer_ring->graph_decode_span_pages : UINT32_MAX; + + struct chunk_descriptor { + int64_t token_begin = 0; + int64_t token_count = 0; + size_t k_row_bytes = 0; + size_t v_row_bytes = 0; + size_t k_head_bytes = 0; + size_t k_bytes = 0; + size_t v_offset = 0; + size_t v_head_bytes = 0; + size_t v_bytes = 0; + char * k_stage = nullptr; + char * v_stage = nullptr; + size_t k_stage_token_stride = 0; + size_t k_stage_head_stride = 0; + size_t v_stage_token_stride = 0; + size_t v_stage_head_stride = 0; + bool upload = true; + bool resident_refresh = false; + bool streamed = false; + uint32_t slot = 0; + size_t request_index = KV_STREAM_NO_REQUEST; + }; + + uint32_t resident_layer = 0; + if (resident_cache != nullptr) { + resident_layer = kv_stream_resident_layer(resident_cache, K->data); + } + const uint32_t resident_layer_pages = resident_cache == nullptr ? + 0 : resident_cache->layer_pages[resident_layer]; + const bool graph_planned = kv_stream_graph_layer_begin( + transfer_ring, K->data, ctx.stream()); + + std::vector chunks(nchunks); + std::vector streamed_chunks; + streamed_chunks.reserve(nchunks); + + for (int chunk = 0; chunk < nchunks; ++chunk) { + auto & desc = chunks[chunk]; + const int64_t token_begin = chunk*block_tokens; + const int64_t token_count = std::min(block_tokens, K->ne[1] - token_begin); + const size_t k_row_bytes = ggml_row_size(K->type, K->ne[0]); + const size_t v_row_bytes = ggml_row_size(V->type, V->ne[0]); + const size_t k_head_bytes = k_row_bytes*token_count; + const size_t k_bytes = k_head_bytes*K->ne[2]; + const size_t v_offset = GGML_PAD(k_bytes, 128); + const size_t v_head_bytes = v_row_bytes*token_count; + const size_t v_bytes = v_head_bytes*V->ne[2]; + GGML_ASSERT(v_offset <= stage_bytes && v_bytes <= stage_bytes - v_offset); + + desc.token_begin = token_begin; + desc.token_count = token_count; + desc.k_row_bytes = k_row_bytes; + desc.v_row_bytes = v_row_bytes; + desc.k_head_bytes = k_head_bytes; + desc.k_bytes = k_bytes; + desc.v_offset = v_offset; + desc.v_head_bytes = v_head_bytes; + desc.v_bytes = v_bytes; + desc.k_stage = static_cast(stage_data); + desc.v_stage = desc.k_stage + v_offset; + desc.k_stage_token_stride = k_row_bytes*K->ne[2]; + desc.k_stage_head_stride = k_row_bytes; + desc.v_stage_token_stride = v_row_bytes*V->ne[2]; + desc.v_stage_head_stride = v_row_bytes; + + if (resident_cache != nullptr) { + GGML_ASSERT(token_count == resident_cache->page_tokens); + GGML_ASSERT(v_offset + v_bytes == resident_cache->page_bytes); + + const uint32_t page = token_begin/resident_cache->page_tokens; + if (page < resident_layer_pages) { + const size_t resident_index = + kv_stream_resident_index(resident_cache, resident_layer, page); + // Keep each layer's resident K and V in separate token-major + // planes so resident pages form one directly consumable span. + char * layer_base = resident_cache->pool_data + resident_cache->scratch_bytes + + resident_cache->layer_offsets[resident_layer]*resident_cache->page_bytes; + const size_t resident_k_plane_bytes = + size_t(resident_layer_pages)*k_bytes; + desc.k_stage = layer_base + size_t(page)*k_bytes; + desc.v_stage = layer_base + resident_k_plane_bytes + size_t(page)*v_bytes; + if (resident_cache->loaded[resident_index]) { + ++resident_cache->stats.resident_hits; + desc.upload = resident_cache->precise_dirty_tracking[resident_layer] ? + resident_cache->dirty[resident_index] : + (dst->src[0]->ne[1] > 1 || chunk == nchunks - 1); + desc.resident_refresh = desc.upload; + } else { + ++resident_cache->stats.resident_misses; + resident_cache->loaded[resident_index] = 1; + } + } else { + ++resident_cache->stats.streamed_pages; + desc.streamed = true; + } + } else { + desc.streamed = true; + } + + if (desc.streamed) { + const size_t stream_index = streamed_chunks.size(); + if (graph_planned) { + desc.request_index = kv_stream_graph_request_index( + transfer_ring, K->data, uint32_t(chunk)); + GGML_ASSERT(desc.request_index != KV_STREAM_NO_REQUEST); + } else { + desc.slot = uint32_t(stream_index%transfer_ring->active_slots); + desc.k_stage = transfer_ring->pool_data + size_t(desc.slot)*desc.k_bytes; + desc.v_stage = transfer_ring->pool_data + + size_t(transfer_ring->active_slots)*desc.k_bytes + + size_t(desc.slot)*desc.v_bytes; + } + streamed_chunks.push_back(chunk); + } + } + + const bool use_mma_prefill = !convert_to_f16 && + Q->ne[1] > 1 && Q->ne[0] == 256 && V->ne[0] == 256 && + mask != nullptr && Q->ne[2] % K->ne[2] == 0 && Q->ne[2]/K->ne[2] <= 8; + const int partial_count = use_mma_prefill ? 1 : kv_stream_parts_per_chunk(); + GGML_ASSERT(partial_count > 0 && partial_count <= KV_STREAM_MAX_PARTS_PER_CHUNK); + GGML_ASSERT(Q->ne[1] > 0 && nrows % Q->ne[1] == 0); + const int64_t rows_per_query = nrows/Q->ne[1]; + const int64_t workspace_queries = use_mma_prefill ? Q->ne[1] : + std::min(Q->ne[1], KV_STREAM_QUERY_WORKSPACE_TOKENS); + const size_t workspace_rows = size_t(workspace_queries*rows_per_query); + const size_t workspace_elements = workspace_rows*dst->ne[0]; + + ggml_cuda_pool & pool = ctx.pool(); + ggml_cuda_pool_alloc parts(pool); + ggml_cuda_pool_alloc meta(pool); + ggml_cuda_pool_alloc accumulator(pool); + ggml_cuda_pool_alloc accumulator_meta(pool); + const bool needs_partial_reduction = convert_to_f16 || (!streamed_chunks.empty() && nchunks > 1); + if (needs_partial_reduction) { + parts.alloc(size_t(partial_count)*workspace_elements); + meta.alloc(size_t(partial_count)*workspace_rows); + accumulator.alloc(ggml_nelements(dst)); + accumulator_meta.alloc(nrows); + } + + auto upload = [&](const chunk_descriptor & desc, cudaStream_t stream) { + if (!desc.upload) { + return; + } + GGML_ASSERT(K->nb[1] == desc.k_stage_token_stride && K->nb[2] == desc.k_stage_head_stride); + GGML_ASSERT(V->nb[1] == desc.v_stage_token_stride && V->nb[2] == desc.v_stage_head_stride); + size_t dirty_row_count = 0; + if (resident_cache != nullptr && desc.resident_refresh && + !resident_cache->all_pages_mutable) { + for (const int64_t row : resident_cache->dirty_rows) { + if (row >= desc.token_begin && row < desc.token_begin + desc.token_count) { + ++dirty_row_count; + } + } + } + const size_t dirty_bytes = dirty_row_count*(desc.k_stage_token_stride + desc.v_stage_token_stride); + if (dirty_row_count > 0 && dirty_bytes < desc.k_bytes + desc.v_bytes) { + for (const int64_t row : resident_cache->dirty_rows) { + if (row < desc.token_begin || row >= desc.token_begin + desc.token_count) { + continue; + } + const size_t page_row = size_t(row - desc.token_begin); + CUDA_CHECK(cudaMemcpyAsync( + desc.k_stage + page_row*desc.k_stage_token_stride, + static_cast(K->data) + row*K->nb[1], + desc.k_stage_token_stride, cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync( + desc.v_stage + page_row*desc.v_stage_token_stride, + static_cast(V->data) + row*V->nb[1], + desc.v_stage_token_stride, cudaMemcpyHostToDevice, stream)); + } + transfer_ring->host_to_device_copy_commands += 2*dirty_row_count; + resident_cache->stats.host_to_device_bytes += dirty_bytes; + return; + } + CUDA_CHECK(cudaMemcpyAsync( + desc.k_stage, static_cast(K->data) + desc.token_begin*K->nb[1], + desc.k_bytes, cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync( + desc.v_stage, static_cast(V->data) + desc.token_begin*V->nb[1], + desc.v_bytes, cudaMemcpyHostToDevice, stream)); + transfer_ring->host_to_device_copy_commands += 2; + if (resident_cache != nullptr) { + resident_cache->stats.host_to_device_bytes += desc.k_bytes + desc.v_bytes; + } + }; + + auto schedule_streamed = [&](size_t stream_index) { + auto & desc = chunks[streamed_chunks[stream_index]]; + const uint32_t slot = desc.slot; + if (transfer_ring->slot_used[slot]) { + CUDA_CHECK(cudaStreamWaitEvent( + transfer_ring->copy_stream, transfer_ring->consumed[slot], 0)); + ++transfer_ring->stage_slot_reuses; + } + upload(desc, transfer_ring->copy_stream); + CUDA_CHECK(cudaEventRecord(transfer_ring->ready[slot], transfer_ring->copy_stream)); + transfer_ring->slot_used[slot] = 1; + ++transfer_ring->asynchronous_page_uploads; + }; + + if (!graph_planned && !streamed_chunks.empty()) { + // SET_ROWS and all other producers for this layer are ordered before + // this marker on the compute stream. The copy stream may then run + // independently while attention consumes previously prepared pages. + CUDA_CHECK(cudaEventRecord(transfer_ring->producer_ready, ctx.stream())); + CUDA_CHECK(cudaStreamWaitEvent( + transfer_ring->copy_stream, transfer_ring->producer_ready, 0)); + const size_t initial = std::min( + transfer_ring->active_slots, streamed_chunks.size()); + for (size_t i = 0; i < initial; ++i) { + schedule_streamed(i); + } + } + + size_t stream_index = 0; + for (int chunk = 0; chunk < nchunks; ++chunk) { + auto & desc = chunks[chunk]; + uint32_t streamed_span_pages = 0; + if (desc.streamed) { + streamed_span_pages = 1; + uint32_t ready_slot = desc.slot; + if (graph_planned) { + auto & request = transfer_ring->graph_requests[desc.request_index]; + GGML_ASSERT(request.scheduled && !request.consumed); + desc.slot = request.slot; + ready_slot = request.ready_slot; + desc.k_stage = transfer_ring->pool_data + size_t(desc.slot)*desc.k_bytes; + desc.v_stage = transfer_ring->pool_data + + size_t(transfer_ring->active_slots)*desc.k_bytes + + size_t(desc.slot)*desc.v_bytes; + if (request.deadline_sample) { + kv_stream_record_deadline<<<1, 1, 0, ctx.stream()>>>( + transfer_ring->ready_flags_device + desc.slot, + transfer_ring->deadline_counters_device + 0, + transfer_ring->deadline_counters_device + 1); + CUDA_CHECK(cudaGetLastError()); + } + } + CUDA_CHECK(cudaStreamWaitEvent(ctx.stream(), transfer_ring->ready[ready_slot], 0)); + ++transfer_ring->compute_stream_waits; + + // Coalesce ready pages that occupy consecutive plane slots. The + // head stride remains the full active-ring plane width, while the + // tensor's token extent grows across adjacent slots. + while (!convert_to_f16 && + streamed_span_pages < maximum_streamed_span_pages && + chunk + int(streamed_span_pages) < nchunks) { + auto & candidate = chunks[chunk + streamed_span_pages]; + uint32_t candidate_ready_slot = candidate.slot; + if (!candidate.streamed) { + break; + } + if (graph_planned) { + auto & request = transfer_ring->graph_requests[candidate.request_index]; + if (!request.scheduled || request.consumed) { + break; + } + candidate.slot = request.slot; + candidate_ready_slot = request.ready_slot; + candidate.k_stage = transfer_ring->pool_data + + size_t(candidate.slot)*candidate.k_bytes; + candidate.v_stage = transfer_ring->pool_data + + size_t(transfer_ring->active_slots)*candidate.k_bytes + + size_t(candidate.slot)*candidate.v_bytes; + } + if (candidate.slot != desc.slot + streamed_span_pages || + candidate.token_begin != desc.token_begin + + int64_t(streamed_span_pages)*block_tokens) { + break; + } + if (graph_planned) { + auto & request = transfer_ring->graph_requests[candidate.request_index]; + if (request.deadline_sample) { + kv_stream_record_deadline<<<1, 1, 0, ctx.stream()>>>( + transfer_ring->ready_flags_device + candidate.slot, + transfer_ring->deadline_counters_device + 0, + transfer_ring->deadline_counters_device + 1); + CUDA_CHECK(cudaGetLastError()); + } + } + if (candidate_ready_slot != ready_slot) { + CUDA_CHECK(cudaStreamWaitEvent( + ctx.stream(), transfer_ring->ready[candidate_ready_slot], 0)); + ++transfer_ring->compute_stream_waits; + ready_slot = candidate_ready_slot; + } + ++streamed_span_pages; + } + desc.token_count = int64_t(streamed_span_pages)*block_tokens; + if (resident_cache != nullptr) { + ++resident_cache->stats.streamed_attention_spans; + resident_cache->stats.streamed_pages_attended += streamed_span_pages; + } + } else { + if (convert_to_f16) { + // Generic quantized K/V is converted one page at a time into + // the bounded workspace. Keep resident pages separate so the + // fallback never creates a context-sized F16 allocation. + upload(desc, ctx.stream()); + if (desc.upload && resident_cache->precise_dirty_tracking[resident_layer]) { + resident_cache->dirty[kv_stream_resident_index( + resident_cache, resident_layer, uint32_t(chunk))] = 0; + } + ++resident_cache->stats.resident_attention_spans; + ++resident_cache->stats.resident_pages_attended; + } else { + // Native kernels consume the resident K/V prefix in one span. + if (chunk > 0) { + continue; + } + const uint32_t resident_span_pages = + std::min(uint32_t(nchunks), resident_layer_pages); + GGML_ASSERT(resident_span_pages > 0); + for (uint32_t page = 0; page < resident_span_pages; ++page) { + upload(chunks[page], ctx.stream()); + if (chunks[page].upload && + resident_cache->precise_dirty_tracking[resident_layer]) { + resident_cache->dirty[ + kv_stream_resident_index(resident_cache, resident_layer, page)] = 0; + } + } + desc.token_count = int64_t(resident_span_pages)*block_tokens; + desc.k_head_bytes = desc.k_row_bytes*desc.token_count; + desc.k_bytes = desc.k_head_bytes*K->ne[2]; + desc.v_head_bytes = desc.v_row_bytes*desc.token_count; + desc.v_bytes = desc.v_head_bytes*V->ne[2]; + ++resident_cache->stats.resident_attention_spans; + resident_cache->stats.resident_pages_attended += resident_span_pages; + } + } + + ggml_tensor staged_k = *K; + ggml_tensor staged_v = *V; + staged_k.data = desc.k_stage; + staged_k.ne[1] = desc.token_count; + staged_k.nb[1] = desc.k_stage_token_stride; + staged_k.nb[2] = desc.k_stage_head_stride; + staged_k.nb[3] = desc.k_stage_token_stride*desc.token_count; + staged_v.data = desc.v_stage; + staged_v.ne[1] = desc.token_count; + staged_v.nb[1] = desc.v_stage_token_stride; + staged_v.nb[2] = desc.v_stage_head_stride; + staged_v.nb[3] = desc.v_stage_token_stride*desc.token_count; + + ggml_tensor converted_k{}; + ggml_tensor converted_v{}; + if (convert_to_f16) { + GGML_ASSERT(transfer_ring->conversion_data != nullptr); + const size_t k_elements = size_t(K->ne[0])*desc.token_count*K->ne[2]; + const size_t v_elements = size_t(V->ne[0])*desc.token_count*V->ne[2]; + const size_t k_f16_bytes = k_elements*sizeof(half); + const size_t v_f16_offset = GGML_PAD(k_f16_bytes, 128); + const size_t v_f16_bytes = v_elements*sizeof(half); + GGML_ASSERT(v_f16_offset <= transfer_ring->conversion_bytes && + v_f16_bytes <= transfer_ring->conversion_bytes - v_f16_offset); + + auto convert_page = [&](const ggml_tensor & src, half * converted, + size_t elements, to_fp16_cuda_t converter) { + if (converter == nullptr) { + CUDA_CHECK(cudaMemcpyAsync( + converted, src.data, elements*sizeof(half), + cudaMemcpyDeviceToDevice, ctx.stream())); + return; + } + converter(src.data, converted, elements, ctx.stream()); + }; + + auto * converted_k_data = reinterpret_cast(transfer_ring->conversion_data); + auto * converted_v_data = reinterpret_cast( + transfer_ring->conversion_data + v_f16_offset); + convert_page(staged_k, converted_k_data, k_elements, converter_k); + convert_page(staged_v, converted_v_data, v_elements, converter_v); + + converted_k = staged_k; + converted_k.type = GGML_TYPE_F16; + converted_k.data = converted_k_data; + converted_k.nb[0] = sizeof(half); + converted_k.nb[1] = size_t(K->ne[0])*K->ne[2]*sizeof(half); + converted_k.nb[2] = size_t(K->ne[0])*sizeof(half); + converted_k.nb[3] = converted_k.nb[1]*converted_k.ne[1]; + converted_v = staged_v; + converted_v.type = GGML_TYPE_F16; + converted_v.data = converted_v_data; + converted_v.nb[0] = sizeof(half); + converted_v.nb[1] = size_t(V->ne[0])*V->ne[2]*sizeof(half); + converted_v.nb[2] = size_t(V->ne[0])*sizeof(half); + converted_v.nb[3] = converted_v.nb[1]*converted_v.ne[1]; + } + ggml_tensor staged_mask{}; + ggml_tensor * staged_mask_ptr = nullptr; + if (mask != nullptr) { + staged_mask = *mask; + staged_mask.data = static_cast(mask->data) + desc.token_begin*mask->nb[0]; + staged_mask.ne[0] = desc.token_count; + staged_mask_ptr = &staged_mask; + } + + ggml_tensor staged_dst = *dst; + staged_dst.src[1] = convert_to_f16 ? &converted_k : &staged_k; + staged_dst.src[2] = convert_to_f16 ? &converted_v : &staged_v; + staged_dst.src[3] = staged_mask_ptr; + + // Preserve normal CUDA flash attention when the active cache is fully resident or fits in one streamed page. + // This avoids a partial reduction and keeps logits identical to a non-streamed cache. + if (!convert_to_f16 && (streamed_chunks.empty() || nchunks == 1)) { + ggml_cuda_flash_attn_ext(ctx, &staged_dst); + if (desc.streamed) { + if (graph_planned) { + kv_stream_graph_release(transfer_ring, desc.request_index, ctx.stream()); + kv_stream_graph_fill_free_slots(transfer_ring); + } else { + CUDA_CHECK(cudaEventRecord(transfer_ring->consumed[desc.slot], ctx.stream())); + } + } + return; + } + + GGML_ASSERT(needs_partial_reduction); + + // Keep each staged KV span alive while all bounded query tiles consume it. + // This bounds vector scratch without issuing another H2D transfer for the span. + for (int64_t query_begin = 0; query_begin < Q->ne[1]; query_begin += workspace_queries) { + const int64_t query_count = std::min( + workspace_queries, Q->ne[1] - query_begin); + ggml_tensor query_q = *Q; + query_q.data = static_cast(Q->data) + query_begin*Q->nb[1]; + query_q.ne[1] = query_count; + + ggml_tensor query_mask{}; + ggml_tensor * query_mask_ptr = nullptr; + if (staged_mask_ptr != nullptr) { + query_mask = *staged_mask_ptr; + query_mask.data = static_cast(staged_mask_ptr->data) + + query_begin*mask->nb[1]; + query_mask.ne[1] = query_count; + query_mask_ptr = &query_mask; + } + + ggml_tensor query_dst = staged_dst; + query_dst.ne[1] = query_count; + query_dst.src[0] = &query_q; + query_dst.src[3] = query_mask_ptr; + + if (use_mma_prefill) { + ggml_cuda_flash_attn_ext_mma_f16_partial_case<256, 256, 8, 8>( + ctx, &query_dst, parts.ptr, meta.ptr); + if (resident_cache != nullptr) { + ++resident_cache->stats.mma_prefill_attention_spans; + } + } else { + if (convert_to_f16) { + ggml_cuda_flash_attn_ext_vec_partial_case< + KV_STREAM_HEAD_DIM, GGML_TYPE_F16, GGML_TYPE_F16>( + ctx, &query_dst, parts.ptr, meta.ptr, partial_count); + } else { +#ifdef GGML_CUDA_FA_ALL_QUANTS + GGML_ASSERT(native_partial != nullptr); + native_partial(ctx, &query_dst, parts.ptr, meta.ptr, partial_count); +#else + GGML_ABORT("native quantized KV streaming requires GGML_CUDA_FA_ALL_QUANTS"); +#endif // GGML_CUDA_FA_ALL_QUANTS + } + } + + const int tile_nrows = int(query_count*rows_per_query); + const size_t row_offset = size_t(query_begin*rows_per_query); + const dim3 blocks(tile_nrows, 1, 1); + const dim3 threads(KV_STREAM_HEAD_DIM, 1, 1); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch( + kv_stream_accumulate_chunk_results, launch_params, + parts.ptr, meta.ptr, accumulator.ptr + row_offset*dst->ne[0], + accumulator_meta.ptr + row_offset, tile_nrows, chunk == 0, partial_count); + CUDA_CHECK(cudaGetLastError()); + } + + if (desc.streamed) { + for (uint32_t page = 0; page < streamed_span_pages; ++page) { + auto & member = chunks[chunk + page]; + if (graph_planned) { + kv_stream_graph_release( + transfer_ring, member.request_index, ctx.stream()); + } else { + CUDA_CHECK(cudaEventRecord( + transfer_ring->consumed[member.slot], ctx.stream())); + const size_t next = stream_index + page + transfer_ring->active_slots; + if (next < streamed_chunks.size()) { + schedule_streamed(next); + } + } + } + stream_index += streamed_span_pages; + if (graph_planned) { + kv_stream_graph_fill_free_slots(transfer_ring); + } + chunk += int(streamed_span_pages) - 1; + } + } + + const dim3 blocks(nrows, 1, 1); + const dim3 threads(KV_STREAM_HEAD_DIM, 1, 1); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(kv_stream_normalize_chunk_results, launch_params, + accumulator.ptr, accumulator_meta.ptr, static_cast(dst->data), nrows); + CUDA_CHECK(cudaGetLastError()); +} template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { diff --git a/ggml/src/ggml-cuda/fattn.cuh b/ggml/src/ggml-cuda/fattn.cuh index f9a7e15fbd62..815d9635bd2c 100644 --- a/ggml/src/ggml-cuda/fattn.cuh +++ b/ggml/src/ggml-cuda/fattn.cuh @@ -5,3 +5,107 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst); size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * dst); + +struct ggml_cuda_kv_stream_resident_cache; +struct ggml_cuda_kv_stream_transfer_ring; + +struct ggml_cuda_kv_stream_resident_stats { + uint64_t resident_hits = 0; + uint64_t resident_misses = 0; + uint64_t streamed_pages = 0; + uint64_t host_to_device_bytes = 0; + uint64_t resident_attention_spans = 0; + uint64_t resident_pages_attended = 0; + uint64_t streamed_attention_spans = 0; + uint64_t streamed_pages_attended = 0; + uint64_t mma_prefill_attention_spans = 0; +}; + +struct ggml_cuda_kv_stream_transfer_stats { + uint64_t asynchronous_page_uploads = 0; + uint64_t host_to_device_copy_commands = 0; + uint64_t compute_stream_waits = 0; + uint64_t stage_slot_reuses = 0; + uint64_t cross_layer_prefetches = 0; + uint64_t deadline_samples = 0; + uint64_t deadline_misses = 0; + uint32_t ring_peak_occupancy = 0; +}; + +ggml_cuda_kv_stream_resident_cache * ggml_cuda_kv_stream_resident_cache_new( + void * pool_data, size_t pool_bytes, size_t scratch_bytes, size_t page_bytes, + uint32_t layer_count, uint32_t page_tokens); +void ggml_cuda_kv_stream_resident_cache_free(ggml_cuda_kv_stream_resident_cache * cache); +void ggml_cuda_kv_stream_resident_cache_reset(ggml_cuda_kv_stream_resident_cache * cache); +bool ggml_cuda_kv_stream_resident_cache_reconfigure( + ggml_cuda_kv_stream_resident_cache * cache, size_t scratch_bytes, + uint32_t active_pages_per_layer); +bool ggml_cuda_kv_stream_resident_cache_repartition( + ggml_cuda_kv_stream_resident_cache * cache, size_t scratch_bytes); +bool ggml_cuda_kv_stream_resident_cache_set_decode_layout( + ggml_cuda_kv_stream_resident_cache * cache, uint32_t active_pages_per_layer); +uint32_t ggml_cuda_kv_stream_resident_cache_pages_per_layer( + const ggml_cuda_kv_stream_resident_cache * cache); +uint32_t ggml_cuda_kv_stream_resident_cache_decode_active_pages( + const ggml_cuda_kv_stream_resident_cache * cache); +ggml_cuda_kv_stream_resident_stats ggml_cuda_kv_stream_resident_cache_get_stats( + const ggml_cuda_kv_stream_resident_cache * cache); +void ggml_cuda_kv_stream_resident_cache_mark_dirty( + ggml_cuda_kv_stream_resident_cache * cache, + const ggml_tensor * target, const ggml_tensor * indices); +bool ggml_cuda_kv_stream_resident_cache_mark_dirty_rows( + ggml_cuda_kv_stream_resident_cache * cache, + const int64_t * rows, size_t count); +bool ggml_cuda_kv_stream_resident_cache_all_layers_fit( + const ggml_cuda_kv_stream_resident_cache * cache, + uint32_t active_pages); +bool ggml_cuda_kv_stream_resident_cache_get_mirror( + ggml_cuda_kv_stream_resident_cache * cache, + const ggml_tensor * target, + void ** data); +void ggml_cuda_kv_stream_resident_cache_mark_mirrored( + ggml_cuda_kv_stream_resident_cache * cache, + const ggml_tensor * target); + +ggml_cuda_kv_stream_transfer_ring * ggml_cuda_kv_stream_transfer_ring_new( + void * pool_data, size_t page_bytes, uint32_t stage_slots, + void * conversion_data, size_t conversion_bytes, + uint32_t forced_decode_span_pages); +void ggml_cuda_kv_stream_transfer_ring_free(ggml_cuda_kv_stream_transfer_ring * ring); +bool ggml_cuda_kv_stream_transfer_ring_set_active_slots( + ggml_cuda_kv_stream_transfer_ring * ring, uint32_t stage_slots); +void ggml_cuda_kv_stream_transfer_ring_reset_span_tuner(ggml_cuda_kv_stream_transfer_ring * ring); +bool ggml_cuda_kv_stream_transfer_ring_observe_decode_latency( + ggml_cuda_kv_stream_transfer_ring * ring, double elapsed_ms); +ggml_cuda_kv_stream_transfer_stats ggml_cuda_kv_stream_transfer_ring_get_stats( + const ggml_cuda_kv_stream_transfer_ring * ring); +void ggml_cuda_kv_stream_graph_begin(ggml_cuda_kv_stream_transfer_ring * ring); +bool ggml_cuda_kv_stream_graph_add_attention( + ggml_cuda_kv_stream_transfer_ring * ring, + ggml_cuda_kv_stream_resident_cache * resident_cache, + const ggml_tensor * dst); +void ggml_cuda_kv_stream_graph_finalize( + ggml_cuda_kv_stream_transfer_ring * ring, cudaStream_t compute_stream); +void ggml_cuda_kv_stream_graph_end( + ggml_cuda_kv_stream_transfer_ring * ring, cudaStream_t compute_stream); +double ggml_cuda_kv_stream_copy_engine_busy_ratio( + ggml_cuda_kv_stream_transfer_ring * ring); +uint32_t ggml_cuda_kv_stream_last_ring_peak_occupancy( + const ggml_cuda_kv_stream_transfer_ring * ring); + +bool ggml_cuda_kv_stream_page_bytes( + ggml_type type_k, ggml_type type_v, + uint32_t head_dim_k, uint32_t head_dim_v, uint32_t head_count, + uint32_t page_tokens, size_t * page_bytes); +bool ggml_cuda_kv_stream_workspace_bytes( + ggml_type type_k, ggml_type type_v, + uint32_t head_dim_k, uint32_t head_dim_v, uint32_t head_count, + uint32_t page_tokens, size_t * workspace_bytes); + +bool ggml_cuda_flash_attn_ext_streamed_supported(const ggml_tensor * dst, size_t stage_bytes); + +void ggml_cuda_flash_attn_ext_streamed( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + ggml_cuda_kv_stream_transfer_ring * transfer_ring, + ggml_cuda_kv_stream_resident_cache * resident_cache); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index dcc80563b6db..afe85be0bfbb 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -89,6 +89,7 @@ #include #include #include +#include #include static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); @@ -765,6 +766,87 @@ static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) { return ctx->dev_ptr; } +static bool ggml_backend_cuda_buffer_set_preferred_device(ggml_backend_buffer_t buffer) { + if (!ggml_backend_buffer_is_cuda(buffer) || getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") == nullptr) { + return false; + } + + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; + cudaMemLocation location = {}; + location.type = cudaMemLocationTypeDevice; + location.id = ggml_cuda_get_physical_device(ctx->device); + + ggml_cuda_set_device(ctx->device); + const cudaError_t err = cudaMemAdvise( + ctx->dev_ptr, + ggml_backend_buffer_get_size(buffer), + cudaMemAdviseSetPreferredLocation, + location); + + if (err != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_WARN("failed to set GPU-preferred placement for %.2f MiB CUDA buffer: %s\n", + ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0, cudaGetErrorString(err)); + return false; + } + + GGML_LOG_INFO("set GPU-preferred placement for %.2f MiB model-weight buffer on device %d\n", + ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0, ctx->device); + return true; +} + + + +static bool ggml_backend_cuda_buffer_set_preferred_host(ggml_backend_buffer_t buffer) { + if (!ggml_backend_buffer_is_cuda(buffer) || getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") == nullptr) { + return false; + } + + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; + cudaMemLocation location = {}; + location.type = cudaMemLocationTypeHost; + + ggml_cuda_set_device(ctx->device); + cudaError_t err = cudaMemAdvise( + ctx->dev_ptr, + ggml_backend_buffer_get_size(buffer), + cudaMemAdviseSetPreferredLocation, + location); + + if (err != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_WARN("failed to set host-preferred placement for %.2f MiB KV buffer: %s\n", + ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0, cudaGetErrorString(err)); + return false; + } + + if (getenv("GGML_CUDA_KV_ACCESSED_BY_GPU") != nullptr) { + cudaMemLocation gpu_location = {}; + gpu_location.type = cudaMemLocationTypeDevice; + gpu_location.id = ggml_cuda_get_physical_device(ctx->device); + + err = cudaMemAdvise( + ctx->dev_ptr, + ggml_backend_buffer_get_size(buffer), + cudaMemAdviseSetAccessedBy, + gpu_location); + + if (err != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_WARN("failed to set GPU AccessedBy mapping for %.2f MiB KV buffer: %s\n", + ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0, cudaGetErrorString(err)); + return false; + } + + GGML_LOG_INFO("set GPU AccessedBy mapping for %.2f MiB KV buffer on device %d\n", + ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0, ctx->device); + } + + GGML_LOG_INFO("set host-preferred placement for %.2f MiB KV buffer on device %d\n", + ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0, ctx->device); + return true; +} + static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; @@ -1336,6 +1418,534 @@ ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { return &ggml_backend_cuda_buffer_type_host; } +struct ggml_backend_cuda_kv_stream_runtime { + int device = 0; + size_t stage_bytes = 0; + uint32_t stage_slots = 0; + size_t pool_bytes = 0; + size_t staging_pool_bytes = 0; + size_t conversion_bytes = 0; + uint32_t resident_layer_count = 0; + void * stage_data = nullptr; + ggml_cuda_kv_stream_transfer_ring * transfer_ring = nullptr; + ggml_cuda_kv_stream_resident_cache * resident_cache = nullptr; + std::vector dirty_rows; + uint32_t dirty_rows_remaining = 0; + uint64_t staged_set_rows = 0; + uint64_t staged_set_rows_bytes = 0; + uint64_t generation = 0; + + std::atomic references{1}; + ggml_backend_buffer_type buffer_type{}; +}; + +struct ggml_backend_cuda_kv_stream_buffer_context { + ggml_backend_cuda_kv_stream_runtime_t runtime = nullptr; + void * host_data = nullptr; +}; + +static void ggml_backend_cuda_kv_stream_runtime_release( + ggml_backend_cuda_kv_stream_runtime_t runtime) { + if (runtime == nullptr || runtime->references.fetch_sub(1, std::memory_order_acq_rel) != 1) { + return; + } + + ggml_cuda_set_device(runtime->device); + ggml_cuda_kv_stream_transfer_ring_free(runtime->transfer_ring); + ggml_cuda_kv_stream_resident_cache_free(runtime->resident_cache); + CUDA_CHECK(cudaFree(runtime->stage_data)); + delete runtime; +} + +static const char * ggml_backend_cuda_kv_stream_buffer_type_name(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return GGML_CUDA_NAME "_KV_Stream_Host"; +} + +static bool ggml_backend_buft_is_cuda_kv_stream(ggml_backend_buffer_type_t buft) { + return buft != nullptr && buft->iface.get_name == ggml_backend_cuda_kv_stream_buffer_type_name; +} + +static void ggml_backend_cuda_kv_stream_buffer_free(ggml_backend_buffer_t buffer) { + auto * context = static_cast(buffer->context); + CUDA_CHECK(cudaFreeHost(context->host_data)); + ggml_backend_cuda_kv_stream_runtime_release(context->runtime); + delete context; +} + +static void * ggml_backend_cuda_kv_stream_buffer_base(ggml_backend_buffer_t buffer) { + auto * context = static_cast(buffer->context); + return context->host_data; +} + +static void ggml_backend_cuda_kv_stream_buffer_memset( + ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + auto * context = static_cast(buffer->context); + ggml_cuda_kv_stream_resident_cache_reset(context->runtime->resident_cache); + ++context->runtime->generation; + memset(static_cast(tensor->data) + offset, value, size); +} + +static void ggml_backend_cuda_kv_stream_buffer_set( + ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + auto * context = static_cast(buffer->context); + ggml_cuda_kv_stream_resident_cache_reset(context->runtime->resident_cache); + ++context->runtime->generation; + memcpy(static_cast(tensor->data) + offset, data, size); +} + +static void ggml_backend_cuda_kv_stream_buffer_get( + ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + GGML_UNUSED(buffer); + memcpy(data, static_cast(tensor->data) + offset, size); +} + +static void ggml_backend_cuda_kv_stream_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + auto * context = static_cast(buffer->context); + ggml_cuda_kv_stream_resident_cache_reset(context->runtime->resident_cache); + ++context->runtime->generation; + memset(context->host_data, value, buffer->size); +} + +static const ggml_backend_buffer_i ggml_backend_cuda_kv_stream_buffer_interface = { + /* .free_buffer = */ ggml_backend_cuda_kv_stream_buffer_free, + /* .get_base = */ ggml_backend_cuda_kv_stream_buffer_base, + /* .init_tensor = */ nullptr, + /* .memset_tensor = */ ggml_backend_cuda_kv_stream_buffer_memset, + /* .set_tensor = */ ggml_backend_cuda_kv_stream_buffer_set, + /* .get_tensor = */ ggml_backend_cuda_kv_stream_buffer_get, + /* .set_tensor_2d = */ nullptr, + /* .get_tensor_2d = */ nullptr, + /* .cpy_tensor = */ nullptr, + /* .clear = */ ggml_backend_cuda_kv_stream_buffer_clear, + /* .reset = */ nullptr, +}; + +static ggml_backend_buffer_t ggml_backend_cuda_kv_stream_buffer_alloc( + ggml_backend_buffer_type_t buft, size_t size) { + auto * runtime = static_cast(buft->context); + ggml_cuda_set_device(runtime->device); + + void * host_data = nullptr; +#if defined(_WIN32) + // Direct decode writes to mapped write-combined host memory have been + // observed to fault under WDDM. Keep the storage mapped, but omit the + // write-combined hint on native Windows. + const unsigned int host_flags = cudaHostAllocMapped; +#else + const unsigned int host_flags = cudaHostAllocMapped | cudaHostAllocWriteCombined; +#endif + const cudaError_t error = cudaHostAlloc(&host_data, size, host_flags); + if (error != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_ERROR("%s: allocating %.2f MiB pinned KV storage failed: %s\n", + __func__, size/1024.0/1024.0, cudaGetErrorString(error)); + return nullptr; + } + + runtime->references.fetch_add(1, std::memory_order_relaxed); + auto * context = new ggml_backend_cuda_kv_stream_buffer_context{runtime, host_data}; + return ggml_backend_buffer_init(buft, ggml_backend_cuda_kv_stream_buffer_interface, context, size); +} + +static size_t ggml_backend_cuda_kv_stream_buffer_alignment(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return ggml_backend_buft_get_alignment(ggml_backend_cpu_buffer_type()); +} + +static size_t ggml_backend_cuda_kv_stream_buffer_alloc_size( + ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + GGML_UNUSED(buft); + return ggml_nbytes(tensor); +} + +static bool ggml_backend_cuda_kv_stream_buffer_is_host(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return true; +} + +static const ggml_backend_buffer_type_i ggml_backend_cuda_kv_stream_buffer_type_interface = { + /* .get_name = */ ggml_backend_cuda_kv_stream_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_cuda_kv_stream_buffer_alloc, + /* .get_alignment = */ ggml_backend_cuda_kv_stream_buffer_alignment, + /* .get_max_size = */ nullptr, + /* .get_alloc_size = */ ggml_backend_cuda_kv_stream_buffer_alloc_size, + /* .is_host = */ ggml_backend_cuda_kv_stream_buffer_is_host, +}; + +ggml_backend_cuda_kv_stream_runtime_t ggml_backend_cuda_kv_stream_runtime_new( + ggml_backend_cuda_kv_stream_params params) { + if (params.device < 0 || params.device >= ggml_backend_cuda_get_device_count() || + params.stage_bytes == 0 || params.stage_slots == 0 || + params.stage_bytes > std::numeric_limits::max()/params.stage_slots) { + return nullptr; + } + + const bool resident_enabled = + params.pool_bytes != 0 || params.resident_layer_count != 0 || params.page_tokens != 0; + const size_t scratch_bytes = params.stage_bytes*params.stage_slots; + if (params.conversion_bytes > std::numeric_limits::max() - scratch_bytes) { + return nullptr; + } + const size_t minimum_pool_bytes = scratch_bytes + params.conversion_bytes; + if (resident_enabled && (params.pool_bytes <= minimum_pool_bytes || + params.resident_layer_count == 0 || params.page_tokens != 256)) { + return nullptr; + } + + auto * runtime = new ggml_backend_cuda_kv_stream_runtime; + runtime->device = params.device; + runtime->stage_bytes = params.stage_bytes; + runtime->stage_slots = params.stage_slots; + runtime->pool_bytes = resident_enabled ? params.pool_bytes : minimum_pool_bytes; + runtime->conversion_bytes = params.conversion_bytes; + runtime->staging_pool_bytes = runtime->pool_bytes - runtime->conversion_bytes; + GGML_ASSERT(runtime->staging_pool_bytes >= scratch_bytes); + runtime->resident_layer_count = params.resident_layer_count; + + const size_t total_stage_bytes = runtime->pool_bytes; + ggml_cuda_set_device(params.device); + const cudaError_t error = cudaMalloc(&runtime->stage_data, total_stage_bytes); + if (error != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_ERROR("%s: allocating %.2f MiB KV staging on device %d failed: %s\n", + __func__, total_stage_bytes/1024.0/1024.0, params.device, cudaGetErrorString(error)); + delete runtime; + return nullptr; + } + + if (resident_enabled) { + runtime->resident_cache = ggml_cuda_kv_stream_resident_cache_new( + runtime->stage_data, runtime->staging_pool_bytes, scratch_bytes, params.stage_bytes, + params.resident_layer_count, params.page_tokens); + if (runtime->resident_cache == nullptr) { + CUDA_CHECK(cudaFree(runtime->stage_data)); + delete runtime; + return nullptr; + } + } + + const size_t ring_capacity = resident_enabled ? + runtime->staging_pool_bytes/runtime->stage_bytes : runtime->stage_slots; + if (ring_capacity == 0 || ring_capacity > UINT32_MAX) { + ggml_cuda_kv_stream_resident_cache_free(runtime->resident_cache); + CUDA_CHECK(cudaFree(runtime->stage_data)); + delete runtime; + return nullptr; + } + runtime->transfer_ring = ggml_cuda_kv_stream_transfer_ring_new( + runtime->stage_data, runtime->stage_bytes, uint32_t(ring_capacity), + runtime->conversion_bytes == 0 ? nullptr : + static_cast(runtime->stage_data) + runtime->staging_pool_bytes, + runtime->conversion_bytes, + params.decode_span_pages); + if (runtime->transfer_ring == nullptr) { + ggml_cuda_kv_stream_resident_cache_free(runtime->resident_cache); + CUDA_CHECK(cudaFree(runtime->stage_data)); + delete runtime; + return nullptr; + } + GGML_ASSERT(ggml_cuda_kv_stream_transfer_ring_set_active_slots( + runtime->transfer_ring, runtime->stage_slots)); + + runtime->buffer_type = { + /* .iface = */ ggml_backend_cuda_kv_stream_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), params.device), + /* .context = */ runtime, + }; + return runtime; +} + +void ggml_backend_cuda_kv_stream_runtime_free(ggml_backend_cuda_kv_stream_runtime_t runtime) { + ggml_backend_cuda_kv_stream_runtime_release(runtime); +} + +ggml_backend_buffer_type_t ggml_backend_cuda_kv_stream_buffer_type( + ggml_backend_cuda_kv_stream_runtime_t runtime) { + return runtime == nullptr ? nullptr : &runtime->buffer_type; +} + +size_t ggml_backend_cuda_kv_stream_stage_bytes(ggml_backend_cuda_kv_stream_runtime_t runtime) { + return runtime == nullptr ? 0 : runtime->stage_bytes; +} + +uint32_t ggml_backend_cuda_kv_stream_stage_slots(ggml_backend_cuda_kv_stream_runtime_t runtime) { + return runtime == nullptr ? 0 : runtime->stage_slots; +} + +uint32_t ggml_backend_cuda_kv_stream_resident_pages_per_layer( + ggml_backend_cuda_kv_stream_runtime_t runtime) { + return runtime == nullptr ? 0 : + ggml_cuda_kv_stream_resident_cache_pages_per_layer(runtime->resident_cache); +} + +bool ggml_backend_cuda_kv_stream_reconfigure( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t active_pages_per_layer, + uint32_t stage_slots) { + if (runtime == nullptr || runtime->resident_cache == nullptr || stage_slots == 0 || + runtime->stage_bytes > std::numeric_limits::max()/stage_slots) { + return false; + } + + const size_t scratch_bytes = runtime->stage_bytes*stage_slots; + if (scratch_bytes > runtime->staging_pool_bytes) { + return false; + } + + const bool ring_changed = runtime->stage_slots != stage_slots; + const bool layout_changed = + ggml_cuda_kv_stream_resident_cache_decode_active_pages( + runtime->resident_cache) != active_pages_per_layer; + if (!ring_changed && !layout_changed) { + return true; + } + + ggml_cuda_set_device(runtime->device); + CUDA_CHECK(cudaDeviceSynchronize()); + if (!ggml_cuda_kv_stream_resident_cache_reconfigure( + runtime->resident_cache, scratch_bytes, active_pages_per_layer) || + !ggml_cuda_kv_stream_transfer_ring_set_active_slots( + runtime->transfer_ring, stage_slots)) { + return false; + } + if (layout_changed && !ring_changed) { + ggml_cuda_kv_stream_transfer_ring_reset_span_tuner(runtime->transfer_ring); + } + + runtime->stage_slots = stage_slots; + runtime->dirty_rows.clear(); + runtime->dirty_rows_remaining = 0; + runtime->staged_set_rows = 0; + runtime->staged_set_rows_bytes = 0; + ++runtime->generation; + return true; +} + + +bool ggml_backend_cuda_kv_stream_repartition( + ggml_backend_cuda_kv_stream_runtime_t runtime, uint32_t stage_slots) { + if (runtime == nullptr || runtime->resident_cache == nullptr || stage_slots == 0 || + runtime->stage_bytes > std::numeric_limits::max()/stage_slots) { + return false; + } + + const size_t scratch_bytes = runtime->stage_bytes*stage_slots; + if (scratch_bytes > runtime->staging_pool_bytes) { + return false; + } + + const bool changed = runtime->stage_slots != stage_slots; + ggml_cuda_set_device(runtime->device); + CUDA_CHECK(cudaDeviceSynchronize()); + if (!ggml_cuda_kv_stream_transfer_ring_set_active_slots(runtime->transfer_ring, stage_slots) || + !ggml_cuda_kv_stream_resident_cache_repartition(runtime->resident_cache, scratch_bytes)) { + return false; + } + runtime->stage_slots = stage_slots; + runtime->dirty_rows.clear(); + runtime->dirty_rows_remaining = 0; + runtime->staged_set_rows = 0; + runtime->staged_set_rows_bytes = 0; + runtime->generation += changed; + return true; +} + +bool ggml_backend_cuda_kv_stream_set_decode_layout( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t active_pages_per_layer) { + if (runtime == nullptr || runtime->resident_cache == nullptr) { + return false; + } + if (ggml_cuda_kv_stream_resident_cache_decode_active_pages( + runtime->resident_cache) == active_pages_per_layer) { + return true; + } + + ggml_cuda_set_device(runtime->device); + CUDA_CHECK(cudaDeviceSynchronize()); + if (!ggml_cuda_kv_stream_resident_cache_set_decode_layout( + runtime->resident_cache, active_pages_per_layer)) { + return false; + } + ggml_cuda_kv_stream_transfer_ring_reset_span_tuner(runtime->transfer_ring); + runtime->dirty_rows.clear(); + runtime->dirty_rows_remaining = 0; + runtime->staged_set_rows = 0; + runtime->staged_set_rows_bytes = 0; + ++runtime->generation; + return true; +} + +bool ggml_backend_cuda_kv_stream_mark_dirty_rows( + ggml_backend_cuda_kv_stream_runtime_t runtime, + const int64_t * rows, size_t count) { + if (runtime == nullptr || rows == nullptr || count == 0 || + !ggml_cuda_kv_stream_resident_cache_mark_dirty_rows( + runtime->resident_cache, rows, count)) { + return false; + } + runtime->dirty_rows.assign(rows, rows + count); + // Every attention layer sharing this runtime updates both K and V. + const uint64_t set_rows_uses = 2ULL*runtime->resident_layer_count; + runtime->dirty_rows_remaining = set_rows_uses > UINT32_MAX ? + UINT32_MAX : uint32_t(set_rows_uses); + return true; +} + +ggml_backend_cuda_kv_stream_stats ggml_backend_cuda_kv_stream_get_stats( + ggml_backend_cuda_kv_stream_runtime_t runtime) { + if (runtime == nullptr) { + return {}; + } + const auto stats = ggml_cuda_kv_stream_resident_cache_get_stats(runtime->resident_cache); + const auto transfer_stats = ggml_cuda_kv_stream_transfer_ring_get_stats(runtime->transfer_ring); + return { + stats.resident_hits, + stats.resident_misses, + stats.streamed_pages, + stats.host_to_device_bytes, + stats.resident_attention_spans, + stats.resident_pages_attended, + stats.streamed_attention_spans, + stats.streamed_pages_attended, + stats.mma_prefill_attention_spans, + transfer_stats.asynchronous_page_uploads, + transfer_stats.host_to_device_copy_commands, + transfer_stats.compute_stream_waits, + transfer_stats.stage_slot_reuses, + transfer_stats.cross_layer_prefetches, + transfer_stats.deadline_samples, + transfer_stats.deadline_misses, + transfer_stats.ring_peak_occupancy, + runtime->staged_set_rows, + runtime->staged_set_rows_bytes, + }; +} + +bool ggml_backend_cuda_kv_stream_stage_upload( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t slot, + size_t offset, + const void * source, + size_t size) { + if (runtime == nullptr || slot >= runtime->stage_slots || + offset > runtime->stage_bytes || size > runtime->stage_bytes - offset || + (size > 0 && source == nullptr)) { + return false; + } + + ggml_cuda_set_device(runtime->device); + void * destination = static_cast(runtime->stage_data) + slot*runtime->stage_bytes + offset; + const cudaError_t error = cudaMemcpy(destination, source, size, cudaMemcpyHostToDevice); + if (error != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_ERROR("%s: staging upload failed: %s\n", __func__, cudaGetErrorString(error)); + return false; + } + return true; +} + +bool ggml_backend_cuda_kv_stream_stage_download( + ggml_backend_cuda_kv_stream_runtime_t runtime, + uint32_t slot, + size_t offset, + void * destination, + size_t size) { + if (runtime == nullptr || slot >= runtime->stage_slots || + offset > runtime->stage_bytes || size > runtime->stage_bytes - offset || + (size > 0 && destination == nullptr)) { + return false; + } + + ggml_cuda_set_device(runtime->device); + const void * source = static_cast(runtime->stage_data) + slot*runtime->stage_bytes + offset; + const cudaError_t error = cudaMemcpy(destination, source, size, cudaMemcpyDeviceToHost); + if (error != cudaSuccess) { + (void) cudaGetLastError(); + GGML_LOG_ERROR("%s: staging download failed: %s\n", __func__, cudaGetErrorString(error)); + return false; + } + return true; +} + +static ggml_backend_cuda_kv_stream_runtime_t ggml_cuda_kv_stream_runtime_from_tensor( + const ggml_tensor * tensor) { + if (tensor == nullptr || tensor->buffer == nullptr || + !ggml_backend_buft_is_cuda_kv_stream(tensor->buffer->buft)) { + return nullptr; + } + + auto * context = static_cast(tensor->buffer->context); + return context->runtime; +} + +static bool ggml_cuda_kv_stream_fattn_fits(const ggml_tensor * dst) { + GGML_ASSERT(dst != nullptr && dst->op == GGML_OP_FLASH_ATTN_EXT); + + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + auto * k_runtime = ggml_cuda_kv_stream_runtime_from_tensor(k); + auto * v_runtime = ggml_cuda_kv_stream_runtime_from_tensor(v); + if (k_runtime == nullptr || v_runtime == nullptr || k_runtime != v_runtime) { + // arm117 diagnostics: why did identity fail? (k/v bufts + buffers + geoms) + GGML_LOG_ERROR("arm117 fits=false: dst=%p ne=[%lld,%lld,%lld,%lld] type=%s " + "k=%p kbuf=%p %s | v=%p vbuf=%p %s | kr=%p vr=%p\n", + (void *) dst, (long long) dst->ne[0], (long long) dst->ne[1], (long long) dst->ne[2], (long long) dst->ne[3], + ggml_type_name(dst->type), + (void *) k, k ? (void *) k->buffer : nullptr, k && k->buffer ? ggml_backend_buffer_name(k->buffer) : "NULL", + (void *) v, v ? (void *) v->buffer : nullptr, v && v->buffer ? ggml_backend_buffer_name(v->buffer) : "NULL", + (void *) k_runtime, (void *) v_runtime); + return false; + } + return ggml_cuda_flash_attn_ext_streamed_supported(dst, k_runtime->stage_bytes); +} + +static void ggml_cuda_kv_stream_fattn( + ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + GGML_ASSERT(ggml_cuda_kv_stream_fattn_fits(dst)); + + const ggml_tensor * k = dst->src[1]; + auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(k); + GGML_ASSERT(runtime->device == ctx.device); + ggml_cuda_flash_attn_ext_streamed( + ctx, dst, runtime->transfer_ring, runtime->resident_cache); +} + +static std::vector ggml_cuda_kv_stream_prepare_graph( + const ggml_cgraph * cgraph, cudaStream_t compute_stream) { + std::vector runtimes; + std::unordered_set seen; + + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (node->op != GGML_OP_FLASH_ATTN_EXT) { + continue; + } + auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(node->src[1]); + if (runtime != nullptr && runtime == ggml_cuda_kv_stream_runtime_from_tensor(node->src[2]) && + seen.insert(runtime).second) { + ggml_cuda_kv_stream_graph_begin(runtime->transfer_ring); + runtimes.push_back(runtime); + } + } + + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (node->op != GGML_OP_FLASH_ATTN_EXT || !ggml_cuda_kv_stream_fattn_fits(node)) { + continue; + } + auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(node->src[1]); + (void) ggml_cuda_kv_stream_graph_add_attention( + runtime->transfer_ring, runtime->resident_cache, node); + } + + for (auto * runtime : runtimes) { + ggml_cuda_kv_stream_graph_finalize(runtime->transfer_ring, compute_stream); + } + return runtimes; +} + //static bool ggml_backend_buffer_is_cuda_host(ggml_backend_buffer_t buffer) { // return buffer->buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; //} @@ -2069,6 +2679,40 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * nb1, nb2, nb3, stream); } +static bool ggml_cuda_kv_stream_staged_set_rows_range( + const ggml_backend_cuda_kv_stream_runtime_t runtime, const ggml_tensor * dst, + int64_t * first_row, int64_t * row_count) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + if (runtime == nullptr || runtime->resident_cache == nullptr || + runtime->dirty_rows_remaining == 0 || runtime->dirty_rows.size() <= 1 || + src0 == nullptr || src1 == nullptr || src0->type != GGML_TYPE_F32 || + src1->type != GGML_TYPE_I64 || + !ggml_backend_cuda_kv_stream_get_type_capabilities( + dst->type).online_write || + src0->ne[1] != int64_t(runtime->dirty_rows.size()) || + ggml_nelements(src1) != int64_t(runtime->dirty_rows.size()) || + src0->ne[2] != 1 || src0->ne[3] != 1 || + dst->nb[1] != ggml_row_size(dst->type, dst->ne[0])) { + return false; + } + + const int64_t first = runtime->dirty_rows.front(); + const int64_t count = int64_t(runtime->dirty_rows.size()); + if (first < 0 || first + count > dst->ne[1]) { + return false; + } + for (int64_t i = 1; i < count; ++i) { + if (runtime->dirty_rows[size_t(i)] != first + i) { + return false; + } + } + + *first_row = first; + *row_count = count; + return true; +} + static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { switch (dst->op) { case GGML_OP_ARGMAX: @@ -2089,9 +2733,53 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_GET_ROWS_BACK: ggml_cuda_op_get_rows_back(ctx, dst); break; - case GGML_OP_SET_ROWS: - ggml_cuda_op_set_rows(ctx, dst); + case GGML_OP_SET_ROWS: { + bool staged = false; + bool mirrored = false; + if (auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(dst); + runtime != nullptr && runtime->resident_cache != nullptr) { + ggml_cuda_kv_stream_resident_cache_mark_dirty( + runtime->resident_cache, dst, dst->src[1]); + int64_t first_row = 0; + int64_t row_count = 0; + staged = ggml_cuda_kv_stream_staged_set_rows_range( + runtime, dst, &first_row, &row_count); + if (runtime->dirty_rows_remaining > 0) { + --runtime->dirty_rows_remaining; + } + if (staged) { + void * mirror_data = nullptr; + ggml_cuda_kv_stream_resident_cache_get_mirror( + runtime->resident_cache, dst, &mirror_data); + ggml_cuda_op_set_rows_staged( + ctx, dst, first_row, row_count, mirror_data); + if (mirror_data != nullptr) { + ggml_cuda_kv_stream_resident_cache_mark_mirrored( + runtime->resident_cache, dst); + } + ++runtime->staged_set_rows; + runtime->staged_set_rows_bytes += size_t(row_count)*dst->nb[1]; + } else if (dst->src[0]->type == GGML_TYPE_F32 && + dst->src[1]->type == GGML_TYPE_I64 && + ggml_backend_cuda_kv_stream_get_type_capabilities(dst->type).online_write) { + void * mirror_data = nullptr; + if (ggml_cuda_kv_stream_resident_cache_get_mirror( + runtime->resident_cache, dst, &mirror_data)) { + ggml_cuda_op_set_rows(ctx, dst); + ggml_tensor mirror = *dst; + mirror.data = mirror_data; + ggml_cuda_op_set_rows(ctx, &mirror); + ggml_cuda_kv_stream_resident_cache_mark_mirrored( + runtime->resident_cache, dst); + mirrored = true; + } + } + } + if (!staged && !mirrored) { + ggml_cuda_op_set_rows(ctx, dst); + } break; + } case GGML_OP_SET: ggml_cuda_op_set(ctx, dst); break; @@ -2370,7 +3058,23 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg ggml_cuda_op_argsort(ctx, dst); break; case GGML_OP_FLASH_ATTN_EXT: - ggml_cuda_flash_attn_ext(ctx, dst); + // arm117 fix: gate the streamed dispatch on the full fits() + // predicate, not buffer membership alone. A multi-stream ubatch + // (parallel slots, K/V 4D views with ne[3] == ns > 1) breaks + // flashed_supported's single-stream geometric assumptions + // (Q/K/V ne[3] == 1); previously such a node still entered + // ggml_cuda_kv_stream_fattn and died on its GGML_ASSERT. + if (ggml_cuda_kv_stream_runtime_from_tensor(dst->src[1]) != nullptr || + ggml_cuda_kv_stream_runtime_from_tensor(dst->src[2]) != nullptr) { + if (ggml_cuda_kv_stream_fattn_fits(dst)) { + ggml_cuda_kv_stream_fattn(ctx, dst); + } else { + GGML_LOG_WARN("kv-stream: FA node not streaming (unsupported geometry, falling back to resident KV path)\n"); + ggml_cuda_flash_attn_ext(ctx, dst); + } + } else { + ggml_cuda_flash_attn_ext(ctx, dst); + } break; case GGML_OP_CROSS_ENTROPY_LOSS: ggml_cuda_cross_entropy_loss(ctx, dst); @@ -2572,6 +3276,40 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { continue; } + auto * node_runtime = ggml_cuda_kv_stream_runtime_from_tensor(node); + bool mirrorable_set_rows = false; + if (node_runtime != nullptr) { + mirrorable_set_rows = node->op == GGML_OP_SET_ROWS && + node_runtime->dirty_rows.size() == 1 && node_runtime->dirty_rows[0] >= 0 && + node->src[0] != nullptr && node->src[0]->type == GGML_TYPE_F32 && + node->src[1] != nullptr && node->src[1]->type == GGML_TYPE_I64 && + ggml_backend_cuda_kv_stream_get_type_capabilities( + node->type).online_write; + void * mirror_data = nullptr; + mirrorable_set_rows = mirrorable_set_rows && + ggml_cuda_kv_stream_resident_cache_get_mirror( + node_runtime->resident_cache, node, &mirror_data); + if (!mirrorable_set_rows) { + use_cuda_graph = false; + } + } + + for (int j = 0; use_cuda_graph && j < GGML_MAX_SRC; ++j) { + auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(node->src[j]); + if (runtime == nullptr) { + continue; + } + const bool resident_attention = mirrorable_set_rows || + (node->op == GGML_OP_FLASH_ATTN_EXT && + node->src[0]->ne[1] == 1 && ggml_cuda_kv_stream_fattn_fits(node) && + ggml_cuda_kv_stream_resident_cache_all_layers_fit( + runtime->resident_cache, + uint32_t((node->src[1]->ne[1] + 255)/256))); + if (!resident_attention) { + use_cuda_graph = false; + } + } + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (node->op == GGML_OP_MUL_MAT_ID) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; @@ -2603,8 +3341,26 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx const void * graph_key = ggml_cuda_graph_get_key(cgraph); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + uint64_t kv_stream_generation = 0; + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(cgraph->nodes[i])) { + kv_stream_generation = std::max(kv_stream_generation, runtime->generation); + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(cgraph->nodes[i]->src[j])) { + kv_stream_generation = std::max(kv_stream_generation, runtime->generation); + } + } + } + uint64_t previous_kv_stream_generation = 0; + for (const auto & prop : graph->node_props) { + previous_kv_stream_generation = std::max( + previous_kv_stream_generation, prop.kv_stream_generation); + } + if (cgraph->uid != 0 && - cgraph->uid == graph->uid) { + cgraph->uid == graph->uid && + kv_stream_generation == previous_kv_stream_generation) { GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); return false; @@ -2622,8 +3378,16 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx ggml_cuda_graph::node_properties prop = {}; memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor)); + if (auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(cgraph->nodes[i])) { + prop.kv_stream_generation = runtime->generation; + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { if (cgraph->nodes[i]->src[j]) { + if (auto * runtime = ggml_cuda_kv_stream_runtime_from_tensor(cgraph->nodes[i]->src[j])) { + prop.kv_stream_generation = std::max( + prop.kv_stream_generation, runtime->generation); + } prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data; memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j])); memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j])); @@ -4285,6 +5049,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud stream_ctx.concurrent_events.clear(); } + // The graph order is now final. Build one shared deadline queue + // per KV runtime before issuing any attention work. + const auto kv_stream_runtimes = ggml_cuda_kv_stream_prepare_graph( + cgraph, cuda_ctx->stream()); + for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -4344,11 +5113,13 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // node's output on the host-visible buffer, which the compute path // handles. Allow that here, mirroring the src-tensor check below. assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || + ggml_backend_buft_is_cuda_kv_stream(node->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->buffer->buft))); for (int j = 0; j < GGML_MAX_SRC; j++) { if (node->src[j] != nullptr) { assert(node->src[j]->buffer); assert(node->src[j]->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || + ggml_backend_buft_is_cuda_kv_stream(node->src[j]->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->src[j]->buffer->buft))); } } @@ -4366,6 +5137,9 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud try_launch_concurrent_event(node); } } + for (auto * runtime : kv_stream_runtimes) { + ggml_cuda_kv_stream_graph_end(runtime->transfer_ring, cuda_ctx->stream()); + } } #ifdef USE_CUDA_GRAPH @@ -5068,6 +5842,7 @@ static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_host_buffer_type( // TODO: move these functions here static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + bool uses_streamed_kv = false; // check if all the sources are allocated on this device for (int i = 0; i < GGML_MAX_SRC; i++) { @@ -5077,6 +5852,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return false; } } + if (op->src[i] && op->src[i]->buffer && + ggml_backend_buft_is_cuda_kv_stream(op->src[i]->buffer->buft)) { + if (op->src[i]->buffer->buft->device != dev) { + return false; + } + uses_streamed_kv = true; + } + } + + if (uses_streamed_kv && op->op != GGML_OP_SET_ROWS && op->op != GGML_OP_FLASH_ATTN_EXT) { + return false; } switch (op->op) { @@ -5506,6 +6292,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_FLASH_ATTN_EXT: + if (uses_streamed_kv) { + return ggml_cuda_kv_stream_fattn_fits(op); + } return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); case GGML_OP_CROSS_ENTROPY_LOSS: case GGML_OP_CROSS_ENTROPY_LOSS_BACK: @@ -5528,7 +6317,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g static bool ggml_backend_cuda_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; const bool integrated = ggml_cuda_info().devices[dev_ctx->device].integrated; - return (ggml_backend_buft_is_cuda(buft) && buft->device == dev) || (integrated && ggml_backend_buft_is_cuda_host(buft)); + return (ggml_backend_buft_is_cuda(buft) && buft->device == dev) || + (ggml_backend_buft_is_cuda_kv_stream(buft) && buft->device == dev) || + (integrated && ggml_backend_buft_is_cuda_host(buft)); } static int64_t get_op_batch_size(const ggml_tensor * op) { @@ -5697,6 +6488,180 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) { return (void *)ggml_backend_cuda_unregister_host_buffer; } + if (strcmp(name, "ggml_backend_cuda_buffer_set_preferred_device") == 0) { + return (void *) ggml_backend_cuda_buffer_set_preferred_device; + } + if (strcmp(name, "ggml_backend_cuda_buffer_set_preferred_host") == 0) { + return (void *) ggml_backend_cuda_buffer_set_preferred_host; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_type_pair_supported") == 0) { + return (void *) +[](ggml_type type_k, ggml_type type_v) { + return ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v) != + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_UNSUPPORTED; + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_page_bytes") == 0) { + return (void *) +[](ggml_type type_k, ggml_type type_v, + uint32_t head_dim_k, uint32_t head_dim_v, uint32_t head_count, + uint32_t page_tokens, size_t * page_bytes) { + return ggml_cuda_kv_stream_page_bytes( + type_k, type_v, head_dim_k, head_dim_v, head_count, page_tokens, page_bytes); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_workspace_bytes") == 0) { + return (void *) +[](ggml_type type_k, ggml_type type_v, + uint32_t head_dim_k, uint32_t head_dim_v, uint32_t head_count, + uint32_t page_tokens, size_t * workspace_bytes) { + return ggml_cuda_kv_stream_workspace_bytes( + type_k, type_v, head_dim_k, head_dim_v, head_count, page_tokens, workspace_bytes); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_runtime_new_for_device") == 0) { + return (void *) +[](ggml_backend_dev_t dev, size_t pool_bytes, size_t page_bytes, + size_t conversion_bytes, + uint32_t resident_layer_count) -> void * { + if (dev == nullptr || dev->iface.get_name != ggml_backend_cuda_device_get_name) { + return nullptr; + } + auto * dev_ctx = static_cast(dev->context); + if (page_bytes == 0 || conversion_bytes > pool_bytes) { + return nullptr; + } + const size_t staging_pool_bytes = pool_bytes - conversion_bytes; + if (staging_pool_bytes/page_bytes <= resident_layer_count) { + return nullptr; + } + const size_t pool_pages = staging_pool_bytes/page_bytes; + size_t minimum_slots = std::min( + 8, pool_pages - resident_layer_count); + // Diagnostic override for reproducible resident:ring sweeps. The + // llama-layer controller treats its presence as a fixed boundary. + if (const char * fixed_slots = getenv("GGML_CUDA_KV_STREAM_FIXED_RING_SLOTS")) { + char * end = nullptr; + const unsigned long requested = strtoul(fixed_slots, &end, 10); + if (end == fixed_slots || *end != '\0' || requested == 0 || requested >= pool_pages) { + GGML_LOG_ERROR("%s: invalid GGML_CUDA_KV_STREAM_FIXED_RING_SLOTS=%s for %zu pool pages\n", + __func__, fixed_slots, pool_pages); + return nullptr; + } + minimum_slots = size_t(requested); + } + const size_t resident_pages_per_layer = + (pool_pages - minimum_slots)/resident_layer_count; + const size_t stage_slots_wide = + pool_pages - resident_pages_per_layer*resident_layer_count; + if (stage_slots_wide == 0 || stage_slots_wide > UINT32_MAX) { + return nullptr; + } + const uint32_t stage_slots = uint32_t(stage_slots_wide); + ggml_backend_cuda_kv_stream_params params{}; + params.device = dev_ctx->device; + params.stage_bytes = page_bytes; + params.stage_slots = stage_slots; + params.pool_bytes = pool_bytes; + params.resident_layer_count = resident_layer_count; + params.page_tokens = 256; + params.conversion_bytes = conversion_bytes; + return ggml_backend_cuda_kv_stream_runtime_new(params); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_runtime_free") == 0) { + return (void *) +[](void * runtime) { + ggml_backend_cuda_kv_stream_runtime_free( + static_cast(runtime)); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_buffer_type") == 0) { + return (void *) +[](void * runtime) -> ggml_backend_buffer_type_t { + return ggml_backend_cuda_kv_stream_buffer_type( + static_cast(runtime)); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_observe_decode_latency") == 0) { + return (void *) +[](void * opaque, double elapsed_ms) -> bool { + auto * runtime = static_cast(opaque); + return runtime != nullptr && + ggml_cuda_kv_stream_transfer_ring_observe_decode_latency( + runtime->transfer_ring, elapsed_ms); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_feedback") == 0) { + return (void *) +[](void * opaque, + uint64_t * deadline_samples, + uint64_t * deadline_misses, + double * copy_engine_busy_ratio, + uint32_t * ring_peak_occupancy, + uint32_t * ring_slots, + uint32_t * resident_pages_per_layer, + uint32_t * controlled_pool_pages) -> bool { + auto * runtime = static_cast(opaque); + if (runtime == nullptr || deadline_samples == nullptr || deadline_misses == nullptr || + copy_engine_busy_ratio == nullptr || ring_peak_occupancy == nullptr || + ring_slots == nullptr || resident_pages_per_layer == nullptr || + controlled_pool_pages == nullptr || runtime->resident_cache == nullptr) { + return false; + } + const auto stats = ggml_cuda_kv_stream_transfer_ring_get_stats(runtime->transfer_ring); + *deadline_samples = stats.deadline_samples; + *deadline_misses = stats.deadline_misses; + *copy_engine_busy_ratio = + ggml_cuda_kv_stream_copy_engine_busy_ratio(runtime->transfer_ring); + *ring_peak_occupancy = + ggml_cuda_kv_stream_last_ring_peak_occupancy(runtime->transfer_ring); + *ring_slots = runtime->stage_slots; + *resident_pages_per_layer = + ggml_cuda_kv_stream_resident_cache_pages_per_layer(runtime->resident_cache); + const uint64_t controlled = uint64_t(*ring_slots) + + uint64_t(*resident_pages_per_layer)*runtime->resident_layer_count; + if (controlled > UINT32_MAX) { + return false; + } + *controlled_pool_pages = uint32_t(controlled); + return true; + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_host_is_write_combined") == 0) { + return (void *) +[](const void * pointer, bool * write_combined) -> bool { + if (pointer == nullptr || write_combined == nullptr) { + return false; + } + unsigned int flags = 0; + const cudaError_t error = cudaHostGetFlags(&flags, const_cast(pointer)); + if (error != cudaSuccess) { + (void) cudaGetLastError(); + return false; + } + *write_combined = (flags & cudaHostAllocWriteCombined) != 0; + return true; + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_reconfigure") == 0) { + return (void *) +[](void * runtime, uint32_t active_pages, uint32_t ring_slots) -> bool { + return ggml_backend_cuda_kv_stream_reconfigure( + static_cast(runtime), + active_pages, ring_slots); + }; + } + + if (strcmp(name, "ggml_backend_cuda_kv_stream_repartition") == 0) { + return (void *) +[](void * runtime, uint32_t ring_slots) -> bool { + return ggml_backend_cuda_kv_stream_repartition( + static_cast(runtime), ring_slots); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_set_decode_layout") == 0) { + return (void *) +[](void * runtime, uint32_t active_pages_per_layer) -> bool { + return ggml_backend_cuda_kv_stream_set_decode_layout( + static_cast(runtime), + active_pages_per_layer); + }; + } + if (strcmp(name, "ggml_backend_cuda_kv_stream_mark_dirty_rows") == 0) { + return (void *) +[](void * runtime, const int64_t * rows, size_t count) -> bool { + return ggml_backend_cuda_kv_stream_mark_dirty_rows( + static_cast(runtime), rows, count); + }; + } if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } diff --git a/ggml/src/ggml-cuda/kv-stream-span-tuner.h b/ggml/src/ggml-cuda/kv-stream-span-tuner.h new file mode 100644 index 000000000000..958c67533cc2 --- /dev/null +++ b/ggml/src/ggml-cuda/kv-stream-span-tuner.h @@ -0,0 +1,90 @@ +#pragma once + +#include +#include +#include + +// Chooses between the ordinary coalesced streamed-attention kernel and a +// bounded-span pipeline using end-to-end decode graph timings. The tuner is +// reset whenever the resident/ring layout changes. +class ggml_cuda_kv_stream_span_tuner { +public: + explicit ggml_cuda_kv_stream_span_tuner( + uint32_t trial_samples = 16, + double minimum_relative_gain = 0.005, + uint32_t warmup_samples = 1) : + trial_samples_(std::max(trial_samples, 1)), + minimum_relative_gain_( + std::isfinite(minimum_relative_gain) ? + std::clamp(minimum_relative_gain, 0.0, 0.999) : 0.005), + warmup_samples_(warmup_samples) { + } + + void reset() { + bounded_ = false; + selected_ = false; + unbounded_samples_ = 0; + bounded_samples_ = 0; + unbounded_warmups_ = 0; + bounded_warmups_ = 0; + unbounded_ms_ = 0.0; + bounded_ms_ = 0.0; + } + + void observe(double elapsed_ms, bool streamed, bool sample_was_bounded) { + if (selected_ || !streamed || !std::isfinite(elapsed_ms) || elapsed_ms <= 0.0) { + return; + } + + uint32_t & warmups = sample_was_bounded ? bounded_warmups_ : unbounded_warmups_; + if (warmups < warmup_samples_) { + ++warmups; + return; + } + + if (sample_was_bounded) { + bounded_ms_ += elapsed_ms; + ++bounded_samples_; + } else { + unbounded_ms_ += elapsed_ms; + ++unbounded_samples_; + } + + if (unbounded_samples_ < trial_samples_) { + bounded_ = false; + return; + } + + if (bounded_samples_ < trial_samples_) { + bounded_ = true; + return; + } + + const double unbounded_average_ms = unbounded_ms_ / unbounded_samples_; + const double bounded_average_ms = bounded_ms_ / bounded_samples_; + bounded_ = bounded_average_ms < unbounded_average_ms * (1.0 - minimum_relative_gain_); + selected_ = true; + } + + bool use_bounded() const { return bounded_; } + bool selected() const { return selected_; } + double unbounded_average_ms() const { + return unbounded_samples_ == 0 ? 0.0 : unbounded_ms_ / unbounded_samples_; + } + double bounded_average_ms() const { + return bounded_samples_ == 0 ? 0.0 : bounded_ms_ / bounded_samples_; + } + +private: + uint32_t trial_samples_ = 16; + double minimum_relative_gain_ = 0.005; + uint32_t warmup_samples_ = 1; + uint32_t unbounded_warmups_ = 0; + uint32_t bounded_warmups_ = 0; + bool bounded_ = false; + bool selected_ = false; + uint32_t unbounded_samples_ = 0; + uint32_t bounded_samples_ = 0; + double unbounded_ms_ = 0.0; + double bounded_ms_ = 0.0; +}; diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu index 4659970651e5..7ea5cd226670 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/ggml/src/ggml-cuda/set-rows.cu @@ -26,7 +26,8 @@ static __global__ void k_set_rows_quant(const float * __restrict__ src0, const uint3 ne01, const uint3 ne02, const uint3 ne11_fd, - const uint3 ne12_fd) { + const uint3 ne12_fd, + const int64_t dst_row_base) { const int64_t i = int64_t(blockDim.x) * blockIdx.x + threadIdx.x; if (i >= ne_total) { @@ -54,7 +55,7 @@ static __global__ void k_set_rows_quant(const float * __restrict__ src0, const int64_t i10 = i01; ggml_cuda_pdl_sync(); - const int64_t dst_row = *(src1 + i10*s10 + i11*s11 + i12*s12); + const int64_t dst_row = *(src1 + i10*s10 + i11*s11 + i12*s12) - dst_row_base; const float * src0_row = src0 + i01*s01 + i02*s02 + i03*s03; block_type * dst_row_ptr = dst + (dst_row*s1 + i02*s2 + i03*s3) / sizeof(block_type); @@ -79,7 +80,8 @@ static void set_rows_cuda_quant( const size_t nb01, const size_t nb02, const size_t nb03, const size_t nb10, const size_t nb11, const size_t nb12, const size_t nb1, const size_t nb2, const size_t nb3, - cudaStream_t stream) { + cudaStream_t stream, + const int64_t dst_row_base) { GGML_ASSERT(ne00 % qk == 0); const int64_t ne_total = (ne00 * ne01 * ne02 * ne03) / qk; @@ -106,7 +108,7 @@ static void set_rows_cuda_quant( k_set_rows_quant<<>>( src0_d, src1_d, dst_d, ne_total, ne10, ne11, ne12, ne13, s01, s02, s03, s10, s11, s12, s1, s2, s3, ne00_fd, - ne01_fd, ne02_fd, ne11_fd, ne12_fd); + ne01_fd, ne02_fd, ne11_fd, ne12_fd, dst_row_base); } } @@ -132,7 +134,8 @@ static __global__ void k_set_rows(const src_t * src0_ptr, const uint3 ne01, const uint3 ne02, const uint3 ne11_fd, - const uint3 ne12_fd) { + const uint3 ne12_fd, + const int64_t dst_row_base) { const src_t * GGML_CUDA_RESTRICT src0 = src0_ptr; const idx_t * GGML_CUDA_RESTRICT src1 = src1_ptr; dst_t * GGML_CUDA_RESTRICT dst = dst_ptr; @@ -162,7 +165,7 @@ static __global__ void k_set_rows(const src_t * src0_ptr, const int64_t i10 = i01; ggml_cuda_pdl_sync(); - const int64_t dst_row = *(src1 + i10*s10 + i11*s11 + i12*s12); + const int64_t dst_row = *(src1 + i10*s10 + i11*s11 + i12*s12) - dst_row_base; ggml_cuda_pdl_lc(); const src_t * src0_row = src0 + i01*s01 + i02*s02 + i03*s03; @@ -184,7 +187,8 @@ static void set_rows_cuda( const size_t nb01, const size_t nb02, const size_t nb03, const size_t nb10, const size_t nb11, const size_t nb12, const size_t nb1, const size_t nb2, const size_t nb3, - cudaStream_t stream) { + cudaStream_t stream, + const int64_t dst_row_base) { const int64_t ne_total = ne00 * ne01 * ne02 * ne03; const int num_blocks = (ne_total + CUDA_SET_ROWS_BLOCK_SIZE - 1) / CUDA_SET_ROWS_BLOCK_SIZE; @@ -213,12 +217,14 @@ static void set_rows_cuda( ggml_cuda_kernel_launch(k_set_rows, launch_params, src0_d, src1_d, dst_d, ne_total, ne10, ne11, ne12, ne13, s01, s02, s03, s10, s11, s12, s1, s2, s3, ne00_fd, ne01_fd, ne02_fd, - ne11_fd, ne12_fd); + ne11_fd, ne12_fd, dst_row_base); } } template -static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +static void set_rows_cuda( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + ggml_tensor * dst, const int64_t dst_row_base = 0) { const src_t * src0_d = (const src_t *)src0->data; const idx_t * src1_d = (const idx_t *)src1->data; @@ -226,7 +232,6 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s cudaStream_t stream = ctx.stream(); - if (dst->type == GGML_TYPE_F32) { set_rows_cuda( src0_d, src1_d, (float*)dst->data, @@ -235,7 +240,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_F16) { set_rows_cuda( @@ -245,7 +250,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_BF16) { set_rows_cuda( @@ -255,7 +260,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_Q4_0) { set_rows_cuda_quant( @@ -265,7 +270,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_Q4_1) { set_rows_cuda_quant( @@ -275,7 +280,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_Q5_0) { set_rows_cuda_quant( @@ -285,7 +290,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_Q5_1) { set_rows_cuda_quant( @@ -295,7 +300,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_Q8_0) { set_rows_cuda_quant( @@ -305,7 +310,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else if (dst->type == GGML_TYPE_IQ4_NL) { set_rows_cuda_quant( @@ -315,7 +320,7 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else { GGML_ABORT("unsupported type %s", ggml_type_name(dst->type)); @@ -323,7 +328,10 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s } template<> -void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +void set_rows_cuda( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + ggml_tensor * dst, const int64_t dst_row_base) { + GGML_ASSERT(dst_row_base == 0); const half * src0_d = (const half *)src0->data; const int32_t * src1_d = (const int32_t *)src1->data; @@ -340,7 +348,7 @@ void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_te nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else { GGML_ABORT("unsupported type %s", ggml_type_name(dst->type)); @@ -348,7 +356,10 @@ void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_te } template<> -void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +void set_rows_cuda( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + ggml_tensor * dst, const int64_t dst_row_base) { + GGML_ASSERT(dst_row_base == 0); const half * src0_d = (const half *)src0->data; const int64_t * src1_d = (const int64_t *)src1->data; @@ -365,7 +376,7 @@ void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_te nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, - stream + stream, dst_row_base ); } else { GGML_ABORT("unsupported type %s", ggml_type_name(dst->type)); @@ -396,3 +407,37 @@ void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ABORT("unsupported type %s", ggml_type_name(src0->type)); } } + +void ggml_cuda_op_set_rows_staged( + ggml_backend_cuda_context & ctx, ggml_tensor * dst, + int64_t first_row, int64_t row_count, void * mirror_data) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_I64); + GGML_ASSERT(first_row >= 0 && row_count > 1 && first_row + row_count <= dst->ne[1]); + + const size_t row_bytes = ggml_row_size(dst->type, dst->ne[0]); + const size_t stage_bytes = row_bytes*size_t(row_count); + ggml_cuda_pool_alloc staging(ctx.pool(), stage_bytes); + + ggml_tensor staged = *dst; + staged.data = staging.get(); + staged.ne[1] = row_count; + staged.nb[1] = row_bytes; + staged.nb[2] = row_bytes*size_t(row_count); + staged.nb[3] = staged.nb[2]*size_t(staged.ne[2]); + + set_rows_cuda(ctx, src0, src1, &staged, first_row); + CUDA_CHECK(cudaMemcpyAsync( + static_cast(dst->data) + size_t(first_row)*row_bytes, + staging.get(), stage_bytes, cudaMemcpyDeviceToHost, ctx.stream())); + if (mirror_data != nullptr) { + CUDA_CHECK(cudaMemcpyAsync( + static_cast(mirror_data) + size_t(first_row)*row_bytes, + staging.get(), stage_bytes, cudaMemcpyDeviceToDevice, ctx.stream())); + } + // Resident uploads are ordered later on this stream. Streamed uploads wait on + // the transfer ring's producer-ready event, also recorded after this copy. +} diff --git a/ggml/src/ggml-cuda/set-rows.cuh b/ggml/src/ggml-cuda/set-rows.cuh index c140c0873c8a..38a9c9be0074 100644 --- a/ggml/src/ggml-cuda/set-rows.cuh +++ b/ggml/src/ggml-cuda/set-rows.cuh @@ -5,3 +5,9 @@ #define CUDA_SET_ROWS_BLOCK_SIZE 256 void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_set_rows_staged( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + int64_t first_row, + int64_t row_count, + void * mirror_data); diff --git a/include/llama.h b/include/llama.h index ef7a012c43a1..451c6036628a 100644 --- a/include/llama.h +++ b/include/llama.h @@ -388,6 +388,7 @@ extern "C" { enum ggml_type type_k; // data type for K cache [EXPERIMENTAL] enum ggml_type type_v; // data type for V cache [EXPERIMENTAL] + uint32_t kv_stream_stage_mib; // block-streaming staging budget, 0 = disabled [EXPERIMENTAL] // Abort callback // if it returns true, execution of llama_decode() will be aborted diff --git a/scripts/arm117/boot-cell.sh b/scripts/arm117/boot-cell.sh new file mode 100644 index 000000000000..a2d580d1336c --- /dev/null +++ b/scripts/arm117/boot-cell.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# arm: adaptive-KV-streaming — boot helper for the 2xRTX RPC split +# Usage: bash scripts/arm117/boot-cell.sh [np] +# Rig safety: caller must have verified rig free (no prod pod, GPUs idle). +set -euo pipefail + +CTX_TOTAL=$1 +STAGE_MIB=$2 +PORT=${3:-8080} +NP=${4:-2} +ROOT=$(cd "$(dirname "$0")/../.." && pwd) + +export GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 +export LLAMA_KV_STREAM_DEVICE=${LLAMA_KV_STREAM_DEVICE:-CUDA0} +[ "${STAGE_MIB}" != "0" ] && export LLAMA_KV_STREAM_ALLOW_MULTISEQ=${LLAMA_KV_STREAM_ALLOW_MULTISEQ:-1} + +echo "[$(date)] boot-cell ctx=$CTX_TOTAL stage=$STAGE_MIB port=$PORT np=$NP stream_dev=$LLAMA_KV_STREAM_DEVICE" | tee -a arm117-boot-log.txt + +# peer (3060, CUDA1) +pkill -f 'ggml-rpc-server.*50052' 2>/dev/null || true +$ROOT/build/bin/ggml-rpc-server --host 127.0.0.1 --port 50052 -d CUDA1 > arm117-rpc-peer.log 2>&1 & +sleep 1 + +# server (5060 Ti, CUDA0) +$ROOT/build/bin/llama-server \ + -m "${MODEL_PATH:-/mnt/SSD/Qwen3.8-27B-UD-Q5_K_M.gguf}" \ + --rpc 127.0.0.1:50052 -ts 27,38 -ngl 99 \ + --rope-scaling yarn --rope-scale 5 --yarn-orig-ctx 32768 \ + -fa on -ctk q8_0 -ctv q5_1 -ctkd q8_0 -ctvd q5_1 \ + --no-kv-unified --cache-prompt --cache-reuse 64 --cache-idle-slots \ + --cache-ram 1024 --ubatch-size 512 --cont-batching \ + -np $NP -c $CTX_TOTAL \ + --parallel-ctx-threshold 100000 --spec-type draft-mtp \ + --prio-batch 1 --kv-stream-stage-mib $STAGE_MIB \ + --jinja --host 0.0.0.0 --port $PORT --metrics --slots --log-verbosity 4 \ + > arm117-server-port$PORT.log 2>&1 & + +SERVER_PID=$! +echo "server pid=$SERVER_PID" +for i in $(seq 1 120); do + code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 2 localhost:$PORT/health || true) + [ "$code" = "200" ] && echo "health 200 after ${i}s" && exit 0 + sleep 1 +done +echo "BOOT FAILED (health never 200)"; tail -40 arm117-server-port$PORT.log; exit 1 diff --git a/scripts/arm117/cell-a-parity.sh b/scripts/arm117/cell-a-parity.sh new file mode 100644 index 000000000000..0c34fcc4117c --- /dev/null +++ b/scripts/arm117/cell-a-parity.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# arm117 cell A: numerical-parity probe — streaming OFF vs ON at a fit context. +# Greedy temp=0 kernel prompts; byte-diff OFF vs ON text. +# Server phases are booted by the runner (boot-cell.sh); this script assumes +# the server on $PORT matches the current phase expectation. +set -euo pipefail +PORT=${1:-8080} +OUT=arm117-artifacts/cell-a +mkdir -p $OUT +LABEL=${2:-run} + +greedy() { # $1 tag, $2 req file + curl -s --max-time 600 localhost:$PORT/v1/chat/completions -H "Content-Type: application/json" -d @"$2" -o "$OUT/resp-$1.json" + python3 -c 'import json,sys; d=json.load(open(sys.argv[1])); c=d["choices"][0]["message"]; print((c.get("reasoning_content") or "")+(c.get("content") or ""))' "$OUT/resp-$1.json" > "$OUT/text-$1.txt" 2>/dev/null || echo RESPONSE_PARSE_FAIL > "$OUT/text-$1.txt" +} + +python3 - > "$OUT/req-p1-$LABEL.json" <<'PY' +import json +filler=" ".join(["Filler sentence number %d for the parity kernel prompt body."%i for i in range(150)]) +body=filler+"\n\nNow reply with exactly this text and nothing else: PARITY-KERNEL-A" +print(json.dumps({"messages":[{"role":"user","content":body}],"max_tokens":96,"temperature":0,"cache_prompt":True})) +PY + +greedy "p1-$LABEL" "$OUT/req-p1-$LABEL.json" +echo "text-p1-$LABEL.txt:"; cat "$OUT/text-p1-$LABEL.txt" diff --git a/scripts/arm117/cell-c-longctx.sh b/scripts/arm117/cell-c-longctx.sh new file mode 100644 index 000000000000..e243f5c47f0b --- /dev/null +++ b/scripts/arm117/cell-c-longctx.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# arm117 long-context capability probe: plant early (M1) + late (M2) markers, +# huge filler, then ask for BOTH. Streaming claims exactness => both must hit. +# Usage: bash scripts/arm117/cell-c-longctx.sh +set -euo pipefail +PORT=${1:-8080} +FILLER_K=${2:-30} +OUT=arm117-artifacts/cell-c +mkdir -p $OUT + +python3 - "$FILLER_K" "$OUT/prompt.json" <<'PY' +import json,sys +k=int(sys.argv[1]) +m1="M1 marker: code word FERRITE-77 guards the northern relay, lease the amber conduit before dawn cycle eleven." +m2="M2 marker: ledger 42 is filed with the onyx ferryman, and the harvest vault answers only to chartreuse." +filler=" ".join("Archive fragment %d records ambient sensor drift in sector %d with no anomalies."%(i,i%97) for i in range(k*22)) +body=m1+"\n\n"+filler+"\n\n"+m2+"\n\nAnswer both questions:\n1) Which code word guards the northern relay?\n2) Which ledger is filed with the onyx ferryman?" +json.dump({"messages":[{"role":"user","content":body}],"max_tokens":160,"temperature":0,"cache_prompt":True}, open(sys.argv[2],"w")) +PY + +curl -s --max-time 1800 localhost:$PORT/v1/chat/completions -H "Content-Type: application/json" -d @"$OUT/prompt.json" -o "$OUT/resp.json" + +python3 - "$OUT/resp.json" <<'PY' +import json,re,sys +d=json.load(open(sys.argv[1])) +t=d['choices'][0]['message']['content'] +print(t) +hits={'M1(FERRITE-77)': bool(re.search(r'FERRITE[- ]?77', t, re.I)), + 'M2(ledger 42)': bool(re.search(r'ledger 42|42.*(onyx|ferryman)|onyx ferryman', t, re.I))} +for k,v in hits.items(): print(("PASS" if v else "FAIL"), k) +PY diff --git a/scripts/arm117/snapshot.sh b/scripts/arm117/snapshot.sh new file mode 100644 index 000000000000..4264c0b13eca --- /dev/null +++ b/scripts/arm117/snapshot.sh @@ -0,0 +1,15 @@ +#!/bin/bash +# arm117 rig snapshot: health, slots, GPUs, key stream metrics — for artifacts +set -euo pipefail +PORT=${1:-8080} +OUT=${2:-arm117-artifacts/snapshot} +mkdir -p $OUT +curl -s --max-time 5 localhost:$PORT/health | tee $OUT/health.json +echo +curl -s --max-time 5 localhost:$PORT/slots | tee $OUT/slots.json | head -c 2000 +echo +nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv,noheader | tee $OUT/gpu-mem.txt +for pid in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do + ps -p $pid -o pid,cmd --no-headers >> $OUT/gpu-procs.txt 2>/dev/null || true +done +grep -h "kv.stream\|kv_stream\|repartition" arm117-server-port$PORT.log > $OUT/log-stream-lines.txt 2>/dev/null || true diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 221e14f7ff23..ed5e4805d584 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -28,6 +28,12 @@ add_library(llama llama-kv-cache-dsa-iswa.cpp llama-kv-cache-msa.cpp llama-kv-cache-dsv4.cpp + llama-kv-stream-config.cpp + llama-kv-stream-config.h + llama-kv-stream-plan.cpp + llama-kv-stream-plan.h + llama-kv-stream-softmax.cpp + llama-kv-stream-softmax.h llama-memory.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3cc27717ece8..edac0cf5f8cd 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -6,7 +6,9 @@ #include "llama-impl.h" #include "llama-batch.h" #include "llama-io.h" +#include "llama-kv-stream-config.h" #include "llama-memory.h" +#include "llama-memory-hybrid.h" #include "llama-mmap.h" #include "llama-model.h" #include "llama-ext.h" @@ -118,6 +120,7 @@ llama_context::llama_context( cparams.embeddings_nextn = false; cparams.embeddings_nextn_masked = false; cparams.offload_kqv = params.offload_kqv; + cparams.kv_stream_stage_mib = params.kv_stream_stage_mib; cparams.no_perf = params.no_perf; cparams.warmup = false; @@ -384,12 +387,45 @@ llama_context::llama_context( // init the memory module if (!hparams.vocab_only) { + const uint64_t kv_stream_stage_bytes = uint64_t(cparams.kv_stream_stage_mib)*1024ULL*1024ULL; + uint64_t kv_stream_minimum_stage_bytes = 0; + if (kv_stream_stage_bytes != 0 && model.arch == LLM_ARCH_QWEN35) { + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + if (hparams.has_kv(il) && !hparams.is_recr(il)) { + kv_stream_minimum_stage_bytes = 256ULL*( + ggml_row_size(params.type_k, hparams.n_embd_k_gqa(il)) + + ggml_row_size(params.type_v, hparams.n_embd_v_gqa(il))); + break; + } + } + } + + const llama_kv_stream_config stream_config = { + /*.stage_bytes =*/ kv_stream_stage_bytes, + /*.minimum_stage_bytes =*/ kv_stream_minimum_stage_bytes, + /*.arch_qwen35 =*/ model.arch == LLM_ARCH_QWEN35, + /*.context_default =*/ cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT, + /*.single_sequence =*/ cparams.n_seq_max == 1, + /*.flash_attention =*/ cparams.flash_attn, + /*.kv_offload =*/ cparams.offload_kqv, + /*.multi_sequence_allowed =*/ getenv("LLAMA_KV_STREAM_ALLOW_MULTISEQ") != nullptr, + }; + const auto stream_validation = llama_kv_stream_config_validate(stream_config); + if (!stream_validation.valid) { + throw std::runtime_error(stream_validation.error); + } + if (stream_validation.enabled) { + LLAMA_LOG_INFO("%s: experimental block KV streaming enabled, pool = %.2f MiB\n", + __func__, kv_stream_stage_bytes/1024.0/1024.0); + } + llama_memory_params params_mem = { - /*.type_k =*/ params.type_k, - /*.type_v =*/ params.type_v, - /*.swa_full =*/ params.swa_full, - /*.ctx_type =*/ cparams.ctx_type, - /*.mem_other =*/ llama_get_memory(cparams.ctx_other), + /*.type_k =*/ params.type_k, + /*.type_v =*/ params.type_v, + /*.kv_stream_stage_bytes =*/ kv_stream_stage_bytes, + /*.swa_full =*/ params.swa_full, + /*.ctx_type =*/ cparams.ctx_type, + /*.mem_other =*/ llama_get_memory(cparams.ctx_other), }; memory.reset(model.create_memory(params_mem, cparams)); @@ -1338,6 +1374,19 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } + // Consume feedback from the previous completed decode before building the + // next graph. Repartition itself synchronizes only when hysteresis selects + // a new boundary, so steady-state token generation stays asynchronous. + if (auto * hybrid_memory = dynamic_cast(memory.get())) { + if (auto * hybrid_context = dynamic_cast(mctx)) { + const llama_kv_cache_context * attn_context = hybrid_context->get_attn(); + if (attn_context != nullptr) { + (void) hybrid_memory->get_mem_attn()->kv_stream_adapt( + attn_context->get_n_kv(), ubatch.n_tokens); + } + } + } + auto * res = gf_res_prev.get(); auto * gf = res->get_gf(); @@ -3640,6 +3689,7 @@ llama_context_params llama_context_default_params() { /*.cb_eval_user_data =*/ nullptr, /*.type_k =*/ GGML_TYPE_F16, /*.type_v =*/ GGML_TYPE_F16, + /*.kv_stream_stage_mib =*/ 0, /*.abort_callback =*/ nullptr, /*.abort_callback_data =*/ nullptr, /*.embeddings =*/ false, diff --git a/src/llama-cparams.h b/src/llama-cparams.h index b592de18c794..c2e51ee42ecc 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -54,6 +54,8 @@ struct llama_cparams { bool kv_unified; bool pipeline_parallel; + uint32_t kv_stream_stage_mib; + std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer enum llama_context_type ctx_type; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a342ee1191d4..7d8a8ff91977 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1,4 +1,5 @@ #include "llama-kv-cache.h" +#include "llama-kv-stream-plan.h" #include "llama-impl.h" #include "llama-io.h" @@ -8,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -79,7 +81,8 @@ llama_kv_cache::llama_kv_cache( const layer_filter_cb & filter, const layer_reuse_cb & reuse, const layer_share_cb & share, - const char * name_tag) : + const char * name_tag, + size_t kv_stream_stage_bytes) : model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), @@ -162,6 +165,44 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); + // fork (arm: adaptive-KV-streaming, Gate R): under the Hydra RPC layer + // split, KV layers hosted by the RPC peer have no CUDA backend in this + // process, so the stock all-layers-on-one-CUDA-device wiring cannot boot. + // LLAMA_KV_STREAM_DEVICE= (e.g. "CUDA0") restricts streaming to the + // KV layers whose model layer device matches; unlisted devices keep their + // ordinary fully-resident KV path. Unset preserves stock behavior. + const char * kv_stream_device_name = nullptr; + if (kv_stream_stage_bytes != 0) { + kv_stream_device_name = getenv("LLAMA_KV_STREAM_DEVICE"); + } + auto kv_stream_wants_dev = [kv_stream_device_name](ggml_backend_dev_t dev) { + if (kv_stream_device_name == nullptr) { + return true; + } + const char * dev_name = ggml_backend_dev_name(dev); + return dev_name != nullptr && strcmp(dev_name, kv_stream_device_name) == 0; + }; + + ggml_backend_dev_t kv_stream_dev = nullptr; + ggml_backend_buffer_type_t kv_stream_buft = nullptr; + uint32_t kv_stream_layer_count = 0; + if (kv_stream_stage_bytes != 0) { + for (uint32_t il = 0; il < n_layer; ++il) { + if (!hparams.has_kv(il) || (filter && !filter(il))) { + continue; + } + if (offload && !hparams.no_alloc && kv_stream_wants_dev(model.dev_layer(il))) { + ++kv_stream_layer_count; + } + } + if (kv_stream_layer_count == 0) { + throw std::runtime_error("block KV streaming enabled but no KV layers on the streaming device"); + } + if (kv_stream_device_name != nullptr) { + LLAMA_LOG_INFO("%s: KV streaming device filter = %s (%u layers)\n", __func__, kv_stream_device_name, kv_stream_layer_count); + } + } + for (uint32_t il = 0; il < n_layer; il++) { if (!hparams.has_kv(il)) { LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il); @@ -218,6 +259,108 @@ llama_kv_cache::llama_kv_cache( buft = ggml_backend_dev_buffer_type(dev); dev_name = ggml_backend_dev_name(dev); + + if (kv_stream_stage_bytes != 0 && !hparams.no_alloc && kv_stream_wants_dev(dev)) { + if (kv_stream_dev != nullptr && kv_stream_dev != dev) { + throw std::runtime_error("block KV streaming requires every attention layer on one CUDA device"); + } + + if (kv_stream_runtime.runtime == nullptr) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + using type_pair_supported_fn_t = bool (*)(ggml_type, ggml_type); + using page_bytes_fn_t = bool (*)( + ggml_type, ggml_type, uint32_t, uint32_t, uint32_t, uint32_t, size_t *); + using runtime_new_fn_t = void * (*)( + ggml_backend_dev_t, size_t, size_t, size_t, uint32_t); + using runtime_free_fn_t = void (*)(void *); + using buffer_type_fn_t = ggml_backend_buffer_type_t (*)(void *); + using feedback_fn_t = kv_stream_runtime_owner::feedback_fn_t; + using span_feedback_fn_t = kv_stream_runtime_owner::span_feedback_fn_t; + using reconfigure_fn_t = kv_stream_runtime_owner::reconfigure_fn_t; + using repartition_fn_t = kv_stream_runtime_owner::repartition_fn_t; + using decode_layout_fn_t = kv_stream_runtime_owner::decode_layout_fn_t; + using mark_dirty_rows_fn_t = kv_stream_runtime_owner::mark_dirty_rows_fn_t; + + auto * type_pair_supported_fn = (type_pair_supported_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_type_pair_supported"); + auto * page_bytes_fn = (page_bytes_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_page_bytes"); + auto * workspace_bytes_fn = (page_bytes_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_workspace_bytes"); + auto * runtime_new_fn = (runtime_new_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_runtime_new_for_device"); + auto * runtime_free_fn = (runtime_free_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_runtime_free"); + auto * buffer_type_fn = (buffer_type_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_buffer_type"); + auto * feedback_fn = (feedback_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_feedback"); + auto * span_feedback_fn = (span_feedback_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_observe_decode_latency"); + auto * reconfigure_fn = (reconfigure_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_reconfigure"); + auto * repartition_fn = (repartition_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_repartition"); + auto * decode_layout_fn = (decode_layout_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_set_decode_layout"); + auto * mark_dirty_rows_fn = (mark_dirty_rows_fn_t) ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_mark_dirty_rows"); + + if (type_pair_supported_fn == nullptr || page_bytes_fn == nullptr || + workspace_bytes_fn == nullptr || runtime_new_fn == nullptr || + runtime_free_fn == nullptr || + buffer_type_fn == nullptr || feedback_fn == nullptr || + span_feedback_fn == nullptr || + repartition_fn == nullptr || decode_layout_fn == nullptr || + reconfigure_fn == nullptr || + mark_dirty_rows_fn == nullptr) { + throw std::runtime_error("block KV streaming requires the CUDA backend"); + } + + if (!type_pair_supported_fn(type_k, type_v)) { + throw std::runtime_error( + "block KV streaming does not support K " + std::string(ggml_type_name(type_k)) + + " and V " + ggml_type_name(type_v)); + } + + size_t page_bytes = 0; + if (!page_bytes_fn( + type_k, type_v, + hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), hparams.n_head_kv(il), + 256, &page_bytes)) { + throw std::runtime_error("invalid block KV streaming page geometry"); + } + size_t conversion_bytes = 0; + if (!workspace_bytes_fn( + type_k, type_v, + hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), hparams.n_head_kv(il), + 256, &conversion_bytes)) { + throw std::runtime_error("invalid block KV streaming conversion workspace geometry"); + } + kv_stream_runtime.runtime = runtime_new_fn( + dev, kv_stream_stage_bytes, page_bytes, conversion_bytes, kv_stream_layer_count); + kv_stream_runtime.free_fn = runtime_free_fn; + kv_stream_runtime.feedback_fn = feedback_fn; + kv_stream_runtime.span_feedback_fn = span_feedback_fn; + kv_stream_runtime.repartition_fn = repartition_fn; + kv_stream_runtime.reconfigure_fn = reconfigure_fn; + kv_stream_runtime.decode_layout_fn = decode_layout_fn; + kv_stream_runtime.mark_dirty_rows_fn = mark_dirty_rows_fn; + kv_stream_runtime.layer_count = kv_stream_layer_count; + if (kv_stream_runtime.runtime == nullptr) { + throw std::runtime_error("failed to create CUDA block KV streaming runtime"); + } + + kv_stream_buft = buffer_type_fn(kv_stream_runtime.runtime); + if (kv_stream_buft == nullptr) { + throw std::runtime_error("failed to obtain CUDA block KV streaming buffer type"); + } + kv_stream_dev = dev; + } + + buft = kv_stream_buft; + dev_name = ggml_backend_buft_name(buft); + } } LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); @@ -291,6 +434,19 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); ggml_backend_buffer_clear(buf, 0); + if (kv_stream_stage_bytes == 0 && getenv("GGML_CUDA_PREFER_KV_HOST") != nullptr && ggml_backend_buffer_get_size(buf) > 0) { + ggml_backend_dev_t dev_kv = ggml_backend_buft_get_device(buft); + if (dev_kv != nullptr) { + ggml_backend_reg_t reg_kv = ggml_backend_dev_backend_reg(dev_kv); + using prefer_host_fn_t = bool (*)(ggml_backend_buffer_t); + auto * prefer_host_fn = (prefer_host_fn_t) ggml_backend_reg_get_proc_address( + reg_kv, "ggml_backend_cuda_buffer_set_preferred_host"); + if (prefer_host_fn != nullptr) { + (void) prefer_host_fn(buf); + } + } + } + ctxs_bufs.emplace_back(std::move(ctx), buf); } @@ -1206,6 +1362,145 @@ uint32_t llama_kv_cache::get_n_stream() const { return n_stream; } +bool llama_kv_cache::kv_stream_adapt(uint32_t active_tokens, uint32_t query_tokens) { + auto & owner = kv_stream_runtime; + if (owner.runtime == nullptr || owner.feedback_fn == nullptr || + owner.span_feedback_fn == nullptr || owner.reconfigure_fn == nullptr || + owner.layer_count == 0) { + return false; + } + + constexpr uint32_t MAX_DECODE_QUERY_TOKENS = 32; + constexpr double MAX_DECODE_INTERVAL_MS = 1000.0; + const int64_t now_us = ggml_time_us(); + if (owner.previous_adapt_us != 0 && + owner.previous_query_tokens <= MAX_DECODE_QUERY_TOKENS && + query_tokens <= MAX_DECODE_QUERY_TOKENS) { + const double elapsed_ms = (now_us - owner.previous_adapt_us)/1000.0; + if (elapsed_ms > 0.0 && elapsed_ms <= MAX_DECODE_INTERVAL_MS) { + (void) owner.span_feedback_fn(owner.runtime, elapsed_ms); + } + } + owner.previous_adapt_us = now_us; + owner.previous_query_tokens = query_tokens; + + uint64_t deadline_samples = 0; + uint64_t deadline_misses = 0; + double copy_busy_ratio = 0.0; + uint32_t peak_occupancy = 0; + uint32_t ring_slots = 0; + uint32_t resident_pages = 0; + uint32_t controlled_pages = 0; + if (!owner.feedback_fn(owner.runtime, + &deadline_samples, &deadline_misses, ©_busy_ratio, + &peak_occupancy, &ring_slots, &resident_pages, &controlled_pages)) { + return false; + } + + const auto delta = llama_kv_stream_feedback_delta_make( + { deadline_samples, deadline_misses }, + { owner.previous_deadline_samples, owner.previous_deadline_misses }); + owner.previous_deadline_samples = deadline_samples; + owner.previous_deadline_misses = deadline_misses; + if (getenv("LLAMA_KV_STREAM_TRACE") != nullptr) { + LLAMA_LOG_WARN("%s: active %u, resident %u, ring %u, samples %llu, misses %llu, copy busy %.1f%%, peak %u\n", + __func__, active_tokens, resident_pages, ring_slots, + (unsigned long long) delta.deadline_samples, + (unsigned long long) delta.deadline_misses, + 100.0*copy_busy_ratio, peak_occupancy); + } + + const uint32_t active_pages = (active_tokens + 255)/256; + // Prompt chunks use the uniform layout because it grows without + // repartitioning. Decode-like microbatches concentrate the same page + // budget into fewer split layers, bounded by the ring working set so copy + // and compute can still overlap. A zero target restores prefill. + const uint32_t decode_layout_pages = + query_tokens <= MAX_DECODE_QUERY_TOKENS && active_pages > resident_pages ? + active_pages : 0; + const bool entering_decode_layout = + decode_layout_pages != 0 && decode_layout_pages != owner.decode_layout_pages; + + if (ring_slots != 0 && owner.minimum_ring_slots == 0) { + owner.minimum_ring_slots = ring_slots; + } + if (!delta.valid) { + owner.starved_evaluations = 0; + owner.overprovisioned_evaluations = 0; + LLAMA_LOG_WARN("%s: ignoring invalid CUDA feedback: %s\n", + __func__, delta.error.c_str()); + } + + uint32_t target_ring_slots = ring_slots; + uint32_t target_resident_pages = resident_pages; + bool partition_changed = false; + const bool fixed_ring = + getenv("GGML_CUDA_KV_STREAM_FIXED_RING_SLOTS") != nullptr; + + // No streaming pressure exists while every active page fits in the + // resident partition; preserve the current boundary without churn. + if (active_pages <= resident_pages) { + owner.starved_evaluations = 0; + owner.overprovisioned_evaluations = 0; + } else if (ring_slots != 0 && !fixed_ring && + (entering_decode_layout || (delta.valid && delta.has_evaluation))) { + if (delta.has_evaluation && owner.evaluations_since_repartition != UINT32_MAX) { + ++owner.evaluations_since_repartition; + } + + llama_kv_stream_partition_params params; + params.total_pool_pages = controlled_pages; + params.layer_count = owner.layer_count; + params.active_pages_per_layer = active_pages; + params.minimum_ring_slots = owner.minimum_ring_slots; + params.previous_resident_pages_per_layer = resident_pages; + params.previous_ring_slots = ring_slots; + params.deadline_miss_ratio = delta.valid ? delta.deadline_miss_ratio : 0.0; + params.copy_engine_busy_ratio = copy_busy_ratio; + params.ring_peak_occupancy_ratio = + std::min(1.0, double(peak_occupancy)/double(ring_slots)); + params.starved_evaluations = owner.starved_evaluations; + params.overprovisioned_evaluations = owner.overprovisioned_evaluations; + params.evaluations_since_repartition = owner.evaluations_since_repartition; + params.entering_decode_layout = entering_decode_layout; + + const auto partition = llama_kv_stream_partition_adapt(params); + if (!partition.valid) { + LLAMA_LOG_WARN("%s: ignoring invalid partition feedback: %s\n", + __func__, partition.error.c_str()); + } else { + owner.starved_evaluations = partition.starved_evaluations; + owner.overprovisioned_evaluations = partition.overprovisioned_evaluations; + partition_changed = partition.changed; + target_ring_slots = partition.ring_slots; + target_resident_pages = partition.resident_pages_per_layer; + } + } + + const bool layout_changed = decode_layout_pages != owner.decode_layout_pages; + if ((layout_changed || partition_changed) && ring_slots != 0) { + if (!owner.reconfigure_fn( + owner.runtime, decode_layout_pages, target_ring_slots)) { + LLAMA_LOG_WARN( + "%s: failed to publish CUDA KV decode layout %u with ring %u\n", + __func__, decode_layout_pages, target_ring_slots); + return false; + } + owner.decode_layout_pages = decode_layout_pages; + } + + if (!partition_changed) { + return false; + } + owner.evaluations_since_repartition = 0; + LLAMA_LOG_WARN("%s: adaptive KV partition: resident pages/layer %u -> %u, ring slots %u -> %u, miss %.1f%%, copy busy %.1f%%\n", + __func__, resident_pages, target_resident_pages, + ring_slots, target_ring_slots, + 100.0*(delta.valid ? delta.deadline_miss_ratio : 0.0), + 100.0*copy_busy_ratio); + return true; +} + bool llama_kv_cache::get_has_shift() const { bool result = false; @@ -1487,6 +1782,12 @@ void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ub data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; } } + + if (kv_stream_runtime.runtime != nullptr) { + GGML_ASSERT(kv_stream_runtime.mark_dirty_rows_fn != nullptr); + GGML_ASSERT(kv_stream_runtime.mark_dirty_rows_fn( + kv_stream_runtime.runtime, data, n_tokens)); + } } void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c4d8699def12..b5d670688aeb 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -114,7 +114,8 @@ class llama_kv_cache : public llama_memory_i { const layer_reuse_cb & reuse, const layer_share_cb & share, // a model can hold more than one cache, so the tensor names have to stay unique - const char * name_tag = ""); + const char * name_tag = "", + size_t kv_stream_stage_bytes = 0); ~llama_kv_cache() = default; @@ -158,6 +159,8 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_size() const; uint32_t get_n_stream() const; + bool kv_stream_adapt(uint32_t active_tokens, uint32_t query_tokens); + bool get_has_shift() const; ggml_type type_k() const; @@ -289,6 +292,45 @@ class llama_kv_cache : public llama_memory_i { const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; // ggml contexts for the KV cache along with the allocated backend buffers: + struct kv_stream_runtime_owner { + using feedback_fn_t = bool (*)( + void *, uint64_t *, uint64_t *, double *, uint32_t *, + uint32_t *, uint32_t *, uint32_t *); + using span_feedback_fn_t = bool (*)(void *, double); + using reconfigure_fn_t = bool (*)(void *, uint32_t, uint32_t); + using repartition_fn_t = bool (*)(void *, uint32_t); + using decode_layout_fn_t = bool (*)(void *, uint32_t); + using mark_dirty_rows_fn_t = bool (*)(void *, const int64_t *, size_t); + + void * runtime = nullptr; + void (*free_fn)(void *) = nullptr; + feedback_fn_t feedback_fn = nullptr; + span_feedback_fn_t span_feedback_fn = nullptr; + reconfigure_fn_t reconfigure_fn = nullptr; + repartition_fn_t repartition_fn = nullptr; + decode_layout_fn_t decode_layout_fn = nullptr; + mark_dirty_rows_fn_t mark_dirty_rows_fn = nullptr; + uint32_t layer_count = 0; + uint32_t minimum_ring_slots = 0; + uint32_t decode_layout_pages = 0; + uint32_t starved_evaluations = 0; + uint32_t overprovisioned_evaluations = 0; + uint32_t evaluations_since_repartition = UINT32_MAX; + uint64_t previous_deadline_samples = 0; + uint64_t previous_deadline_misses = 0; + int64_t previous_adapt_us = 0; + uint32_t previous_query_tokens = UINT32_MAX; + + ~kv_stream_runtime_owner() { + if (runtime != nullptr) { + free_fn(runtime); + } + } + }; + + // Declared before ctxs_bufs so the custom buffers release their runtime + // references before this owner releases the initial reference. + kv_stream_runtime_owner kv_stream_runtime; std::vector> ctxs_bufs; // the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot()) diff --git a/src/llama-kv-stream-config.cpp b/src/llama-kv-stream-config.cpp new file mode 100644 index 000000000000..dd3709ab8798 --- /dev/null +++ b/src/llama-kv-stream-config.cpp @@ -0,0 +1,80 @@ +#include "llama-kv-stream-config.h" + +#include + +#include "llama-impl.h" + +llama_kv_stream_config_result llama_kv_stream_config_validate(const llama_kv_stream_config & config) { + if (config.stage_bytes == 0) { + return { true, false, {} }; + } + + auto invalid = [](const char * error) { + return llama_kv_stream_config_result { false, false, error }; + }; + + if (!config.arch_qwen35) { + return invalid("block KV streaming currently supports only Qwen3.5"); + } + if (!config.context_default) { + return invalid("block KV streaming currently supports only the target context, not MTP/draft contexts"); + } + if (!config.single_sequence && !config.multi_sequence_allowed) { + return invalid("block KV streaming requires exactly one sequence (-np 1); " + "set LLAMA_KV_STREAM_ALLOW_MULTISEQ=1 to opt in to parallel-slot streaming (upstream-untested)"); + } + if (!config.single_sequence && config.multi_sequence_allowed) { + // surfaced at the validator so the operator sees it on every boot + LLAMA_LOG_WARN("block KV streaming: parallel slots enabled via LLAMA_KV_STREAM_ALLOW_MULTISEQ; " + "upstream validated single-slot only — concurrent correctness/perf is under test\n"); + } + if (!config.flash_attention) { + return invalid("block KV streaming requires Flash Attention"); + } + if (!config.kv_offload) { + return invalid("block KV streaming requires GPU KV offload"); + } + if (config.minimum_stage_bytes == 0 || config.stage_bytes < config.minimum_stage_bytes) { + return invalid("block KV streaming stage is too small for one 256-token cache page"); + } + + return { true, true, {} }; +} + +llama_kv_stream_pool_layout llama_kv_stream_pool_layout_make( + const llama_kv_stream_pool_layout_params & params) { + llama_kv_stream_pool_layout result; + + auto invalid = [&](const char * error) { + result.error = error; + return result; + }; + + if (params.pool_bytes == 0 || params.page_bytes == 0 || + params.layer_count == 0 || params.scratch_pages == 0) { + return invalid("pool, page, layer, and scratch counts must be nonzero"); + } + if (params.page_bytes > std::numeric_limits::max()/params.scratch_pages) { + return invalid("scratch byte count overflow"); + } + result.scratch_bytes = params.page_bytes*params.scratch_pages; + if (result.scratch_bytes >= params.pool_bytes) { + return invalid("pool has no resident capacity after reserving scratch pages"); + } + if (params.page_bytes > std::numeric_limits::max()/params.layer_count) { + return invalid("per-layer partition byte count overflow"); + } + + const uint64_t bytes_per_round = params.page_bytes*params.layer_count; + const uint64_t pages = (params.pool_bytes - result.scratch_bytes)/bytes_per_round; + if (pages == 0 || pages > std::numeric_limits::max()/256U) { + return invalid("pool cannot hold one resident page per layer"); + } + + result.resident_pages_per_layer = pages; + result.resident_tokens_per_layer = pages*256U; + result.resident_bytes = pages*bytes_per_round; + result.unused_bytes = params.pool_bytes - result.scratch_bytes - result.resident_bytes; + result.valid = true; + return result; +} diff --git a/src/llama-kv-stream-config.h b/src/llama-kv-stream-config.h new file mode 100644 index 000000000000..aa6ed68924bb --- /dev/null +++ b/src/llama-kv-stream-config.h @@ -0,0 +1,50 @@ +#pragma once + +#include +#include + +struct llama_kv_stream_config { + uint64_t stage_bytes = 0; + uint64_t minimum_stage_bytes = 0; + + bool arch_qwen35 = false; + bool context_default = false; + bool single_sequence = false; + bool flash_attention = false; + bool kv_offload = false; + + // fork (arm: adaptive-KV-streaming, Gate R): upstream hard-rejects + // parallel slots (untested there). LLAMA_KV_STREAM_ALLOW_MULTISEQ=1 + // opts in to -np>1 streaming so the concurrency hypothesis is measurable + // on the rig instead of being asserted unbootable. + bool multi_sequence_allowed = false; +}; + +struct llama_kv_stream_config_result { + bool valid = false; + bool enabled = false; + std::string error; +}; + +llama_kv_stream_config_result llama_kv_stream_config_validate(const llama_kv_stream_config & config); + +struct llama_kv_stream_pool_layout_params { + uint64_t pool_bytes = 0; + uint64_t page_bytes = 0; + uint32_t layer_count = 0; + uint32_t scratch_pages = 0; +}; + +struct llama_kv_stream_pool_layout { + bool valid = false; + std::string error; + + uint32_t resident_pages_per_layer = 0; + uint32_t resident_tokens_per_layer = 0; + uint64_t scratch_bytes = 0; + uint64_t resident_bytes = 0; + uint64_t unused_bytes = 0; +}; + +llama_kv_stream_pool_layout llama_kv_stream_pool_layout_make( + const llama_kv_stream_pool_layout_params & params); diff --git a/src/llama-kv-stream-plan.cpp b/src/llama-kv-stream-plan.cpp new file mode 100644 index 000000000000..17ecaf967442 --- /dev/null +++ b/src/llama-kv-stream-plan.cpp @@ -0,0 +1,510 @@ +#include "llama-kv-stream-plan.h" + +#include +#include +#include +#include +#include + +namespace { + +struct logical_interval { + llama_kv_stream_region_role role; + int32_t layer_id; + uint64_t begin; + uint64_t end; +}; + +bool checked_add(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(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; +} + +llama_kv_stream_plan fail(llama_kv_stream_plan result, const char * message) { + result.valid = false; + result.error = message; + return result; +} + +} // namespace + +llama_kv_stream_plan llama_kv_stream_plan_make(const llama_kv_stream_plan_params & params) { + llama_kv_stream_plan result; + result.pool_bytes = params.pool_bytes; + + if (params.pool_bytes == 0) { + return fail(std::move(result), "KV stream pool must be non-zero"); + } + + if (params.stage_slots > 0 && params.stage_slot_bytes == 0) { + return fail(std::move(result), "non-zero stage count requires a non-zero stage size"); + } + + if (!checked_mul(params.stage_slots, params.stage_slot_bytes, result.reserved_stage_bytes)) { + return fail(std::move(result), "stage reservation size overflow"); + } + + if (result.reserved_stage_bytes > params.pool_bytes) { + return fail(std::move(result), "stage reservations exceed the KV stream pool"); + } + + std::vector optional_regions; + std::vector logical_intervals; + optional_regions.reserve(params.regions.size()); + logical_intervals.reserve(params.regions.size()); + result.resident_regions.reserve(params.regions.size()); + result.streamed_regions.reserve(params.regions.size()); + + uint64_t pinned_bytes = 0; + + for (size_t index = 0; index < params.regions.size(); ++index) { + const auto & region = params.regions[index]; + + if (region.layer_id < 0) { + return fail(std::move(result), "KV stream region has an invalid layer id"); + } + + if (region.token_count == 0 || region.bytes == 0) { + return fail(std::move(result), "KV stream region must contain tokens and bytes"); + } + + const uint64_t token_end = uint64_t(region.token_begin) + region.token_count; + if (token_end > uint64_t(std::numeric_limits::max()) + 1) { + return fail(std::move(result), "KV stream region token range overflow"); + } + + logical_intervals.push_back({ region.role, region.layer_id, region.token_begin, token_end }); + + if (region.pinned) { + if (!checked_add(pinned_bytes, region.bytes, pinned_bytes)) { + return fail(std::move(result), "pinned KV region size overflow"); + } + result.resident_regions.push_back(index); + } else { + optional_regions.push_back(index); + } + } + + std::sort(logical_intervals.begin(), logical_intervals.end(), [](const auto & lhs, const auto & rhs) { + return std::tie(lhs.role, lhs.layer_id, lhs.begin, lhs.end) < + std::tie(rhs.role, rhs.layer_id, rhs.begin, rhs.end); + }); + + for (size_t i = 1; i < logical_intervals.size(); ++i) { + const auto & previous = logical_intervals[i - 1]; + const auto & current = logical_intervals[i]; + + if (previous.role == current.role && + previous.layer_id == current.layer_id && + current.begin < previous.end) { + return fail(std::move(result), "KV stream regions overlap in one logical layer"); + } + } + + uint64_t required_bytes = 0; + if (!checked_add(result.reserved_stage_bytes, pinned_bytes, required_bytes)) { + return fail(std::move(result), "mandatory KV stream pool size overflow"); + } + + if (required_bytes > params.pool_bytes) { + return fail(std::move(result), "pinned KV regions and stages exceed the pool"); + } + + result.resident_bytes = pinned_bytes; + uint64_t remaining_bytes = params.pool_bytes - required_bytes; + + std::sort(optional_regions.begin(), optional_regions.end(), [&](size_t a, size_t b) { + const auto & lhs = params.regions[a]; + const auto & rhs = params.regions[b]; + + return std::tie(lhs.residency_priority, lhs.role, lhs.layer_id, lhs.token_begin, a) < + std::tie(rhs.residency_priority, rhs.role, rhs.layer_id, rhs.token_begin, b); + }); + + for (size_t index : optional_regions) { + const auto & region = params.regions[index]; + + if (region.bytes <= remaining_bytes) { + result.resident_regions.push_back(index); + result.resident_bytes += region.bytes; + remaining_bytes -= region.bytes; + continue; + } + + if (params.stage_slots == 0) { + return fail(std::move(result), "streaming is required but no stage slots are configured"); + } + + if (region.bytes > params.stage_slot_bytes) { + return fail(std::move(result), "a streamed KV region exceeds the stage slot size"); + } + + uint64_t streamed_bytes = 0; + if (!checked_add(result.streamed_bytes, region.bytes, streamed_bytes)) { + return fail(std::move(result), "streamed KV region size overflow"); + } + + result.streamed_bytes = streamed_bytes; + result.streamed_regions.push_back(index); + } + + result.unused_bytes = remaining_bytes; + result.valid = true; + return result; +} + +llama_kv_stream_extent llama_kv_stream_extent_make(const llama_kv_stream_extent_params & params) { + llama_kv_stream_extent result; + + auto fail_extent = [&](const char * message) { + result.valid = false; + result.error = message; + return result; + }; + + if (params.page_tokens == 0) { + return fail_extent("KV stream page size must be non-zero"); + } + + if (params.maximum_tokens == 0 || params.maximum_tokens%params.page_tokens != 0) { + return fail_extent("KV stream maximum must be non-zero and page aligned"); + } + + if (params.previous_extent > params.maximum_tokens || + params.previous_extent%params.page_tokens != 0) { + return fail_extent("previous KV stream extent is invalid"); + } + + const uint64_t requested_tokens = uint64_t(params.live_tokens) + params.reserve_tokens; + if (requested_tokens > params.maximum_tokens) { + return fail_extent("live and reserved KV tokens exceed the configured maximum"); + } + + uint64_t desired_tokens = 0; + if (requested_tokens > 0) { + desired_tokens = ((requested_tokens + params.page_tokens - 1)/params.page_tokens)*params.page_tokens; + } + + if (desired_tokens > params.maximum_tokens) { + return fail_extent("padded KV stream extent exceeds the configured maximum"); + } + + uint32_t selected_tokens = uint32_t(desired_tokens); + if (!params.force_shrink && params.previous_extent > selected_tokens) { + const uint32_t released_tokens = params.previous_extent - selected_tokens; + if (released_tokens < params.shrink_hysteresis_tokens) { + selected_tokens = params.previous_extent; + } + } + + result.valid = true; + result.tokens = selected_tokens; + result.grew = selected_tokens > params.previous_extent; + result.shrunk = selected_tokens < params.previous_extent; + return result; +} + +llama_kv_stream_regions llama_kv_stream_regions_make(const llama_kv_stream_regions_params & params) { + llama_kv_stream_regions result; + + auto fail_regions = [&](const char * message) { + result.valid = false; + result.error = message; + result.regions.clear(); + result.total_bytes = 0; + return result; + }; + + if (params.page_tokens == 0) { + return fail_regions("KV stream region page size must be non-zero"); + } + + std::set> logical_layers; + + for (const auto & layer : params.layers) { + if (layer.layer_id < 0) { + return fail_regions("KV stream layer layout has an invalid layer id"); + } + + if (!logical_layers.emplace(layer.role, layer.layer_id).second) { + return fail_regions("KV stream layer layout is duplicated"); + } + + if (layer.n_tokens > 0 && layer.bytes_per_token == 0) { + return fail_regions("non-empty KV stream layer must have a non-zero token size"); + } + + for (uint64_t token_begin = 0; token_begin < layer.n_tokens; token_begin += params.page_tokens) { + const uint32_t token_count = uint32_t(std::min( + params.page_tokens, uint64_t(layer.n_tokens) - token_begin)); + + llama_kv_stream_region region; + region.role = layer.role; + region.layer_id = layer.layer_id; + region.token_begin = uint32_t(token_begin); + region.token_count = token_count; + region.pinned = layer.pin_all || + (layer.pin_tail && token_begin + token_count == layer.n_tokens); + + if (!checked_mul(token_count, layer.bytes_per_token, region.bytes)) { + return fail_regions("KV stream region byte size overflow"); + } + + uint64_t total_bytes = 0; + if (!checked_add(result.total_bytes, region.bytes, total_bytes)) { + return fail_regions("KV stream layout total byte size overflow"); + } + result.total_bytes = total_bytes; + + const uint64_t page_index = token_begin/params.page_tokens; + region.residency_priority = (page_index << 32) | layer.layer_priority; + result.regions.push_back(region); + } + } + + result.valid = true; + return result; +} + +llama_kv_stream_prefetch_dispatch_result llama_kv_stream_prefetch_dispatch( + const llama_kv_stream_prefetch_params & params) { + llama_kv_stream_prefetch_dispatch_result result; + + auto fail_dispatch = [&](const char * message) { + result.valid = false; + result.error = message; + result.assignments.clear(); + return result; + }; + + if (params.stage_slot_bytes == 0) { + return fail_dispatch("KV prefetch stage slot size must be non-zero"); + } + if (params.states.size() != params.requests.size()) { + return fail_dispatch("KV prefetch request states do not match requests"); + } + + std::vector slots = params.free_slots; + std::sort(slots.begin(), slots.end()); + if (std::adjacent_find(slots.begin(), slots.end()) != slots.end()) { + return fail_dispatch("KV prefetch free slot list contains duplicates"); + } + + const uint64_t window_end = params.adaptive_lookahead ? + std::numeric_limits::max() : + uint64_t(params.current_attention) + params.lookahead_layers; + std::vector eligible; + eligible.reserve(params.requests.size()); + + for (size_t index = 0; index < params.requests.size(); ++index) { + const auto & request = params.requests[index]; + if (request.layer_id < 0 || request.bytes == 0) { + return fail_dispatch("KV prefetch request is invalid"); + } + if (request.bytes > params.stage_slot_bytes) { + return fail_dispatch("KV prefetch request exceeds a stage slot"); + } + if (params.states[index] != llama_kv_stream_prefetch_state::pending) { + continue; + } + if (request.attention_index < params.current_attention || + uint64_t(request.attention_index) > window_end) { + continue; + } + + if (request.producer_attention_index >= 0) { + const uint32_t producer = uint32_t(request.producer_attention_index); + if (producer > params.current_attention || + (producer == params.current_attention && !params.current_producer_complete)) { + continue; + } + } + eligible.push_back(index); + } + + std::sort(eligible.begin(), eligible.end(), [&](size_t a, size_t b) { + const auto & lhs = params.requests[a]; + const auto & rhs = params.requests[b]; + return std::tie(lhs.attention_index, lhs.page_index, lhs.layer_id, a) < + std::tie(rhs.attention_index, rhs.page_index, rhs.layer_id, b); + }); + + const size_t count = std::min(slots.size(), eligible.size()); + result.assignments.reserve(count); + for (size_t i = 0; i < count; ++i) { + result.assignments.push_back({ eligible[i], slots[i] }); + } + + result.valid = true; + return result; +} + +llama_kv_stream_partition llama_kv_stream_partition_adapt( + const llama_kv_stream_partition_params & params) { + llama_kv_stream_partition result; + result.resident_pages_per_layer = params.previous_resident_pages_per_layer; + result.ring_slots = params.previous_ring_slots; + + auto fail_partition = [&](const char * message) { + result.valid = false; + result.error = message; + return result; + }; + + if (params.total_pool_pages == 0 || params.layer_count == 0) { + return fail_partition("KV stream partition geometry must be non-zero"); + } + if (params.grow_hysteresis_evaluations == 0 || + params.shrink_hysteresis_evaluations == 0 || + params.repartition_cooldown_evaluations == 0) { + return fail_partition("KV stream partition hysteresis must be non-zero"); + } + const auto ratio_valid = [](double value) { + return value >= 0.0 && value <= 1.0; + }; + if (!ratio_valid(params.deadline_miss_ratio) || + !ratio_valid(params.copy_engine_busy_ratio) || + !ratio_valid(params.ring_peak_occupancy_ratio)) { + return fail_partition("KV stream partition metrics must be ratios"); + } + if (!std::isfinite(params.target_ring_working_set_ratio) || + params.target_ring_working_set_ratio <= 0.0) { + return fail_partition("KV stream target working-set ratio must be positive"); + } + if (params.previous_resident_pages_per_layer > params.active_pages_per_layer) { + return fail_partition("resident KV pages exceed active pages"); + } + + uint64_t resident_pages = 0; + if (!checked_mul(params.previous_resident_pages_per_layer, params.layer_count, resident_pages) || + resident_pages + params.previous_ring_slots != params.total_pool_pages) { + return fail_partition("previous KV stream partition does not cover the fixed pool"); + } + if (params.previous_ring_slots < params.minimum_ring_slots) { + return fail_partition("previous KV stream ring is below its minimum"); + } + + constexpr double MISS_THRESHOLD = 0.01; + // A repartition invalidates compact resident addresses. Do not pay that + // transition cost when the copy engine already has too little headroom + // for a larger lookahead ring to increase sustained throughput. + constexpr double COPY_SATURATED = 0.80; + constexpr double COPY_LIGHT = 0.50; + constexpr double RING_LIGHT = 0.50; + constexpr uint32_t FEEDBACK_GROWTH_EPOCHS = 1; + + uint32_t target_resident_pages = 0; + const uint32_t maximum_resident_pages = std::min( + params.active_pages_per_layer, + (params.total_pool_pages - params.minimum_ring_slots)/params.layer_count); + for (uint32_t resident = maximum_resident_pages;; --resident) { + const uint32_t ring = params.total_pool_pages - resident*params.layer_count; + const uint32_t streamed = params.active_pages_per_layer - resident; + if (streamed == 0 || double(ring) >= + params.target_ring_working_set_ratio*double(streamed)) { + target_resident_pages = resident; + break; + } + if (resident == 0) { + break; + } + } + const uint32_t feedback_resident_floor = target_resident_pages > + FEEDBACK_GROWTH_EPOCHS ? target_resident_pages - FEEDBACK_GROWTH_EPOCHS : 0; + const bool overgrown_ring = + params.previous_resident_pages_per_layer < feedback_resident_floor; + + const bool below_overlap_target = + params.previous_resident_pages_per_layer > target_resident_pages; + const bool feedback_starved = + params.previous_resident_pages_per_layer > feedback_resident_floor && + params.deadline_miss_ratio > MISS_THRESHOLD && + params.copy_engine_busy_ratio < COPY_SATURATED; + const bool starved = below_overlap_target || feedback_starved; + const bool overprovisioned = params.deadline_miss_ratio <= MISS_THRESHOLD && + params.copy_engine_busy_ratio < COPY_LIGHT && + params.ring_peak_occupancy_ratio < RING_LIGHT; + + result.starved_evaluations = starved ? params.starved_evaluations + 1 : 0; + result.overprovisioned_evaluations = overprovisioned ? + params.overprovisioned_evaluations + 1 : 0; + + const bool cooldown_complete = params.evaluations_since_repartition >= + params.repartition_cooldown_evaluations; + if (params.entering_decode_layout && below_overlap_target) { + result.resident_pages_per_layer = target_resident_pages; + result.ring_slots = params.total_pool_pages - + result.resident_pages_per_layer*params.layer_count; + result.starved_evaluations = 0; + result.overprovisioned_evaluations = 0; + result.changed = true; + } else if (cooldown_complete && overgrown_ring) { + result.resident_pages_per_layer = feedback_resident_floor; + result.ring_slots = params.total_pool_pages - + result.resident_pages_per_layer*params.layer_count; + result.starved_evaluations = 0; + result.overprovisioned_evaluations = 0; + result.changed = true; + } else if (cooldown_complete && starved && + result.starved_evaluations >= params.grow_hysteresis_evaluations && + result.resident_pages_per_layer > 0) { + result.resident_pages_per_layer = below_overlap_target ? + target_resident_pages : result.resident_pages_per_layer - 1; + result.ring_slots = params.total_pool_pages - + result.resident_pages_per_layer*params.layer_count; + result.starved_evaluations = 0; + result.overprovisioned_evaluations = 0; + result.changed = true; + } else if (cooldown_complete && overprovisioned && + result.overprovisioned_evaluations >= params.shrink_hysteresis_evaluations && + result.resident_pages_per_layer < target_resident_pages && + result.ring_slots >= params.minimum_ring_slots + params.layer_count) { + ++result.resident_pages_per_layer; + result.ring_slots -= params.layer_count; + result.starved_evaluations = 0; + result.overprovisioned_evaluations = 0; + result.changed = true; + } + + result.valid = true; + return result; +} + +llama_kv_stream_feedback_delta llama_kv_stream_feedback_delta_make( + const llama_kv_stream_feedback_counters & current, + const llama_kv_stream_feedback_counters & previous) { + llama_kv_stream_feedback_delta result; + if (current.deadline_samples < previous.deadline_samples || + current.deadline_misses < previous.deadline_misses) { + result.error = "KV stream feedback counters moved backwards"; + return result; + } + + result.deadline_samples = current.deadline_samples - previous.deadline_samples; + result.deadline_misses = current.deadline_misses - previous.deadline_misses; + if (result.deadline_misses > result.deadline_samples) { + result.error = "KV stream deadline misses exceed samples"; + return result; + } + + result.valid = true; + result.has_evaluation = result.deadline_samples != 0; + if (result.has_evaluation) { + result.deadline_miss_ratio = + double(result.deadline_misses)/double(result.deadline_samples); + } + return result; +} diff --git a/src/llama-kv-stream-plan.h b/src/llama-kv-stream-plan.h new file mode 100644 index 000000000000..9d065e96b49f --- /dev/null +++ b/src/llama-kv-stream-plan.h @@ -0,0 +1,220 @@ +#pragma once + +#include +#include +#include +#include + +enum class llama_kv_stream_region_role : uint8_t { + target, + mtp, +}; + +// Smallest independently resident piece of a logical KV cache. The CUDA copy +// layer may coalesce adjacent regions into a larger transfer slot. +struct llama_kv_stream_region { + llama_kv_stream_region_role role = llama_kv_stream_region_role::target; + + int32_t layer_id = -1; + uint32_t token_begin = 0; + uint32_t token_count = 0; + + uint64_t bytes = 0; + + // Lower values have stronger residency preference. The ordering must be + // stable across replans to avoid unnecessary promotion/demotion churn. + uint64_t residency_priority = 0; + + // Pinned is a logical pool policy, not cudaMalloc pinning. A pinned region + // must receive permanent pool space or planning fails. + bool pinned = false; +}; + +struct llama_kv_stream_plan_params { + uint64_t pool_bytes = 0; + uint64_t stage_slot_bytes = 0; + uint32_t stage_slots = 0; + + std::vector regions; +}; + +struct llama_kv_stream_plan { + bool valid = false; + std::string error; + + uint64_t pool_bytes = 0; + uint64_t reserved_stage_bytes = 0; + uint64_t resident_bytes = 0; + uint64_t streamed_bytes = 0; + uint64_t unused_bytes = 0; + + std::vector resident_regions; + std::vector streamed_regions; +}; + +llama_kv_stream_plan llama_kv_stream_plan_make(const llama_kv_stream_plan_params & params); + +struct llama_kv_stream_extent_params { + uint32_t live_tokens = 0; + uint32_t reserve_tokens = 0; + uint32_t page_tokens = 256; + uint32_t previous_extent = 0; + uint32_t shrink_hysteresis_tokens = 0; + uint32_t maximum_tokens = 0; + + bool force_shrink = false; +}; + +struct llama_kv_stream_extent { + bool valid = false; + std::string error; + + uint32_t tokens = 0; + bool grew = false; + bool shrunk = false; +}; + +llama_kv_stream_extent llama_kv_stream_extent_make(const llama_kv_stream_extent_params & params); + +struct llama_kv_stream_layer_layout { + llama_kv_stream_region_role role = llama_kv_stream_region_role::target; + + int32_t layer_id = -1; + uint32_t n_tokens = 0; + uint64_t bytes_per_token = 0; + + // Lower values spread residency preference within the same token page. + uint32_t layer_priority = 0; + + bool pin_all = false; + bool pin_tail = false; +}; + +struct llama_kv_stream_regions_params { + uint32_t page_tokens = 256; + std::vector layers; +}; + +struct llama_kv_stream_regions { + bool valid = false; + std::string error; + + uint64_t total_bytes = 0; + std::vector regions; +}; + +llama_kv_stream_regions llama_kv_stream_regions_make(const llama_kv_stream_regions_params & params); + +enum class llama_kv_stream_prefetch_state : uint8_t { + pending, + scheduled, + consumed, +}; + +struct llama_kv_stream_prefetch_request { + int32_t layer_id = -1; + uint32_t attention_index = 0; + uint32_t page_index = 0; + uint64_t bytes = 0; + + // -1 means immutable at graph start. A non-negative value identifies the + // attention step whose SET_ROWS operation must finish before this page may + // be copied (normally the mutable tail page). + int32_t producer_attention_index = -1; +}; + +struct llama_kv_stream_prefetch_assignment { + size_t request_index = 0; + uint32_t slot = 0; +}; + +struct llama_kv_stream_prefetch_params { + uint32_t current_attention = 0; + uint32_t lookahead_layers = 0; + uint64_t stage_slot_bytes = 0; + + bool adaptive_lookahead = false; + bool current_producer_complete = false; + + std::vector free_slots; + std::vector requests; + std::vector states; +}; + +struct llama_kv_stream_prefetch_dispatch_result { + bool valid = false; + std::string error; + std::vector assignments; +}; + +// Selects work for the currently free transfer slots. The caller owns the +// asynchronous lifecycle and changes request states only after recording the +// corresponding CUDA ready/consumed events. +llama_kv_stream_prefetch_dispatch_result llama_kv_stream_prefetch_dispatch( + const llama_kv_stream_prefetch_params & params); + +struct llama_kv_stream_partition_params { + uint32_t total_pool_pages = 0; + uint32_t layer_count = 0; + uint32_t active_pages_per_layer = 0; + uint32_t minimum_ring_slots = 0; + + uint32_t previous_resident_pages_per_layer = 0; + uint32_t previous_ring_slots = 0; + + double deadline_miss_ratio = 0.0; + double copy_engine_busy_ratio = 0.0; + double ring_peak_occupancy_ratio = 0.0; + + // Ring slots relative to the streamed pages consumed by one attention + // layer. A value above one leaves room to begin the next layer early. + double target_ring_working_set_ratio = 1.10; + + uint32_t starved_evaluations = 0; + uint32_t overprovisioned_evaluations = 0; + uint32_t grow_hysteresis_evaluations = 3; + uint32_t shrink_hysteresis_evaluations = 8; + uint32_t evaluations_since_repartition = UINT32_MAX; + uint32_t repartition_cooldown_evaluations = 64; + + // The first decode graph has no useful streaming feedback yet, but its + // active working set is already known. Select the deterministic overlap + // target immediately instead of spending several tokens in an undersized + // prefill ring before hysteresis can react. + bool entering_decode_layout = false; +}; + +struct llama_kv_stream_partition { + bool valid = false; + std::string error; + bool changed = false; + + uint32_t resident_pages_per_layer = 0; + uint32_t ring_slots = 0; + uint32_t starved_evaluations = 0; + uint32_t overprovisioned_evaluations = 0; +}; + +// Adjusts the boundary inside a fixed page pool. A growth step demotes exactly +// one page from every layer, turning those addresses into immediately reusable +// ring slots. Promotion performs the inverse operation after longer hysteresis. +llama_kv_stream_partition llama_kv_stream_partition_adapt( + const llama_kv_stream_partition_params & params); + +struct llama_kv_stream_feedback_counters { + uint64_t deadline_samples = 0; + uint64_t deadline_misses = 0; +}; + +struct llama_kv_stream_feedback_delta { + bool valid = false; + bool has_evaluation = false; + std::string error; + uint64_t deadline_samples = 0; + uint64_t deadline_misses = 0; + double deadline_miss_ratio = 0.0; +}; + +llama_kv_stream_feedback_delta llama_kv_stream_feedback_delta_make( + const llama_kv_stream_feedback_counters & current, + const llama_kv_stream_feedback_counters & previous); diff --git a/src/llama-kv-stream-softmax.cpp b/src/llama-kv-stream-softmax.cpp new file mode 100644 index 000000000000..a666058d9703 --- /dev/null +++ b/src/llama-kv-stream-softmax.cpp @@ -0,0 +1,75 @@ +#include "llama-kv-stream-softmax.h" + +#include +#include + +llama_kv_stream_softmax_result llama_kv_stream_softmax_merge( + const std::vector & parts) { + llama_kv_stream_softmax_result result; + + auto fail = [&](const char * message) { + result.valid = false; + result.error = message; + result.numerator.clear(); + result.value.clear(); + return result; + }; + + if (parts.empty()) { + return fail("softmax block merge requires at least one partial result"); + } + + const size_t value_width = parts.front().numerator.size(); + if (value_width == 0) { + return fail("softmax block merge requires a non-empty value vector"); + } + + result.max_logit = parts.front().max_logit; + + for (const auto & part : parts) { + if (part.numerator.size() != value_width) { + return fail("softmax block partial results have different value dimensions"); + } + + if (!std::isfinite(part.max_logit) || + !std::isfinite(part.normalizer) || part.normalizer <= 0.0f) { + return fail("softmax block partial result has invalid metadata"); + } + + if (!std::all_of(part.numerator.begin(), part.numerator.end(), [](float value) { + return std::isfinite(value); + })) { + return fail("softmax block partial result has a non-finite numerator"); + } + + result.max_logit = std::max(result.max_logit, part.max_logit); + } + + std::vector numerator(value_width, 0.0); + double normalizer = 0.0; + + for (const auto & part : parts) { + const double scale = std::exp(double(part.max_logit) - result.max_logit); + normalizer += scale*part.normalizer; + + for (size_t i = 0; i < value_width; ++i) { + numerator[i] += scale*part.numerator[i]; + } + } + + if (!std::isfinite(normalizer) || normalizer <= 0.0) { + return fail("combined softmax normalizer is invalid"); + } + + result.normalizer = float(normalizer); + result.numerator.resize(value_width); + result.value.resize(value_width); + + for (size_t i = 0; i < value_width; ++i) { + result.numerator[i] = float(numerator[i]); + result.value[i] = float(numerator[i]/normalizer); + } + + result.valid = true; + return result; +} diff --git a/src/llama-kv-stream-softmax.h b/src/llama-kv-stream-softmax.h new file mode 100644 index 000000000000..aa7ee8606fa3 --- /dev/null +++ b/src/llama-kv-stream-softmax.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include + +// Partial FlashAttention result for one logical KV block. The numerator is +// deliberately unnormalized so independently computed blocks can be combined +// with the same max-rescaling used by the CUDA FlashAttention fixup kernel. +struct llama_kv_stream_softmax_part { + float max_logit = 0.0f; + float normalizer = 0.0f; + + std::vector numerator; +}; + +struct llama_kv_stream_softmax_result { + bool valid = false; + std::string error; + + float max_logit = 0.0f; + float normalizer = 0.0f; + + std::vector numerator; + std::vector value; +}; + +llama_kv_stream_softmax_result llama_kv_stream_softmax_merge( + const std::vector & parts); diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..c9b0c2947c95 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -29,7 +29,8 @@ llama_memory_hybrid::llama_memory_hybrid( bool unified, /* layer filters */ const layer_filter_cb & filter_attn, - const layer_filter_cb & filter_recr) : + const layer_filter_cb & filter_recr, + size_t kv_stream_stage_bytes) : hparams(model.hparams), mem_attn(new llama_kv_cache( model, @@ -49,7 +50,9 @@ llama_memory_hybrid::llama_memory_hybrid( [&](int32_t il) { return !hparams.is_recr(il); } : filter_attn, nullptr, - nullptr + nullptr, + "", + kv_stream_stage_bytes )), mem_recr(new llama_memory_recurrent( model, diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..a5436502f790 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -39,7 +39,8 @@ class llama_memory_hybrid : public llama_memory_i { bool unified, /* layer filters */ const layer_filter_cb & filter_attn = nullptr, - const layer_filter_cb & filter_recr = nullptr); + const layer_filter_cb & filter_recr = nullptr, + size_t kv_stream_stage_bytes = 0); ~llama_memory_hybrid() = default; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..bc27cbcc0bb7 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -18,6 +18,7 @@ struct llama_memory_params { // kv cache ggml_type type_k; ggml_type type_v; + uint64_t kv_stream_stage_bytes; // use full-size SWA cache bool swa_full; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 9e2e0c893675..b1bbe354602a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -1777,6 +1778,20 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // indicate that this buffer contains weights // this is used by ggml_backend_sched to improve op scheduling: ops that use a weight are preferably scheduled to the backend that contains the weight ggml_backend_buffer_set_usage(buf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + if (getenv("GGML_CUDA_PREFER_MODEL_WEIGHTS") != nullptr) { + ggml_backend_buffer_type_t buft_weights = ggml_backend_buffer_get_type(buf.get()); + ggml_backend_dev_t dev_weights = ggml_backend_buft_get_device(buft_weights); + if (dev_weights != nullptr) { + ggml_backend_reg_t reg_weights = ggml_backend_dev_backend_reg(dev_weights); + using prefer_device_fn_t = bool (*)(ggml_backend_buffer_t); + auto * prefer_device_fn = (prefer_device_fn_t) ggml_backend_reg_get_proc_address( + reg_weights, "ggml_backend_cuda_buffer_set_preferred_device"); + if (prefer_device_fn != nullptr) { + (void) prefer_device_fn(buf.get()); + } + } + } } pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs)); @@ -2559,7 +2574,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* offload */ cparams.offload_kqv, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), - /* filter_recr */ std::move(filter_recr)); + /* filter_recr */ std::move(filter_recr), + /* kv stream stage */ arch == LLM_ARCH_QWEN35 ? params.kv_stream_stage_bytes : 0); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c46377c7623d..dfdf0b1a9bc4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -289,6 +289,21 @@ llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p " set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model) llama_build_and_test(test-arg-parser.cpp) +llama_build_and_test(test-kv-stream-plan.cpp) +target_include_directories(test-kv-stream-plan PRIVATE ${PROJECT_SOURCE_DIR}/src) +llama_build_and_test(test-kv-stream-config.cpp) +target_include_directories(test-kv-stream-config PRIVATE ${PROJECT_SOURCE_DIR}/src) +llama_build_and_test(test-kv-stream-softmax.cpp) +target_include_directories(test-kv-stream-softmax PRIVATE ${PROJECT_SOURCE_DIR}/src) +llama_build_and_test( + test-kv-stream-bench-config.cpp + ${PROJECT_SOURCE_DIR}/tools/kv-stream-bench/kv-stream-bench-config.cpp) +target_include_directories(test-kv-stream-bench-config PRIVATE ${PROJECT_SOURCE_DIR}/tools/kv-stream-bench) +if (GGML_CUDA) + llama_build_and_test(test-kv-stream-cuda-buffer.cpp LABEL cuda) + llama_build_and_test(test-kv-stream-cuda-attn.cpp LABEL cuda) + llama_build_and_test(test-kv-stream-cuda-set-rows.cpp LABEL cuda) +endif() llama_build_and_test(test-model-resolution.cpp) # the test serves its repos from an httplib server, and the library links it privately target_link_libraries(test-model-resolution PRIVATE cpp-httplib) diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index e0907631abd8..81ea5afc1d75 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -95,10 +95,12 @@ static void test(void) { common_params base; base.n_parallel = 4; base.n_outputs_max_per_seq = 8; + base.kv_stream_stage_mib = 64; const auto draft = common_base_params_to_speculative(base); assert(draft.n_outputs_max == 4); assert(draft.n_outputs_max_per_seq == 1); + assert(draft.kv_stream_stage_mib == 0); } printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n"); @@ -186,6 +188,18 @@ static void test(void) { argv = {"binary_name", "-sm", "hello"}; assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); + { + common_params stream_params; + assert(stream_params.kv_stream_stage_mib == 0); + + argv = {"binary_name", "-m", "model_file.gguf", "--kv-stream-stage-mib", "64"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), stream_params, LLAMA_EXAMPLE_COMMON)); + assert(stream_params.kv_stream_stage_mib == 64); + + argv = {"binary_name", "-m", "model_file.gguf", "--kv-stream-stage-mib", "-1"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), stream_params, LLAMA_EXAMPLE_COMMON)); + } + { common_params penalty_params; assert(penalty_params.sampling.penalty_last_n == 64); diff --git a/tests/test-kv-stream-bench-config.cpp b/tests/test-kv-stream-bench-config.cpp new file mode 100644 index 000000000000..fdc6a04551f4 --- /dev/null +++ b/tests/test-kv-stream-bench-config.cpp @@ -0,0 +1,80 @@ +#include "kv-stream-bench-config.h" +#include "testing.h" + +#include +#include +#include + +namespace { + +constexpr uint64_t MIB = 1024ULL*1024ULL; + +llama_kv_stream_bench_config parse(std::initializer_list args) { + std::vector argv(args); + return llama_kv_stream_bench_config_parse(int(argv.size()), argv.data()); +} + +} // namespace + +int main() { + testing t; + + t.test("defaults are dry-run and cover observed transfer sizes", [](testing & t) { + const auto config = parse({ "llama-kv-stream-bench" }); + + if (!t.assert_true("configuration is valid", config.valid)) { + return; + } + + t.assert_true("execution requires opt in", !config.execute); + t.assert_equal(0, config.device); + t.assert_equal(5, config.iterations); + t.assert_equal(uint64_t(1024)*MIB, config.reserve_bytes); + t.assert_true("default sizes are exact", config.transfer_bytes == std::vector({ + 26*MIB, 197*MIB, 312*MIB, 416*MIB, + })); + }); + + t.test("explicit safe options are parsed", [](testing & t) { + const auto config = parse({ + "llama-kv-stream-bench", + "--execute", + "--device", "1", + "--iterations", "3", + "--reserve-mib", "2048", + "--sizes-mib", "4,26,416", + }); + + if (!t.assert_true("configuration is valid", config.valid)) { + return; + } + + t.assert_true("execution is enabled", config.execute); + t.assert_equal(1, config.device); + t.assert_equal(3, config.iterations); + t.assert_equal(uint64_t(2048)*MIB, config.reserve_bytes); + t.assert_true("explicit sizes are exact", config.transfer_bytes == std::vector({ + 4*MIB, 26*MIB, 416*MIB, + })); + }); + + t.test("unsafe or malformed options are rejected", [](testing & t) { + t.assert_true("zero iterations are rejected", !parse({ + "bench", "--iterations", "0", + }).valid); + t.assert_true("zero transfer is rejected", !parse({ + "bench", "--sizes-mib", "26,0,416", + }).valid); + t.assert_true("negative device is rejected", !parse({ + "bench", "--device", "-1", + }).valid); + t.assert_true("unknown option is rejected", !parse({ + "bench", "--surprise", + }).valid); + t.assert_true("missing value is rejected", !parse({ + "bench", "--reserve-mib", + }).valid); + }); + + return t.summary(); +} diff --git a/tests/test-kv-stream-config.cpp b/tests/test-kv-stream-config.cpp new file mode 100644 index 000000000000..c7b343ecca51 --- /dev/null +++ b/tests/test-kv-stream-config.cpp @@ -0,0 +1,92 @@ +#include "llama-kv-stream-config.h" +#include "testing.h" + +int main() { + testing t; + + t.test("streaming is opt-in", [](testing & t) { + llama_kv_stream_config config; + const auto result = llama_kv_stream_config_validate(config); + t.assert_true("disabled config is valid", result.valid); + t.assert_true("disabled config remains disabled", !result.enabled); + }); + + t.test("supported target configuration is accepted", [](testing & t) { + llama_kv_stream_config config; + config.stage_bytes = 64ULL*1024ULL*1024ULL; + config.minimum_stage_bytes = 1664ULL*256ULL; + config.arch_qwen35 = true; + config.context_default = true; + config.single_sequence = true; + config.flash_attention = true; + config.kv_offload = true; + + const auto result = llama_kv_stream_config_validate(config); + t.assert_true("config is valid", result.valid); + t.assert_true("config is enabled", result.enabled); + }); + + t.test("each unsupported condition fails loudly", [](testing & t) { + llama_kv_stream_config base; + base.stage_bytes = 64ULL*1024ULL*1024ULL; + base.minimum_stage_bytes = 1664ULL*256ULL; + base.arch_qwen35 = true; + base.context_default = true; + base.single_sequence = true; + base.flash_attention = true; + base.kv_offload = true; + + auto expect_invalid = [&](const char * name, const llama_kv_stream_config & config) { + const auto result = llama_kv_stream_config_validate(config); + t.assert_true(name, !result.valid && !result.enabled && !result.error.empty()); + }; + + auto config = base; + config.arch_qwen35 = false; + expect_invalid("non-Qwen architecture", config); + config = base; + config.context_default = false; + expect_invalid("draft/MTP context", config); + config = base; + config.single_sequence = false; + expect_invalid("parallel sequences", config); + config = base; + config.flash_attention = false; + expect_invalid("Flash Attention disabled", config); + config = base; + config.kv_offload = false; + expect_invalid("KV offload disabled", config); + config.stage_bytes = config.minimum_stage_bytes - 1; + expect_invalid("stage smaller than one page", config); + }); + + t.test("pool is partitioned evenly across layers with one scratch page", [](testing & t) { + const auto layout = llama_kv_stream_pool_layout_make({ + /*.pool_bytes =*/ 64ULL*1024ULL*1024ULL, + /*.page_bytes =*/ 1664ULL*256ULL, + /*.layer_count =*/ 16, + /*.scratch_pages=*/ 1, + }); + + t.assert_true("layout is valid", layout.valid); + t.assert_equal(uint32_t(9), layout.resident_pages_per_layer); + t.assert_equal(uint32_t(9*256), layout.resident_tokens_per_layer); + t.assert_equal(1664ULL*256ULL, layout.scratch_bytes); + t.assert_equal( + 64ULL*1024ULL*1024ULL, + layout.scratch_bytes + layout.resident_bytes + layout.unused_bytes); + }); + + t.test("pool rejects missing scratch or resident capacity", [](testing & t) { + auto layout = llama_kv_stream_pool_layout_make({ 0, 1664ULL*256ULL, 16, 1 }); + t.assert_true("zero pool", !layout.valid); + + layout = llama_kv_stream_pool_layout_make({ 1664ULL*256ULL, 1664ULL*256ULL, 16, 1 }); + t.assert_true("scratch-only pool", !layout.valid); + + layout = llama_kv_stream_pool_layout_make({ 64ULL*1024ULL*1024ULL, 0, 16, 1 }); + t.assert_true("zero page", !layout.valid); + }); + + return t.summary(); +} diff --git a/tests/test-kv-stream-cuda-attn.cpp b/tests/test-kv-stream-cuda-attn.cpp new file mode 100644 index 000000000000..fab8287163b9 --- /dev/null +++ b/tests/test-kv-stream-cuda-attn.cpp @@ -0,0 +1,1627 @@ +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpp.h" +#include "ggml-cuda.h" +#include "../ggml/src/ggml-impl.h" +#include "../ggml/src/ggml-cuda/kv-stream-span-tuner.h" +#include "ggml.h" +#include "testing.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int64_t HEAD_DIM = 256; +constexpr int64_t N_KV_HEAD = 4; +constexpr int64_t N_Q_HEAD = 24; + +size_t align_up(size_t value, size_t alignment) { + return (value + alignment - 1)/alignment*alignment; +} + +struct attention_inputs { + ggml_type type_k = GGML_TYPE_Q8_0; + ggml_type type_v = GGML_TYPE_Q4_0; + std::vector q; + std::vector k; + std::vector v; + std::vector mask; +}; + +attention_inputs make_inputs( + int64_t n_kv, int64_t n_batch, int64_t query_start = 0, + ggml_type type_k = GGML_TYPE_Q8_0, ggml_type type_v = GGML_TYPE_Q4_0) { + attention_inputs result; + result.type_k = type_k; + result.type_v = type_v; + + result.q.resize(HEAD_DIM*n_batch*N_Q_HEAD); + for (size_t i = 0; i < result.q.size(); ++i) { + result.q[i] = 0.15f*std::sin(float(i)*0.03125f) + 0.05f*std::cos(float(i)*0.0078125f); + } + + const int64_t nrows = n_kv*N_KV_HEAD; + std::vector source(HEAD_DIM*nrows); + for (size_t i = 0; i < source.size(); ++i) { + source[i] = 0.4f*std::sin(float(i)*0.001953125f) + 0.2f*std::cos(float(i)*0.00048828125f); + } + + result.k.resize(ggml_row_size(type_k, HEAD_DIM)*nrows); + const size_t k_written = ggml_quantize_chunk( + type_k, source.data(), result.k.data(), 0, nrows, HEAD_DIM, nullptr); + GGML_ASSERT(k_written == result.k.size()); + + for (size_t i = 0; i < source.size(); ++i) { + source[i] = 0.35f*std::cos(float(i)*0.00146484375f) - 0.1f*std::sin(float(i)*0.00390625f); + } + result.v.resize(ggml_row_size(type_v, HEAD_DIM)*nrows); + const size_t v_written = ggml_quantize_chunk( + type_v, source.data(), result.v.data(), 0, nrows, HEAD_DIM, nullptr); + GGML_ASSERT(v_written == result.v.size()); + + result.mask.resize(n_kv*n_batch); + for (int64_t batch = 0; batch < n_batch; ++batch) { + for (int64_t token = 0; token < n_kv; ++token) { + // Vary both token blocks and both query rows. This catches a + // streamed implementation that offsets the first mask row but + // accidentally uses the compact block width as the next-row pitch. + const float bias = token <= query_start + batch ? + -0.015625f*float((token + 73*batch) % 127) : -INFINITY; + result.mask[batch*n_kv + token] = ggml_fp32_to_fp16(bias); + } + } + return result; +} + +attention_inputs make_f16_reference(const attention_inputs & inputs, int64_t n_kv) { + attention_inputs result; + result.type_k = GGML_TYPE_F16; + result.type_v = GGML_TYPE_F16; + result.q = inputs.q; + result.mask = inputs.mask; + + const int64_t nrows = n_kv*N_KV_HEAD; + result.k.resize(ggml_row_size(GGML_TYPE_F16, HEAD_DIM)*nrows); + result.v.resize(ggml_row_size(GGML_TYPE_F16, HEAD_DIM)*nrows); + std::vector row(HEAD_DIM); + + auto convert = [&](ggml_type type, const std::vector & source, + std::vector & destination) { + const size_t source_row_bytes = ggml_row_size(type, HEAD_DIM); + const size_t destination_row_bytes = ggml_row_size(GGML_TYPE_F16, HEAD_DIM); + const auto * traits = ggml_get_type_traits(type); + for (int64_t i = 0; i < nrows; ++i) { + const void * source_row = source.data() + size_t(i)*source_row_bytes; + if (type == GGML_TYPE_F32) { + std::memcpy(row.data(), source_row, HEAD_DIM*sizeof(float)); + } else { + GGML_ASSERT(traits->to_float != nullptr); + traits->to_float(source_row, row.data(), HEAD_DIM); + } + ggml_fp32_to_fp16_row( + row.data(), + reinterpret_cast( + destination.data() + size_t(i)*destination_row_bytes), + HEAD_DIM); + } + }; + + convert(inputs.type_k, inputs.k, result.k); + convert(inputs.type_v, inputs.v, result.v); + return result; +} + +std::vector run_attention( + ggml_backend_t backend, + const attention_inputs & inputs, + ggml_backend_buffer_type_t kv_buft, + int64_t n_kv, + int64_t n_batch, + int repeats = 1, + int64_t update_rows = 1, + bool change_updates = false, + ggml_type index_type = GGML_TYPE_I32, + bool index_on_host = true, + ggml_backend_cuda_kv_stream_runtime_t dirty_runtime = nullptr, + bool change_indices = false, + bool replace_cache = false, + uint64_t graph_uid = 0) { + constexpr size_t N_TENSORS = 32; + const size_t context_bytes = ggml_tensor_overhead()*N_TENSORS + ggml_graph_overhead_custom(N_TENSORS, false); + + ggml_init_params params{ + /* .mem_size = */ context_bytes, + /* .mem_buffer = */ nullptr, + /* .no_alloc = */ true, + }; + ggml_context_ptr compute_ctx(ggml_init(params)); + ggml_context_ptr kv_ctx(ggml_init(params)); + ggml_context_ptr index_ctx(ggml_init(params)); + GGML_ASSERT(compute_ctx && kv_ctx && index_ctx); + + ggml_tensor * q = ggml_new_tensor_4d( + compute_ctx.get(), GGML_TYPE_F32, HEAD_DIM, n_batch, N_Q_HEAD, 1); + ggml_tensor * mask = ggml_new_tensor_4d( + compute_ctx.get(), GGML_TYPE_F16, n_kv, n_batch, 1, 1); + ggml_tensor * k_storage = ggml_new_tensor_2d( + kv_ctx.get(), inputs.type_k, HEAD_DIM*N_KV_HEAD, n_kv); + ggml_tensor * v_storage = ggml_new_tensor_2d( + kv_ctx.get(), inputs.type_v, HEAD_DIM*N_KV_HEAD, n_kv); + ggml_tensor * k_cache = ggml_view_4d( + kv_ctx.get(), k_storage, HEAD_DIM, N_KV_HEAD, n_kv, 1, + ggml_row_size(inputs.type_k, HEAD_DIM), + ggml_row_size(inputs.type_k, HEAD_DIM*N_KV_HEAD), + ggml_row_size(inputs.type_k, HEAD_DIM*N_KV_HEAD)*n_kv, 0); + ggml_tensor * v_cache = ggml_view_4d( + kv_ctx.get(), v_storage, HEAD_DIM, N_KV_HEAD, n_kv, 1, + ggml_row_size(inputs.type_v, HEAD_DIM), + ggml_row_size(inputs.type_v, HEAD_DIM*N_KV_HEAD), + ggml_row_size(inputs.type_v, HEAD_DIM*N_KV_HEAD)*n_kv, 0); + ggml_tensor * k = ggml_permute(kv_ctx.get(), k_cache, 0, 2, 1, 3); + ggml_tensor * v = ggml_permute(kv_ctx.get(), v_cache, 0, 2, 1, 3); + + ggml_tensor * k_update = ggml_new_tensor_2d( + compute_ctx.get(), GGML_TYPE_F32, HEAD_DIM*N_KV_HEAD, update_rows); + ggml_tensor * v_update = ggml_new_tensor_2d( + compute_ctx.get(), GGML_TYPE_F32, HEAD_DIM*N_KV_HEAD, update_rows); + ggml_tensor * update_index = ggml_new_tensor_1d(index_ctx.get(), index_type, update_rows); + ggml_tensor * updated_k = ggml_set_rows(compute_ctx.get(), k_storage, k_update, update_index); + ggml_tensor * updated_v = ggml_set_rows(compute_ctx.get(), v_storage, v_update, update_index); + + ggml_tensor * out = ggml_flash_attn_ext( + compute_ctx.get(), q, k, v, mask, 1.0f/std::sqrt(float(HEAD_DIM)), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_set_name(out, "streamed-attention-output"); + + ggml_backend_buffer_ptr kv_buffer( + ggml_backend_alloc_ctx_tensors_from_buft(kv_ctx.get(), kv_buft)); + ggml_backend_buffer_ptr compute_buffer( + ggml_backend_alloc_ctx_tensors(compute_ctx.get(), backend)); + ggml_backend_buffer_ptr index_buffer( + ggml_backend_alloc_ctx_tensors_from_buft(index_ctx.get(), index_on_host ? + ggml_backend_cuda_host_buffer_type() : ggml_backend_get_default_buffer_type(backend))); + GGML_ASSERT(kv_buffer && compute_buffer && index_buffer); + + ggml_backend_tensor_set(q, inputs.q.data(), 0, inputs.q.size()*sizeof(float)); + ggml_backend_tensor_set(k_storage, inputs.k.data(), 0, inputs.k.size()); + ggml_backend_tensor_set(v_storage, inputs.v.data(), 0, inputs.v.size()); + ggml_backend_tensor_set(mask, inputs.mask.data(), 0, inputs.mask.size()*sizeof(uint16_t)); + + std::vector k_update_data(HEAD_DIM*N_KV_HEAD*update_rows); + std::vector v_update_data(HEAD_DIM*N_KV_HEAD*update_rows); + for (size_t i = 0; i < k_update_data.size(); ++i) { + k_update_data[i] = 0.6f*std::sin(float(i)*0.0234375f); + v_update_data[i] = 0.4f*std::cos(float(i)*0.017578125f); + } + ggml_backend_tensor_set(k_update, k_update_data.data(), 0, k_update_data.size()*sizeof(float)); + ggml_backend_tensor_set(v_update, v_update_data.data(), 0, v_update_data.size()*sizeof(float)); + std::vector dirty_rows(update_rows); + for (int64_t row = 0; row < update_rows; ++row) { dirty_rows[row] = row; } + if (index_type == GGML_TYPE_I32) { + std::vector update_index_data(update_rows); + for (int64_t row = 0; row < update_rows; ++row) { update_index_data[row] = int32_t(row); } + ggml_backend_tensor_set( + update_index, update_index_data.data(), 0, update_index_data.size()*sizeof(int32_t)); + } else { + GGML_ASSERT(index_type == GGML_TYPE_I64); + ggml_backend_tensor_set( + update_index, dirty_rows.data(), 0, dirty_rows.size()*sizeof(int64_t)); + } + + ggml_cgraph * graph = ggml_new_graph_custom(compute_ctx.get(), N_TENSORS, false); + graph->uid = graph_uid; + ggml_build_forward_expand(graph, updated_k); + ggml_build_forward_expand(graph, updated_v); + ggml_build_forward_expand(graph, out); + GGML_ASSERT(ggml_backend_supports_op(backend, updated_k)); + GGML_ASSERT(ggml_backend_supports_op(backend, updated_v)); + if (!ggml_backend_supports_op(backend, out)) { + std::fprintf(stderr, "unsupported attention K=%s V=%s buft=%s\n", + ggml_type_name(inputs.type_k), ggml_type_name(inputs.type_v), + ggml_backend_buft_name(kv_buft)); + } + GGML_ASSERT(ggml_backend_supports_op(backend, out)); + for (int repeat = 0; repeat < repeats; ++repeat) { + if (replace_cache && repeat == 3) { + std::vector zero_k(inputs.k.size(), 0); + std::vector zero_v(inputs.v.size(), 0); + ggml_backend_tensor_set(k_storage, zero_k.data(), 0, zero_k.size()); + ggml_backend_tensor_set(v_storage, zero_v.data(), 0, zero_v.size()); + } + if (change_indices) { + GGML_ASSERT(index_type == GGML_TYPE_I64); + for (int64_t row = 0; row < update_rows; ++row) { + dirty_rows[row] = (int64_t(repeat)*update_rows + row)%n_kv; + } + ggml_backend_tensor_set( + update_index, dirty_rows.data(), 0, dirty_rows.size()*sizeof(int64_t)); + } + if (dirty_runtime != nullptr) { + GGML_ASSERT(ggml_backend_cuda_kv_stream_mark_dirty_rows( + dirty_runtime, dirty_rows.data(), dirty_rows.size())); + } + if (repeat > 0 && change_updates) { + for (float & value : k_update_data) { value = -2.0f*value; } + for (float & value : v_update_data) { value = -2.0f*value; } + ggml_backend_tensor_set(k_update, k_update_data.data(), 0, k_update_data.size()*sizeof(float)); + ggml_backend_tensor_set(v_update, v_update_data.data(), 0, v_update_data.size()*sizeof(float)); + } + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + } + + std::vector result(ggml_nelements(out)); + ggml_backend_tensor_get(out, result.data(), 0, result.size()*sizeof(float)); + return result; +} + +std::vector run_attention_layers( + ggml_backend_t backend, + const std::vector & layers, + ggml_backend_buffer_type_t kv_buft, + int64_t n_kv, + int64_t n_batch, + int repeats = 1, + int64_t update_rows = 1, + ggml_type index_type = GGML_TYPE_I32, + ggml_backend_cuda_kv_stream_runtime_t dirty_runtime = nullptr, + uint32_t layout_after_first = 0) { + constexpr size_t N_TENSORS = 256; + const size_t context_bytes = ggml_tensor_overhead()*N_TENSORS + + ggml_graph_overhead_custom(N_TENSORS, false); + + ggml_init_params params{ + /* .mem_size = */ context_bytes, + /* .mem_buffer = */ nullptr, + /* .no_alloc = */ true, + }; + ggml_context_ptr compute_ctx(ggml_init(params)); + ggml_context_ptr kv_ctx(ggml_init(params)); + GGML_ASSERT(compute_ctx && kv_ctx); + + struct layer_tensors { + ggml_tensor * q; + ggml_tensor * mask; + ggml_tensor * k_storage; + ggml_tensor * v_storage; + ggml_tensor * k_update; + ggml_tensor * v_update; + ggml_tensor * update_index; + ggml_tensor * updated_k; + ggml_tensor * updated_v; + ggml_tensor * out; + }; + std::vector tensors; + tensors.reserve(layers.size()); + + for (size_t layer = 0; layer < layers.size(); ++layer) { + layer_tensors current{}; + current.q = ggml_new_tensor_4d( + compute_ctx.get(), GGML_TYPE_F32, HEAD_DIM, n_batch, N_Q_HEAD, 1); + current.mask = ggml_new_tensor_4d( + compute_ctx.get(), GGML_TYPE_F16, n_kv, n_batch, 1, 1); + current.k_storage = ggml_new_tensor_2d( + kv_ctx.get(), GGML_TYPE_Q8_0, HEAD_DIM*N_KV_HEAD, n_kv); + current.v_storage = ggml_new_tensor_2d( + kv_ctx.get(), GGML_TYPE_Q4_0, HEAD_DIM*N_KV_HEAD, n_kv); + ggml_tensor * k_cache = ggml_view_4d( + kv_ctx.get(), current.k_storage, HEAD_DIM, N_KV_HEAD, n_kv, 1, + ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM), + ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM*N_KV_HEAD), + ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM*N_KV_HEAD)*n_kv, 0); + ggml_tensor * v_cache = ggml_view_4d( + kv_ctx.get(), current.v_storage, HEAD_DIM, N_KV_HEAD, n_kv, 1, + ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM), + ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM*N_KV_HEAD), + ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM*N_KV_HEAD)*n_kv, 0); + ggml_tensor * k = ggml_permute(kv_ctx.get(), k_cache, 0, 2, 1, 3); + ggml_tensor * v = ggml_permute(kv_ctx.get(), v_cache, 0, 2, 1, 3); + + current.k_update = ggml_new_tensor_2d( + compute_ctx.get(), GGML_TYPE_F32, HEAD_DIM*N_KV_HEAD, update_rows); + current.v_update = ggml_new_tensor_2d( + compute_ctx.get(), GGML_TYPE_F32, HEAD_DIM*N_KV_HEAD, update_rows); + current.update_index = ggml_new_tensor_1d(compute_ctx.get(), index_type, update_rows); + current.updated_k = ggml_set_rows( + compute_ctx.get(), current.k_storage, current.k_update, current.update_index); + current.updated_v = ggml_set_rows( + compute_ctx.get(), current.v_storage, current.v_update, current.update_index); + current.out = ggml_flash_attn_ext( + compute_ctx.get(), current.q, k, v, current.mask, + 1.0f/std::sqrt(float(HEAD_DIM)), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(current.out, GGML_PREC_F32); + tensors.push_back(current); + } + + ggml_backend_buffer_ptr kv_buffer( + ggml_backend_alloc_ctx_tensors_from_buft(kv_ctx.get(), kv_buft)); + ggml_backend_buffer_ptr compute_buffer( + ggml_backend_alloc_ctx_tensors(compute_ctx.get(), backend)); + GGML_ASSERT(kv_buffer && compute_buffer); + + std::vector dirty_rows(update_rows); + for (int64_t row = 0; row < update_rows; ++row) { dirty_rows[row] = row; } + for (size_t layer = 0; layer < layers.size(); ++layer) { + const auto & input = layers[layer]; + auto & current = tensors[layer]; + ggml_backend_tensor_set(current.q, input.q.data(), 0, input.q.size()*sizeof(float)); + ggml_backend_tensor_set(current.k_storage, input.k.data(), 0, input.k.size()); + ggml_backend_tensor_set(current.v_storage, input.v.data(), 0, input.v.size()); + ggml_backend_tensor_set(current.mask, input.mask.data(), 0, input.mask.size()*sizeof(uint16_t)); + + std::vector k_update_data(HEAD_DIM*N_KV_HEAD*update_rows); + std::vector v_update_data(HEAD_DIM*N_KV_HEAD*update_rows); + for (size_t i = 0; i < k_update_data.size(); ++i) { + k_update_data[i] = (0.5f + 0.03f*layer)*std::sin(float(i)*0.0234375f); + v_update_data[i] = (0.3f + 0.02f*layer)*std::cos(float(i)*0.017578125f); + } + ggml_backend_tensor_set(current.k_update, k_update_data.data(), 0, + k_update_data.size()*sizeof(float)); + ggml_backend_tensor_set(current.v_update, v_update_data.data(), 0, + v_update_data.size()*sizeof(float)); + if (index_type == GGML_TYPE_I32) { + std::vector rows_i32(update_rows); + for (int64_t row = 0; row < update_rows; ++row) { rows_i32[row] = int32_t(row); } + ggml_backend_tensor_set( + current.update_index, rows_i32.data(), 0, rows_i32.size()*sizeof(int32_t)); + } else { + GGML_ASSERT(index_type == GGML_TYPE_I64); + ggml_backend_tensor_set( + current.update_index, dirty_rows.data(), 0, dirty_rows.size()*sizeof(int64_t)); + } + } + + ggml_cgraph * graph = ggml_new_graph_custom(compute_ctx.get(), N_TENSORS, false); + for (auto & current : tensors) { + ggml_build_forward_expand(graph, current.updated_k); + ggml_build_forward_expand(graph, current.updated_v); + ggml_build_forward_expand(graph, current.out); + } + for (int repeat = 0; repeat < repeats; ++repeat) { + if (dirty_runtime != nullptr) { + GGML_ASSERT(ggml_backend_cuda_kv_stream_mark_dirty_rows( + dirty_runtime, dirty_rows.data(), dirty_rows.size())); + } + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + if (repeat == 0 && layout_after_first != 0) { + GGML_ASSERT(ggml_backend_cuda_kv_stream_set_decode_layout( + dirty_runtime, layout_after_first)); + } + } + + std::vector result; + for (auto & current : tensors) { + const size_t begin = result.size(); + result.resize(begin + ggml_nelements(current.out)); + ggml_backend_tensor_get(current.out, result.data() + begin, 0, + ggml_nbytes(current.out)); + } + return result; +} + +} // namespace + +int main() { + testing t; + + t.test("decode span tuner selects the faster measured mode per layout", [](testing & t) { + ggml_cuda_kv_stream_span_tuner production_tuner; + production_tuner.observe(100.0, /* streamed = */ true, /* bounded = */ false); + for (uint32_t sample = 0; sample < 4; ++sample) { + production_tuner.observe(10.0, /* streamed = */ true, /* bounded = */ false); + } + t.assert_true("production tuner does not decide from four samples", + !production_tuner.use_bounded() && !production_tuner.selected()); + + ggml_cuda_kv_stream_span_tuner tuner(/* trial_samples = */ 2, 0.005, /* warmup_samples = */ 1); + + tuner.observe(100.0, /* streamed = */ false, /* bounded = */ false); + t.assert_true("non-streamed graphs do not start a trial", !tuner.use_bounded()); + t.assert_true("non-streamed graphs do not select a mode", !tuner.selected()); + + tuner.observe(100.0, /* streamed = */ true, /* bounded = */ false); + t.assert_true("unbounded warmup is not measured", !tuner.use_bounded()); + + tuner.observe(10.0, /* streamed = */ true, /* bounded = */ false); + tuner.observe(10.2, /* streamed = */ true, /* bounded = */ false); + t.assert_true("tuner advances to bounded trials", tuner.use_bounded()); + t.assert_true("both modes are measured before selection", !tuner.selected()); + + tuner.observe(100.0, /* streamed = */ true, /* bounded = */ true); + t.assert_true("bounded warmup is not measured", !tuner.selected()); + + tuner.observe(8.0, /* streamed = */ true, /* bounded = */ true); + tuner.observe(8.2, /* streamed = */ true, /* bounded = */ true); + t.assert_true("bounded mode is selected when materially faster", tuner.selected()); + t.assert_true("bounded mode remains active after selection", tuner.use_bounded()); + + tuner.reset(); + tuner.observe(100.0, /* streamed = */ true, /* bounded = */ false); + tuner.observe(10.0, /* streamed = */ true, /* bounded = */ false); + tuner.observe(10.0, /* streamed = */ true, /* bounded = */ false); + tuner.observe(100.0, /* streamed = */ true, /* bounded = */ true); + tuner.observe(10.0, /* streamed = */ true, /* bounded = */ true); + tuner.observe(10.0, /* streamed = */ true, /* bounded = */ true); + t.assert_true("tuner selects after both trials", tuner.selected()); + t.assert_true("noise does not displace the ordinary kernel", !tuner.use_bounded()); + + tuner.observe(1.0, /* streamed = */ true, /* bounded = */ true); + t.assert_true("selection remains stable until layout reset", !tuner.use_bounded()); + }); + + t.test("all native CUDA KV pairs preserve streamed prefill results", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 4; + const ggml_type native_types[] = { + GGML_TYPE_F16, + GGML_TYPE_Q4_0, + GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, + GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, + GGML_TYPE_BF16, + }; + + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + for (const ggml_type type_k : native_types) { + for (const ggml_type type_v : native_types) { + const attention_inputs inputs = + make_inputs(n_kv, n_batch, n_kv - n_batch, type_k, type_v); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch); + + const size_t k_page_bytes = + ggml_row_size(type_k, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = + ggml_row_size(type_v, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + if (!std::isfinite(max_abs) || max_abs > 5e-4f || + stats.asynchronous_page_uploads == 0) { + std::fprintf(stderr, + "native pair K=%s V=%s max_abs=%g async_uploads=%llu\n", + ggml_type_name(type_k), ggml_type_name(type_v), max_abs, + (unsigned long long) stats.asynchronous_page_uploads); + } + t.assert_true("native pair executes streamed attention", + stats.asynchronous_page_uploads > 0); + t.assert_true("native pair remains numerically equivalent", + std::isfinite(max_abs) && max_abs <= 5e-4f); + } + } + }); + + t.test("all bounded-fallback KV pairs preserve streamed prefill results", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 4; + const ggml_type kv_types[] = { + GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, + }; + + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + for (const ggml_type type_k : kv_types) { + for (const ggml_type type_v : kv_types) { + if (ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v) != + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16) { + continue; + } + + const attention_inputs inputs = + make_inputs(n_kv, n_batch, n_kv - n_batch, type_k, type_v); + const attention_inputs reference = make_f16_reference(inputs, n_kv); + const std::vector expected = run_attention( + backend.get(), reference, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch); + + const size_t k_page_bytes = + ggml_row_size(type_k, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = + ggml_row_size(type_v, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + const size_t f16_k_page_bytes = + ggml_row_size(GGML_TYPE_F16, HEAD_DIM)*N_KV_HEAD*256; + const size_t f16_v_page_bytes = + ggml_row_size(GGML_TYPE_F16, HEAD_DIM)*N_KV_HEAD*256; + const size_t conversion_bytes = + align_up(f16_k_page_bytes, 128) + f16_v_page_bytes; + + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + params.conversion_bytes = conversion_bytes; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + if (!std::isfinite(max_abs) || max_abs > 2e-3f || + stats.asynchronous_page_uploads == 0) { + std::fprintf(stderr, + "fallback pair K=%s V=%s max_abs=%g async_uploads=%llu\n", + ggml_type_name(type_k), ggml_type_name(type_v), max_abs, + (unsigned long long) stats.asynchronous_page_uploads); + } + t.assert_true("fallback pair executes streamed attention", + stats.asynchronous_page_uploads > 0); + t.assert_true("fallback pair remains numerically equivalent", + std::isfinite(max_abs) && max_abs <= 2e-3f); + } + } + }); + + t.test("wide generic fallback remains equivalent across query workspace tiles", [](testing & t) { + constexpr int64_t n_kv = 1024; + constexpr int64_t n_batch = 513; + constexpr ggml_type type_k = GGML_TYPE_IQ4_NL; + constexpr ggml_type type_v = GGML_TYPE_IQ4_NL; + + t.assert_equal( + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16, + ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v)); + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = + make_inputs(n_kv, n_batch, n_kv - n_batch, type_k, type_v); + const attention_inputs reference = make_f16_reference(inputs, n_kv); + const std::vector expected = run_attention( + backend.get(), reference, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(type_k, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(type_v, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + const size_t f16_page_bytes = + ggml_row_size(GGML_TYPE_F16, HEAD_DIM)*N_KV_HEAD*256; + const size_t conversion_bytes = align_up(f16_page_bytes, 128) + f16_page_bytes; + + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 2; + params.conversion_bytes = conversion_bytes; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + std::fprintf(stderr, + "wide-fallback n_batch=%lld max_abs=%g async_uploads=%llu\n", + (long long) n_batch, max_abs, + (unsigned long long) stats.asynchronous_page_uploads); + t.assert_true("wide fallback executes streamed attention", + stats.asynchronous_page_uploads > 0); + t.assert_equal(uint64_t(8), stats.host_to_device_copy_commands); + t.assert_true("wide fallback remains numerically equivalent", + std::isfinite(max_abs) && max_abs <= 2e-3f); + }); + + t.test("server-shaped causal prefill pipelines four Q8/Q4 blocks through two slots", [](testing & t) { + constexpr int64_t n_kv = 1024; + constexpr int64_t n_batch = 83; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + // Exercise the final causal block. With query_start == 0, every block + // after the first is masked and corrupted streamed pages are invisible. + const attention_inputs inputs = make_inputs(n_kv, n_batch, n_kv - 256); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch, 2, 256, true); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 2; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + std::vector cleared(params.stage_bytes, 0); + GGML_ASSERT(ggml_backend_cuda_kv_stream_stage_upload( + runtime, 0, 0, cleared.data(), cleared.size())); + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), n_kv, n_batch, 2, 256, true); + + + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + t.assert_true("multi-token streamed spans use MMA partial attention", + stats.mma_prefill_attention_spans > 0); + t.assert_equal(uint64_t(6), stats.asynchronous_page_uploads); + t.assert_equal(uint64_t(16), stats.host_to_device_copy_commands); + t.assert_equal(uint64_t(6), stats.compute_stream_waits); + t.assert_equal(uint64_t(4), stats.stage_slot_reuses); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + t.assert_equal(uint64_t(4), stats.streamed_attention_spans); + t.assert_equal(uint64_t(6), stats.streamed_pages_attended); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + + float max_abs = 0.0f; + float max_rel = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + max_rel = std::max(max_rel, std::abs(expected[i] - actual[i])/(std::abs(expected[i]) + 1e-6f)); + } + std::fprintf(stderr, "streamed attention max_abs=%g max_rel=%g\n", max_abs, max_rel); + t.assert_true("outputs remain finite", std::isfinite(max_abs) && std::isfinite(max_rel)); + t.assert_true("streamed output is numerically equivalent", max_abs <= 3e-4f); + }); + + t.test("one-page causal prefill stays bit-identical to ordinary CUDA attention", [](testing & t) { + constexpr int64_t n_kv = 256; + constexpr int64_t n_batch = 5; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), n_kv, n_batch); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(0), stats.staged_set_rows); + t.assert_equal(expected.size(), actual.size()); + t.assert_true("one-page outputs are bit-identical", expected == actual); + }); + + t.test("fully resident multi-page prefill stays bit-identical to ordinary CUDA attention", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 4; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch, 340); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("resident runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), n_kv, n_batch); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(0), stats.streamed_pages); + t.assert_equal(expected.size(), actual.size()); + t.assert_true("fully resident outputs are bit-identical", expected == actual); + }); + + t.test("four-query page-boundary prefill remains finite and equivalent", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 4; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch, 340); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 2; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("resident runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), n_kv, n_batch); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + bool all_finite = true; + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + all_finite = all_finite && std::isfinite(actual[i]); + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + std::fprintf(stderr, "four-query page-boundary max_abs=%g\n", max_abs); + t.assert_true("page-boundary output remains finite", all_finite); + t.assert_true("page-boundary output remains equivalent", max_abs <= 3e-4f); + }); + + t.test("wide causal prefills remain equivalent across the 256-query boundary", [](testing & t) { + constexpr int64_t n_kv = 1024; + const int64_t query_counts[] = { 257, 512, 513, 1024 }; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + + for (const int64_t n_batch : query_counts) { + const attention_inputs inputs = make_inputs(n_kv, n_batch, n_kv - n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch, 1, n_batch); + + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 2; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 1, n_batch, false, GGML_TYPE_I64, false, runtime); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + bool all_finite = true; + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + all_finite = all_finite && std::isfinite(actual[i]); + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + std::fprintf(stderr, + "wide-query n_batch=%lld max_abs=%g streamed_pages=%llu\n", + (long long) n_batch, max_abs, + (unsigned long long) stats.streamed_pages); + t.assert_true("wide-query output remains finite", all_finite); + t.assert_true("wide-query output remains equivalent", max_abs <= 3e-4f); + t.assert_true("wide-query test exercises streamed pages", stats.streamed_pages > 0); + } + }); + + t.test("resident pages survive between evaluations while the tail is refreshed", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 4; + params.pool_bytes = 6*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("resident runtime initializes", runtime != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(512, 256); + (void) run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + 512, 256, 2, 256, false, GGML_TYPE_I64, false, runtime); + + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + t.assert_equal(uint64_t(2), stats.resident_misses); + t.assert_equal(uint64_t(2), stats.resident_hits); + t.assert_equal(uint64_t(0), stats.streamed_pages); + t.assert_equal(uint64_t(2*page_bytes), stats.host_to_device_bytes); + t.assert_equal(uint64_t(2), stats.resident_attention_spans); + t.assert_equal(uint64_t(4), stats.resident_pages_attended); + t.assert_equal(uint64_t(4), stats.staged_set_rows); + t.assert_equal(uint64_t(2*page_bytes), stats.staged_set_rows_bytes); + + t.assert_true("one resident page is demoted into the ring", + ggml_backend_cuda_kv_stream_repartition(runtime, 5)); + t.assert_equal(uint32_t(1), + ggml_backend_cuda_kv_stream_resident_pages_per_layer(runtime)); + (void) run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), 512, 256); + + const auto repartitioned = ggml_backend_cuda_kv_stream_get_stats(runtime); + t.assert_equal(uint64_t(1), repartitioned.resident_misses); + t.assert_equal(uint64_t(1), repartitioned.streamed_pages); + t.assert_equal(uint64_t(1), repartitioned.asynchronous_page_uploads); + t.assert_equal(uint64_t(1), repartitioned.resident_attention_spans); + t.assert_equal(uint64_t(1), repartitioned.resident_pages_attended); + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + t.test("one-row decode refreshes only the changed resident K and V rows", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch, 2, 1, true, GGML_TYPE_I32, true); + + const size_t k_token_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD; + const size_t v_token_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD; + const size_t k_page_bytes = k_token_bytes*256; + const size_t v_page_bytes = v_token_bytes*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("resident runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 2, 1, true, GGML_TYPE_I32, true, runtime); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(2), stats.resident_misses); + t.assert_equal(uint64_t(2), stats.resident_hits); + t.assert_equal( + uint64_t(2*page_bytes + k_token_bytes + v_token_bytes), + stats.host_to_device_bytes); + t.assert_equal(expected.size(), actual.size()); + t.assert_true("one-row resident refresh remains bit-identical", expected == actual); + }); + + t.test("resident SET_ROWS mirrors changing decode slots without page refreshes", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch, 5, 1, true, GGML_TYPE_I64, false, nullptr, true); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("resident runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 5, 1, true, GGML_TYPE_I64, false, runtime, true); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + + using feedback_fn_t = bool (*)( + void *, uint64_t *, uint64_t *, double *, uint32_t *, + uint32_t *, uint32_t *, uint32_t *); + ggml_backend_dev_t device = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); + auto feedback_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_kv_stream_feedback")); + uint64_t deadline_samples = 0; + uint64_t deadline_misses = 0; + double copy_busy_ratio = -1.0; + uint32_t peak_occupancy = 0; + uint32_t ring_slots = 0; + uint32_t resident_pages = 0; + uint32_t controlled_pages = 0; + t.assert_true("feedback remains readable after resident graph replay", + feedback_fn != nullptr && feedback_fn( + runtime, &deadline_samples, &deadline_misses, ©_busy_ratio, + &peak_occupancy, &ring_slots, &resident_pages, &controlled_pages)); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(2*page_bytes), stats.host_to_device_bytes); + t.assert_equal(expected.size(), actual.size()); + t.assert_true("mirrored changing slots remain bit-identical", expected == actual); + }); + + t.test("resident staged writes support every exposed KV-cache format", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 1; + const ggml_type kv_types[] = { + GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, + }; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + for (const ggml_type type : kv_types) { + const attention_inputs inputs = + make_inputs(n_kv, n_batch, n_kv - 1, type, type); + const size_t typed_page_bytes = + ggml_row_size(type, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(typed_page_bytes, 128) + typed_page_bytes; + const bool fallback = + ggml_backend_cuda_kv_stream_get_attention_mode(type, type) == + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16; + const size_t f16_page_bytes = + ggml_row_size(GGML_TYPE_F16, HEAD_DIM)*N_KV_HEAD*256; + const size_t conversion_bytes = fallback ? + align_up(f16_page_bytes, 128) + f16_page_bytes : 0; + + ggml_backend_cuda_kv_stream_params baseline_params{}; + baseline_params.device = 0; + baseline_params.stage_bytes = page_bytes; + baseline_params.stage_slots = 1; + baseline_params.conversion_bytes = conversion_bytes; + auto baseline_runtime = ggml_backend_cuda_kv_stream_runtime_new(baseline_params); + if (!t.assert_true("baseline runtime initializes", baseline_runtime != nullptr)) { + return; + } + const std::vector expected = run_attention( + backend.get(), inputs, + ggml_backend_cuda_kv_stream_buffer_type(baseline_runtime), + n_kv, n_batch, 5, 3, true, GGML_TYPE_I64, false, nullptr, true); + ggml_backend_cuda_kv_stream_runtime_free(baseline_runtime); + + ggml_backend_cuda_kv_stream_params resident_params{}; + resident_params.device = 0; + resident_params.stage_bytes = page_bytes; + resident_params.stage_slots = 1; + resident_params.pool_bytes = 3*page_bytes + conversion_bytes; + resident_params.conversion_bytes = conversion_bytes; + resident_params.resident_layer_count = 1; + resident_params.page_tokens = 256; + auto resident_runtime = ggml_backend_cuda_kv_stream_runtime_new(resident_params); + if (!t.assert_true("resident runtime initializes", resident_runtime != nullptr)) { + return; + } + const std::vector actual = run_attention( + backend.get(), inputs, + ggml_backend_cuda_kv_stream_buffer_type(resident_runtime), + n_kv, n_batch, 5, 3, true, GGML_TYPE_I64, false, + resident_runtime, true); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(resident_runtime); + ggml_backend_cuda_kv_stream_runtime_free(resident_runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + if (!t.assert_true(ggml_type_name(type), + std::isfinite(max_abs) && max_abs <= 5e-4f)) { + std::fprintf(stderr, "resident staged write type=%s max_abs=%g\n", + ggml_type_name(type), max_abs); + return; + } + t.assert_equal(uint64_t(2*page_bytes), stats.host_to_device_bytes); + t.assert_equal(uint64_t(10), stats.staged_set_rows); + t.assert_equal( + uint64_t(10*3*ggml_row_size(type, HEAD_DIM*N_KV_HEAD)), + stats.staged_set_rows_bytes); + } + }); + + t.test("resident graph reloads after authoritative cache replacement", [](testing & t) { + constexpr int64_t n_kv = 512; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch, 5, 1, true, GGML_TYPE_I64, false, nullptr, true, true, 1); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + params.pool_bytes = 3*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("resident runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 5, 1, true, GGML_TYPE_I64, false, runtime, true, true, 1); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(2), stats.resident_misses); + t.assert_equal(expected.size(), actual.size()); + t.assert_true("reloaded resident output is bit-identical", expected == actual); + }); + + t.test("decode batches contiguous streamed pages without changing logits", [](testing & t) { + constexpr int64_t n_kv = 1536; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch, 2); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 4; + params.pool_bytes = 5*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("batched runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 2); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(10), stats.asynchronous_page_uploads); + t.assert_equal(uint64_t(12), stats.host_to_device_copy_commands); + t.assert_equal(uint64_t(4), stats.compute_stream_waits); + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + t.assert_true("batched streamed logits remain equivalent", max_abs <= 3e-4f); + }); + + t.test("decode pipelines bounded attention chunks and samples every immutable copy batch", [](testing & t) { + constexpr int64_t n_kv = 41*256; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const attention_inputs inputs = make_inputs(n_kv, n_batch); + const std::vector expected = run_attention( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 40; + params.pool_bytes = 41*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + params.decode_span_pages = 32; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("chunk-pipelined runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(40), stats.streamed_pages); + t.assert_equal(uint64_t(2), stats.streamed_attention_spans); + // The 39 immutable pages form one full 32-page transfer and one + // seven-page transfer. The mutable tail is deliberately excluded from + // adaptive prefetch feedback because its producer runs in this graph. + t.assert_equal(uint64_t(2), stats.deadline_samples); + t.assert_true("chunk deadline misses cannot exceed samples", + stats.deadline_misses <= stats.deadline_samples); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + t.assert_true("chunk-pipelined logits remain equivalent", max_abs <= 3e-4f); + }); + + t.test("sixteen attention layers share one resident/ring pool during causal prefill", [](testing & t) { + constexpr int64_t n_kv = 1024; + constexpr int64_t n_batch = 83; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + std::vector inputs(16, make_inputs(n_kv, n_batch, n_kv - 256)); + for (size_t layer = 1; layer < inputs.size(); ++layer) { + for (size_t i = 0; i < inputs[layer].q.size(); ++i) { + inputs[layer].q[i] += 0.025f*layer*std::sin(float(i)*0.015625f); + } + } + const std::vector expected = run_attention_layers( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), + n_kv, n_batch, 1, 2, GGML_TYPE_I64); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 8; + params.pool_bytes = 24*page_bytes; + params.resident_layer_count = 16; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("shared runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention_layers( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 1, 2, GGML_TYPE_I64, runtime); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + t.assert_equal(uint64_t(32), stats.staged_set_rows); + t.assert_equal( + uint64_t(16*(2*ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM*N_KV_HEAD) + + 2*ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM*N_KV_HEAD))), + stats.staged_set_rows_bytes); + t.assert_equal(uint64_t(48), stats.asynchronous_page_uploads); + t.assert_equal(uint64_t(48), stats.compute_stream_waits); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(expected.size(), actual.size()); + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + t.assert_true("sixteen-layer prefill remains equivalent", max_abs <= 3e-4f); + }); + + t.test("one shared ring prefetches across attention layers", [](testing & t) { + constexpr int64_t n_kv = 768; + // Cross-layer prefetch is a decode optimization. A one-token batch + // also guarantees that only the mutable tail page is changed by the + // SET_ROWS producers in this graph. + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + std::vector inputs{ + make_inputs(n_kv, n_batch), + make_inputs(n_kv, n_batch), + make_inputs(n_kv, n_batch), + }; + for (size_t layer = 1; layer < inputs.size(); ++layer) { + for (size_t i = 0; i < inputs[layer].q.size(); ++i) { + inputs[layer].q[i] += 0.025f*layer*std::sin(float(i)*0.015625f); + } + } + const std::vector expected = run_attention_layers( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 2; + params.pool_bytes = 5*page_bytes; + params.resident_layer_count = 3; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("shared runtime initializes", runtime != nullptr)) { + return; + } + + const std::vector actual = run_attention_layers( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), n_kv, n_batch, 2); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + t.assert_equal(uint64_t(12), stats.asynchronous_page_uploads); + t.assert_equal(uint64_t(12), stats.compute_stream_waits); + t.assert_equal(uint64_t(4), stats.cross_layer_prefetches); + t.assert_equal(uint64_t(6), stats.deadline_samples); + t.assert_true("deadline misses cannot exceed samples", + stats.deadline_misses <= stats.deadline_samples); + t.assert_equal(uint32_t(2), stats.ring_peak_occupancy); + t.assert_equal(uint64_t(6), stats.streamed_attention_spans); + t.assert_equal(uint64_t(12), stats.streamed_pages_attended); + + using feedback_fn_t = bool (*)( + void *, uint64_t *, uint64_t *, double *, uint32_t *, + uint32_t *, uint32_t *, uint32_t *); + ggml_backend_dev_t device = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); + auto feedback_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_feedback")); + if (!t.assert_true("dynamic feedback API is exported", feedback_fn != nullptr)) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + return; + } + uint64_t deadline_samples = 0; + uint64_t deadline_misses = 0; + double copy_busy_ratio = -1.0; + uint32_t peak_occupancy = 0; + uint32_t ring_slots = 0; + uint32_t resident_pages = 0; + uint32_t controlled_pages = 0; + t.assert_true("dynamic feedback is readable", feedback_fn( + runtime, &deadline_samples, &deadline_misses, ©_busy_ratio, + &peak_occupancy, &ring_slots, &resident_pages, &controlled_pages)); + t.assert_equal(uint64_t(6), deadline_samples); + t.assert_true("copy busy ratio is normalized", + copy_busy_ratio >= 0.0 && copy_busy_ratio <= 1.0); + t.assert_equal(uint32_t(2), peak_occupancy); + t.assert_equal(uint32_t(2), ring_slots); + t.assert_equal(uint32_t(1), resident_pages); + t.assert_equal(uint32_t(5), controlled_pages); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + std::fprintf(stderr, "multi-layer streamed attention max_abs=%g\n", max_abs); + t.assert_true("multi-layer streamed logits remain equivalent", max_abs <= 3e-4f); + }); + + t.test("decode layout permits multiple streaming waves through a small ring", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 1; + params.pool_bytes = 5*page_bytes; + params.resident_layer_count = 4; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("small-ring runtime initializes", runtime != nullptr)) { + return; + } + + // Each of four layers has one resident and two streamed pages. The + // one-slot ring services each layer in two waves, so requiring one + // split layer per ring-sized portion would incorrectly demand eight + // split layers from a four-layer model. + t.assert_true("small ring selects a valid decode layout", + ggml_backend_cuda_kv_stream_set_decode_layout(runtime, 3)); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + t.test("decode layout and ring boundary publish as one reconfiguration", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 8; + params.pool_bytes = 12*page_bytes; + params.resident_layer_count = 4; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("shared runtime initializes", runtime != nullptr)) { + return; + } + + t.assert_true("combined reconfiguration succeeds", + ggml_backend_cuda_kv_stream_reconfigure( + runtime, /* active pages per layer = */ 3, /* ring slots = */ 4)); + t.assert_equal(uint32_t(4), + ggml_backend_cuda_kv_stream_stage_slots(runtime)); + t.assert_equal(uint32_t(2), + ggml_backend_cuda_kv_stream_resident_pages_per_layer(runtime)); + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + + t.test("decode layout spreads an oversized streamed deficit across enough layers", [](testing & t) { + constexpr int64_t n_kv = 768; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + std::vector inputs{ + make_inputs(n_kv, n_batch, n_kv - 1), + make_inputs(n_kv, n_batch, n_kv - 1), + make_inputs(n_kv, n_batch, n_kv - 1), + make_inputs(n_kv, n_batch, n_kv - 1), + }; + const std::vector expected = run_attention_layers( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 8; + params.pool_bytes = 12*page_bytes; + params.resident_layer_count = 4; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("shared runtime initializes", runtime != nullptr)) { + return; + } + + // Four layers with three active pages and one uniformly resident page + // have an eight-page streamed deficit. An eight-slot ring can hold the + // entire deficit, but one layer contains only three pages, so the layout + // must distribute the deficit over at least ceil(8/3) layers. + if (!t.assert_true("decode layout respects per-layer active-page capacity", + ggml_backend_cuda_kv_stream_set_decode_layout(runtime, 3))) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + return; + } + + const std::vector actual = run_attention_layers( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 1, 1, GGML_TYPE_I32, runtime); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(8), stats.streamed_pages); + t.assert_equal(uint64_t(4), stats.resident_pages_attended); + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + max_abs = std::max(max_abs, std::abs(expected[i] - actual[i])); + } + std::fprintf(stderr, "oversized-ring decode max_abs=%g\n", max_abs); + t.assert_true("oversized-ring decode remains equivalent", max_abs <= 3e-4f); + }); + + t.test("decode layout bounds layer concentration by the transfer ring", [](testing & t) { + constexpr int64_t n_kv = 768; + constexpr int64_t n_batch = 1; + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + std::vector inputs{ + make_inputs(n_kv, n_batch, n_kv - 1), + make_inputs(n_kv, n_batch, n_kv - 1), + make_inputs(n_kv, n_batch, n_kv - 1), + }; + const std::vector expected = run_attention_layers( + backend.get(), inputs, ggml_backend_get_default_buffer_type(backend.get()), n_kv, n_batch); + + const size_t k_page_bytes = ggml_row_size(GGML_TYPE_Q8_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t v_page_bytes = ggml_row_size(GGML_TYPE_Q4_0, HEAD_DIM)*N_KV_HEAD*256; + const size_t page_bytes = align_up(k_page_bytes, 128) + v_page_bytes; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 3; + params.pool_bytes = 6*page_bytes; + params.resident_layer_count = 3; + params.page_tokens = 256; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("shared runtime initializes", runtime != nullptr)) { + return; + } + const std::vector uniform = run_attention_layers( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 1, 1, GGML_TYPE_I32, runtime); + float uniform_max_abs = 0.0f; + for (size_t i = 0; i < expected.size(); ++i) { + uniform_max_abs = + std::max(uniform_max_abs, std::abs(expected[i] - uniform[i])); + } + std::fprintf(stderr, "uniform decode max_abs=%g\n", uniform_max_abs); + t.assert_true("uniform decode control remains equivalent", uniform_max_abs <= 3e-4f); + if (!t.assert_true("decode residency becomes ring-bounded", + ggml_backend_cuda_kv_stream_set_decode_layout(runtime, 3))) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + return; + } + + const std::vector actual = run_attention_layers( + backend.get(), inputs, ggml_backend_cuda_kv_stream_buffer_type(runtime), + n_kv, n_batch, 2, 1, GGML_TYPE_I32, runtime, 4); + const auto stats = ggml_backend_cuda_kv_stream_get_stats(runtime); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + + t.assert_equal(uint64_t(6), stats.resident_misses); + t.assert_equal(uint64_t(0), stats.resident_hits); + t.assert_equal(uint64_t(12), stats.streamed_pages); + t.assert_equal(uint64_t(4), stats.resident_attention_spans); + t.assert_equal(uint64_t(6), stats.resident_pages_attended); + if (!t.assert_equal(expected.size(), actual.size())) { + return; + } + float max_abs = 0.0f; + std::vector layer_max_abs(inputs.size(), 0.0f); + const size_t layer_elements = expected.size()/inputs.size(); + for (size_t i = 0; i < expected.size(); ++i) { + const float error = std::abs(expected[i] - actual[i]); + max_abs = std::max(max_abs, error); + layer_max_abs[i/layer_elements] = std::max(layer_max_abs[i/layer_elements], error); + } + std::fprintf(stderr, "ring-bounded decode max_abs=%g layers=%g,%g,%g\n", + max_abs, layer_max_abs[0], layer_max_abs[1], layer_max_abs[2]); + t.assert_true("ring-bounded decode remains equivalent", max_abs <= 3e-4f); + }); + + ggml_quantize_free(); + return t.summary(); +} diff --git a/tests/test-kv-stream-cuda-buffer.cpp b/tests/test-kv-stream-cuda-buffer.cpp new file mode 100644 index 000000000000..80e73cbb55ad --- /dev/null +++ b/tests/test-kv-stream-cuda-buffer.cpp @@ -0,0 +1,237 @@ +#include "ggml-cuda.h" +#include "testing.h" + +#include +#include + +int main() { + testing t; + + t.test("runtime exposes exact staging geometry", [](testing & t) { + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = 1024*1024; + params.stage_slots = 2; + + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("runtime allocation succeeds", runtime != nullptr)) { + return; + } + + t.assert_equal(params.stage_bytes, ggml_backend_cuda_kv_stream_stage_bytes(runtime)); + t.assert_equal(params.stage_slots, ggml_backend_cuda_kv_stream_stage_slots(runtime)); + auto buft = ggml_backend_cuda_kv_stream_buffer_type(runtime); + if (!t.assert_true("streamed buffer type exists", buft != nullptr)) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + return; + } + t.assert_true( + "owning CUDA device accepts streamed buffers", + ggml_backend_dev_supports_buft(ggml_backend_buft_get_device(buft), buft)); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + t.test("authoritative buffer is pinned-host accessible", [](testing & t) { + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = 1024*1024; + params.stage_slots = 1; + + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("runtime allocation succeeds", runtime != nullptr)) { + return; + } + + auto buft = ggml_backend_cuda_kv_stream_buffer_type(runtime); + auto buffer = ggml_backend_buft_alloc_buffer(buft, 256*1024); + if (!t.assert_true("host buffer allocation succeeds", buffer != nullptr)) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + return; + } + + t.assert_true("buffer reports host accessibility", ggml_backend_buffer_is_host(buffer)); + t.assert_equal(size_t(256*1024), ggml_backend_buffer_get_size(buffer)); + void * base = ggml_backend_buffer_get_base(buffer); + t.assert_true("buffer base exists", base != nullptr); + + using query_fn_t = bool (*)(const void *, bool *); + ggml_backend_dev_t device = ggml_backend_buft_get_device(buft); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); + auto query_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cuda_kv_stream_host_is_write_combined")); + if (!t.assert_true("host allocation query is exported", query_fn != nullptr)) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + ggml_backend_buffer_free(buffer); + return; + } + bool write_combined = false; + t.assert_true("host allocation flags are readable", + query_fn(base, &write_combined)); +#if defined(_WIN32) + t.assert_true("authoritative KV storage is not write-combined on Windows", !write_combined); +#else + t.assert_true("authoritative KV storage is write-combined", write_combined); +#endif + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + ggml_backend_buffer_free(buffer); + }); + + t.test("invalid runtime geometry is rejected", [](testing & t) { + ggml_backend_cuda_kv_stream_params params{}; + params.device = -1; + params.stage_bytes = 1024*1024; + params.stage_slots = 1; + t.assert_true("negative device is rejected", ggml_backend_cuda_kv_stream_runtime_new(params) == nullptr); + + params.device = 0; + params.stage_bytes = 0; + t.assert_true("zero stage size is rejected", ggml_backend_cuda_kv_stream_runtime_new(params) == nullptr); + + params.stage_bytes = 1024*1024; + params.stage_slots = 0; + t.assert_true("zero stage slots are rejected", ggml_backend_cuda_kv_stream_runtime_new(params) == nullptr); + }); + + t.test("unimplemented operations cannot consume streamed storage", [](testing & t) { + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = 1024*1024; + params.stage_slots = 1; + + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("runtime allocation succeeds", runtime != nullptr)) { + return; + } + + auto buft = ggml_backend_cuda_kv_stream_buffer_type(runtime); + auto buffer = ggml_backend_buft_alloc_buffer(buft, 4096); + if (!t.assert_true("host buffer allocation succeeds", buffer != nullptr)) { + ggml_backend_cuda_kv_stream_runtime_free(runtime); + return; + } + + ggml_tensor source{}; + source.type = GGML_TYPE_F32; + source.buffer = buffer; + source.data = ggml_backend_buffer_get_base(buffer); + source.ne[0] = 1; + source.ne[1] = 1; + source.ne[2] = 1; + source.ne[3] = 1; + source.nb[0] = sizeof(float); + source.nb[1] = sizeof(float); + source.nb[2] = sizeof(float); + source.nb[3] = sizeof(float); + + ggml_tensor fill{}; + fill.op = GGML_OP_FILL; + fill.type = GGML_TYPE_F32; + fill.ne[0] = fill.ne[1] = fill.ne[2] = fill.ne[3] = 1; + fill.nb[0] = fill.nb[1] = fill.nb[2] = fill.nb[3] = sizeof(float); + + auto device = ggml_backend_buft_get_device(buft); + t.assert_true("ordinary fill is supported", ggml_backend_dev_supports_op(device, &fill)); + fill.src[0] = &source; + t.assert_true("fill from streamed storage is rejected", !ggml_backend_dev_supports_op(device, &fill)); + + ggml_backend_buffer_free(buffer); + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + t.test("stage slots round-trip bounded byte ranges", [](testing & t) { + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = 64*1024; + params.stage_slots = 2; + + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("runtime allocation succeeds", runtime != nullptr)) { + return; + } + + std::vector source(8192); + for (size_t i = 0; i < source.size(); ++i) { + source[i] = uint8_t((i*37 + 11) & 0xff); + } + std::vector destination(source.size(), 0); + + t.assert_true("bounded upload succeeds", ggml_backend_cuda_kv_stream_stage_upload( + runtime, 1, 123, source.data(), source.size())); + t.assert_true("bounded download succeeds", ggml_backend_cuda_kv_stream_stage_download( + runtime, 1, 123, destination.data(), destination.size())); + t.assert_true("round-trip bytes are exact", source == destination); + + t.assert_true("invalid slot is rejected", !ggml_backend_cuda_kv_stream_stage_upload( + runtime, 2, 0, source.data(), source.size())); + t.assert_true("cross-slot range is rejected", !ggml_backend_cuda_kv_stream_stage_upload( + runtime, 0, params.stage_bytes - 16, source.data(), source.size())); + t.assert_true("null source is rejected", !ggml_backend_cuda_kv_stream_stage_upload( + runtime, 0, 0, nullptr, source.size())); + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + t.test("fixed pool moves pages from balanced residency into the active ring", [](testing & t) { + constexpr size_t page_bytes = 64*1024; + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = page_bytes; + params.stage_slots = 2; + params.pool_bytes = 6*page_bytes; + params.resident_layer_count = 1; + params.page_tokens = 256; + + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("runtime allocation succeeds", runtime != nullptr)) { + return; + } + + t.assert_equal(uint32_t(2), ggml_backend_cuda_kv_stream_stage_slots(runtime)); + t.assert_equal(uint32_t(4), ggml_backend_cuda_kv_stream_resident_pages_per_layer(runtime)); + t.assert_true("ring grows inside the existing pool", + ggml_backend_cuda_kv_stream_repartition(runtime, 3)); + t.assert_equal(uint32_t(3), ggml_backend_cuda_kv_stream_stage_slots(runtime)); + t.assert_equal(uint32_t(3), ggml_backend_cuda_kv_stream_resident_pages_per_layer(runtime)); + t.assert_true("ring cannot exceed the fixed pool", + !ggml_backend_cuda_kv_stream_repartition(runtime, 7)); + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + t.test("device factory assigns unusable resident remainder pages to the ring", [](testing & t) { + constexpr size_t page_bytes = 64*1024; + constexpr size_t pool_pages = 160; + constexpr uint32_t layers = 16; + using new_fn_t = void * (*)(ggml_backend_dev_t, size_t, size_t, size_t, uint32_t); + + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + ggml_backend_dev_t device = ggml_backend_get_device(backend); + auto new_fn = reinterpret_cast(ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(device), + "ggml_backend_cuda_kv_stream_runtime_new_for_device")); + if (!t.assert_true("device factory is exported", new_fn != nullptr)) { + ggml_backend_free(backend); + return; + } + auto runtime = static_cast( + new_fn(device, pool_pages*page_bytes, page_bytes, 0, layers)); + if (!t.assert_true("factory runtime initializes", runtime != nullptr)) { + ggml_backend_free(backend); + return; + } + + t.assert_equal(uint32_t(16), ggml_backend_cuda_kv_stream_stage_slots(runtime)); + t.assert_equal(uint32_t(9), + ggml_backend_cuda_kv_stream_resident_pages_per_layer(runtime)); + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + ggml_backend_free(backend); + }); + + return t.summary(); +} diff --git a/tests/test-kv-stream-cuda-set-rows.cpp b/tests/test-kv-stream-cuda-set-rows.cpp new file mode 100644 index 000000000000..ff909f380746 --- /dev/null +++ b/tests/test-kv-stream-cuda-set-rows.cpp @@ -0,0 +1,311 @@ +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpp.h" +#include "ggml-cuda.h" +#include "ggml.h" +#include "testing.h" + +#include +#include +#include +#include + +namespace { + +constexpr int64_t CACHE_WIDTH = 256*4; +constexpr int64_t CACHE_ROWS = 512; +constexpr int64_t UPDATE_ROWS = 3; + +std::vector run_set_rows( + ggml_backend_t backend, + ggml_backend_buffer_type_t cache_buft, + ggml_type cache_type) { + constexpr size_t N_TENSORS = 16; + const size_t context_bytes = ggml_tensor_overhead()*N_TENSORS + + ggml_graph_overhead_custom(N_TENSORS, false); + const ggml_init_params params{context_bytes, nullptr, true}; + + ggml_context_ptr cache_ctx(ggml_init(params)); + ggml_context_ptr compute_ctx(ggml_init(params)); + GGML_ASSERT(cache_ctx && compute_ctx); + + ggml_tensor * cache = ggml_new_tensor_2d(cache_ctx.get(), cache_type, CACHE_WIDTH, CACHE_ROWS); + ggml_tensor * values = ggml_new_tensor_2d(compute_ctx.get(), GGML_TYPE_F32, CACHE_WIDTH, UPDATE_ROWS); + ggml_tensor * indices = ggml_new_tensor_1d(compute_ctx.get(), GGML_TYPE_I32, UPDATE_ROWS); + ggml_tensor * updated = ggml_set_rows(compute_ctx.get(), cache, values, indices); + + ggml_backend_buffer_ptr cache_buffer( + ggml_backend_alloc_ctx_tensors_from_buft(cache_ctx.get(), cache_buft)); + ggml_backend_buffer_ptr compute_buffer( + ggml_backend_alloc_ctx_tensors(compute_ctx.get(), backend)); + GGML_ASSERT(cache_buffer && compute_buffer); + ggml_backend_buffer_clear(cache_buffer.get(), 0); + + std::vector source(CACHE_WIDTH*UPDATE_ROWS); + for (size_t i = 0; i < source.size(); ++i) { + source[i] = 0.5f*std::sin(float(i)*0.01953125f) - 0.25f*std::cos(float(i)*0.00390625f); + } + const std::vector rows{1, 257, 511}; + ggml_backend_tensor_set(values, source.data(), 0, source.size()*sizeof(float)); + ggml_backend_tensor_set(indices, rows.data(), 0, rows.size()*sizeof(int32_t)); + + ggml_cgraph * graph = ggml_new_graph_custom(compute_ctx.get(), N_TENSORS, false); + ggml_build_forward_expand(graph, updated); + GGML_ASSERT(ggml_backend_supports_op(backend, updated)); + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + + std::vector result(ggml_nbytes(cache)); + ggml_backend_tensor_get(cache, result.data(), 0, result.size()); + return result; +} + +} // namespace + +int main() { + testing t; + + t.test("KV stream quant types are classified", [](testing & t) { + for (int type = 0; type < GGML_TYPE_COUNT; ++type) { + const ggml_type ggml_type_value = (ggml_type) type; + if (!ggml_is_quantized(ggml_type_value) && + ggml_type_value != GGML_TYPE_F32 && + ggml_type_value != GGML_TYPE_F16 && + ggml_type_value != GGML_TYPE_BF16) { + continue; + } + + const auto capabilities = + ggml_backend_cuda_kv_stream_get_type_capabilities(ggml_type_value); + t.assert_true(ggml_type_name(ggml_type_value), capabilities.classified); + } + }); + + t.test("Q8 K and Q4 V retain direct streamed attention", [](testing & t) { + const auto k = ggml_backend_cuda_kv_stream_get_type_capabilities(GGML_TYPE_Q8_0); + const auto v = ggml_backend_cuda_kv_stream_get_type_capabilities(GGML_TYPE_Q4_0); + + t.assert_true("Q8 supports direct attention", k.direct_attention); + t.assert_true("Q4 supports direct attention", v.direct_attention); + t.assert_equal( + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT, + ggml_backend_cuda_kv_stream_get_attention_mode(GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); + }); + + t.test("all native CUDA flash-attention KV pairs use direct streaming", [](testing & t) { + const ggml_type native_types[] = { + GGML_TYPE_F16, + GGML_TYPE_Q4_0, + GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, + GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, + GGML_TYPE_BF16, + }; + + for (const ggml_type type_k : native_types) { + for (const ggml_type type_v : native_types) { + t.assert_equal( + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT, + ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v)); + } + } + }); + + t.test("all exposed KV-cache pairs select an optimized execution class", [](testing & t) { + const ggml_type kv_types[] = { + GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, + }; + size_t direct_pairs = 0; + size_t fallback_pairs = 0; + for (const ggml_type type_k : kv_types) { + const auto k = ggml_backend_cuda_kv_stream_get_type_capabilities(type_k); + for (const ggml_type type_v : kv_types) { + const auto v = ggml_backend_cuda_kv_stream_get_type_capabilities(type_v); + const auto expected = k.direct_attention && v.direct_attention ? + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT : + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16; + const auto actual = + ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v); + if (!t.assert_equal(expected, actual)) { + std::fprintf(stderr, "mode mismatch K=%s V=%s\n", + ggml_type_name(type_k), ggml_type_name(type_v)); + return; + } + direct_pairs += actual == GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT; + fallback_pairs += actual == GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_F16; + } + } + t.assert_equal(size_t(49), direct_pairs); + t.assert_equal(size_t(32), fallback_pairs); + }); + + t.test("every exposed KV-cache type has a GPU online writer", [](testing & t) { + const ggml_type kv_types[] = { + GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, + }; + for (const ggml_type type : kv_types) { + const auto capabilities = + ggml_backend_cuda_kv_stream_get_type_capabilities(type); + t.assert_true(ggml_type_name(type), capabilities.online_write); + t.assert_true(ggml_type_name(type), capabilities.decode_f16); + } + const auto internal = + ggml_backend_cuda_kv_stream_get_type_capabilities(GGML_TYPE_Q2_K); + t.assert_true("internal Q2_K tensor format is classified", internal.classified); + t.assert_true("internal Q2_K is not advertised as online writable", !internal.online_write); + t.assert_equal( + GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_UNSUPPORTED, + ggml_backend_cuda_kv_stream_get_attention_mode( + GGML_TYPE_Q2_K, GGML_TYPE_Q4_0)); + }); + + t.test("CUDA backend reports executable KV stream pairs", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + using supported_fn_t = bool (*)(ggml_type, ggml_type); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get())); + auto supported_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_kv_stream_type_pair_supported")); + + if (!t.assert_true("type-pair capability query is available", supported_fn != nullptr)) { + return; + } + t.assert_true("Q8 K and Q4 V are executable", + supported_fn(GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); + t.assert_true("auxiliary Q8_1 pair is rejected", + !supported_fn(GGML_TYPE_Q8_1, GGML_TYPE_Q8_1)); + }); + + t.test("CUDA backend computes block-safe KV page geometry", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + using page_bytes_fn_t = bool (*)( + ggml_type, ggml_type, uint32_t, uint32_t, uint32_t, uint32_t, size_t *); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get())); + auto page_bytes_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_kv_stream_page_bytes")); + if (!t.assert_true("page geometry query is available", page_bytes_fn != nullptr)) { + return; + } + + size_t page_bytes = 0; + t.assert_true("Q8/Q4 geometry is valid", page_bytes_fn( + GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, 256, 256, 4, 256, &page_bytes)); + t.assert_equal( + 256*(4*ggml_row_size(GGML_TYPE_Q8_0, 256) + 4*ggml_row_size(GGML_TYPE_Q4_0, 256)), + page_bytes); + + size_t reversed_page_bytes = 0; + t.assert_true("reversed geometry is valid", page_bytes_fn( + GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, 256, 256, 4, 256, &reversed_page_bytes)); + t.assert_equal(page_bytes, reversed_page_bytes); + + t.assert_true("K quant rejects a partial superblock", !page_bytes_fn( + GGML_TYPE_Q2_K, GGML_TYPE_Q4_0, 128, 256, 4, 256, &page_bytes)); + t.assert_true("K quant accepts a complete superblock", page_bytes_fn( + GGML_TYPE_Q2_K, GGML_TYPE_Q4_0, 256, 256, 4, 256, &page_bytes)); + t.assert_true("overflow is rejected", !page_bytes_fn( + GGML_TYPE_F32, GGML_TYPE_F32, UINT32_MAX, UINT32_MAX, + UINT32_MAX, UINT32_MAX, &page_bytes)); + }); + + t.test("conversion workspace is bounded to one page for every exposed KV pair", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + using workspace_fn_t = bool (*)( + ggml_type, ggml_type, uint32_t, uint32_t, uint32_t, uint32_t, size_t *); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get())); + auto workspace_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_kv_stream_workspace_bytes")); + if (!t.assert_true("workspace query is available", workspace_fn != nullptr)) { + return; + } + + constexpr uint32_t head_dim = 256; + constexpr uint32_t head_count = 4; + constexpr uint32_t page_tokens = 256; + const size_t f16_k_page = + ggml_row_size(GGML_TYPE_F16, head_dim)*head_count*page_tokens; + const size_t f16_v_page = f16_k_page; + const size_t expected_fallback = ((f16_k_page + 127) & ~size_t(127)) + f16_v_page; + const ggml_type kv_types[] = { + GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, + }; + + for (const ggml_type type_k : kv_types) { + for (const ggml_type type_v : kv_types) { + size_t workspace_bytes = SIZE_MAX; + if (!t.assert_true("exposed pair has a workspace policy", workspace_fn( + type_k, type_v, head_dim, head_dim, head_count, + page_tokens, &workspace_bytes))) { + std::fprintf(stderr, "workspace query failed K=%s V=%s\n", + ggml_type_name(type_k), ggml_type_name(type_v)); + return; + } + const auto mode = + ggml_backend_cuda_kv_stream_get_attention_mode(type_k, type_v); + const size_t expected = mode == GGML_BACKEND_CUDA_KV_STREAM_ATTENTION_DIRECT ? + 0 : expected_fallback; + if (!t.assert_equal(expected, workspace_bytes)) { + std::fprintf(stderr, "workspace mismatch K=%s V=%s\n", + ggml_type_name(type_k), ggml_type_name(type_v)); + return; + } + } + } + }); + + t.test("mapped authoritative cache SET_ROWS matches CUDA quantization", [](testing & t) { + ggml_backend_ptr backend(ggml_backend_cuda_init(0)); + if (!t.assert_true("CUDA backend initializes", backend != nullptr)) { + return; + } + + ggml_backend_cuda_kv_stream_params params{}; + params.device = 0; + params.stage_bytes = 512*1024; + params.stage_slots = 1; + auto runtime = ggml_backend_cuda_kv_stream_runtime_new(params); + if (!t.assert_true("stream runtime initializes", runtime != nullptr)) { + return; + } + + const ggml_type kv_types[] = { + GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, + }; + for (const ggml_type type : kv_types) { + const std::vector expected = run_set_rows( + backend.get(), ggml_backend_get_default_buffer_type(backend.get()), type); + const std::vector actual = run_set_rows( + backend.get(), ggml_backend_cuda_kv_stream_buffer_type(runtime), type); + if (!t.assert_true(ggml_type_name(type), expected == actual)) { + return; + } + } + + ggml_backend_cuda_kv_stream_runtime_free(runtime); + }); + + return t.summary(); +} diff --git a/tests/test-kv-stream-plan.cpp b/tests/test-kv-stream-plan.cpp new file mode 100644 index 000000000000..b71306502808 --- /dev/null +++ b/tests/test-kv-stream-plan.cpp @@ -0,0 +1,894 @@ +#include "llama-kv-stream-plan.h" +#include "testing.h" + +#include +#include +#include +#include + +namespace { + +constexpr uint64_t MIB = 1024ULL*1024ULL; + +constexpr uint32_t N_TARGET_LAYERS = 16; +constexpr uint32_t N_MTP_LAYERS = 1; +constexpr uint32_t PAGE_TOKENS = 256; +constexpr uint32_t STAGE_TOKENS = 16*1024; +constexpr uint64_t BYTES_PER_TOKEN = 1664; + +std::vector make_qwen38_regions(uint32_t n_kv) { + llama_kv_stream_regions_params params; + params.page_tokens = PAGE_TOKENS; + + for (uint32_t il = 0; il < N_TARGET_LAYERS; ++il) { + llama_kv_stream_layer_layout layer; + layer.role = llama_kv_stream_region_role::target; + layer.layer_id = 3 + 4*il; + layer.n_tokens = n_kv; + layer.bytes_per_token = BYTES_PER_TOKEN; + layer.layer_priority = il; + layer.pin_tail = true; + params.layers.push_back(layer); + } + + llama_kv_stream_layer_layout mtp; + mtp.role = llama_kv_stream_region_role::mtp; + mtp.layer_id = 64; + mtp.n_tokens = n_kv; + mtp.bytes_per_token = BYTES_PER_TOKEN; + mtp.pin_all = true; + params.layers.push_back(mtp); + + return llama_kv_stream_regions_make(params).regions; +} + +llama_kv_stream_plan_params make_params(uint32_t n_kv, uint64_t pool_mib = 2526) { + llama_kv_stream_plan_params params; + params.pool_bytes = pool_mib*MIB; + params.stage_slot_bytes = STAGE_TOKENS*BYTES_PER_TOKEN; + params.stage_slots = 2; + params.regions = make_qwen38_regions(n_kv); + return params; +} + +size_t count_role( + const llama_kv_stream_plan_params & params, + const std::vector & region_indices, + llama_kv_stream_region_role role) { + return std::count_if(region_indices.begin(), region_indices.end(), [&](size_t index) { + return params.regions.at(index).role == role; + }); +} + +} // namespace + +int main() { + testing t; + + t.test("small context remains fully resident while transfer slots stay reserved", [](testing & t) { + const auto params = make_params(1024); + const auto plan = llama_kv_stream_plan_make(params); + + t.assert_true("plan is valid", plan.valid); + t.assert_equal(2*STAGE_TOKENS*BYTES_PER_TOKEN, plan.reserved_stage_bytes); + t.assert_equal(uint64_t(0), plan.streamed_bytes); + t.assert_equal(size_t(0), plan.streamed_regions.size()); + t.assert_equal(params.regions.size(), plan.resident_regions.size()); + t.assert_equal( + plan.pool_bytes, + plan.reserved_stage_bytes + plan.resident_bytes + plan.unused_bytes); + }); + + t.test("pinned MTP pages remain resident at 196K", [](testing & t) { + const auto params = make_params(196608); + const auto plan = llama_kv_stream_plan_make(params); + + t.assert_true("plan is valid", plan.valid); + t.assert_equal(size_t(0), count_role(params, plan.streamed_regions, llama_kv_stream_region_role::mtp)); + t.assert_true( + "some target pages stream after the pool fills", + count_role(params, plan.streamed_regions, llama_kv_stream_region_role::target) > 0); + t.assert_true( + "resident bytes never exceed the pool", + plan.reserved_stage_bytes + plan.resident_bytes <= plan.pool_bytes); + }); + + t.test("one padded context step increases streaming smoothly", [](testing & t) { + const auto params_before = make_params(196608); + const auto params_after = make_params(196608 + PAGE_TOKENS); + + const auto before = llama_kv_stream_plan_make(params_before); + const auto after = llama_kv_stream_plan_make(params_after); + + t.assert_true("before plan is valid", before.valid); + t.assert_true("after plan is valid", after.valid); + t.assert_equal( + uint64_t((N_TARGET_LAYERS + N_MTP_LAYERS)*PAGE_TOKENS)*BYTES_PER_TOKEN, + after.streamed_bytes - before.streamed_bytes); + }); + + t.test("planning fails when pinned pages and stages do not fit", [](testing & t) { + auto params = make_params(196608); + params.pool_bytes = 300*MIB; + + const auto plan = llama_kv_stream_plan_make(params); + t.assert_true("plan is rejected", !plan.valid); + t.assert_true("error is reported", !plan.error.empty()); + }); + + t.test("a streamed region must fit one transfer slot", [](testing & t) { + llama_kv_stream_plan_params params; + params.pool_bytes = 1*MIB; + params.stage_slot_bytes = 1*MIB; + params.stage_slots = 1; + + llama_kv_stream_region region; + region.layer_id = 3; + region.token_count = 1; + region.bytes = 2*MIB; + params.regions.push_back(region); + + const auto plan = llama_kv_stream_plan_make(params); + t.assert_true("plan is rejected", !plan.valid); + t.assert_true("error is reported", !plan.error.empty()); + }); + + t.test("every region is assigned exactly once and accounting closes", [](testing & t) { + for (uint32_t n_kv : { 1024U, 92160U, 196608U, 262144U }) { + const auto params = make_params(n_kv); + const auto plan = llama_kv_stream_plan_make(params); + + t.assert_true("plan is valid", plan.valid); + + std::vector seen(params.regions.size(), 0); + uint64_t logical_bytes = 0; + + for (const auto & region : params.regions) { + logical_bytes += region.bytes; + } + + bool indices_in_range = true; + for (size_t index : plan.resident_regions) { + if (index < seen.size()) { + ++seen[index]; + } else { + indices_in_range = false; + } + } + + for (size_t index : plan.streamed_regions) { + if (index < seen.size()) { + ++seen[index]; + } else { + indices_in_range = false; + } + } + + t.assert_true("all assignment indices are in range", indices_in_range); + t.assert_true("every region is assigned once", std::all_of(seen.begin(), seen.end(), [](uint8_t count) { + return count == 1; + })); + t.assert_equal(logical_bytes, plan.resident_bytes + plan.streamed_bytes); + t.assert_equal( + plan.pool_bytes, + plan.reserved_stage_bytes + plan.resident_bytes + plan.unused_bytes); + } + }); + + t.test("streamed bytes grow monotonically without a context cliff", [](testing & t) { + std::vector contexts = { 1024, 65536, 89088 }; + for (uint32_t n_kv = 89344; n_kv <= 94208; n_kv += PAGE_TOKENS) { + contexts.push_back(n_kv); + } + contexts.insert(contexts.end(), { 131072, 196608, 196864, 262144 }); + + uint64_t previous_streamed = 0; + uint32_t previous_context = 0; + + for (uint32_t n_kv : contexts) { + const auto plan = llama_kv_stream_plan_make(make_params(n_kv)); + + if (!t.assert_true("plan is valid", plan.valid)) { + break; + } + + t.assert_true("streamed bytes are monotonic", plan.streamed_bytes >= previous_streamed); + if (previous_context != 0) { + const uint64_t maximum_step = + uint64_t(N_TARGET_LAYERS + N_MTP_LAYERS)*(n_kv - previous_context)*BYTES_PER_TOKEN; + t.assert_true( + "context growth cannot stream more bytes than the added logical KV", + plan.streamed_bytes - previous_streamed <= maximum_step); + } + + previous_streamed = plan.streamed_bytes; + previous_context = n_kv; + } + }); + + t.test("prefetch window prioritizes imminent pages and stops at lookahead", [](testing & t) { + llama_kv_stream_prefetch_params params; + params.current_attention = 4; + params.lookahead_layers = 3; + params.stage_slot_bytes = 4096; + params.free_slots = { 2, 0, 1, 3, 4, 5, 6, 7 }; + + for (uint32_t attention = 4; attention <= 8; ++attention) { + for (uint32_t page = 0; page < 2; ++page) { + llama_kv_stream_prefetch_request request; + request.layer_id = int32_t(3 + 4*attention); + request.attention_index = attention; + request.page_index = page; + request.bytes = 4096; + params.requests.push_back(request); + } + } + params.states.resize(params.requests.size(), llama_kv_stream_prefetch_state::pending); + + const auto dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_true("dispatch is valid", dispatch.valid); + t.assert_equal(size_t(8), dispatch.assignments.size()); + t.assert_true("free slots are normalized", dispatch.assignments[0].slot == 0); + t.assert_true("current attention is first", dispatch.assignments[0].request_index == 0); + t.assert_true("x+3 is inside the window", dispatch.assignments.back().request_index == 7); + t.assert_true("x+4 is outside the window", std::none_of( + dispatch.assignments.begin(), dispatch.assignments.end(), [](const auto & assignment) { + return assignment.request_index >= 8; + })); + }); + + t.test("prefetch slot reuse waits for release and selects the earliest deadline", [](testing & t) { + llama_kv_stream_prefetch_params params; + params.current_attention = 0; + params.lookahead_layers = 3; + params.stage_slot_bytes = 4096; + params.free_slots = { 0, 1 }; + + for (uint32_t attention = 0; attention < 4; ++attention) { + llama_kv_stream_prefetch_request request; + request.layer_id = int32_t(3 + 4*attention); + request.attention_index = attention; + request.page_index = 0; + request.bytes = 4096; + params.requests.push_back(request); + } + params.states.resize(params.requests.size(), llama_kv_stream_prefetch_state::pending); + + auto dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_equal(size_t(2), dispatch.assignments.size()); + t.assert_equal(size_t(0), dispatch.assignments[0].request_index); + t.assert_equal(size_t(1), dispatch.assignments[1].request_index); + + params.states[0] = llama_kv_stream_prefetch_state::consumed; + params.states[1] = llama_kv_stream_prefetch_state::scheduled; + params.free_slots = { dispatch.assignments[0].slot }; + dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_equal(size_t(1), dispatch.assignments.size()); + t.assert_equal(size_t(2), dispatch.assignments[0].request_index); + t.assert_equal(uint32_t(0), dispatch.assignments[0].slot); + }); + + t.test("future mutable tails are not prefetched before their SET_ROWS producer", [](testing & t) { + llama_kv_stream_prefetch_params params; + params.current_attention = 2; + params.lookahead_layers = 3; + params.stage_slot_bytes = 4096; + params.free_slots = { 0, 1, 2 }; + + llama_kv_stream_prefetch_request current_tail; + current_tail.layer_id = 11; + current_tail.attention_index = 2; + current_tail.page_index = 9; + current_tail.bytes = 4096; + current_tail.producer_attention_index = 2; + + auto future_tail = current_tail; + future_tail.layer_id = 15; + future_tail.attention_index = 3; + future_tail.producer_attention_index = 3; + + auto future_stable = future_tail; + future_stable.page_index = 8; + future_stable.producer_attention_index = -1; + + params.requests = { current_tail, future_tail, future_stable }; + params.states.resize(params.requests.size(), llama_kv_stream_prefetch_state::pending); + + auto dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_equal(size_t(1), dispatch.assignments.size()); + t.assert_equal(size_t(2), dispatch.assignments[0].request_index); + + params.current_producer_complete = true; + dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_equal(size_t(2), dispatch.assignments.size()); + t.assert_equal(size_t(0), dispatch.assignments[0].request_index); + t.assert_equal(size_t(2), dispatch.assignments[1].request_index); + }); + + t.test("adaptive prefetch considers every future attention deadline", [](testing & t) { + llama_kv_stream_prefetch_params params; + params.current_attention = 0; + params.adaptive_lookahead = true; + params.stage_slot_bytes = 4096; + params.free_slots = { 0, 1, 2, 3 }; + + for (uint32_t attention = 0; attention < N_TARGET_LAYERS; ++attention) { + llama_kv_stream_prefetch_request request; + request.layer_id = int32_t(3 + 4*attention); + request.attention_index = attention; + request.page_index = 7; + request.bytes = 4096; + params.requests.push_back(request); + } + params.states.resize(params.requests.size(), llama_kv_stream_prefetch_state::pending); + + auto dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_true("dispatch is valid", dispatch.valid); + t.assert_equal(size_t(4), dispatch.assignments.size()); + t.assert_equal(size_t(3), dispatch.assignments.back().request_index); + + params.states[0] = llama_kv_stream_prefetch_state::consumed; + for (size_t i = 1; i < 4; ++i) { + params.states[i] = llama_kv_stream_prefetch_state::scheduled; + } + params.free_slots = { 0 }; + dispatch = llama_kv_stream_prefetch_dispatch(params); + t.assert_equal(size_t(1), dispatch.assignments.size()); + t.assert_equal(size_t(4), dispatch.assignments[0].request_index); + }); + + t.test("repeated deadline misses grow the ring by one balanced layer epoch", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 160; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 12; + params.minimum_ring_slots = 16; + params.previous_resident_pages_per_layer = 9; + params.previous_ring_slots = 16; + params.deadline_miss_ratio = 0.12; + params.copy_engine_busy_ratio = 0.70; + params.starved_evaluations = 2; + params.grow_hysteresis_evaluations = 3; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition changed", partition.changed); + t.assert_equal(uint32_t(8), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(32), partition.ring_slots); + t.assert_equal(uint32_t(0), partition.starved_evaluations); + }); + + t.test("feedback cannot demote beyond one epoch past the overlap target", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 7010; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 716; + params.minimum_ring_slots = 18; + params.previous_resident_pages_per_layer = 416; + params.previous_ring_slots = 354; + params.deadline_miss_ratio = 0.10; + params.copy_engine_busy_ratio = 0.80; + params.starved_evaluations = 2; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition remains stable", !partition.changed); + t.assert_equal(uint32_t(416), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(354), partition.ring_slots); + }); + + t.test("runaway feedback partition heals to the bounded floor", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 7010; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 716; + params.minimum_ring_slots = 18; + params.previous_resident_pages_per_layer = 368; + params.previous_ring_slots = 1122; + params.deadline_miss_ratio = 0.10; + params.copy_engine_busy_ratio = 0.80; + params.starved_evaluations = 2; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition repairs the resident boundary", partition.changed); + t.assert_equal(uint32_t(416), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(354), partition.ring_slots); + t.assert_equal(uint32_t(0), partition.starved_evaluations); + }); + + t.test("PCIe saturation does not sacrifice more resident pages", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 160; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 12; + params.minimum_ring_slots = 16; + params.previous_resident_pages_per_layer = 9; + params.previous_ring_slots = 16; + params.deadline_miss_ratio = 0.20; + params.copy_engine_busy_ratio = 0.99; + params.starved_evaluations = 10; + params.grow_hysteresis_evaluations = 3; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition remains stable", !partition.changed); + t.assert_equal(uint32_t(9), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(16), partition.ring_slots); + }); + + t.test("undersized ring jumps to one-layer overlap target despite copy pressure", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 630; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 257; + params.minimum_ring_slots = 22; + params.previous_resident_pages_per_layer = 38; + params.previous_ring_slots = 22; + params.deadline_miss_ratio = 0.40; + params.copy_engine_busy_ratio = 0.99; + params.starved_evaluations = 2; + params.grow_hysteresis_evaluations = 3; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition changed", partition.changed); + t.assert_equal(uint32_t(23), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(262), partition.ring_slots); + }); + + t.test("decode transition selects its overlap target before feedback exists", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 6971; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 673; + params.minimum_ring_slots = 11; + params.previous_resident_pages_per_layer = 435; + params.previous_ring_slots = 11; + params.evaluations_since_repartition = 0; + params.entering_decode_layout = true; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("decode transition changes immediately", partition.changed); + t.assert_equal(uint32_t(418), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(283), partition.ring_slots); + }); + + t.test("copy pressure stops demotion after overlap target is reached", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 630; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 257; + params.minimum_ring_slots = 22; + params.previous_resident_pages_per_layer = 23; + params.previous_ring_slots = 262; + params.deadline_miss_ratio = 0.20; + params.copy_engine_busy_ratio = 0.90; + params.starved_evaluations = 10; + params.grow_hysteresis_evaluations = 3; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition remains stable", !partition.changed); + t.assert_equal(uint32_t(23), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(262), partition.ring_slots); + }); + t.test("near-saturated copy traffic does not trigger a disruptive feedback epoch", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 6971; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 673; + params.minimum_ring_slots = 11; + params.previous_resident_pages_per_layer = 418; + params.previous_ring_slots = 283; + params.deadline_miss_ratio = 0.053; + params.copy_engine_busy_ratio = 0.848; + params.ring_peak_occupancy_ratio = 1.0; + params.starved_evaluations = 2; + params.evaluations_since_repartition = 128; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("near-saturated partition remains stable", !partition.changed); + t.assert_equal(uint32_t(418), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(283), partition.ring_slots); + t.assert_equal(uint32_t(0), partition.starved_evaluations); + }); + + + t.test("light utilization does not promote above overlap target", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 630; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 257; + params.minimum_ring_slots = 22; + params.previous_resident_pages_per_layer = 23; + params.previous_ring_slots = 262; + params.deadline_miss_ratio = 0.0; + params.copy_engine_busy_ratio = 0.25; + params.ring_peak_occupancy_ratio = 0.20; + params.overprovisioned_evaluations = 7; + params.shrink_hysteresis_evaluations = 8; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition remains at target", !partition.changed); + t.assert_equal(uint32_t(23), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(262), partition.ring_slots); + }); + + t.test("constrained pool demotes all resident pages when target is unreachable", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 39; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 65; + params.minimum_ring_slots = 23; + params.previous_resident_pages_per_layer = 1; + params.previous_ring_slots = 23; + params.deadline_miss_ratio = 0.50; + params.copy_engine_busy_ratio = 0.99; + params.starved_evaluations = 2; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition changed", partition.changed); + t.assert_equal(uint32_t(0), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(39), partition.ring_slots); + }); + + t.test("non-positive overlap target is rejected", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 160; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 12; + params.minimum_ring_slots = 16; + params.previous_resident_pages_per_layer = 9; + params.previous_ring_slots = 16; + params.target_ring_working_set_ratio = 0.0; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is rejected", !partition.valid); + }); + + t.test("partition cooldown accumulates pressure without repeatedly resetting residency", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 160; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 12; + params.minimum_ring_slots = 16; + params.previous_resident_pages_per_layer = 9; + params.previous_ring_slots = 16; + params.deadline_miss_ratio = 0.12; + params.copy_engine_busy_ratio = 0.70; + params.starved_evaluations = 2; + params.evaluations_since_repartition = 2; + + auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("cooldown preserves the resident boundary", !partition.changed); + t.assert_equal(uint32_t(3), partition.starved_evaluations); + + params.evaluations_since_repartition = params.repartition_cooldown_evaluations; + partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition changes after cooldown", partition.changed); + t.assert_equal(uint32_t(8), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(32), partition.ring_slots); + }); + + t.test("sustained overprovision promotes one balanced resident epoch", [](testing & t) { + llama_kv_stream_partition_params params; + params.total_pool_pages = 160; + params.layer_count = N_TARGET_LAYERS; + params.active_pages_per_layer = 12; + params.minimum_ring_slots = 16; + params.previous_resident_pages_per_layer = 8; + params.previous_ring_slots = 32; + params.deadline_miss_ratio = 0.0; + params.copy_engine_busy_ratio = 0.25; + params.ring_peak_occupancy_ratio = 0.20; + params.overprovisioned_evaluations = 7; + params.shrink_hysteresis_evaluations = 8; + + const auto partition = llama_kv_stream_partition_adapt(params); + t.assert_true("partition is valid", partition.valid); + t.assert_true("partition changed", partition.changed); + t.assert_equal(uint32_t(9), partition.resident_pages_per_layer); + t.assert_equal(uint32_t(16), partition.ring_slots); + t.assert_equal(uint32_t(0), partition.overprovisioned_evaluations); + }); + + t.test("cumulative CUDA feedback becomes one bounded evaluation delta", [](testing & t) { + const auto delta = llama_kv_stream_feedback_delta_make( + { 145, 17 }, { 120, 12 }); + t.assert_true("feedback delta is valid", delta.valid); + t.assert_true("feedback contains a new evaluation", delta.has_evaluation); + t.assert_equal(uint64_t(25), delta.deadline_samples); + t.assert_equal(uint64_t(5), delta.deadline_misses); + t.assert_true("deadline ratio is exact", + std::abs(delta.deadline_miss_ratio - 0.20) < 1e-12); + + const auto unchanged = llama_kv_stream_feedback_delta_make( + { 145, 17 }, { 145, 17 }); + t.assert_true("unchanged counters are valid", unchanged.valid); + t.assert_true("unchanged counters do not invent an evaluation", + !unchanged.has_evaluation); + + const auto reset = llama_kv_stream_feedback_delta_make( + { 3, 1 }, { 145, 17 }); + t.assert_true("counter reset is rejected", !reset.valid); + + const auto impossible = llama_kv_stream_feedback_delta_make( + { 150, 30 }, { 145, 17 }); + t.assert_true("more misses than samples are rejected", !impossible.valid); + }); + + t.test("duplicate logical regions are rejected", [](testing & t) { + llama_kv_stream_plan_params params; + params.pool_bytes = 64*MIB; + params.stage_slot_bytes = 1*MIB; + params.stage_slots = 1; + + llama_kv_stream_region region; + region.layer_id = 3; + region.token_begin = 0; + region.token_count = PAGE_TOKENS; + region.bytes = PAGE_TOKENS*BYTES_PER_TOKEN; + + params.regions.push_back(region); + params.regions.push_back(region); + + const auto plan = llama_kv_stream_plan_make(params); + t.assert_true("duplicate plan is rejected", !plan.valid); + }); + + t.test("overlapping regions in one logical layer are rejected", [](testing & t) { + llama_kv_stream_plan_params params; + params.pool_bytes = 64*MIB; + params.stage_slot_bytes = 1*MIB; + params.stage_slots = 1; + + llama_kv_stream_region first; + first.layer_id = 3; + first.token_begin = 0; + first.token_count = PAGE_TOKENS; + first.bytes = PAGE_TOKENS*BYTES_PER_TOKEN; + + auto second = first; + second.token_begin = PAGE_TOKENS/2; + + params.regions.push_back(first); + params.regions.push_back(second); + + const auto plan = llama_kv_stream_plan_make(params); + t.assert_true("overlapping plan is rejected", !plan.valid); + }); + + t.test("overflowing token ranges are rejected", [](testing & t) { + llama_kv_stream_plan_params params; + params.pool_bytes = 64*MIB; + params.stage_slot_bytes = 1*MIB; + params.stage_slots = 1; + + llama_kv_stream_region region; + region.layer_id = 3; + region.token_begin = std::numeric_limits::max() - 127; + region.token_count = PAGE_TOKENS; + region.bytes = PAGE_TOKENS*BYTES_PER_TOKEN; + params.regions.push_back(region); + + const auto plan = llama_kv_stream_plan_make(params); + t.assert_true("overflowing plan is rejected", !plan.valid); + }); + + t.test("extent includes pending speculative positions and page padding", [](testing & t) { + llama_kv_stream_extent_params params; + params.live_tokens = 196607; + params.reserve_tokens = 3; + params.page_tokens = PAGE_TOKENS; + params.maximum_tokens = 262144; + + const auto extent = llama_kv_stream_extent_make(params); + t.assert_true("extent is valid", extent.valid); + t.assert_equal(uint32_t(196864), extent.tokens); + t.assert_true("initial extent is growth", extent.grew); + t.assert_true("initial extent is not shrink", !extent.shrunk); + }); + + t.test("growth is immediate", [](testing & t) { + llama_kv_stream_extent_params params; + params.live_tokens = 131072; + params.reserve_tokens = 3; + params.page_tokens = PAGE_TOKENS; + params.previous_extent = 131072; + params.shrink_hysteresis_tokens = 4096; + params.maximum_tokens = 262144; + + const auto extent = llama_kv_stream_extent_make(params); + t.assert_true("extent is valid", extent.valid); + t.assert_equal(uint32_t(131328), extent.tokens); + t.assert_true("extent grew", extent.grew); + }); + + t.test("small rollback retains the previous extent", [](testing & t) { + llama_kv_stream_extent_params params; + params.live_tokens = 196605; + params.reserve_tokens = 3; + params.page_tokens = PAGE_TOKENS; + params.previous_extent = 196864; + params.shrink_hysteresis_tokens = 4096; + params.maximum_tokens = 262144; + + const auto extent = llama_kv_stream_extent_make(params); + t.assert_true("extent is valid", extent.valid); + t.assert_equal(uint32_t(196864), extent.tokens); + t.assert_true("extent did not shrink", !extent.shrunk); + }); + + t.test("large shrink and forced reset release residency", [](testing & t) { + llama_kv_stream_extent_params params; + params.live_tokens = 180000; + params.page_tokens = PAGE_TOKENS; + params.previous_extent = 196864; + params.shrink_hysteresis_tokens = 4096; + params.maximum_tokens = 262144; + + auto extent = llama_kv_stream_extent_make(params); + t.assert_true("large-shrink extent is valid", extent.valid); + t.assert_equal(uint32_t(180224), extent.tokens); + t.assert_true("large shrink is reported", extent.shrunk); + + params.live_tokens = 0; + params.force_shrink = true; + extent = llama_kv_stream_extent_make(params); + t.assert_true("reset extent is valid", extent.valid); + t.assert_equal(uint32_t(0), extent.tokens); + t.assert_true("reset shrink is reported", extent.shrunk); + }); + + t.test("extent rejects invalid alignment and capacity overflow", [](testing & t) { + llama_kv_stream_extent_params params; + params.live_tokens = 262143; + params.reserve_tokens = 3; + params.page_tokens = PAGE_TOKENS; + params.maximum_tokens = 262144; + + auto extent = llama_kv_stream_extent_make(params); + t.assert_true("capacity overflow is rejected", !extent.valid); + + params.live_tokens = 1024; + params.reserve_tokens = 0; + params.page_tokens = 0; + extent = llama_kv_stream_extent_make(params); + t.assert_true("zero page size is rejected", !extent.valid); + + params.page_tokens = PAGE_TOKENS; + params.maximum_tokens = 262143; + extent = llama_kv_stream_extent_make(params); + t.assert_true("unaligned maximum is rejected", !extent.valid); + }); + + t.test("region builder reproduces exact Qwen3.8 maximum-context geometry", [](testing & t) { + llama_kv_stream_regions_params params; + params.page_tokens = PAGE_TOKENS; + + for (uint32_t il = 0; il < N_TARGET_LAYERS; ++il) { + llama_kv_stream_layer_layout layer; + layer.role = llama_kv_stream_region_role::target; + layer.layer_id = 3 + 4*il; + layer.n_tokens = 262144; + layer.bytes_per_token = BYTES_PER_TOKEN; + layer.layer_priority = il; + layer.pin_tail = true; + params.layers.push_back(layer); + } + + llama_kv_stream_layer_layout mtp; + mtp.role = llama_kv_stream_region_role::mtp; + mtp.layer_id = 64; + mtp.n_tokens = 262144; + mtp.bytes_per_token = BYTES_PER_TOKEN; + mtp.pin_all = true; + params.layers.push_back(mtp); + + const auto built = llama_kv_stream_regions_make(params); + t.assert_true("regions are valid", built.valid); + t.assert_equal(size_t(17*1024), built.regions.size()); + t.assert_equal(uint64_t(7072)*MIB, built.total_bytes); + + const size_t pinned = std::count_if(built.regions.begin(), built.regions.end(), [](const auto & region) { + return region.pinned; + }); + t.assert_equal(size_t(1024 + N_TARGET_LAYERS), pinned); + t.assert_true("every full page has exact bytes", std::all_of( + built.regions.begin(), built.regions.end(), [](const auto & region) { + return region.token_count == PAGE_TOKENS && + region.bytes == PAGE_TOKENS*BYTES_PER_TOKEN; + })); + }); + + t.test("region builder represents a partial tail without reserving a full page", [](testing & t) { + llama_kv_stream_regions_params params; + params.page_tokens = PAGE_TOKENS; + + llama_kv_stream_layer_layout layer; + layer.layer_id = 3; + layer.n_tokens = 196609; + layer.bytes_per_token = BYTES_PER_TOKEN; + layer.pin_tail = true; + params.layers.push_back(layer); + + const auto built = llama_kv_stream_regions_make(params); + if (!t.assert_true("regions are valid", built.valid)) { + return; + } + t.assert_equal(size_t(769), built.regions.size()); + + if (built.regions.empty()) { + return; + } + + const auto & tail = built.regions.back(); + t.assert_equal(uint32_t(196608), tail.token_begin); + t.assert_equal(uint32_t(1), tail.token_count); + t.assert_equal(uint64_t(BYTES_PER_TOKEN), tail.bytes); + t.assert_true("tail is pinned", tail.pinned); + }); + + t.test("region priority spreads the same context page across layers", [](testing & t) { + llama_kv_stream_regions_params params; + params.page_tokens = PAGE_TOKENS; + + for (uint32_t il = 0; il < 4; ++il) { + llama_kv_stream_layer_layout layer; + layer.layer_id = 3 + 4*il; + layer.n_tokens = 2*PAGE_TOKENS; + layer.bytes_per_token = BYTES_PER_TOKEN; + layer.layer_priority = il; + params.layers.push_back(layer); + } + + const auto built = llama_kv_stream_regions_make(params); + t.assert_true("regions are valid", built.valid); + + uint64_t maximum_page_zero_priority = 0; + uint64_t minimum_page_one_priority = std::numeric_limits::max(); + for (const auto & region : built.regions) { + if (region.token_begin == 0) { + maximum_page_zero_priority = std::max(maximum_page_zero_priority, region.residency_priority); + } else { + minimum_page_one_priority = std::min(minimum_page_one_priority, region.residency_priority); + } + } + + t.assert_true( + "all layers of one context page are preferred before the next page", + maximum_page_zero_priority < minimum_page_one_priority); + }); + + t.test("region builder rejects invalid and overflowing layouts", [](testing & t) { + llama_kv_stream_regions_params params; + params.page_tokens = 0; + auto built = llama_kv_stream_regions_make(params); + t.assert_true("zero page size is rejected", !built.valid); + + params.page_tokens = PAGE_TOKENS; + llama_kv_stream_layer_layout layer; + layer.layer_id = 3; + layer.n_tokens = 2; + layer.bytes_per_token = std::numeric_limits::max(); + params.layers.push_back(layer); + built = llama_kv_stream_regions_make(params); + t.assert_true("byte overflow is rejected", !built.valid); + + params.layers.clear(); + layer.n_tokens = PAGE_TOKENS; + layer.bytes_per_token = BYTES_PER_TOKEN; + params.layers.push_back(layer); + params.layers.push_back(layer); + built = llama_kv_stream_regions_make(params); + t.assert_true("duplicate logical layers are rejected", !built.valid); + }); + + return t.summary(); +} diff --git a/tests/test-kv-stream-softmax.cpp b/tests/test-kv-stream-softmax.cpp new file mode 100644 index 000000000000..807e49e12533 --- /dev/null +++ b/tests/test-kv-stream-softmax.cpp @@ -0,0 +1,110 @@ +#include "llama-kv-stream-softmax.h" +#include "testing.h" + +#include +#include +#include +#include + +namespace { + +llama_kv_stream_softmax_part make_part( + const std::vector & logits, + const std::vector> & values, + size_t begin, + size_t end) { + llama_kv_stream_softmax_part result; + result.max_logit = *std::max_element(logits.begin() + begin, logits.begin() + end); + result.numerator.assign(values.front().size(), 0.0f); + + for (size_t i = begin; i < end; ++i) { + const float weight = std::exp(logits[i] - result.max_logit); + result.normalizer += weight; + for (size_t j = 0; j < result.numerator.size(); ++j) { + result.numerator[j] += weight*values[i][j]; + } + } + + return result; +} + +void assert_vector_near(testing & t, const std::vector & expected, const std::vector & actual, float tolerance) { + if (!t.assert_equal("vector dimensions match", expected.size(), actual.size())) { + return; + } + + for (size_t i = 0; i < expected.size(); ++i) { + t.assert_true("component is within tolerance", std::abs(expected[i] - actual[i]) <= tolerance); + } +} + +} // namespace + +int main() { + testing t; + + const std::vector logits = { -1000.0f, -5.0f, 1.0f, 7.0f, 999.0f, 1000.0f, 3.0f }; + const std::vector> values = { + { 1.0f, 2.0f, 3.0f }, + { -2.0f, 0.5f, 4.0f }, + { 8.0f, -1.0f, 0.0f }, + { 0.0f, 3.0f, -7.0f }, + { 11.0f, -9.0f, 2.0f }, + { 17.0f, 5.0f, -3.0f }, + { -4.0f, 12.0f, 6.0f }, + }; + + t.test("partitioned merge matches monolithic stable softmax", [&](testing & t) { + const auto monolithic = llama_kv_stream_softmax_merge({ make_part(logits, values, 0, logits.size()) }); + const auto partitioned = llama_kv_stream_softmax_merge({ + make_part(logits, values, 0, 2), + make_part(logits, values, 2, 5), + make_part(logits, values, 5, 7), + }); + + if (!t.assert_true("monolithic merge is valid", monolithic.valid) || + !t.assert_true("partitioned merge is valid", partitioned.valid)) { + return; + } + + assert_vector_near(t, monolithic.value, partitioned.value, 1e-5f); + t.assert_true("global max is preserved", std::abs(partitioned.max_logit - 1000.0f) <= 1e-6f); + }); + + t.test("merge is invariant to block order", [&](testing & t) { + const auto first = make_part(logits, values, 0, 3); + const auto second = make_part(logits, values, 3, 5); + const auto third = make_part(logits, values, 5, 7); + + const auto forward = llama_kv_stream_softmax_merge({ first, second, third }); + const auto reverse = llama_kv_stream_softmax_merge({ third, second, first }); + + if (!t.assert_true("forward merge is valid", forward.valid) || + !t.assert_true("reverse merge is valid", reverse.valid)) { + return; + } + + assert_vector_near(t, forward.value, reverse.value, 1e-5f); + assert_vector_near(t, forward.numerator, reverse.numerator, 1e-4f); + t.assert_true("normalizer is invariant", std::abs(forward.normalizer - reverse.normalizer) <= 1e-6f); + }); + + t.test("invalid partial results are rejected", [](testing & t) { + llama_kv_stream_softmax_part first; + first.max_logit = 1.0f; + first.normalizer = 1.0f; + first.numerator = { 1.0f, 2.0f }; + + auto second = first; + second.numerator.push_back(3.0f); + t.assert_true("dimension mismatch is rejected", !llama_kv_stream_softmax_merge({ first, second }).valid); + + second = first; + second.normalizer = 0.0f; + t.assert_true("zero normalizer is rejected", !llama_kv_stream_softmax_merge({ first, second }).valid); + + t.assert_true("empty merge is rejected", !llama_kv_stream_softmax_merge({}).valid); + }); + + return t.summary(); +} diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index c8ad1db43623..4cf738398585 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -17,6 +17,9 @@ else() add_subdirectory(batched-bench) add_subdirectory(gguf-split) add_subdirectory(imatrix) + if (GGML_CUDA) + add_subdirectory(kv-stream-bench) + endif() add_subdirectory(llama-bench) add_subdirectory(completion) add_subdirectory(perplexity) diff --git a/tools/kv-stream-bench/CMakeLists.txt b/tools/kv-stream-bench/CMakeLists.txt new file mode 100644 index 000000000000..0cdeca14ead1 --- /dev/null +++ b/tools/kv-stream-bench/CMakeLists.txt @@ -0,0 +1,15 @@ +set(TARGET llama-kv-stream-bench) + +find_package(CUDAToolkit REQUIRED) + +add_executable(${TARGET} + main.cpp + kv-stream-bench-config.cpp + kv-stream-bench-config.h) + +target_compile_features(${TARGET} PRIVATE cxx_std_17) +target_link_libraries(${TARGET} PRIVATE CUDA::cudart) + +if (LLAMA_TOOLS_INSTALL) + install(TARGETS ${TARGET} RUNTIME) +endif() diff --git a/tools/kv-stream-bench/kv-stream-bench-config.cpp b/tools/kv-stream-bench/kv-stream-bench-config.cpp new file mode 100644 index 000000000000..35d9e823dd22 --- /dev/null +++ b/tools/kv-stream-bench/kv-stream-bench-config.cpp @@ -0,0 +1,137 @@ +#include "kv-stream-bench-config.h" + +#include +#include +#include + +namespace { + +constexpr uint64_t MIB = 1024ULL*1024ULL; + +bool parse_u64(std::string_view text, uint64_t & value) { + if (text.empty()) { + return false; + } + + const char * begin = text.data(); + const char * end = begin + text.size(); + const auto parsed = std::from_chars(begin, end, value); + return parsed.ec == std::errc() && parsed.ptr == end; +} + +bool mib_to_bytes(uint64_t mib, uint64_t & bytes) { + if (mib > std::numeric_limits::max()/MIB) { + return false; + } + + bytes = mib*MIB; + return true; +} + +} // namespace + +llama_kv_stream_bench_config llama_kv_stream_bench_config_parse(int argc, const char * const * argv) { + llama_kv_stream_bench_config result; + result.transfer_bytes = { 26*MIB, 197*MIB, 312*MIB, 416*MIB }; + + auto fail = [&](const std::string & message) { + result.valid = false; + result.error = message; + return result; + }; + + auto next_value = [&](int & index, std::string_view option, std::string_view & value) { + if (index + 1 >= argc || argv[index + 1] == nullptr) { + result.error = std::string(option) + " requires a value"; + return false; + } + + value = argv[++index]; + return true; + }; + + for (int i = 1; i < argc; ++i) { + if (argv[i] == nullptr) { + return fail("benchmark option must not be null"); + } + + const std::string_view option(argv[i]); + if (option == "--execute") { + result.execute = true; + continue; + } + + std::string_view value; + if (option == "--device") { + if (!next_value(i, option, value)) { + return fail(result.error); + } + + uint64_t parsed = 0; + if (!parse_u64(value, parsed) || parsed > uint64_t(std::numeric_limits::max())) { + return fail("--device must be a non-negative integer"); + } + result.device = int(parsed); + continue; + } + + if (option == "--iterations") { + if (!next_value(i, option, value)) { + return fail(result.error); + } + + uint64_t parsed = 0; + if (!parse_u64(value, parsed) || parsed == 0 || parsed > uint64_t(std::numeric_limits::max())) { + return fail("--iterations must be a positive integer"); + } + result.iterations = int(parsed); + continue; + } + + if (option == "--reserve-mib") { + if (!next_value(i, option, value)) { + return fail(result.error); + } + + uint64_t parsed = 0; + if (!parse_u64(value, parsed) || !mib_to_bytes(parsed, result.reserve_bytes)) { + return fail("--reserve-mib must be a valid non-negative integer"); + } + continue; + } + + if (option == "--sizes-mib") { + if (!next_value(i, option, value)) { + return fail(result.error); + } + + std::vector transfer_bytes; + size_t begin = 0; + while (begin <= value.size()) { + const size_t comma = value.find(',', begin); + const size_t end = comma == std::string_view::npos ? value.size() : comma; + + uint64_t parsed = 0; + uint64_t bytes = 0; + if (!parse_u64(value.substr(begin, end - begin), parsed) || + parsed == 0 || !mib_to_bytes(parsed, bytes)) { + return fail("--sizes-mib must be a comma-separated list of positive integers"); + } + transfer_bytes.push_back(bytes); + + if (comma == std::string_view::npos) { + break; + } + begin = comma + 1; + } + + result.transfer_bytes = std::move(transfer_bytes); + continue; + } + + return fail("unknown benchmark option: " + std::string(option)); + } + + result.valid = true; + return result; +} diff --git a/tools/kv-stream-bench/kv-stream-bench-config.h b/tools/kv-stream-bench/kv-stream-bench-config.h new file mode 100644 index 000000000000..924658be0dc3 --- /dev/null +++ b/tools/kv-stream-bench/kv-stream-bench-config.h @@ -0,0 +1,20 @@ +#pragma once + +#include +#include +#include + +struct llama_kv_stream_bench_config { + bool valid = false; + bool execute = false; + + std::string error; + + int device = 0; + int iterations = 5; + + uint64_t reserve_bytes = 1024ULL*1024ULL*1024ULL; + std::vector transfer_bytes; +}; + +llama_kv_stream_bench_config llama_kv_stream_bench_config_parse(int argc, const char * const * argv); diff --git a/tools/kv-stream-bench/main.cpp b/tools/kv-stream-bench/main.cpp new file mode 100644 index 000000000000..b8aaf6079481 --- /dev/null +++ b/tools/kv-stream-bench/main.cpp @@ -0,0 +1,202 @@ +#include "kv-stream-bench-config.h" + +#include + +#include +#include +#include +#include + +namespace { + +constexpr double MIB = 1024.0*1024.0; +constexpr double GIB = 1024.0*1024.0*1024.0; + +bool cuda_ok(cudaError_t status, const char * operation) { + if (status == cudaSuccess) { + return true; + } + + std::cerr << operation << " failed: " << cudaGetErrorString(status) << '\n'; + return false; +} + +cudaError_t prefetch_async(void * pointer, size_t bytes, int device_id, cudaStream_t stream) { +#if CUDART_VERSION >= 13000 + cudaMemLocation location{}; + location.type = device_id == cudaCpuDeviceId ? cudaMemLocationTypeHost : cudaMemLocationTypeDevice; + location.id = device_id == cudaCpuDeviceId ? 0 : device_id; + return cudaMemPrefetchAsync(pointer, bytes, location, 0, stream); +#else + return cudaMemPrefetchAsync(pointer, bytes, device_id, stream); +#endif +} + +bool has_budget(uint64_t transfer_bytes, uint64_t reserve_bytes) { + size_t free_bytes = 0; + size_t total_bytes = 0; + if (!cuda_ok(cudaMemGetInfo(&free_bytes, &total_bytes), "cudaMemGetInfo")) { + return false; + } + + const bool fits = transfer_bytes <= free_bytes && reserve_bytes <= free_bytes - transfer_bytes; + if (!fits) { + std::cout << "skip " << transfer_bytes/MIB << " MiB: free=" << free_bytes/MIB + << " MiB reserve=" << reserve_bytes/MIB << " MiB\n"; + } + return fits; +} + +void print_rate(const char * name, uint64_t bytes, int iterations, float milliseconds) { + const double seconds = milliseconds/1000.0; + const double gib_per_second = (double(bytes)*iterations/GIB)/seconds; + + std::cout << std::fixed << std::setprecision(2) + << name << " size=" << bytes/MIB << " MiB" + << " iterations=" << iterations + << " elapsed=" << milliseconds << " ms" + << " bandwidth=" << gib_per_second << " GiB/s\n"; +} + +bool benchmark_pinned(uint64_t bytes, int iterations) { + void * host = nullptr; + void * device = nullptr; + cudaStream_t stream = nullptr; + cudaEvent_t begin = nullptr; + cudaEvent_t end = nullptr; + + bool ok = cuda_ok(cudaMallocHost(&host, size_t(bytes)), "cudaMallocHost") && + cuda_ok(cudaMalloc(&device, size_t(bytes)), "cudaMalloc") && + cuda_ok(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), "cudaStreamCreateWithFlags") && + cuda_ok(cudaEventCreate(&begin), "cudaEventCreate(begin)") && + cuda_ok(cudaEventCreate(&end), "cudaEventCreate(end)"); + + if (ok) { + std::memset(host, 0x5a, size_t(bytes)); + ok = cuda_ok(cudaMemcpyAsync(device, host, size_t(bytes), cudaMemcpyHostToDevice, stream), "warmup cudaMemcpyAsync") && + cuda_ok(cudaStreamSynchronize(stream), "warmup cudaStreamSynchronize") && + cuda_ok(cudaEventRecord(begin, stream), "cudaEventRecord(begin)"); + } + + for (int i = 0; ok && i < iterations; ++i) { + ok = cuda_ok(cudaMemcpyAsync(device, host, size_t(bytes), cudaMemcpyHostToDevice, stream), "cudaMemcpyAsync"); + } + + float milliseconds = 0.0f; + if (ok) { + ok = cuda_ok(cudaEventRecord(end, stream), "cudaEventRecord(end)") && + cuda_ok(cudaEventSynchronize(end), "cudaEventSynchronize") && + cuda_ok(cudaEventElapsedTime(&milliseconds, begin, end), "cudaEventElapsedTime"); + } + if (ok) { + print_rate("pinned-h2d", bytes, iterations, milliseconds); + } + + if (end) cudaEventDestroy(end); + if (begin) cudaEventDestroy(begin); + if (stream) cudaStreamDestroy(stream); + if (device) cudaFree(device); + if (host) cudaFreeHost(host); + return ok; +} + +bool benchmark_managed(uint64_t bytes, int iterations, int device_id) { + void * managed = nullptr; + cudaStream_t stream = nullptr; + cudaEvent_t begin = nullptr; + cudaEvent_t end = nullptr; + + bool ok = cuda_ok(cudaMallocManaged(&managed, size_t(bytes)), "cudaMallocManaged") && + cuda_ok(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), "cudaStreamCreateWithFlags") && + cuda_ok(cudaEventCreate(&begin), "cudaEventCreate(begin)") && + cuda_ok(cudaEventCreate(&end), "cudaEventCreate(end)"); + + if (ok) { + std::memset(managed, 0x6b, size_t(bytes)); + } + + double total_milliseconds = 0.0; + for (int i = 0; ok && i < iterations; ++i) { + ok = cuda_ok(prefetch_async(managed, size_t(bytes), cudaCpuDeviceId, stream), "prefetch to CPU") && + cuda_ok(cudaStreamSynchronize(stream), "CPU prefetch synchronize") && + cuda_ok(cudaEventRecord(begin, stream), "cudaEventRecord(begin)") && + cuda_ok(prefetch_async(managed, size_t(bytes), device_id, stream), "prefetch to GPU"); + + float milliseconds = 0.0f; + if (ok) { + ok = cuda_ok(cudaEventRecord(end, stream), "cudaEventRecord(end)") && + cuda_ok(cudaEventSynchronize(end), "cudaEventSynchronize") && + cuda_ok(cudaEventElapsedTime(&milliseconds, begin, end), "cudaEventElapsedTime"); + total_milliseconds += milliseconds; + } + } + + if (ok) { + print_rate("managed-prefetch-h2d", bytes, iterations, float(total_milliseconds)); + } + + if (end) cudaEventDestroy(end); + if (begin) cudaEventDestroy(begin); + if (stream) cudaStreamDestroy(stream); + if (managed) cudaFree(managed); + return ok; +} + +void print_plan(const llama_kv_stream_bench_config & config) { + std::cout << "KV stream CUDA transfer benchmark\n" + << " mode: " << (config.execute ? "execute" : "dry-run") << '\n' + << " device: " << config.device << '\n' + << " iterations: " << config.iterations << '\n' + << " free-VRAM reserve: " << config.reserve_bytes/MIB << " MiB\n" + << " transfer sizes:"; + for (uint64_t bytes : config.transfer_bytes) { + std::cout << ' ' << bytes/MIB << " MiB"; + } + std::cout << "\nUse --execute to allocate and measure.\n"; +} + +} // namespace + +int main(int argc, char ** argv) { + const auto config = llama_kv_stream_bench_config_parse(argc, argv); + if (!config.valid) { + std::cerr << config.error << '\n'; + return 2; + } + + print_plan(config); + if (!config.execute) { + return 0; + } + + int device_count = 0; + if (!cuda_ok(cudaGetDeviceCount(&device_count), "cudaGetDeviceCount") || + config.device >= device_count || + !cuda_ok(cudaSetDevice(config.device), "cudaSetDevice")) { + return 3; + } + + cudaDeviceProp properties{}; + if (!cuda_ok(cudaGetDeviceProperties(&properties, config.device), "cudaGetDeviceProperties")) { + return 3; + } + std::cout << "GPU: " << properties.name << '\n'; + + for (uint64_t bytes : config.transfer_bytes) { + if (!has_budget(bytes, config.reserve_bytes)) { + continue; + } + if (!benchmark_pinned(bytes, config.iterations)) { + return 4; + } + + if (!has_budget(bytes, config.reserve_bytes)) { + continue; + } + if (!benchmark_managed(bytes, config.iterations, config.device)) { + return 5; + } + } + + return 0; +}