From a2e6ab871582a3884f966c56047b28b21451c3f9 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 23 Jul 2026 09:21:49 +0000 Subject: [PATCH 1/3] feat(qwen35): add paged-attention foundation Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QwBCUqGbQV4LSm6qUkZS3H --- README.md | 1 + optimizations/paged_attention/README.md | 158 +++ server/CMakeLists.txt | 42 + server/deps/llama.cpp/ggml/include/ggml-rpc.h | 4 +- server/deps/llama.cpp/ggml/include/ggml.h | 23 + .../llama.cpp/ggml/src/ggml-backend-meta.cpp | 25 + .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c | 5 + .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp | 2 + .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 6 + .../ggml/src/ggml-cuda/paged-attn.cu | 338 ++++++ .../ggml/src/ggml-cuda/paged-attn.cuh | 7 + server/deps/llama.cpp/ggml/src/ggml.c | 56 +- server/src/common/backend_factory.cpp | 38 + server/src/common/backend_factory.h | 1 + server/src/common/paged_attention_config.h | 83 ++ server/src/common/paged_kv_pool.cpp | 350 ++++++ server/src/common/paged_kv_pool.h | 235 ++++ server/src/common/step_graph.h | 11 +- server/src/internal.h | 19 +- server/src/qwen35/graph_builders.cpp | 28 +- server/src/qwen35/graph_builders.h | 3 +- server/src/qwen35/qwen35_backend.cpp | 346 +++++- server/src/qwen35/qwen35_backend.h | 22 +- server/src/qwen35/qwen35_target_graph.cpp | 94 +- server/src/server/server_main.cpp | 48 + server/test/bench_paged_attention.cpp | 1052 +++++++++++++++++ server/test/test_paged_attention.cpp | 257 ++++ server/test/test_paged_kv_pool.cpp | 291 +++++ 28 files changed, 3491 insertions(+), 54 deletions(-) create mode 100644 optimizations/paged_attention/README.md create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cuh create mode 100644 server/src/common/paged_attention_config.h create mode 100644 server/src/common/paged_kv_pool.cpp create mode 100644 server/src/common/paged_kv_pool.h create mode 100644 server/test/bench_paged_attention.cpp create mode 100644 server/test/test_paged_attention.cpp create mode 100644 server/test/test_paged_kv_pool.cpp diff --git a/README.md b/README.md index bc524efde..f0edf1575 100644 --- a/README.md +++ b/README.md @@ -348,6 +348,7 @@ When compression is on, the request path picks one of three modes automatically, | `--prefix-cache-slots N` | — | Live prefix-cache slot count | | `--kv-cache-dir ` | — | Persist prefix cache to disk | | `--kv-cache-budget N` | — | On-disk cache size cap | +| `--paged-attention` | off | Exact 16-token block-table decode for single-device Qwen3.6-27B AR; see [paged attention](optimizations/paged_attention/README.md) | **Bounded KV residency (KVFlash)** diff --git a/optimizations/paged_attention/README.md b/optimizations/paged_attention/README.md new file mode 100644 index 000000000..501af73ce --- /dev/null +++ b/optimizations/paged_attention/README.md @@ -0,0 +1,158 @@ +

+ ← lucebox-hub +

+ +

Luce Paged Attention

+ +

+ Exact fixed-block K/V management for concurrent long-context serving.
+ Decode reads reusable physical K/V blocks directly through per-sequence block tables. +

+ +--- + +Lucebox's paged-attention path stores full-attention K/V rows in fixed +16-token blocks. Each active sequence owns a logical block table, while the +K/V tensors are addressed by reusable physical block IDs. Decode attention +walks that table directly instead of requiring a sequence's cache to occupy +one contiguous range. + +This is separate from [KVFlash](../kvflash/README.md). Paged attention is an +exact layout and allocation mechanism: live blocks remain on the device and +no tokens are evicted. KVFlash is a bounded-residency policy that can evict +cold chunks and changes which context is resident. + +The initial integration targets dense Qwen3.5/Qwen3.6 27B (`qwen35`) on one +CUDA or HIP device: + +```bash +./server/build/dflash_server model.gguf \ + --paged-attention \ + --max-ctx 131072 +``` + +The prompt is prefetched with the existing exact, chunked attention path. +Because the first sequence receives blocks from a freshly reset pool, its +prefill layout is identity-mapped. Single-token autoregressive decode then +uses the block table and context length for every full-attention layer. + +## Current compatibility + +- One local Qwen3.5/Qwen3.6 dense CUDA target. +- Autoregressive decode with full attention (`--fa-window 0`). +- Fixed 16-token blocks. +- Exact block allocation with transactional exhaustion handling and stale + sequence-handle detection. +- A multi-sequence GPU attention primitive and host allocator, ready for the + scheduler integration described below. + +Startup rejects speculative drafts, DDTree, target layer splitting, and +KVFlash when paged attention is selected. Prefix-cache snapshots are also +disabled in this mode because the existing snapshot format assumes contiguous +K/V rows. + +Qwen3.6 is a hybrid model: 48 of its 64 layers use recurrent Gated DeltaNet +state rather than an attention cache. That state is currently stored once per +backend, so admitting multiple live requests would mix their recurrent state +even though their K/V blocks were isolated. Continuous batching therefore +requires two additional pieces before it is safe: + +1. make DeltaNet and convolution state sequence-indexed, including state + routing on every scheduled step; +2. replace the HTTP server's run-to-completion worker with a token-step + scheduler that emits block tables, valid K/V sequence lengths, and write + slots for each batch. + +Page-aware prefix snapshots, copy-on-write prefix sharing, paged prefill, +speculative verification, and layer-split execution remain follow-up work. +The implementation follows the fixed-block/block-table shape described in +the [llama.cpp paged-attention discussion][llama-discussion], while keeping +these unsupported combinations explicit. + +## Recommended roadmap + +1. **Current implementation:** decode-only `ggml_paged_attn`, block manager, + and contiguous-vs-paged benchmarks. +2. **Baseline continuous batching:** iteration-level scheduler, admission + control, cancellation, and dynamic decode batches. +3. **Batched prefill:** process multiple prompts in one forward pass, + preferably with variable-length batching. +4. **Chunked prefill:** add a token budget, decode prioritization, and + prefill/decode interleaving. +5. **Paged-aware prefill:** read prefixes and preceding chunks directly from + the paged K/V pool. +6. **Prefix caching/CoW:** share reusable prefix blocks with reference counting + and copy-on-write. + +## GPU attention comparison microbenchmark + +`bench_paged_attention` compares one native batched contiguous-attention +operation with one paged-attention operation at `n_seq=1,2,4,8`, using Qwen's +D=256, Hq=24, and Hkv=4 dimensions: + +The paged kernel, Qwen integration, numerical test, and benchmark build on +both CUDA and HIP. HIP numerical coverage has been validated on gfx1151. + +```bash +cmake --build "$BUILD_DIR" --target bench_paged_attention +"$BUILD_DIR/bench_paged_attention" \ + --context 4096 \ + --k-type q4_0 \ + --v-type q4_0 +``` + +Both layouts receive the same logical queries and identical quantized K/V +rows. The paged inputs use private, interleaved physical blocks for every +sequence, and an untimed output comparison must pass before timing starts. +Each sample window queues repeated graph executions and synchronizes once. + +The CSV reports median and p95 window-averaged step time, paged overhead +relative to contiguous attention, aggregate attention-query throughput, +scaling efficiency, K/V memory, block-table metadata, and output error. +Automatic calibration gives each layout an approximately equal-duration +sample window. All sequences have the same context length, so this is a direct +kernel/layout comparison. + +To expose the allocation trade-off separately, pass a comma-separated set of +live context lengths: + +```bash +"$BUILD_DIR/bench_paged_attention" \ + --ragged-contexts 131072,8192,8192,8192,8192,8192,8192,8192 \ + --k-type q4_0 \ + --v-type q4_0 +``` + +This runs one ragged decode step for each sequence. The native baseline is +specifically a **padded per-sequence contiguous** layout: every sequence owns +`align256(max(contexts))` K/V slots and an F16 mask hides its unused suffix. +The paged layout allocates only `sum(align16(context_i))` slots and interleaves +the sequences' unique physical blocks in that compact pool. Both paths receive +identical logical Q/K/V values, and output parity is mandatory before either +path is timed. The CUDA contiguous kernel derives each live prefix from the +mask and skips fully masked tail tiles, so the timing comparison does not +charge it for all padded tokens even though its K/V allocation remains padded. + +The ragged CSV reports live, contiguous-padded, and paged-pool token counts; +kernel timings and output error; exact one-layer K/V bytes; and a K/V-only +projection across Qwen3.6's 16 full-attention layers. The projection excludes +model weights, activations, recurrent-layer state, allocator workspace, and +other runtime memory, so it is a capacity comparison rather than a claim that +one layout will or will not fit on a particular GPU. + +For Q4_0 K/V, the eight-user profile above projects to 18.000 GiB for padded +per-sequence contiguous K/V and 3.234 GiB for paged K/V: a 5.57x reduction, +or 82.03% less K/V storage. Uniform `8 x 128K` contexts would use 18.000 GiB +in both layouts, apart from small paging metadata and block rounding. + +A second GGML design can use one unified contiguous arena plus a +sequence-aware mask. It can approach sum-sized storage while compact, but +each query attends over a shared high-water span, and request churn can leave +holes that must be reused or defragmented. That alternative is not measured +here. Paging targets sum-sized stable allocation while each sequence follows +only its own block table. + +Neither mode measures the full Qwen graph, HTTP scheduling, or +continuous-batching throughput. + +[llama-discussion]: https://github.com/ggml-org/llama.cpp/discussions/21961 diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index a07701382..6d444d579 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -284,6 +284,7 @@ add_library(dflash_common STATIC src/common/dflash_draft_graph.cpp src/common/dflash_draft_kv.cpp src/common/dflash_spec_decode.cpp + src/common/paged_kv_pool.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp src/qwen35/graph_builders.cpp @@ -1031,6 +1032,15 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) add_test(NAME kvflash_pool_sizing COMMAND test_kvflash_pool_sizing) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_paged_kv_pool.cpp") + # Pure host-side allocator test: no model, ggml, or GPU is required. + add_executable(test_paged_kv_pool + test/test_paged_kv_pool.cpp + src/common/paged_kv_pool.cpp) + target_include_directories(test_paged_kv_pool PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + add_test(NAME paged_kv_pool COMMAND test_paged_kv_pool) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_restore_delta.cpp") add_executable(test_restore_delta test/test_restore_delta.cpp) target_include_directories(test_restore_delta PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) @@ -1169,6 +1179,38 @@ if(DFLASH27B_TESTS) ${DFLASH27B_SRC_INCLUDE_DIRS}) endif() endif() + + if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR + DFLASH27B_GPU_BACKEND STREQUAL "hip") + AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_paged_attention.cpp") + add_executable(test_paged_attention test/test_paged_attention.cpp) + target_link_libraries(test_paged_attention PRIVATE + ggml ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(test_paged_attention PRIVATE CUDA::cudart) + else() + target_link_libraries(test_paged_attention PRIVATE hip::host) + endif() + target_include_directories(test_paged_attention PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) + add_test(NAME paged_attention COMMAND test_paged_attention) + endif() + if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR + DFLASH27B_GPU_BACKEND STREQUAL "hip") + AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_paged_attention.cpp") + add_executable(bench_paged_attention test/bench_paged_attention.cpp) + target_link_libraries(bench_paged_attention PRIVATE + ggml ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_paged_attention PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_paged_attention PRIVATE hip::host) + endif() + target_include_directories(bench_paged_attention PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) + endif() endif() if(DFLASH27B_SERVER) diff --git a/server/deps/llama.cpp/ggml/include/ggml-rpc.h b/server/deps/llama.cpp/ggml/include/ggml-rpc.h index 398f614c6..21fdf726b 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-rpc.h +++ b/server/deps/llama.cpp/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 3 #define RPC_PROTO_MINOR_VERSION 6 -#define RPC_PROTO_PATCH_VERSION 4 +#define RPC_PROTO_PATCH_VERSION 5 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 6710b2cca..809564704 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -613,6 +613,8 @@ extern "C" { // remain stable within a protocol generation. GGML_OP_MUL_MAT_GROUPED_SRC, + GGML_OP_PAGED_ATTN, + GGML_OP_COUNT, }; @@ -2458,6 +2460,27 @@ extern "C" { float scale, float alpha); + // Decode-only attention over independently allocated physical KV blocks. + // q: [D, n_seq, n_head] F32 (one query per sequence) + // k/v: [D, pool_tokens, n_head_kv] + // block_table: [max_blocks, n_seq] I32 + // kv_seq_lens: [n_seq] I32 (valid cached K/V tokens per sequence) + // res: [D, n_seq, n_head] F32 + // + // A logical token t is read from physical row + // block_table[t/block_size, seq]*block_size + t%block_size. + // The initial CUDA/HIP implementation supports D=256 and independently + // typed F16, Q4_0, or Q8_0 K/V pools. + GGML_API struct ggml_tensor * ggml_paged_attn( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * block_table, + struct ggml_tensor * kv_seq_lens, + float scale, + int block_size); + // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). // Used for KV cache rotation in TurboQuant quantization types (TQ3_0). diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index bf974d722..494b0ed05 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -730,6 +730,28 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; }; + auto handle_paged_attn = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + const bool mirrored = + src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED; + if (mirrored) { + return src_ss[0]; + } + + // Head-sharded Q/K/V are valid when every device preserves the same + // GQA ratio. The generic ratio validation below enforces that; routing + // metadata itself must be mirrored. + GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + return src_ss[0]; + }; + auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { @@ -948,6 +970,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co // multi-device split (all sources on same device as output). split_state = handle_flash_attn_ext(src_ss); } break; + case GGML_OP_PAGED_ATTN: { + split_state = handle_paged_attn(src_ss); + } break; case GGML_OP_FLASH_ATTN_BACK: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); } break; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c index 4193c8081..98e4bf4fb 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c @@ -2187,6 +2187,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { GGML_ABORT("GGML_OP_FLASH_ATTN_SPARSE is only supported on the CUDA backend"); } + case GGML_OP_PAGED_ATTN: + { + GGML_ABORT("GGML_OP_PAGED_ATTN is only supported on the CUDA backend"); + } case GGML_OP_FLASH_ATTN_BACK: { int32_t t = ggml_get_op_params_i32(tensor, 0); @@ -2577,6 +2581,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_DS4_INDEXER_MASK: case GGML_OP_FLASH_ATTN_EXT: case GGML_OP_FLASH_ATTN_SPARSE: + case GGML_OP_PAGED_ATTN: case GGML_OP_FLASH_ATTN_BACK: case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp index 26bab55b9..d31c7e985 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -470,6 +470,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st // DS4 attention carries sparse-layout and fused-RoPE semantics. // The generic CPU kernel does not implement that contract. return !ggml_flash_attn_ext_is_ds4(op); + case GGML_OP_PAGED_ATTN: + return false; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 14e464ea8..bb15697bb 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -24,6 +24,7 @@ #include "ggml-cuda/diag.cuh" #include "ggml-cuda/fattn.cuh" #include "ggml-cuda/fattn-sparse.cuh" +#include "ggml-cuda/paged-attn.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/turbo-wht.cuh" #include "ggml-cuda/im2col.cuh" @@ -3144,6 +3145,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_FLASH_ATTN_SPARSE: ggml_cuda_flash_attn_sparse(ctx, dst); break; + case GGML_OP_PAGED_ATTN: + ggml_cuda_paged_attn(ctx, dst); + break; case GGML_OP_CROSS_ENTROPY_LOSS: ggml_cuda_cross_entropy_loss(ctx, dst); break; @@ -5575,6 +5579,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); case GGML_OP_FLASH_ATTN_SPARSE: return true; // Always supported on CUDA + case GGML_OP_PAGED_ATTN: + return ggml_cuda_paged_attn_supported(op); case GGML_OP_CROSS_ENTROPY_LOSS: case GGML_OP_CROSS_ENTROPY_LOSS_BACK: case GGML_OP_OPT_STEP_ADAMW: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu new file mode 100644 index 000000000..3de091832 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -0,0 +1,338 @@ +#include "paged-attn.cuh" + +#include "fattn-common.cuh" + +#include +#include + +// Decode uses one logical warp per (sequence, query head). Keeping Q in the +// same register formats as fattn-vec lets the quantized K dot products and V +// dequantizers share their well-tested F16/Q4_0/Q8_0 implementations. +template +static __global__ void paged_attn_decode( + const char * __restrict__ q, + const char * __restrict__ k, + const char * __restrict__ v, + const char * __restrict__ block_table, + const char * __restrict__ kv_seq_lens, + char * __restrict__ dst, + int64_t q_nb1, int64_t q_nb2, + int64_t k_nb1, int64_t k_nb2, + int64_t v_nb1, int64_t v_nb2, + int64_t bt_nb0, int64_t bt_nb1, + int64_t ksl_nb0, + int64_t dst_nb1, int64_t dst_nb2, + int32_t n_seq, + int32_t n_head, + int32_t n_head_kv, + int32_t pool_tokens, + int32_t max_blocks, + int32_t block_size, + float scale) { + constexpr int nthreads = WARP_SIZE; + constexpr int values_per_load = 4; + constexpr int values_per_lane = D / nthreads; + static_assert(D % (nthreads * values_per_load) == 0, "unsupported head size"); + + const int head = blockIdx.x; + const int seq = blockIdx.y; + const int lane = threadIdx.x; + if (seq >= n_seq || head >= n_head) { + return; + } + + const float * q_row = (const float *) (q + (int64_t) seq * q_nb1 + (int64_t) head * q_nb2); + float * o_row = (float *) (dst + (int64_t) seq * dst_nb1 + (int64_t) head * dst_nb2); + + constexpr bool quantize_q = + type_K != GGML_TYPE_F16; + constexpr int q_registers = (D / 2) / nthreads; + +#ifdef V_DOT2_F32_F16_AVAILABLE + half2 q_reg[q_registers]; +#else + float2 q_reg[q_registers]; +#endif + int q_i32[D / (sizeof(int) * nthreads)]; + float2 q_ds [D / (sizeof(int) * nthreads)]; + + __shared__ int q_i32_shared[D / sizeof(int)]; + __shared__ float2 q_ds_shared[D / QK8_1]; + + if constexpr (quantize_q) { + // Quantizing scale*Q matches fattn-vec and avoids an extra multiply for + // every cached K row. +#pragma unroll + for (int i0 = 0; i0 < D / (int) sizeof(int); i0 += nthreads) { + quantize_q8_1_to_shared( + q_row + i0 * sizeof(int), scale, + q_i32_shared + i0, q_ds_shared + i0 / QI8_1); + } + __syncthreads(); + +#pragma unroll + for (int i0 = 0; i0 < D / (int) sizeof(int); i0 += nthreads) { + const int i = i0 + lane; + q_i32[i0 / nthreads] = q_i32_shared[i]; + q_ds [i0 / nthreads] = q_ds_shared[i / QI8_1]; + } + } else { + constexpr int cpy_nb = ggml_cuda_get_max_cpy_bytes(); + constexpr int cpy_ne = cpy_nb / sizeof(float); + const float2 * q_f2 = (const float2 *) q_row; + +#pragma unroll + for (int i0 = 0; i0 < D / 2; i0 += nthreads * cpy_ne) { + const int i = i0 + lane * cpy_ne; + __align__(16) float2 tmp[cpy_ne]; + ggml_cuda_memcpy_1( + tmp, q_f2 + i); + ggml_cuda_memcpy_1( + tmp + cpy_ne / 2, q_f2 + i + cpy_ne / 2); +#pragma unroll + for (int j = 0; j < cpy_ne; ++j) { +#ifdef V_DOT2_F32_F16_AVAILABLE + q_reg[i0 / nthreads + j] = + make_half2(tmp[j].x * scale, tmp[j].y * scale); +#else + q_reg[i0 / nthreads + j] = + make_float2(tmp[j].x * scale, tmp[j].y * scale); +#endif + } + } + } + + constexpr vec_dot_KQ_t vec_dot_kq = + get_vec_dot_KQ(); + constexpr dequantize_V_t dequantize_v = + get_dequantize_V(); + + float acc[values_per_lane] = { 0.0f }; + float qk_max = -FLT_MAX; + float qk_sum = 0.0f; + + const int32_t kv_seq_len_raw = + *(const int32_t *) (kv_seq_lens + (int64_t) seq * ksl_nb0); + const int64_t table_capacity = + (int64_t) max_blocks * block_size; + const int32_t kv_seq_len = kv_seq_len_raw <= 0 + ? 0 + : (kv_seq_len_raw < table_capacity + ? kv_seq_len_raw + : (int32_t) table_capacity); + const int32_t n_logical_blocks = + (kv_seq_len + block_size - 1) / block_size; + const int32_t n_physical_blocks = pool_tokens / block_size; + const int32_t kv_head = head / (n_head / n_head_kv); + + for (int32_t logical_block = 0; logical_block < n_logical_blocks; ++logical_block) { + const int32_t physical_block = + *(const int32_t *) (block_table + + (int64_t) logical_block * bt_nb0 + + (int64_t) seq * bt_nb1); + + // Invalid entries are never expected from the allocator, but skipping + // them prevents stale metadata from becoming an out-of-bounds read. + if (physical_block < 0 || physical_block >= n_physical_blocks) { + continue; + } + + const int32_t block_begin = logical_block * block_size; + const int32_t block_end = + kv_seq_len < block_begin + block_size + ? kv_seq_len + : block_begin + block_size; + + for (int32_t token = block_begin; token < block_end; ++token) { + const int32_t physical_token = + physical_block * block_size + token - block_begin; + const char * k_row = + k + (int64_t) physical_token * k_nb1 + + (int64_t) kv_head * k_nb2; + const char * v_row = + v + (int64_t) physical_token * v_nb1 + + (int64_t) kv_head * v_nb2; + + float score = vec_dot_kq(k_row, q_reg, q_i32, q_ds); + score = warp_reduce_sum(score); + + const float qk_max_new = fmaxf(qk_max, score); + const float old_scale = expf(qk_max - qk_max_new); + const float weight = expf(score - qk_max_new); + qk_sum = qk_sum * old_scale + weight; + qk_max = qk_max_new; + +#pragma unroll + for (int segment = 0; + segment < D / (nthreads * values_per_load); + ++segment) { + float values[values_per_load]; + const int value0 = + segment * nthreads * values_per_load + + lane * values_per_load; + dequantize_v(v_row, values, value0); +#pragma unroll + for (int j = 0; j < values_per_load; ++j) { + const int ai = segment * values_per_load + j; + acc[ai] = acc[ai] * old_scale + weight * values[j]; + } + } + } + } + + const float inv_sum = qk_sum > 0.0f ? 1.0f / qk_sum : 0.0f; +#pragma unroll + for (int segment = 0; + segment < D / (nthreads * values_per_load); + ++segment) { + const int value0 = + segment * nthreads * values_per_load + + lane * values_per_load; +#pragma unroll + for (int j = 0; j < values_per_load; ++j) { + o_row[value0 + j] = + acc[segment * values_per_load + j] * inv_sum; + } + } +} + +static bool paged_attn_type_supported(ggml_type type) { + return type == GGML_TYPE_F16 || + type == GGML_TYPE_Q4_0 || + type == GGML_TYPE_Q8_0; +} + +bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * block_table = dst->src[3]; + const ggml_tensor * kv_seq_lens = dst->src[4]; + + if (!q || !k || !v || !block_table || !kv_seq_lens) { + return false; + } + + const int32_t block_size = ggml_get_op_params_i32(dst, 1); + return dst->type == GGML_TYPE_F32 && + q->type == GGML_TYPE_F32 && + paged_attn_type_supported(k->type) && + paged_attn_type_supported(v->type) && + block_table->type == GGML_TYPE_I32 && + kv_seq_lens->type == GGML_TYPE_I32 && + q->nb[0] == sizeof(float) && + k->nb[0] == ggml_type_size(k->type) && + v->nb[0] == ggml_type_size(v->type) && + block_table->nb[0] == sizeof(int32_t) && + kv_seq_lens->nb[0] == sizeof(int32_t) && + dst->nb[0] == sizeof(float) && + dst->ne[0] == q->ne[0] && + dst->ne[1] == q->ne[1] && + dst->ne[2] == q->ne[2] && + dst->ne[3] == q->ne[3] && + q->ne[0] == 256 && + q->ne[0] == k->ne[0] && + q->ne[0] == v->ne[0] && + k->ne[1] == v->ne[1] && + k->ne[2] > 0 && + k->ne[2] == v->ne[2] && + q->ne[2] % k->ne[2] == 0 && + q->ne[3] == 1 && + k->ne[3] == 1 && + v->ne[3] == 1 && + block_table->ne[0] > 0 && + block_table->ne[1] == q->ne[1] && + block_table->ne[2] == 1 && + block_table->ne[3] == 1 && + kv_seq_lens->ne[0] == q->ne[1] && + kv_seq_lens->ne[1] == 1 && + kv_seq_lens->ne[2] == 1 && + kv_seq_lens->ne[3] == 1 && + block_size > 0 && + k->ne[1] % block_size == 0; +} + +template +static void launch_paged_attn( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + float scale, + int32_t block_size) { + constexpr int D = 256; + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * block_table = dst->src[3]; + const ggml_tensor * kv_seq_lens = dst->src[4]; + + const dim3 grid((unsigned int) q->ne[2], (unsigned int) q->ne[1], 1); + const dim3 block(WARP_SIZE, 1, 1); + paged_attn_decode<<>>( + (const char *) q->data, + (const char *) k->data, + (const char *) v->data, + (const char *) block_table->data, + (const char *) kv_seq_lens->data, + (char *) dst->data, + q->nb[1], q->nb[2], + k->nb[1], k->nb[2], + v->nb[1], v->nb[2], + block_table->nb[0], block_table->nb[1], + kv_seq_lens->nb[0], + dst->nb[1], dst->nb[2], + (int32_t) q->ne[1], + (int32_t) q->ne[2], + (int32_t) k->ne[2], + (int32_t) k->ne[1], + (int32_t) block_table->ne[0], + block_size, + scale); +} + +template +static void launch_paged_attn_v( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + float scale, + int32_t block_size) { + switch (dst->src[2]->type) { + case GGML_TYPE_F16: + launch_paged_attn(ctx, dst, scale, block_size); + break; + case GGML_TYPE_Q4_0: + launch_paged_attn(ctx, dst, scale, block_size); + break; + case GGML_TYPE_Q8_0: + launch_paged_attn(ctx, dst, scale, block_size); + break; + default: + GGML_ABORT("unsupported paged-attention V type: %s", + ggml_type_name(dst->src[2]->type)); + } +} + +void ggml_cuda_paged_attn( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + GGML_ASSERT(ggml_cuda_paged_attn_supported(dst)); + + float scale; + memcpy(&scale, dst->op_params, sizeof(scale)); + const int32_t block_size = ggml_get_op_params_i32(dst, 1); + + switch (dst->src[1]->type) { + case GGML_TYPE_F16: + launch_paged_attn_v(ctx, dst, scale, block_size); + break; + case GGML_TYPE_Q4_0: + launch_paged_attn_v(ctx, dst, scale, block_size); + break; + case GGML_TYPE_Q8_0: + launch_paged_attn_v(ctx, dst, scale, block_size); + break; + default: + GGML_ABORT("unsupported paged-attention K type: %s", + ggml_type_name(dst->src[1]->type)); + } + CUDA_CHECK(cudaGetLastError()); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cuh new file mode 100644 index 000000000..4e3bf2527 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cuh @@ -0,0 +1,7 @@ +#pragma once + +#include "common.cuh" + +bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst); + +void ggml_cuda_paged_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index cee7dc7d0..c1a16f29e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -1136,9 +1136,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "DS4_INDEXER_MASK", "MUL_MAT_GROUPED_SRC", + + "PAGED_ATTN", }; -static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); +static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1261,9 +1263,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "ds4_indexer_mask(x,topk)", "X*grouped(Y)", + + "paged_attn(q,k,v)", }; -static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); +static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -5604,6 +5608,54 @@ struct ggml_tensor * ggml_flash_attn_sparse( return result; } +// ggml_paged_attn + +struct ggml_tensor * ggml_paged_attn( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * block_table, + struct ggml_tensor * kv_seq_lens, + float scale, + int block_size) { + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_Q8_0); + GGML_ASSERT(v->type == GGML_TYPE_F16 || v->type == GGML_TYPE_Q4_0 || v->type == GGML_TYPE_Q8_0); + GGML_ASSERT(block_table->type == GGML_TYPE_I32); + GGML_ASSERT(kv_seq_lens->type == GGML_TYPE_I32); + + GGML_ASSERT(q->ne[0] == k->ne[0] && q->ne[0] == v->ne[0]); + GGML_ASSERT(k->ne[1] == v->ne[1]); + GGML_ASSERT(k->ne[2] > 0); + GGML_ASSERT(k->ne[2] == v->ne[2]); + GGML_ASSERT(q->ne[2] % k->ne[2] == 0); + GGML_ASSERT(q->ne[3] == 1 && k->ne[3] == 1 && v->ne[3] == 1); + + GGML_ASSERT(block_table->ne[0] > 0); + GGML_ASSERT(block_table->ne[1] == q->ne[1]); + GGML_ASSERT(block_table->ne[2] == 1 && block_table->ne[3] == 1); + GGML_ASSERT(kv_seq_lens->ne[0] == q->ne[1]); + GGML_ASSERT(kv_seq_lens->ne[1] == 1 && kv_seq_lens->ne[2] == 1 && kv_seq_lens->ne[3] == 1); + + GGML_ASSERT(block_size > 0); + GGML_ASSERT(k->ne[1] % block_size == 0); + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, q->ne); + + ggml_set_op_params_f32(result, 0, scale); + ggml_set_op_params_i32(result, 1, block_size); + + result->op = GGML_OP_PAGED_ATTN; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = block_table; + result->src[4] = kv_seq_lens; + + return result; +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 9348cd8e0..34b181cb7 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -2,6 +2,7 @@ #include "backend_factory.h" #include "gguf_inspect.h" +#include "paged_attention_config.h" #include "qwen35_backend.h" #include "qwen35moe_backend.h" @@ -20,6 +21,19 @@ namespace dflash::common { +namespace { + +constexpr bool has_compiled_accelerator_backend() { +#if defined(DFLASH27B_BACKEND_CUDA) || defined(GGML_USE_CUDA) || \ + defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + return true; +#else + return false; +#endif +} + +} // namespace + std::string detect_arch(const char * model_path) { auto info = inspect_gguf_model_info(model_path); return info.arch; @@ -55,6 +69,29 @@ std::unique_ptr create_backend(const BackendArgs & args) { return nullptr; } + const PagedAttentionOptions paged_options{ + args.paged_attention, + arch.c_str(), + has_compiled_accelerator_backend(), + args.device.is_layer_split(), + args.remote_target_shard.enabled(), + args.draft_path != nullptr, + args.remote_draft.enabled(), + args.ddtree_mode, + args.fa_window, + false, + false, + args.device.max_ctx, + }; + const std::string paged_error = + validate_paged_attention_options(paged_options); + if (!paged_error.empty()) { + std::fprintf(stderr, + "[backend_factory] --paged-attention %s\n", + paged_error.c_str()); + return nullptr; + } + if (arch == "qwen35") { if (args.device.is_layer_split()) { Qwen35LayerSplitAdapterConfig cfg; @@ -91,6 +128,7 @@ std::unique_ptr create_backend(const BackendArgs & args) { cfg.remote_draft = args.remote_draft; cfg.stream_fd = args.stream_fd; cfg.fa_window = args.fa_window; + cfg.paged_attention = args.paged_attention; cfg.kq_stride_pad = args.kq_stride_pad; cfg.draft_swa_window = args.draft_swa_window; cfg.draft_ctx_max = args.draft_ctx_max; diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index 7b468000c..b462bea2b 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -54,6 +54,7 @@ struct BackendArgs { // qwen35-specific speculative decode options int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. + bool paged_attention = false; // qwen35 single-device AR decode (16-token pages) int kq_stride_pad = 32; int draft_swa_window = 0; int draft_ctx_max = 4096; diff --git a/server/src/common/paged_attention_config.h b/server/src/common/paged_attention_config.h new file mode 100644 index 000000000..c2e1a2857 --- /dev/null +++ b/server/src/common/paged_attention_config.h @@ -0,0 +1,83 @@ +// Shared paged-attention sizing and compatibility rules. + +#pragma once + +#include +#include + +namespace dflash::common { + +constexpr int PAGED_BLOCK_SIZE = 16; + +constexpr int paged_block_count(int max_ctx) { + return (max_ctx + PAGED_BLOCK_SIZE - 1) / PAGED_BLOCK_SIZE; +} + +constexpr int paged_token_capacity(int max_ctx) { + return paged_block_count(max_ctx) * PAGED_BLOCK_SIZE; +} + +// TODO(#558): once the cross-feature gate lands in main, move these rules +// into check_feature_compatibility() and delete both this struct and the +// function below; the geometry helpers above stay. Six of the eight rules +// read straight from BackendArgs plus the gate's `arch` parameter. Three +// things need doing by hand: +// - `pflash` and `kvflash` need new BackendArgs fields. They are the +// reason both non-server callers construct this struct with them +// hardcoded to false — neither layer can see ServerConfig or the env. +// - `accelerator_build` needs has_compiled_accelerator_backend() moved +// out of backend_factory.cpp into placement/placement_backend.h, next +// to compiled_placement_backend(). +// - the prefix/prefill/disk-cache rule in server_main stays where it is: +// it is pure ServerConfig, and naming one feature is not enough to earn +// a place in the gate. +// The 256-wide K/V head check in Qwen35Backend::init() does not move either +// — it needs loaded tensor dims, which GgufModelInfo does not carry. +struct PagedAttentionOptions { + bool enabled = false; + const char * architecture = nullptr; // null when not known yet + bool accelerator_build = true; + bool layer_split = false; + bool remote_target_shard = false; + bool draft = false; + bool remote_draft = false; + bool ddtree = false; + int fa_window = 0; + bool pflash = false; + bool kvflash = false; + int max_ctx = 0; +}; + +inline std::string validate_paged_attention_options( + const PagedAttentionOptions & options) { + if (!options.enabled) return {}; + if (options.architecture && + std::string(options.architecture) != "qwen35") { + return "supports only Qwen3.5/Qwen3.6 dense targets"; + } + if (!options.accelerator_build) { + return "requires a CUDA or HIP build"; + } + if (options.layer_split || options.remote_target_shard) { + return "requires one local target device"; + } + if (options.draft || options.remote_draft || options.ddtree) { + return "requires autoregressive decode without a draft or DDTree"; + } + if (options.fa_window != 0) { + return "requires full attention (--fa-window 0)"; + } + if (options.pflash) { + return "cannot be combined with PFlash prefill compression"; + } + if (options.kvflash) { + return "cannot be combined with KVFlash"; + } + if (options.max_ctx <= 0 || + options.max_ctx > INT_MAX - PAGED_BLOCK_SIZE + 1) { + return "requires max_ctx in [1, INT_MAX - PAGED_BLOCK_SIZE + 1]"; + } + return {}; +} + +} // namespace dflash::common diff --git a/server/src/common/paged_kv_pool.cpp b/server/src/common/paged_kv_pool.cpp new file mode 100644 index 000000000..8a437f412 --- /dev/null +++ b/server/src/common/paged_kv_pool.cpp @@ -0,0 +1,350 @@ +#include "paged_kv_pool.h" + +#include +#include +#include + +namespace dflash::common { + +const char * paged_kv_status_string(PagedKvStatus status) { + switch (status) { + case PagedKvStatus::Ok: + return "ok"; + case PagedKvStatus::InvalidArgument: + return "invalid argument"; + case PagedKvStatus::DuplicateRequest: + return "duplicate request"; + case PagedKvStatus::SequenceSlotsExhausted: + return "sequence slots exhausted"; + case PagedKvStatus::BlocksExhausted: + return "physical blocks exhausted"; + case PagedKvStatus::RequestNotFound: + return "request not found"; + case PagedKvStatus::StaleHandle: + return "stale sequence handle"; + } + return "unknown paged KV status"; +} + +uint32_t PagedKvPool::take_lowest(std::vector & free_list) { + const uint32_t index = free_list.front(); + std::pop_heap(free_list.begin(), free_list.end(), std::greater<>()); + free_list.pop_back(); + return index; +} + +void PagedKvPool::give_back(std::vector & free_list, + uint32_t index) { + free_list.push_back(index); + std::push_heap(free_list.begin(), free_list.end(), std::greater<>()); +} + +void PagedKvPool::refill(std::vector & free_list, uint32_t count) { + // Ascending order is already a valid min-heap, so the refilled list needs + // no heapify pass. + free_list.resize(count); + for (uint32_t i = 0; i < count; ++i) { + free_list[i] = i; + } +} + +PagedKvPool::PagedKvPool(uint32_t physical_block_count, + uint32_t max_sequences, + uint32_t block_size) + : block_size_(block_size), + physical_block_count_(physical_block_count) { + const uint64_t token_capacity = + static_cast(physical_block_count) * block_size; + if (physical_block_count == 0 || max_sequences == 0 || block_size == 0 || + token_capacity > std::numeric_limits::max()) { + throw std::invalid_argument("invalid paged KV pool dimensions"); + } + + sequences_.resize(max_sequences); + refill(free_sequence_slots_, max_sequences); + refill(free_blocks_, physical_block_count); +} + +PagedKvStatus PagedKvPool::acquire(PagedKvRequestId request_id, + PagedKvSequenceHandle & out_handle) { + if (request_to_slot_.find(request_id) != request_to_slot_.end()) { + return PagedKvStatus::DuplicateRequest; + } + if (free_sequence_slots_.empty()) { + return PagedKvStatus::SequenceSlotsExhausted; + } + + // Peek rather than take: the request-map insert below can throw, and the + // slot must stay free if it does. + const uint32_t slot = free_sequence_slots_.front(); + SequenceState & sequence = sequences_[slot]; + uint64_t generation = sequence.generation + 1; + if (generation == 0) generation = 1; + + // A free slot owns no blocks and has no live request mapping. + sequence.request_id = request_id; + sequence.generation = generation; + sequence.kv_seq_len = 0; + sequence.active = true; + sequence.block_table.clear(); + request_to_slot_.emplace(request_id, slot); + take_lowest(free_sequence_slots_); + + out_handle = {slot, generation}; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::lookup(PagedKvRequestId request_id, + PagedKvSequenceHandle & out_handle) const { + const auto it = request_to_slot_.find(request_id); + if (it == request_to_slot_.end()) { + return PagedKvStatus::RequestNotFound; + } + const SequenceState & sequence = sequences_[it->second]; + out_handle = {it->second, sequence.generation}; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::reserve(PagedKvSequenceHandle handle, + uint32_t token_capacity) { + SequenceState * sequence = nullptr; + const PagedKvStatus status = validate(handle, sequence); + if (status != PagedKvStatus::Ok) return status; + return extend_block_table(*sequence, blocks_for_tokens(token_capacity)); +} + +PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, + uint32_t token_count) { + PagedKvAppendResult result; + SequenceState * sequence = nullptr; + result.status = validate(handle, sequence); + if (result.status != PagedKvStatus::Ok || token_count == 0) { + return result; + } + if (token_count > std::numeric_limits::max() - + sequence->kv_seq_len) { + result.status = PagedKvStatus::InvalidArgument; + return result; + } + + const uint32_t old_kv_seq_len = sequence->kv_seq_len; + const uint32_t new_kv_seq_len = old_kv_seq_len + token_count; + const uint32_t required_blocks = blocks_for_tokens(new_kv_seq_len); + const uint32_t current_blocks = + static_cast(sequence->block_table.size()); + const uint32_t additional_blocks = + required_blocks > current_blocks ? required_blocks - current_blocks : 0; + if (additional_blocks > free_blocks_.size()) { + result.status = PagedKvStatus::BlocksExhausted; + return result; + } + + std::vector allocated_blocks; + allocated_blocks.reserve(additional_blocks); + auto free_it = free_blocks_.begin(); + for (uint32_t i = 0; i < additional_blocks; ++i, ++free_it) { + allocated_blocks.push_back(*free_it); + } + + result.write_slots.reserve(token_count); + for (uint32_t i = 0; i < token_count; ++i) { + const uint32_t logical_position = old_kv_seq_len + i; + const uint32_t logical_block = logical_position / block_size_; + const uint32_t block_offset = logical_position % block_size_; + const uint32_t physical_block = + logical_block < current_blocks + ? sequence->block_table[logical_block] + : allocated_blocks[logical_block - current_blocks]; + result.write_slots.push_back({ + logical_position, + physical_block, + block_offset, + static_cast(physical_block) * block_size_ + block_offset, + }); + } + + // Allocation state changes only after every write slot has been built. + for (uint32_t block : allocated_blocks) free_blocks_.erase(block); + sequence->block_table.insert(sequence->block_table.end(), + allocated_blocks.begin(), + allocated_blocks.end()); + sequence->kv_seq_len = new_kv_seq_len; + result.status = PagedKvStatus::Ok; + return result; +} + +PagedKvAppendSpan PagedKvPool::append_compact( + PagedKvSequenceHandle handle, + uint32_t token_count) { + PagedKvAppendSpan result; + SequenceState * sequence = nullptr; + result.status = validate(handle, sequence); + if (result.status != PagedKvStatus::Ok || token_count == 0) { + return result; + } + if (token_count > std::numeric_limits::max() - + sequence->kv_seq_len) { + result.status = PagedKvStatus::InvalidArgument; + return result; + } + + const uint32_t old_kv_seq_len = sequence->kv_seq_len; + const uint32_t new_kv_seq_len = old_kv_seq_len + token_count; + result.status = + extend_block_table(*sequence, blocks_for_tokens(new_kv_seq_len)); + if (result.status != PagedKvStatus::Ok) return result; + + const auto make_slot = [&](uint32_t logical_position) { + const uint32_t logical_block = logical_position / block_size_; + const uint32_t block_offset = logical_position % block_size_; + const uint32_t physical_block = + sequence->block_table[logical_block]; + return PagedKvWriteSlot{ + logical_position, + physical_block, + block_offset, + static_cast(physical_block) * block_size_ + block_offset, + }; + }; + result.token_count = token_count; + result.first = make_slot(old_kv_seq_len); + result.last = make_slot(new_kv_seq_len - 1); + sequence->kv_seq_len = new_kv_seq_len; + return result; +} + +PagedKvStatus PagedKvPool::release(PagedKvSequenceHandle handle) { + SequenceState * sequence = nullptr; + const PagedKvStatus status = validate(handle, sequence); + if (status != PagedKvStatus::Ok) return status; + release_sequence(handle.slot); + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::cancel(PagedKvRequestId request_id) { + const auto it = request_to_slot_.find(request_id); + if (it == request_to_slot_.end()) { + return PagedKvStatus::RequestNotFound; + } + release_sequence(it->second); + return PagedKvStatus::Ok; +} + +void PagedKvPool::reset() { + request_to_slot_.clear(); + + for (SequenceState & sequence : sequences_) { + sequence.request_id = 0; + sequence.kv_seq_len = 0; + sequence.active = false; + sequence.block_table.clear(); + } + refill(free_sequence_slots_, static_cast(sequences_.size())); + refill(free_blocks_, physical_block_count_); +} + +PagedKvStatus PagedKvPool::sequence( + PagedKvSequenceHandle handle, + PagedKvSequenceSnapshot & out_sequence) const { + const SequenceState * sequence = nullptr; + const PagedKvStatus status = validate(handle, sequence); + if (status != PagedKvStatus::Ok) return status; + + PagedKvSequenceSnapshot snapshot; + snapshot.request_id = sequence->request_id; + snapshot.handle = handle; + snapshot.kv_seq_len = sequence->kv_seq_len; + snapshot.block_table = sequence->block_table; + out_sequence = std::move(snapshot); + return PagedKvStatus::Ok; +} + +PagedKvMetadataSnapshot PagedKvPool::metadata_snapshot() const { + PagedKvMetadataSnapshot snapshot; + snapshot.block_size = block_size_; + snapshot.physical_block_count = physical_block_count_; + snapshot.sequences.reserve(request_to_slot_.size()); + snapshot.block_table.reserve( + physical_block_count_ - static_cast(free_blocks_.size())); + + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & sequence = sequences_[slot]; + if (!sequence.active) continue; + + PagedKvSequenceMetadata metadata; + metadata.request_id = sequence.request_id; + metadata.generation = sequence.generation; + metadata.sequence_slot = slot; + metadata.kv_seq_len = sequence.kv_seq_len; + metadata.block_table_offset = + static_cast(snapshot.block_table.size()); + metadata.block_count = + static_cast(sequence.block_table.size()); + snapshot.sequences.push_back(metadata); + snapshot.block_table.insert(snapshot.block_table.end(), + sequence.block_table.begin(), + sequence.block_table.end()); + } + return snapshot; +} + +uint32_t PagedKvPool::blocks_for_tokens(uint32_t token_count) const { + if (token_count == 0) return 0; + return 1 + (token_count - 1) / block_size_; +} + +PagedKvStatus PagedKvPool::validate( + PagedKvSequenceHandle handle, + const SequenceState *& out_state) const { + if (handle.slot >= sequences_.size()) { + return PagedKvStatus::StaleHandle; + } + const SequenceState & sequence = sequences_[handle.slot]; + if (!sequence.active || sequence.generation != handle.generation) { + return PagedKvStatus::StaleHandle; + } + out_state = &sequence; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::validate(PagedKvSequenceHandle handle, + SequenceState *& out_state) { + const SequenceState * state = nullptr; + const PagedKvStatus status = + static_cast(this)->validate(handle, state); + out_state = const_cast(state); + return status; +} + +PagedKvStatus PagedKvPool::extend_block_table(SequenceState & sequence, + uint32_t required_blocks) { + const uint32_t current_blocks = + static_cast(sequence.block_table.size()); + if (required_blocks <= current_blocks) return PagedKvStatus::Ok; + + const uint32_t additional_blocks = required_blocks - current_blocks; + if (additional_blocks > free_blocks_.size()) { + return PagedKvStatus::BlocksExhausted; + } + + sequence.block_table.reserve(required_blocks); + for (uint32_t i = 0; i < additional_blocks; ++i) { + sequence.block_table.push_back(take_lowest(free_blocks_)); + } + return PagedKvStatus::Ok; +} + +void PagedKvPool::release_sequence(uint32_t slot) { + SequenceState & sequence = sequences_[slot]; + request_to_slot_.erase(sequence.request_id); + for (uint32_t block : sequence.block_table) { + give_back(free_blocks_, block); + } + sequence.request_id = 0; + sequence.kv_seq_len = 0; + sequence.active = false; + sequence.block_table.clear(); + give_back(free_sequence_slots_, slot); +} + +} // namespace dflash::common diff --git a/server/src/common/paged_kv_pool.h b/server/src/common/paged_kv_pool.h new file mode 100644 index 000000000..38e9cc875 --- /dev/null +++ b/server/src/common/paged_kv_pool.h @@ -0,0 +1,235 @@ +// PagedKvPool — block-table bookkeeping for paged KV-cache attention. +// +// Splits each sequence's KV cache into fixed-size blocks (vLLM-style) drawn +// from one shared physical pool, so sequences grow in block granularity +// instead of reserving max-context capacity up front. The pool tracks +// indices only — it owns no K/V storage. Callers translate the returned +// block/slot indices into offsets within their own pooled K/V tensors and +// upload metadata_snapshot() to drive the GPU paged-attention kernels. +// +// Not thread-safe; callers must serialize access. + +#pragma once + +#include +#include +#include +#include + +namespace dflash::common { + +// Caller-chosen identifier for one inference request (one KV sequence). +using PagedKvRequestId = uint64_t; + +// Result code of every pool operation. Any value other than Ok means the +// call left the pool state unchanged. +enum class PagedKvStatus : uint8_t { + Ok = 0, + InvalidArgument, // size argument would overflow the sequence length + DuplicateRequest, // acquire() for a request id that is still active + SequenceSlotsExhausted, // all max_sequences slots are in use + BlocksExhausted, // not enough free physical blocks for the growth + RequestNotFound, // lookup()/cancel() for an unknown request id + StaleHandle, // handle refers to a released or reused slot +}; + +// Human-readable status name for logs and error messages. +const char * paged_kv_status_string(PagedKvStatus status); + +// Ticket for one active sequence. `slot` indexes the pool's internal +// sequence array; `generation` is bumped every time the slot is re-acquired, +// so a handle kept past release()/cancel() is rejected as StaleHandle +// instead of silently aliasing the slot's next owner. +struct PagedKvSequenceHandle { + uint32_t slot = std::numeric_limits::max(); + uint64_t generation = 0; +}; + +constexpr bool operator==(const PagedKvSequenceHandle & lhs, + const PagedKvSequenceHandle & rhs) { + return lhs.slot == rhs.slot && lhs.generation == rhs.generation; +} + +constexpr bool operator!=(const PagedKvSequenceHandle & lhs, + const PagedKvSequenceHandle & rhs) { + return !(lhs == rhs); +} + +// Cache destination for one appended token. +struct PagedKvWriteSlot { + uint32_t logical_position = 0; // 0-based position in the sequence + uint32_t physical_block = 0; // block index in the shared pool + uint32_t block_offset = 0; // token slot within that block + // Flat row into the pooled K/V buffer: + // physical_block * block_size + block_offset. + uint64_t physical_token_index = 0; +}; + +// Outcome of append(). On success, `write_slots` holds one entry per +// appended token in logical order; on failure it is empty. +struct PagedKvAppendResult { + PagedKvStatus status = PagedKvStatus::Ok; + std::vector write_slots; + + explicit operator bool() const { return status == PagedKvStatus::Ok; } +}; + +// Allocation summary for callers that only need the range endpoints. Unlike +// append(), this does not materialize one PagedKvWriteSlot per token. +struct PagedKvAppendSpan { + PagedKvStatus status = PagedKvStatus::Ok; + uint32_t token_count = 0; + PagedKvWriteSlot first; + PagedKvWriteSlot last; + + explicit operator bool() const { return status == PagedKvStatus::Ok; } +}; + +// Copy of one sequence's bookkeeping state, as returned by sequence(). +struct PagedKvSequenceSnapshot { + PagedKvRequestId request_id = 0; + PagedKvSequenceHandle handle; + uint32_t kv_seq_len = 0; + std::vector block_table; +}; + +// Per-sequence entry of a metadata snapshot. Each entry addresses its own +// range [block_table_offset, block_table_offset + block_count) in +// PagedKvMetadataSnapshot::block_table. +struct PagedKvSequenceMetadata { + PagedKvRequestId request_id = 0; + uint64_t generation = 0; + uint32_t sequence_slot = 0; + uint32_t kv_seq_len = 0; + uint32_t block_table_offset = 0; + uint32_t block_count = 0; +}; + +// Flattened view of every active sequence, ordered by slot index and laid +// out for upload to the GPU paged-attention kernels: `block_table` is the +// concatenation of all per-sequence block tables. +struct PagedKvMetadataSnapshot { + uint32_t block_size = 0; + uint32_t physical_block_count = 0; + std::vector sequences; + std::vector block_table; +}; + +// Allocator front-end: hands out sequence slots and physical block indices. +class PagedKvPool { +public: + // `block_size` is the number of tokens per physical block. Throws + // std::invalid_argument if any dimension is zero or the total token + // capacity (physical_block_count * block_size) overflows uint32_t. + explicit PagedKvPool(uint32_t physical_block_count, + uint32_t max_sequences, + uint32_t block_size = 16); + + uint32_t block_size() const { return block_size_; } + uint32_t physical_block_count() const { return physical_block_count_; } + uint32_t max_sequences() const { return static_cast(sequences_.size()); } + uint32_t active_sequence_count() const { + return static_cast(request_to_slot_.size()); + } + uint32_t free_block_count() const { + return static_cast(free_blocks_.size()); + } + + // Claim a free sequence slot for `request_id`. The new sequence starts + // empty; no blocks are allocated until reserve()/append(). + PagedKvStatus acquire(PagedKvRequestId request_id, + PagedKvSequenceHandle & out_handle); + + // Fetch the live handle of an already-acquired request. + PagedKvStatus lookup(PagedKvRequestId request_id, + PagedKvSequenceHandle & out_handle) const; + + // Grow the sequence's block table until it can hold `token_capacity` + // tokens in total (never shrinks). This does not advance kv_seq_len; a + // later append() fills the pre-allocated blocks first. Lets callers fail + // fast before prefill instead of mid-append. + PagedKvStatus reserve(PagedKvSequenceHandle handle, + uint32_t token_capacity); + + // Advance kv_seq_len by `token_count`, allocating new blocks as needed, + // and return the cache destination of every appended token. All-or- + // nothing: on failure (stale handle, length overflow, BlocksExhausted) + // no state changes. `token_count == 0` returns Ok with no write slots. + PagedKvAppendResult append(PagedKvSequenceHandle handle, + uint32_t token_count); + + // Advance by `token_count` with the same allocation semantics as append(), + // but return only the first and last destinations. This is the preferred + // path for large contiguous prefills and single-token decode. + PagedKvAppendSpan append_compact(PagedKvSequenceHandle handle, + uint32_t token_count); + + // Return the sequence's blocks and slot to the pool; the handle (and + // any copy of it) becomes stale. + PagedKvStatus release(PagedKvSequenceHandle handle); + + // release() by request id, for callers that no longer hold the handle + // (e.g. client disconnect). + PagedKvStatus cancel(PagedKvRequestId request_id); + + // Drop every sequence and reclaim all blocks. Every outstanding handle + // becomes stale. + void reset(); + + // Copy one sequence's current length and block table. + PagedKvStatus sequence(PagedKvSequenceHandle handle, + PagedKvSequenceSnapshot & out_sequence) const; + + // Flattened copy of all active sequences, ready for device upload. + PagedKvMetadataSnapshot metadata_snapshot() const; + +private: + // Bookkeeping for one sequence slot. `generation` survives release so + // the next acquire on this slot invalidates old handles. + struct SequenceState { + PagedKvRequestId request_id = 0; + uint64_t generation = 0; + uint32_t kv_seq_len = 0; + bool active = false; + std::vector block_table; + }; + + // Blocks needed to hold `token_count` tokens (ceiling division). + uint32_t blocks_for_tokens(uint32_t token_count) const; + + // Resolve a handle to its live state; StaleHandle when the slot is out + // of range, inactive, or from an older generation. + PagedKvStatus validate(PagedKvSequenceHandle handle, + const SequenceState *& out_state) const; + PagedKvStatus validate(PagedKvSequenceHandle handle, + SequenceState *& out_state); + + // Grow `sequence` to own `required_blocks` blocks (no-op if it already + // does). All-or-nothing on BlocksExhausted. + PagedKvStatus extend_block_table(SequenceState & sequence, + uint32_t required_blocks); + + // Return one slot's blocks and request mapping to the free lists. + void release_sequence(uint32_t slot); + + // Take the lowest free index off `free_list`, which must be non-empty. + static uint32_t take_lowest(std::vector & free_list); + // Return one index to `free_list`. + static void give_back(std::vector & free_list, uint32_t index); + // Refill `free_list` with every index in [0, count). + static void refill(std::vector & free_list, uint32_t count); + + uint32_t block_size_ = 0; + uint32_t physical_block_count_ = 0; + std::vector sequences_; + // Min-heaps (std::greater) used as lowest-index-first free lists, so a + // released index is handed out again before any higher untouched one. A + // heap keeps release at O(k log n) in the blocks actually returned, where + // a sorted vector would re-sort the whole pool on every teardown. Their + // capacities are fixed at construction, so reset/release do not allocate. + std::vector free_sequence_slots_; + std::vector free_blocks_; + std::unordered_map request_to_slot_; +}; + +} // namespace dflash::common diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index b09a91341..a3d3d4609 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -47,8 +47,15 @@ struct StepGraph { // reused by the copy-mode persistent fast path. bool built_view = false; ggml_tensor * hidden_input = nullptr; // lm-head projection only - // [n_tokens,n_head_kv] i64; step-invariant KV write (carries kv_start). Null on non-graph paths. + // [n_tokens,n_head_kv] i64 physical destination rows for ggml_set_rows. + // Used by contiguous replay, KVFlash, and paged attention; null when the + // graph uses the legacy contiguous ggml_cpy write. ggml_tensor * kv_write_rows = nullptr; + // Paged decode metadata. The block table is [max_blocks,n_seqs], while + // kv_seq_lens is [n_seqs] and stores valid cached K/V tokens per sequence. + // Both are null on the dense attention path. + ggml_tensor * paged_block_table = nullptr; + ggml_tensor * paged_kv_seq_lens = nullptr; // Output ggml_tensor * logits = nullptr; @@ -79,6 +86,8 @@ inline void step_graph_free(StepGraph & sg) { sg.hidden_input = nullptr; sg.parent_ids = nullptr; sg.kv_write_rows = nullptr; + sg.paged_block_table = nullptr; + sg.paged_kv_seq_lens = nullptr; sg.logits = nullptr; sg.hidden_states = nullptr; sg.argmax_tokens = nullptr; diff --git a/server/src/internal.h b/server/src/internal.h index 6f5a37405..b474bc836 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -25,6 +25,7 @@ #include "gguf.h" #include "dflash27b.h" +#include "common/paged_attention_config.h" namespace dflash::common { @@ -527,14 +528,18 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, // tensors. When smaller than max_ctx, a KvFlashPager maps logical positions to // pool slots and pages cold chunks to host (bounded KV residency); the // logical context bound stays max_ctx. Recurrent (DeltaNet) state is -// unaffected. +// unaffected. When `paged_attention` is true, ctx_alloc may instead be +// max_ctx rounded up to PAGED_BLOCK_SIZE so the last partial page has physical +// rows. Non-paged callers retain the legacy rule that ctx_alloc cannot grow +// the allocation beyond max_ctx. bool create_target_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, ggml_backend_t backend, TargetCache & out, bool prefill_only = false, - int ctx_alloc = 0); + int ctx_alloc = 0, + bool paged_attention = false); bool create_target_cache_partial(const TargetWeights & w, int max_ctx, @@ -545,7 +550,8 @@ bool create_target_cache_partial(const TargetWeights & w, int layer_begin, int layer_end, bool allocate_target_feat, - int ctx_alloc = 0); + int ctx_alloc = 0, + bool paged_attention = false); void free_target_cache(TargetCache & c); @@ -602,8 +608,13 @@ struct QwenGraphInputs { int fa_window = 0; // sliding window for FA layers: 0 = full attention bool last_token_logits_only = false; // if true, only compute logits for last token (prefill optimization) ggml_tensor * parent_ids = nullptr; // [n_tokens] i32; tree mode when non-null - // [n_tokens,n_head_kv] i64; non-null = step-invariant KV write via ggml_set_rows (carries kv_start). + // [n_tokens,n_head_kv] i64 physical destination rows; non-null selects the + // step-invariant ggml_set_rows KV write. ggml_tensor * kv_write_rows = nullptr; + ggml_tensor * paged_block_table = nullptr; // [max_blocks,n_seqs] i32 + // [n_seqs] i32; valid cached K/V tokens per sequence. + ggml_tensor * paged_kv_seq_lens = nullptr; + int paged_block_size = 0; // Capture the LAST token's post-RoPE/post-rotation Q per full-attention // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 4e4fdf912..0c919620f 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -279,7 +279,8 @@ bool build_target_step( int kq_stride_pad, bool capture_moe_router, bool kvflash_mask, - bool capture_qk) { + bool capture_qk, + bool paged_attention) { step_graph_free(sg); // Persistent thread_local arena: rebuilt step graphs land at identical @@ -323,6 +324,19 @@ bool build_target_step( ggml_set_input(sg.attn_mask); } + if (paged_attention) { + if (n_tokens != 1 || with_mask || fa_window != 0) return false; + const int max_blocks = paged_block_count(cache.max_ctx); + sg.paged_block_table = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I32, max_blocks, 1); + ggml_set_name(sg.paged_block_table, "paged_block_table"); + ggml_set_input(sg.paged_block_table); + sg.paged_kv_seq_lens = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); + ggml_set_name(sg.paged_kv_seq_lens, "paged_kv_seq_lens"); + ggml_set_input(sg.paged_kv_seq_lens); + } + sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); // Step-invariant KV write: only when topology can't vary per step. @@ -335,10 +349,11 @@ bool build_target_step( // physical slots, so the slot-mapped write stays active for masked, // multi-token, and feature-capturing forwards (decode AND spec verify). const bool use_kv_write_rows = - !g_no_kvpad && !capture_delta_intermediate && - (kvflash_mask - ? (fa_window == 0) - : (n_tokens == 1 && fa_window == 0 && !with_mask && !capture)); + paged_attention || + (!g_no_kvpad && !capture_delta_intermediate && + (kvflash_mask + ? (fa_window == 0) + : (n_tokens == 1 && fa_window == 0 && !with_mask && !capture))); if (use_kv_write_rows) { sg.kv_write_rows = ggml_new_tensor_2d(sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv); @@ -358,6 +373,9 @@ bool build_target_step( gi.fa_window = fa_window; gi.last_token_logits_only = last_token_logits_only; gi.kv_write_rows = sg.kv_write_rows; + gi.paged_block_table = sg.paged_block_table; + gi.paged_kv_seq_lens = sg.paged_kv_seq_lens; + gi.paged_block_size = paged_attention ? PAGED_BLOCK_SIZE : 0; gi.q_capture = capture_qk; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index 2cff7804a..4f76bb802 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -94,7 +94,8 @@ bool build_target_step( int kq_stride_pad = KQ_MASK_PAD, bool capture_moe_router = false, bool kvflash_mask = false, - bool capture_qk = false); + bool capture_qk = false, + bool paged_attention = false); // Full target forward: DDTree tree-verify mode. bool build_target_step_tree( diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 357fef97c..88cfa5129 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -176,7 +176,33 @@ KvFlashAutoBudget Qwen35Backend::make_kvflash_budget(const TargetWeights & w, bool Qwen35Backend::init() { const bool use_remote_draft = cfg_.remote_draft.enabled(); - split_gpus_ = !use_remote_draft && (cfg_.device.gpu != cfg_.draft_gpu); + split_gpus_ = !use_remote_draft && + (cfg_.device.gpu != cfg_.draft_gpu); + + const bool paged_kvflash = + cfg_.paged_attention && + kvflash_pool_from_env(cfg_.device.max_ctx, KvFlashConfig{}) > 0; + const PagedAttentionOptions paged_options{ + cfg_.paged_attention, + "qwen35", + true, + cfg_.device.is_layer_split(), + false, + cfg_.draft_path != nullptr, + use_remote_draft, + cfg_.ddtree_mode, + cfg_.fa_window, + false, + paged_kvflash, + cfg_.device.max_ctx, + }; + const std::string paged_error = + validate_paged_attention_options(paged_options); + if (!paged_error.empty()) { + std::fprintf(stderr, "[paged-attention] %s\n", paged_error.c_str()); + set_last_error("invalid paged-attention configuration: " + paged_error); + return false; + } target_backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!target_backend_) { @@ -209,6 +235,15 @@ bool Qwen35Backend::init() { return false; } std::printf("[target] %s\n", dflash27b_last_error()); + if (cfg_.paged_attention && + (w_.n_embd_head_k != 256 || w_.n_embd_head_v != 256)) { + std::fprintf(stderr, + "[paged-attention] unsupported attention head dimensions " + "K=%d V=%d; this kernel requires K=V=256\n", + w_.n_embd_head_k, w_.n_embd_head_v); + set_last_error("paged attention requires 256-wide K/V heads"); + return false; + } // Load draft if (cfg_.draft_path && use_remote_draft) { @@ -279,11 +314,50 @@ bool Qwen35Backend::init() { // Subclass gate (e.g. MoE all-hot): may zero kvflash_tokens_ before the KV // cache is sized, so create_target_cache allocates full max_ctx KV. if (!post_kvflash_init_gate()) return false; + // Paged mode sizes the KV cache to whole blocks; otherwise KVFlash + // decides the allocation (0 = full max_ctx). + const int ctx_alloc = cfg_.paged_attention + ? paged_token_capacity(cfg_.device.max_ctx) + : kvflash_tokens_; if (!create_target_cache(w_, cfg_.device.max_ctx, max_verify_tokens, target_backend_, cache_, - /*prefill_only=*/true, /*ctx_alloc=*/kvflash_tokens_)) { + /*prefill_only=*/true, ctx_alloc, + cfg_.paged_attention)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } + if (cfg_.paged_attention) { + const auto supported_type = [](ggml_type type) { + return type == GGML_TYPE_F16 || + type == GGML_TYPE_Q4_0 || + type == GGML_TYPE_Q8_0; + }; + if (!supported_type(cache_.kv_k_type) || + !supported_type(cache_.kv_v_type)) { + std::fprintf(stderr, + "[paged-attention] unsupported KV types K=%s V=%s; use " + "f16, q4_0, or q8_0\n", + ggml_type_name(cache_.kv_k_type), + ggml_type_name(cache_.kv_v_type)); + return false; + } + try { + paged_kv_pool_ = std::make_unique( + (uint32_t)paged_block_count(cfg_.device.max_ctx), + /*max_sequences=*/1, PAGED_BLOCK_SIZE); + paged_kv_write_rows_.resize((size_t)w_.n_head_kv); + paged_block_table_upload_.resize( + (size_t)paged_block_count(cfg_.device.max_ctx)); + } catch (const std::exception & e) { + std::fprintf(stderr, "[paged-attention] pool init failed: %s\n", + e.what()); + return false; + } + std::printf("[paged-attention] %u physical blocks x %u tokens " + "(%d logical tokens)\n", + paged_kv_pool_->physical_block_count(), + paged_kv_pool_->block_size(), cfg_.device.max_ctx); + std::fflush(stdout); + } if (kvflash_active()) { KvFlashConfig pc; pc.pool_tokens = kvflash_tokens_; @@ -339,6 +413,185 @@ bool Qwen35Backend::run_ar_decode_path(int committed, int n_gen, return do_ar_decode(committed, n_gen, out_tokens, io); } +bool Qwen35Backend::begin_paged_sequence(uint32_t prompt_tokens) { + if (!cfg_.paged_attention) return true; + if (!paged_kv_pool_) return false; + + paged_kv_pool_->reset(); + paged_sequence_.reset(); + paged_block_table_size_ = 0; + if (++paged_request_id_ == 0) ++paged_request_id_; + + PagedKvSequenceHandle handle; + PagedKvStatus status = paged_kv_pool_->acquire(paged_request_id_, handle); + if (status != PagedKvStatus::Ok) { + std::fprintf(stderr, "[paged-attention] acquire failed: %s\n", + paged_kv_status_string(status)); + return false; + } + paged_sequence_ = handle; + + PagedKvAppendSpan append = + paged_kv_pool_->append_compact(handle, prompt_tokens); + if (!append) { + std::fprintf(stderr, + "[paged-attention] prompt allocation (%u tokens) failed: %s\n", + prompt_tokens, paged_kv_status_string(append.status)); + set_last_error("paged attention prompt allocation failed"); + end_paged_sequence(); + return false; + } + + // Dense prefill writes logical row p directly. A freshly reset pool hands + // out the lowest blocks, so checking the compact range endpoints is enough + // to verify that those rows and the initial page table coincide. + if (prompt_tokens > 0 && + (append.first.physical_token_index != append.first.logical_position || + append.last.physical_token_index != append.last.logical_position)) { + std::fprintf(stderr, + "[paged-attention] non-identity prefill allocation in [%u, %u]\n", + append.first.logical_position, append.last.logical_position); + set_last_error("paged attention prefill allocation is not contiguous"); + end_paged_sequence(); + return false; + } + return true; +} + +bool Qwen35Backend::append_paged_gap(uint32_t logical_position) { + if (!cfg_.paged_attention) return true; + if (!paged_kv_pool_ || !paged_sequence_) return false; + PagedKvAppendSpan append = + paged_kv_pool_->append_compact(*paged_sequence_, /*token_count=*/1); + if (!append || append.token_count != 1 || + append.last.logical_position != logical_position) { + std::fprintf(stderr, + "[paged-attention] logical gap allocation failed at %u (%s)\n", + logical_position, paged_kv_status_string(append.status)); + set_last_error("paged attention logical gap allocation failed"); + return false; + } + return true; +} + +bool Qwen35Backend::prepare_paged_decode_step(uint32_t logical_position) { + if (!cfg_.paged_attention) return true; + if (!paged_kv_pool_ || !paged_sequence_ || !sg_.kv_write_rows || + !sg_.paged_block_table || !sg_.paged_kv_seq_lens) { + std::fprintf(stderr, + "[paged-attention] decode graph is missing paging metadata\n"); + set_last_error("paged attention decode graph is missing metadata"); + return false; + } + + PagedKvAppendSpan append = + paged_kv_pool_->append_compact(*paged_sequence_, /*token_count=*/1); + if (!append || append.token_count != 1 || + append.last.logical_position != logical_position) { + std::fprintf(stderr, + "[paged-attention] append failed at %u: %s\n", + logical_position, paged_kv_status_string(append.status)); + set_last_error("paged attention decode allocation failed"); + return false; + } + + const PagedKvWriteSlot & slot = append.last; + if (slot.physical_token_index > (uint64_t)INT64_MAX) { + std::fprintf(stderr, + "[paged-attention] physical row exceeds INT64_MAX at %u\n", + logical_position); + set_last_error("paged attention physical row is out of range"); + return false; + } + const int64_t physical_row = (int64_t)slot.physical_token_index; + std::fill(paged_kv_write_rows_.begin(), + paged_kv_write_rows_.end(), physical_row); + ggml_backend_tensor_set( + sg_.kv_write_rows, paged_kv_write_rows_.data(), 0, + paged_kv_write_rows_.size() * sizeof(paged_kv_write_rows_[0])); + + if (logical_position >= (uint32_t)INT32_MAX) { + std::fprintf(stderr, + "[paged-attention] sequence length exceeds INT32_MAX at %u\n", + logical_position); + set_last_error("paged attention sequence length is out of range"); + return false; + } + const size_t logical_block = logical_position / PAGED_BLOCK_SIZE; + const size_t table_capacity = (size_t)sg_.paged_block_table->ne[0]; + if (logical_block >= table_capacity) { + std::fprintf(stderr, + "[paged-attention] block index %zu exceeds table capacity %zu\n", + logical_block, table_capacity); + set_last_error("paged attention block table capacity exceeded"); + return false; + } + if (slot.physical_block > (uint32_t)INT32_MAX) { + std::fprintf(stderr, + "[paged-attention] physical block %u exceeds INT32_MAX\n", + slot.physical_block); + set_last_error("paged attention physical block is out of range"); + return false; + } + + // Graph inputs are not persistent allocations: gallocr may reuse this + // tensor's buffer after its last attention consumer. Keep the stable table + // on the host, but upload every live entry before every compute. + if (paged_block_table_size_ == 0 || slot.block_offset == 0) { + PagedKvSequenceSnapshot sequence; + const PagedKvStatus status = + paged_kv_pool_->sequence(*paged_sequence_, sequence); + if (status != PagedKvStatus::Ok) { + std::fprintf(stderr, + "[paged-attention] sequence snapshot failed: %s\n", + paged_kv_status_string(status)); + set_last_error("paged attention sequence snapshot failed"); + return false; + } + if (sequence.block_table.size() > table_capacity) { + std::fprintf(stderr, + "[paged-attention] block table needs %zu entries, capacity is %zu\n", + sequence.block_table.size(), table_capacity); + set_last_error("paged attention block table capacity exceeded"); + return false; + } + for (size_t i = 0; i < sequence.block_table.size(); ++i) { + if (sequence.block_table[i] > (uint32_t)INT32_MAX) { + std::fprintf(stderr, + "[paged-attention] physical block %u exceeds INT32_MAX " + "at table entry %zu\n", + sequence.block_table[i], i); + set_last_error("paged attention physical block is out of range"); + return false; + } + paged_block_table_upload_[i] = + (int32_t)sequence.block_table[i]; + } + paged_block_table_size_ = sequence.block_table.size(); + } + ggml_backend_tensor_set( + sg_.paged_block_table, paged_block_table_upload_.data(), 0, + paged_block_table_size_ * sizeof(paged_block_table_upload_[0])); + + const int32_t kv_seq_len = (int32_t)logical_position + 1; + ggml_backend_tensor_set(sg_.paged_kv_seq_lens, &kv_seq_len, 0, + sizeof(kv_seq_len)); + return true; +} + +void Qwen35Backend::end_paged_sequence() { + if (paged_kv_pool_ && paged_sequence_) { + const PagedKvStatus status = + paged_kv_pool_->release(*paged_sequence_); + if (status != PagedKvStatus::Ok && + status != PagedKvStatus::StaleHandle) { + std::fprintf(stderr, "[paged-attention] release failed: %s\n", + paged_kv_status_string(status)); + } + } + paged_sequence_.reset(); +} + // ── print_ready_banner ────────────────────────────────────────────────── void Qwen35Backend::print_ready_banner() const { @@ -432,6 +685,16 @@ bool Qwen35Backend::unpark(ParkTarget target) { // ── Snapshots ─────────────────────────────────────────────────────────── bool Qwen35Backend::snapshot_save(int slot) { + if (cfg_.paged_attention) { + static bool warned = false; + if (!warned) { + std::fprintf(stderr, + "[paged-attention] prefix snapshots are disabled until the " + "snapshot format stores block tables\n"); + warned = true; + } + return false; + } if (slot < 0 || slot >= PREFIX_SLOTS) return false; // kvflash: snapshots right-size to cur_pos, which is a LOGICAL position // that can exceed the physical pool once decode has paged, and they copy @@ -459,22 +722,26 @@ void Qwen35Backend::snapshot_free(int slot) { } bool Qwen35Backend::snapshot_used(int slot) const { + if (cfg_.paged_attention) return false; if (slot < 0 || slot >= PREFIX_SLOTS) return false; return prefix_snapshots_[slot].ctx != nullptr; } bool Qwen35Backend::restore_target_cache_from_snapshot(int slot) { + if (cfg_.paged_attention) return false; if (slot < 0 || slot >= PREFIX_SLOTS || !prefix_snapshots_[slot].ctx) return false; return restore_target_cache(prefix_snapshots_[slot], cache_); } int Qwen35Backend::snapshot_cur_pos(int slot) const { + if (cfg_.paged_attention) return 0; if (slot < 0 || slot >= PREFIX_SLOTS) return 0; return prefix_snapshots_[slot].cur_pos; } ModelBackend::SnapshotRef Qwen35Backend::snapshot_ref(int slot) const { SnapshotRef ref; + if (cfg_.paged_attention) return ref; if (slot < 0 || slot >= PREFIX_SLOTS) return ref; const auto & snap = prefix_snapshots_[slot]; if (!snap.ctx) return ref; @@ -488,6 +755,7 @@ ModelBackend::SnapshotRef Qwen35Backend::snapshot_ref(int slot) const { bool Qwen35Backend::snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, int cur_pos, int32_t last_tok) { + if (cfg_.paged_attention) return false; if (slot < 0 || slot >= PREFIX_SLOTS) return false; snapshot_free(slot); @@ -690,6 +958,12 @@ void Qwen35Backend::free_drafter() { bool Qwen35Backend::try_handle_command(const std::string & line, const DaemonIO & io) { // SNAPSHOT_THIN — lightweight snapshot (SSM state only, no KV copy) if (line.compare(0, 14, "SNAPSHOT_THIN ") == 0) { + if (cfg_.paged_attention) { + std::fprintf(stderr, + "[paged-attention] SNAPSHOT_THIN is not page-table aware\n"); + io.emit(-1); + return true; + } int slot = std::atoi(line.c_str() + 14); if (slot >= 0 && slot < PREFIX_SLOTS) { snapshot_free(slot); @@ -709,6 +983,7 @@ bool Qwen35Backend::try_handle_command(const std::string & line, const DaemonIO // ── DFlash spec decode target ──────────────────────────────────────────── DFlashTarget * Qwen35Backend::dflash_target() { + if (cfg_.paged_attention) return nullptr; if (!dflash_target_) { dflash_target_ = std::make_unique( w_, cache_, target_backend_, sg_, @@ -726,6 +1001,7 @@ DFlashTarget * Qwen35Backend::dflash_target() { void Qwen35Backend::shutdown() { const bool use_remote_draft = cfg_.remote_draft.enabled(); + end_paged_sequence(); free_drafter(); step_graph_destroy(sg_); step_graph_destroy(draft_sg_); @@ -795,10 +1071,27 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, // position-addressed and will be overwritten during prefill. reset_recurrent_state(cache_); + if (cfg_.paged_attention) { + // The pool's block rounding can hold up to PAGED_BLOCK_SIZE-1 tokens + // beyond max_ctx; gate on max_ctx itself so no logical (RoPE) + // position ever exceeds the configured context. + if (req.prompt.size() > (uint64_t)cfg_.device.max_ctx) { + result.fail(GenerateErrorCode::ContextOverflow, + "prompt exceeds max_ctx for paged attention"); + return result; + } + if (!begin_paged_sequence((uint32_t)req.prompt.size())) { + result.fail(GenerateErrorCode::BackendSpecific, + "paged KV sequence initialization failed"); + return result; + } + } + // Prefill auto t_prefill_start = std::chrono::steady_clock::now(); const int committed = do_prefill(req.prompt, out_io, req.snap_pos, req.snap_slot); if (committed < 0) { + end_paged_sequence(); result.fail(GenerateErrorCode::PrefillFailed); return result; } @@ -815,8 +1108,32 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, // generation. Most requests never hit the tail because the // model closes naturally well before the budget edge. bool decode_ok = false; - if (req.force_ar_decode) { - decode_ok = do_ar_decode(committed, req.n_gen, result.tokens, out_io, + int ar_n_gen = req.n_gen; + if (cfg_.paged_attention) { + // The paged pool ends at the block-rounded max_ctx: an unclamped + // budget would fail a mid-stream block allocation (and push RoPE + // positions past max_ctx) instead of finishing cleanly. The HTTP + // admission gate normally guarantees prompt+n_gen <= max_ctx, but + // daemon-command callers are not required to. + const int max_ar_n_gen = cfg_.device.max_ctx - committed; + if (max_ar_n_gen <= 0) { + // do_ar_decode would return true on a non-positive count and + // report an empty completion as success. + end_paged_sequence(); + result.fail(GenerateErrorCode::ContextOverflow, + "no context left for generation after prefill"); + return result; + } + if (ar_n_gen > max_ar_n_gen) { + ar_n_gen = max_ar_n_gen; + std::fprintf(stderr, + "[paged-attention] clamping n_gen %d -> %d " + "(committed=%d, max_ctx=%d)\n", + req.n_gen, ar_n_gen, committed, cfg_.device.max_ctx); + } + } + if (cfg_.paged_attention || req.force_ar_decode) { + decode_ok = do_ar_decode(committed, ar_n_gen, result.tokens, out_io, req.budget_hook, &result.budget_forced_close, &result.degenerate_decode_close); @@ -837,6 +1154,7 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, } } if (!decode_ok) { + end_paged_sequence(); result.fail(GenerateErrorCode::DecodeFailed); return result; } @@ -844,6 +1162,7 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, std::chrono::steady_clock::now() - t_decode_start).count(); } + end_paged_sequence(); result.succeed(); return result; } @@ -855,6 +1174,12 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, const DaemonIO & io) { GenerateResult result; DaemonIO out_io = io.with_token_callback(req.on_token); + if (cfg_.paged_attention) { + result.fail(GenerateErrorCode::BackendSpecific, + "paged-attention snapshots are not yet supported"); + out_io.emit(-1); + return result; + } if (slot < 0 || slot >= PREFIX_SLOTS || !prefix_snapshots_[slot].ctx) { result.fail(GenerateErrorCode::InvalidSnapshotSlot); out_io.emit(-1); @@ -1534,6 +1859,11 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, io.emit(first_tok); if (kvflash_active()) kvflash_history_.push_back(first_tok); if (IS_EOS_TOK(first_tok, w_)) return true; + // Preserve the existing AR position semantics: the first sampled + // token advances cur_pos without a cache write. The page table must + // include that zero-initialized logical row so subsequent attention + // sees exactly the same context as the dense path. + if (!append_paged_gap((uint32_t)committed)) return false; committed++; cache_.cur_pos = committed; } @@ -1550,6 +1880,7 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, // kvflash: graph carries a slot-validity mask alongside the // step-invariant set_rows write; the FA span clamps to the pool. const bool pool = kvflash_active(); + const bool paged = cfg_.paged_attention; if (!build_target_step(sg_, w_, cache_, target_backend_, /*kv_start=*/committed, /*n_tokens=*/1, /*with_mask=*/pool, /*capture=*/false, @@ -1559,13 +1890,16 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, cfg_.kq_stride_pad, should_capture_moe_router(), /*kvflash_mask=*/pool, - /*capture_qk=*/pool && kvflash_qk_policy_)) { + /*capture_qk=*/pool && kvflash_qk_policy_, + /*paged_attention=*/paged)) { return false; } // Fill kv_write_rows with this step's cache slot for set_rows: // the logical position directly, or its pool slot in kvflash mode. - if (sg_.kv_write_rows) { + if (paged) { + if (!prepare_paged_decode_step((uint32_t)committed)) return false; + } else if (sg_.kv_write_rows) { const int n_head_kv = w_.n_head_kv; const int64_t slot = pool ? (int64_t)kvflash_pager_.slot_for(committed) : (int64_t)committed; diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 0a82f2ef6..58bf06456 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -20,6 +20,7 @@ #include "ddtree.h" #include "dflash_feature_ring.h" #include "common/dflash_draft_kv.h" +#include "common/paged_kv_pool.h" #include "internal.h" // TargetWeights, TargetCache, DraftWeights, PrefixSnapshot #include "qwen3/qwen3_drafter.h" // DrafterContext, load_drafter, free_drafter, drafter_score_and_compress #include "kvflash_pager.h" // bounded KV residency pool @@ -30,6 +31,7 @@ #include "ggml-backend.h" #include +#include #include #include #include @@ -48,6 +50,7 @@ struct Qwen35Config { // FA/KV int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. + bool paged_attention = false; int kq_stride_pad = 32; // KQ_MASK_PAD or 256 for TBQ // Draft @@ -118,7 +121,7 @@ class Qwen35Backend : public ModelBackend { bool try_handle_command(const std::string & line, const DaemonIO & io) override; - bool supports_dflash_spec_decode() const override { return true; } + bool supports_dflash_spec_decode() const override { return !cfg_.paged_attention; } DFlashTarget * dflash_target() override; bool supports_remote_draft() const override { return true; } @@ -260,6 +263,18 @@ class Qwen35Backend : public ModelBackend { std::size_t prefill_last_logits_offset_ = 0; bool prefill_last_logits_valid_ = false; + // The HTTP worker still runs one request at a time. The allocator and + // attention op are sequence-aware; this first integration owns one live + // sequence until recurrent DeltaNet state is also sequence-indexed. + // Page size comes from PAGED_BLOCK_SIZE (internal.h), shared with the + // graph builder and the cache's block-aligned physical sizing. + std::unique_ptr paged_kv_pool_; + std::optional paged_sequence_; + PagedKvRequestId paged_request_id_ = 0; + std::vector paged_kv_write_rows_; + std::vector paged_block_table_upload_; + size_t paged_block_table_size_ = 0; + // ── DFlashTarget adapter (lazy-built) ──────────────────────────── std::unique_ptr dflash_target_; @@ -315,6 +330,11 @@ class Qwen35Backend : public ModelBackend { bool * forced_close_out = nullptr, bool * degenerate_close_out = nullptr); + bool begin_paged_sequence(uint32_t prompt_tokens); + bool append_paged_gap(uint32_t logical_position); + bool prepare_paged_decode_step(uint32_t logical_position); + void end_paged_sequence(); + bool sync_remote_draft_features(int start_pos, int n_tokens); // Chain-mode verify (single batch of q_len tokens). diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 24808177a..edcb705a2 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -78,10 +78,12 @@ bool create_target_cache(const TargetWeights & w, ggml_backend_t backend, TargetCache & out, bool prefill_only, - int ctx_alloc) { + int ctx_alloc, + bool paged_attention) { return create_target_cache_partial(w, max_ctx, max_verify_tokens, backend, out, prefill_only, - 0, w.n_layer, true, ctx_alloc); + 0, w.n_layer, true, ctx_alloc, + paged_attention); } bool create_target_cache_partial(const TargetWeights & w, @@ -93,7 +95,8 @@ bool create_target_cache_partial(const TargetWeights & w, int layer_begin, int layer_end, bool allocate_target_feat, - int ctx_alloc) { + int ctx_alloc, + bool paged_attention) { if (layer_begin < 0) layer_begin = 0; if (layer_end < 0 || layer_end > w.n_layer) layer_end = w.n_layer; if (layer_begin > layer_end) { @@ -140,7 +143,14 @@ bool create_target_cache_partial(const TargetWeights & w, // physical pool capacity; logical positions are mapped to pool slots // by KvFlashPager. The 256-stride rounding applies to whichever capacity // is in effect. - const int ctx_phys = (ctx_alloc > 0 && ctx_alloc < max_ctx) ? ctx_alloc : max_ctx; + // KVFlash may shrink the physical pool. Only an explicitly paged caller + // may grow it to the next whole block; other oversized overrides preserve + // the legacy max_ctx fallback. + const bool bounded_pool = ctx_alloc > 0 && ctx_alloc < max_ctx; + const bool paged_padding = + paged_attention && ctx_alloc == paged_token_capacity(max_ctx); + const int ctx_phys = + (bounded_pool || paged_padding) ? ctx_alloc : max_ctx; const int max_ctx_alloc = needs_256_stride ? ((ctx_phys + 255) / 256) * 256 : ctx_phys; @@ -566,7 +576,10 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * q_tail_capture = nullptr, int q_tail_start = 0, ggml_tensor * kv_write_rows = nullptr, - ggml_tensor ** q_fa_out = nullptr // post-RoPE/post-rotation Q [head_dim, n_tokens, n_head] + ggml_tensor ** q_fa_out = nullptr, // post-RoPE/post-rotation Q [head_dim, n_tokens, n_head] + ggml_tensor * paged_block_table = nullptr, + ggml_tensor * paged_kv_seq_lens = nullptr, + int paged_block_size = 0 ) { const int head_dim = w.n_embd_head_k; const int n_head = w.n_head; @@ -655,7 +668,9 @@ static ggml_tensor * build_full_attn_block( } if (kv_write_rows) { - // Step-invariant: constant dst pointer, idx carries kv_start. set_rows needs contiguous src. + // Step-invariant: the destination tensor stays fixed while the input + // indices carry contiguous, KVFlash, or paged physical rows. + // ggml_set_rows requires a contiguous source. ggml_tensor * Kcur_cont = ggml_is_contiguous(Kcur_T) ? Kcur_T : ggml_cont(ctx, Kcur_T); ggml_tensor * Vcur_cont = ggml_is_contiguous(Vcur_T) ? Vcur_T : ggml_cont(ctx, Vcur_T); ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache_k, Kcur_cont, kv_write_rows)); @@ -675,12 +690,7 @@ static ggml_tensor * build_full_attn_block( } // ── Flash attention over the valid slice - // fa_window > 0: attend only to the last fa_window positions (cuts FA cost - // during spec-decode verify at long contexts). - const int win_start = (fa_window > 0 && kv_start > fa_window) - ? (kv_start - fa_window) : 0; const int kv_len = kv_start + n_tokens; - const int win_len = kv_len - win_start; // Stride-256 FA span when (a) TQ3_0 requires it, or (b) the step-invariant // set_rows KV write is active (kv_write_rows): a fixed span within each @@ -691,11 +701,6 @@ static ggml_tensor * build_full_attn_block( const bool step_invariant = (kv_write_rows != nullptr); const int fattn_stride = (kv_k_type == GGML_TYPE_TQ3_0 || kv_v_type == GGML_TYPE_TQ3_0 || step_invariant) ? 256 : 1; - int win_len_padded = ((win_len + fattn_stride - 1) / fattn_stride) * fattn_stride; - if (step_invariant) { - // Never view past the cache tensor (max_ctx may not be 256-aligned). - win_len_padded = std::min(win_len_padded, (int)cache_k->ne[1]); - } ggml_tensor * Qfa = ggml_permute(ctx, Q, 0, 2, 1, 3); // When K is rotated (TQ3_0 or explicit FWHT), Q needs forward rotation too. @@ -713,22 +718,44 @@ static ggml_tensor * build_full_attn_block( // cosine (orthogonal transform). if (q_fa_out) *q_fa_out = Qfa; - // K and V from cache: a windowed view starting at win_start. - ggml_tensor * Kfa = ggml_view_3d(ctx, cache_k, - head_dim, win_len_padded, n_head_kv, - cache_k->nb[1], cache_k->nb[2], cache_k->nb[1] * win_start); - ggml_tensor * Vfa = ggml_view_3d(ctx, cache_v, - head_dim, win_len_padded, n_head_kv, - cache_v->nb[1], cache_v->nb[2], cache_v->nb[1] * win_start); - - // Causal mask: for n_tokens==1 we don't need one (a single query attending - // to all keys is trivially causal). For n_tokens>1 the caller must provide - // a mask shaped [kv_len, n_tokens] with 0 for attendable positions and - // -inf for positions beyond the causal boundary. const float kq_scale = 1.0f / std::sqrt((float)head_dim); - ggml_tensor * attn = ggml_flash_attn_ext(ctx, Qfa, Kfa, Vfa, attn_mask, - kq_scale, 0.0f, 0.0f); - // attn: [head_dim, n_head, n_tokens] (permuted) + ggml_tensor * attn = nullptr; + if (paged_block_table) { + GGML_ASSERT(paged_kv_seq_lens && paged_block_size > 0); + attn = ggml_paged_attn(ctx, Qfa, cache_k, cache_v, + paged_block_table, paged_kv_seq_lens, + kq_scale, paged_block_size); + } else { + // fa_window > 0: attend only to the last fa_window positions (cuts FA + // cost during spec-decode verify at long contexts). Paged attention + // ignores the window entirely — it walks the block table for the whole + // sequence — which is why build_target_step rejects the combination. + const int win_start = (fa_window > 0 && kv_start > fa_window) + ? (kv_start - fa_window) : 0; + const int win_len = kv_len - win_start; + int win_len_padded = + ((win_len + fattn_stride - 1) / fattn_stride) * fattn_stride; + if (step_invariant) { + // Never view past the cache tensor (max_ctx may not be 256-aligned). + win_len_padded = std::min(win_len_padded, (int)cache_k->ne[1]); + } + + // K and V from cache: a windowed view starting at win_start. + ggml_tensor * Kfa = ggml_view_3d(ctx, cache_k, + head_dim, win_len_padded, n_head_kv, + cache_k->nb[1], cache_k->nb[2], cache_k->nb[1] * win_start); + ggml_tensor * Vfa = ggml_view_3d(ctx, cache_v, + head_dim, win_len_padded, n_head_kv, + cache_v->nb[1], cache_v->nb[2], cache_v->nb[1] * win_start); + + // A single query needs no causal mask. Multi-token callers supply one. + attn = ggml_flash_attn_ext(ctx, Qfa, Kfa, Vfa, attn_mask, + kq_scale, 0.0f, 0.0f); + } + // Dense output is [D,Hq,n_tokens]; paged output is [D,n_seq,Hq]. + // They are layout-equivalent here because the paged path is decode-only + // and n_tokens/n_seq is exactly one. A future batched integration must + // permute the paged result before this reshape. // Un-rotate the FA output from FWHT-rotated V space (only when V is TQ3). if (out_rotate) { @@ -1221,7 +1248,10 @@ QwenGraphOutputs build_qwen35_graph( /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, in.kv_write_rows, - want_q_cap ? &q_fa : nullptr); + want_q_cap ? &q_fa : nullptr, + in.paged_block_table, + in.paged_kv_seq_lens, + in.paged_block_size); if (want_q_cap && q_fa) { // Last token's Q, all heads: src [head_dim, 1, n_head] view of // [head_dim, n_tokens, n_head]; dst = q_cap plane fa_idx diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index bc5dfe33b..2887b5101 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -16,6 +16,7 @@ #include "model_card.h" #include "common/backend_factory.h" #include "common/gguf_inspect.h" +#include "common/paged_attention_config.h" #include "common/layer_split_utils.h" #include "common/spark_corpus.h" #include "common/moe_routing_collector.h" @@ -23,6 +24,7 @@ #include "common/peer_access.h" #include "placement/pflash_placement.h" #include "placement/draft_residency.h" +#include "kvflash_pager.h" #include "gguf.h" @@ -227,6 +229,8 @@ static void print_usage(const char * prog) { " --fa-window Flash-attention sliding window (default: 0=full).\n" " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" + " --paged-attention Use 16-token paged KV blocks for Qwen3.6-27B\n" + " autoregressive decode (experimental)\n" " --model-name Model name for /v1/models (default: dflash)\n" " --prefix-cache-slots Prefix cache slots (default: 32, 0 disables)\n" " --prefill-cache-slots Full prompt/prefill cache slots (default: 0)\n" @@ -336,6 +340,8 @@ int main(int argc, char ** argv) { std::string cache_type_k; // explicit --cache-type-k override std::string cache_type_v; // explicit --cache-type-v override bool target_device_seen = false; + bool prefix_cache_slots_seen = false; + bool prefill_cache_slots_seen = false; bool target_devices_seen = false; bool fast_rollback_forced_off = false; @@ -454,12 +460,16 @@ int main(int argc, char ** argv) { } } else if (std::strcmp(argv[i], "--fa-window") == 0 && i + 1 < argc) { bargs.fa_window = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--paged-attention") == 0) { + bargs.paged_attention = true; } else if (std::strcmp(argv[i], "--model-name") == 0 && i + 1 < argc) { sconfig.model_name = argv[++i]; } else if (std::strcmp(argv[i], "--prefix-cache-slots") == 0 && i + 1 < argc) { sconfig.prefix_cache_cap = std::atoi(argv[++i]); + prefix_cache_slots_seen = true; } else if (std::strcmp(argv[i], "--prefill-cache-slots") == 0 && i + 1 < argc) { sconfig.prefill_cache_cap = std::atoi(argv[++i]); + prefill_cache_slots_seen = true; } else if (std::strcmp(argv[i], "--fast-rollback") == 0) { bargs.fast_rollback = true; } else if (std::strcmp(argv[i], "--ddtree") == 0) { @@ -672,6 +682,44 @@ int main(int argc, char ** argv) { } if (fast_rollback_forced_off) bargs.fast_rollback = false; + if (bargs.paged_attention) { + const PlacementBackend compiled_backend = compiled_placement_backend(); + const PagedAttentionOptions paged_options{ + true, + nullptr, + compiled_backend == PlacementBackend::Cuda || + compiled_backend == PlacementBackend::Hip, + bargs.device.is_layer_split(), + bargs.remote_target_shard.enabled(), + bargs.draft_path != nullptr, + bargs.remote_draft.enabled(), + bargs.ddtree_mode, + bargs.fa_window, + sconfig.pflash_mode != ServerConfig::PflashMode::OFF, + kvflash_pool_from_env( + bargs.device.max_ctx, KvFlashConfig{}) > 0, + bargs.device.max_ctx, + }; + const std::string paged_error = + validate_paged_attention_options(paged_options); + if (!paged_error.empty()) { + std::fprintf(stderr, "[server] --paged-attention %s\n", + paged_error.c_str()); + return 2; + } + if ((prefix_cache_slots_seen && sconfig.prefix_cache_cap > 0) || + (prefill_cache_slots_seen && sconfig.prefill_cache_cap > 0) || + !sconfig.disk_cache_dir.empty()) { + std::fprintf(stderr, + "[server] paged attention disables prefix/prefill snapshots " + "until their format stores page tables\n"); + } + sconfig.prefix_cache_cap = 0; + sconfig.prefill_cache_cap = 0; + sconfig.disk_cache_dir.clear(); + sconfig.disk_cache_policy.mode = DiskPrefixCacheMode::Off; + } + if (!validate_server_placement(bargs, sconfig)) return 2; if (bargs.remote_draft.enabled() && bargs.draft_path) { diff --git a/server/test/bench_paged_attention.cpp b/server/test/bench_paged_attention.cpp new file mode 100644 index 000000000..d4a9cfe50 --- /dev/null +++ b/server/test/bench_paged_attention.cpp @@ -0,0 +1,1052 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int D = 256; +constexpr int N_HEAD = 24; +constexpr int N_HEAD_KV = 4; +constexpr int BLOCK_SIZE = 16; +constexpr int PROTOTYPE_ROWS = 256; +constexpr int CONTIGUOUS_CONTEXT_ALIGNMENT = 256; +constexpr int QWEN_FULL_ATTENTION_LAYERS = 16; +constexpr size_t MAX_RAGGED_SEQUENCES = 64; +constexpr double TARGET_SAMPLE_SECONDS = 0.20; + +struct Options { + int context = 4096; + int warmup = 10; + int samples = 7; + int iterations = 0; + ggml_type k_type = GGML_TYPE_Q4_0; + ggml_type v_type = GGML_TYPE_Q4_0; + std::vector ragged_contexts; +}; + +struct TimingStats { + int iterations = 0; + double median_us = 0.0; + double p95_us = 0.0; + double aggregate_queries_s = 0.0; +}; + +struct BenchResult { + int n_seq = 0; + int padded_context = 0; + int64_t live_tokens = 0; + int64_t padded_tokens = 0; + int64_t paged_pool_tokens = 0; + TimingStats contiguous; + TimingStats paged; + uint64_t contiguous_kv_bytes = 0; + uint64_t paged_kv_bytes = 0; + uint64_t paged_metadata_bytes = 0; + double contiguous_kv_mib = 0.0; + double paged_kv_mib = 0.0; + double paged_metadata_kib = 0.0; + double max_abs_error = 0.0; + double max_tolerance_ratio = 0.0; + double rmse = 0.0; +}; + +struct CachePair { + std::vector contiguous; + std::vector paged; +}; + +struct ContextGuard { + ggml_context * value = nullptr; + ~ContextGuard() { + if (value) ggml_free(value); + } +}; + +struct AllocatorGuard { + ggml_gallocr_t value = nullptr; + ~AllocatorGuard() { + if (value) ggml_gallocr_free(value); + } +}; + +void print_usage(const char * program) { + std::fprintf(stderr, + "Usage: %s [--context N] [--warmup N] [--samples N]\n" + " [--iterations N] [--k-type TYPE] [--v-type TYPE]\n" + " [--ragged-contexts N1,N2,...]\n" + " TYPE: f16, q4_0, or q8_0\n" + " --ragged-contexts runs one ragged case against a per-sequence\n" + " contiguous cache padded to a 256-token maximum context.\n" + " Between 2 and %zu sequence lengths are accepted.\n" + " --iterations 0 automatically targets %.0f ms timing windows.\n", + program, MAX_RAGGED_SEQUENCES, + TARGET_SAMPLE_SECONDS * 1000.0); +} + +bool parse_positive_int(const char * text, int & value, bool allow_zero) { + char * end = nullptr; + const long parsed = std::strtol(text, &end, 10); + if (!text[0] || !end || *end != '\0' || + parsed < (allow_zero ? 0 : 1) || parsed > INT32_MAX) { + return false; + } + value = static_cast(parsed); + return true; +} + +bool parse_type(const char * text, ggml_type & type) { + if (std::strcmp(text, "f16") == 0) { + type = GGML_TYPE_F16; + return true; + } + if (std::strcmp(text, "q4_0") == 0) { + type = GGML_TYPE_Q4_0; + return true; + } + if (std::strcmp(text, "q8_0") == 0) { + type = GGML_TYPE_Q8_0; + return true; + } + return false; +} + +bool parse_context_list( + const char * text, + std::vector & contexts) { + contexts.clear(); + const char * cursor = text; + while (cursor && cursor[0]) { + int context = 0; + char * end = nullptr; + const long parsed = std::strtol(cursor, &end, 10); + if (end == cursor || parsed < 1 || parsed > INT32_MAX || + (*end != ',' && *end != '\0') || + contexts.size() >= MAX_RAGGED_SEQUENCES) { + contexts.clear(); + return false; + } + context = static_cast(parsed); + contexts.push_back(context); + if (*end == '\0') { + if (contexts.size() >= 2) return true; + contexts.clear(); + return false; + } + cursor = end + 1; + } + contexts.clear(); + return false; +} + +bool parse_options(int argc, char ** argv, Options & options) { + for (int i = 1; i < argc; ++i) { + const auto next = [&]() -> const char * { + return i + 1 < argc ? argv[++i] : nullptr; + }; + if (std::strcmp(argv[i], "--context") == 0) { + const char * value = next(); + if (!value || + !parse_positive_int(value, options.context, false)) { + return false; + } + } else if (std::strcmp(argv[i], "--warmup") == 0) { + const char * value = next(); + if (!value || + !parse_positive_int(value, options.warmup, true)) { + return false; + } + } else if (std::strcmp(argv[i], "--samples") == 0) { + const char * value = next(); + if (!value || + !parse_positive_int(value, options.samples, false)) { + return false; + } + } else if (std::strcmp(argv[i], "--iterations") == 0) { + const char * value = next(); + if (!value || + !parse_positive_int(value, options.iterations, true)) { + return false; + } + } else if (std::strcmp(argv[i], "--k-type") == 0) { + const char * value = next(); + if (!value || !parse_type(value, options.k_type)) return false; + } else if (std::strcmp(argv[i], "--v-type") == 0) { + const char * value = next(); + if (!value || !parse_type(value, options.v_type)) return false; + } else if (std::strcmp(argv[i], "--ragged-contexts") == 0) { + const char * value = next(); + if (!value || + !parse_context_list(value, options.ragged_contexts)) { + return false; + } + } else if (std::strcmp(argv[i], "--help") == 0 || + std::strcmp(argv[i], "-h") == 0) { + print_usage(argv[0]); + std::exit(0); + } else { + return false; + } + } + return true; +} + +std::vector make_prototype_rows( + ggml_type type, + float phase) { + const size_t row_bytes = ggml_row_size(type, D); + if (row_bytes == 0) return {}; + + std::vector source( + static_cast(PROTOTYPE_ROWS) * D); + for (size_t i = 0; i < source.size(); ++i) { + source[i] = + std::sin(static_cast(i) * 0.011f + phase) * 0.25f; + } + + std::vector result( + static_cast(PROTOTYPE_ROWS) * row_bytes); + const size_t written = ggml_quantize_chunk( + type, source.data(), result.data(), 0, PROTOTYPE_ROWS, D, nullptr); + return written == result.size() ? result : std::vector{}; +} + +bool cache_size_fits(ggml_type type, int64_t rows) { + const size_t row_bytes = ggml_row_size(type, D); + return rows > 0 && row_bytes > 0 && + static_cast(rows) <= SIZE_MAX / row_bytes; +} + +CachePair make_cache_pair( + ggml_type type, + const std::vector & kv_seq_lens, + int contiguous_context, + int max_blocks, + int64_t pool_tokens, + const std::vector & block_table, + float phase) { + const int n_seq = static_cast(kv_seq_lens.size()); + const int64_t contiguous_rows = + static_cast(contiguous_context) * N_HEAD_KV * n_seq; + const int64_t paged_rows = pool_tokens * N_HEAD_KV; + if (!cache_size_fits(type, contiguous_rows) || + !cache_size_fits(type, paged_rows)) { + return {}; + } + + const size_t row_bytes = ggml_row_size(type, D); + const std::vector prototype = + make_prototype_rows(type, phase); + if (prototype.empty()) return {}; + + CachePair result; + result.contiguous.resize( + static_cast(contiguous_rows) * row_bytes); + result.paged.resize( + static_cast(paged_rows) * row_bytes); + + // Generate one canonical logical row and place the identical quantized + // bytes in both the contiguous tensor and the paged physical pool. + for (int seq = 0; seq < n_seq; ++seq) { + for (int kv_head = 0; kv_head < N_HEAD_KV; ++kv_head) { + for (int token = 0; token < kv_seq_lens[seq]; ++token) { + const int64_t logical_row = + (static_cast(seq) * N_HEAD_KV + kv_head) * + contiguous_context + + token; + const int32_t physical_block = + block_table[ + static_cast(seq) * max_blocks + + token / BLOCK_SIZE]; + const int64_t physical_token = + static_cast(physical_block) * BLOCK_SIZE + + token % BLOCK_SIZE; + const int64_t paged_row = + static_cast(kv_head) * pool_tokens + + physical_token; + // Keep sequence and head identity in the source pattern. + // Using logical_row % PROTOTYPE_ROWS would collapse both at + // common contexts such as 4096. + const size_t prototype_row = + static_cast( + (static_cast(token) * 17 + + (seq * N_HEAD_KV + kv_head) * 7) % + PROTOTYPE_ROWS); + + std::memcpy( + result.contiguous.data() + + static_cast(logical_row) * row_bytes, + prototype.data() + prototype_row * row_bytes, + row_bytes); + std::memcpy( + result.paged.data() + + static_cast(paged_row) * row_bytes, + prototype.data() + prototype_row * row_bytes, + row_bytes); + } + } + } + return result; +} + +bool make_block_table( + const std::vector & kv_seq_lens, + int max_blocks, + int64_t & pool_tokens, + std::vector & block_table) { + const int n_seq = static_cast(kv_seq_lens.size()); + if (n_seq <= 0 || + static_cast(max_blocks) * n_seq > + SIZE_MAX / sizeof(int32_t)) { + return false; + } + + std::vector sequence_blocks(n_seq); + int64_t physical_blocks = 0; + for (int seq = 0; seq < n_seq; ++seq) { + sequence_blocks[seq] = static_cast( + (static_cast(kv_seq_lens[seq]) + + BLOCK_SIZE - 1) / + BLOCK_SIZE); + physical_blocks += sequence_blocks[seq]; + } + pool_tokens = physical_blocks * BLOCK_SIZE; + if (physical_blocks <= 0 || physical_blocks > INT32_MAX || + pool_tokens <= 0 || pool_tokens > INT32_MAX) { + return false; + } + + block_table.assign( + static_cast(max_blocks) * n_seq, -1); + int32_t physical = 0; + for (int logical = 0; logical < max_blocks; ++logical) { + // Round-robin by logical block keeps live blocks from different + // sequences interleaved while assigning a compact, unique pool. + for (int rank = 0; rank < n_seq; ++rank) { + const int seq = (logical + rank) % n_seq; + if (logical >= sequence_blocks[seq]) continue; + block_table[ + static_cast(seq) * max_blocks + logical] = + physical++; + } + } + return physical == physical_blocks; +} + +bool compute_iterations( + ggml_backend_t backend, + ggml_cgraph * graph, + int iterations, + double & elapsed_seconds) { + ggml_backend_synchronize(backend); + const auto begin = std::chrono::steady_clock::now(); + for (int i = 0; i < iterations; ++i) { + if (ggml_backend_graph_compute_async(backend, graph) != + GGML_STATUS_SUCCESS) { + ggml_backend_synchronize(backend); + return false; + } + } + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + elapsed_seconds = + std::chrono::duration(end - begin).count(); + return true; +} + +bool warm_graph( + ggml_backend_t backend, + ggml_cgraph * graph, + int warmup) { + if (warmup <= 0) return true; + double ignored_seconds = 0.0; + return compute_iterations( + backend, graph, warmup, ignored_seconds); +} + +bool probe_graph( + ggml_backend_t backend, + ggml_cgraph * graph, + double & seconds_per_step) { + constexpr int PROBE_ITERATIONS = 5; + double elapsed_seconds = 0.0; + if (!compute_iterations( + backend, graph, PROBE_ITERATIONS, elapsed_seconds)) { + return false; + } + seconds_per_step = + std::max(elapsed_seconds / PROBE_ITERATIONS, 1.0e-6); + return true; +} + +bool collect_timing( + ggml_backend_t backend, + ggml_cgraph * graph, + const Options & options, + int iterations, + int n_seq, + TimingStats & result) { + if (!warm_graph(backend, graph, options.warmup)) return false; + + std::vector sample_us; + sample_us.reserve(options.samples); + for (int sample = 0; sample < options.samples; ++sample) { + double elapsed_seconds = 0.0; + if (!compute_iterations( + backend, graph, iterations, elapsed_seconds)) { + return false; + } + sample_us.push_back( + elapsed_seconds * 1.0e6 / static_cast(iterations)); + } + + std::sort(sample_us.begin(), sample_us.end()); + const size_t p95_index = std::min( + sample_us.size() - 1, + static_cast( + std::ceil(sample_us.size() * 0.95) - 1.0)); + result.iterations = iterations; + const size_t upper_middle = sample_us.size() / 2; + result.median_us = + sample_us.size() % 2 != 0 + ? sample_us[upper_middle] + : (sample_us[upper_middle - 1] + + sample_us[upper_middle]) / + 2.0; + result.p95_us = sample_us[p95_index]; + result.aggregate_queries_s = + static_cast(n_seq) * 1.0e6 / result.median_us; + return true; +} + +bool compare_outputs( + ggml_tensor * paged_output, + ggml_tensor * contiguous_output, + int n_seq, + BenchResult & result) { + std::vector paged(ggml_nelements(paged_output)); + std::vector contiguous(ggml_nelements(contiguous_output)); + ggml_backend_tensor_get( + paged_output, paged.data(), 0, + paged.size() * sizeof(paged[0])); + ggml_backend_tensor_get( + contiguous_output, contiguous.data(), 0, + contiguous.size() * sizeof(contiguous[0])); + + double squared_error = 0.0; + bool ok = true; + for (int seq = 0; seq < n_seq; ++seq) { + for (int head = 0; head < N_HEAD; ++head) { + for (int d = 0; d < D; ++d) { + const size_t paged_index = + (static_cast(head) * n_seq + seq) * D + d; + const size_t contiguous_index = + (static_cast(seq) * N_HEAD + head) * D + d; + const float paged_value = paged[paged_index]; + const float contiguous_value = + contiguous[contiguous_index]; + if (!std::isfinite(paged_value) || + !std::isfinite(contiguous_value)) { + return false; + } + + const double abs_error = + std::fabs( + static_cast(paged_value) - + contiguous_value); + result.max_abs_error = + std::max(result.max_abs_error, abs_error); + squared_error += abs_error * abs_error; + + const double tolerance = + 1.0e-4 + + 1.0e-3 * + std::fabs(static_cast(contiguous_value)); + const double tolerance_ratio = abs_error / tolerance; + result.max_tolerance_ratio = + std::max( + result.max_tolerance_ratio, + tolerance_ratio); + ok = ok && tolerance_ratio <= 1.0; + } + } + } + result.rmse = + std::sqrt( + squared_error / + static_cast(D * N_HEAD * n_seq)); + return ok; +} + +bool run_case( + ggml_backend_t backend, + const Options & options, + const std::vector & kv_seq_lens, + int contiguous_context, + bool use_padding_mask, + bool contiguous_first, + BenchResult & result) { + if (kv_seq_lens.empty() || + kv_seq_lens.size() > MAX_RAGGED_SEQUENCES) { + return false; + } + const int n_seq = static_cast(kv_seq_lens.size()); + const int max_context = + *std::max_element(kv_seq_lens.begin(), kv_seq_lens.end()); + const int64_t max_blocks_64 = + (static_cast(max_context) + BLOCK_SIZE - 1) / + BLOCK_SIZE; + if (max_blocks_64 <= 0 || max_blocks_64 > INT32_MAX || + contiguous_context < max_context) { + std::fprintf( + stderr, "context layout exceeds attention kernel limits\n"); + return false; + } + const int max_blocks = static_cast(max_blocks_64); + + int64_t pool_tokens = 0; + std::vector block_table; + if (!make_block_table( + kv_seq_lens, max_blocks, pool_tokens, block_table)) { + std::fprintf( + stderr, "paged block table exceeds attention kernel limits\n"); + return false; + } + + const CachePair k_data = + make_cache_pair( + options.k_type, kv_seq_lens, contiguous_context, max_blocks, + pool_tokens, block_table, 0.1f); + const CachePair v_data = + make_cache_pair( + options.v_type, kv_seq_lens, contiguous_context, max_blocks, + pool_tokens, block_table, 0.7f); + if (k_data.contiguous.empty() || k_data.paged.empty() || + v_data.contiguous.empty() || v_data.paged.empty()) { + return false; + } + + std::vector paged_q( + static_cast(D) * n_seq * N_HEAD); + std::vector contiguous_q(paged_q.size()); + for (int seq = 0; seq < n_seq; ++seq) { + for (int head = 0; head < N_HEAD; ++head) { + for (int d = 0; d < D; ++d) { + const size_t canonical_index = + (static_cast(seq) * N_HEAD + head) * D + d; + const float value = + std::cos( + static_cast(canonical_index) * 0.013f) * + 0.20f; + contiguous_q[canonical_index] = value; + paged_q[ + (static_cast(head) * n_seq + seq) * D + d] = + value; + } + } + } + + ggml_init_params params{}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ContextGuard ctx{ggml_init(params)}; + if (!ctx.value) return false; + + ggml_tensor * q_paged = + ggml_new_tensor_3d( + ctx.value, GGML_TYPE_F32, D, n_seq, N_HEAD); + ggml_tensor * k_paged = + ggml_new_tensor_3d( + ctx.value, options.k_type, D, pool_tokens, N_HEAD_KV); + ggml_tensor * v_paged = + ggml_new_tensor_3d( + ctx.value, options.v_type, D, pool_tokens, N_HEAD_KV); + ggml_tensor * table = + ggml_new_tensor_2d( + ctx.value, GGML_TYPE_I32, max_blocks, n_seq); + ggml_tensor * kv_seq_lens_tensor = + ggml_new_tensor_1d(ctx.value, GGML_TYPE_I32, n_seq); + + ggml_tensor * q_contiguous = + ggml_new_tensor_4d( + ctx.value, GGML_TYPE_F32, D, 1, N_HEAD, n_seq); + ggml_tensor * k_contiguous = + ggml_new_tensor_4d( + ctx.value, options.k_type, D, contiguous_context, + N_HEAD_KV, n_seq); + ggml_tensor * v_contiguous = + ggml_new_tensor_4d( + ctx.value, options.v_type, D, contiguous_context, + N_HEAD_KV, n_seq); + ggml_tensor * padding_mask = + use_padding_mask + ? ggml_new_tensor_4d( + ctx.value, GGML_TYPE_F16, contiguous_context, 1, 1, + n_seq) + : nullptr; + + for (ggml_tensor * input : { + q_paged, k_paged, v_paged, table, kv_seq_lens_tensor, + q_contiguous, k_contiguous, v_contiguous}) { + ggml_set_input(input); + } + if (padding_mask) ggml_set_input(padding_mask); + + ggml_tensor * paged_output = ggml_paged_attn( + ctx.value, q_paged, k_paged, v_paged, table, kv_seq_lens_tensor, + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE); + ggml_tensor * contiguous_output = ggml_flash_attn_ext( + ctx.value, q_contiguous, k_contiguous, v_contiguous, + padding_mask, 1.0f / std::sqrt(static_cast(D)), + 0.0f, 0.0f); + ggml_set_output(paged_output); + ggml_set_output(contiguous_output); + + if (!ggml_backend_supports_op(backend, paged_output) || + !ggml_backend_supports_op(backend, contiguous_output)) { + std::fprintf( + stderr, + "GPU backend does not support K=%s V=%s for both layouts\n", + ggml_type_name(options.k_type), + ggml_type_name(options.v_type)); + return false; + } + + ggml_cgraph * paged_graph = ggml_new_graph(ctx.value); + ggml_build_forward_expand(paged_graph, paged_output); + ggml_cgraph * contiguous_graph = ggml_new_graph(ctx.value); + ggml_build_forward_expand(contiguous_graph, contiguous_output); + ggml_cgraph * allocation_graph = ggml_new_graph(ctx.value); + ggml_build_forward_expand(allocation_graph, paged_output); + ggml_build_forward_expand(allocation_graph, contiguous_output); + + AllocatorGuard allocator{ + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend))}; + if (!allocator.value || + !ggml_gallocr_alloc_graph(allocator.value, allocation_graph)) { + return false; + } + + ggml_backend_tensor_set( + q_paged, paged_q.data(), 0, + paged_q.size() * sizeof(paged_q[0])); + ggml_backend_tensor_set( + k_paged, k_data.paged.data(), 0, k_data.paged.size()); + ggml_backend_tensor_set( + v_paged, v_data.paged.data(), 0, v_data.paged.size()); + ggml_backend_tensor_set( + table, block_table.data(), 0, + block_table.size() * sizeof(block_table[0])); + ggml_backend_tensor_set( + kv_seq_lens_tensor, kv_seq_lens.data(), 0, + kv_seq_lens.size() * sizeof(kv_seq_lens[0])); + ggml_backend_tensor_set( + q_contiguous, contiguous_q.data(), 0, + contiguous_q.size() * sizeof(contiguous_q[0])); + ggml_backend_tensor_set( + k_contiguous, k_data.contiguous.data(), 0, + k_data.contiguous.size()); + ggml_backend_tensor_set( + v_contiguous, v_data.contiguous.data(), 0, + v_data.contiguous.size()); + if (padding_mask) { + const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f); + const ggml_fp16_t negative_infinity = + ggml_fp32_to_fp16(-INFINITY); + std::vector mask_data( + static_cast(contiguous_context) * n_seq, + negative_infinity); + for (int seq = 0; seq < n_seq; ++seq) { + std::fill_n( + mask_data.begin() + + static_cast(seq) * contiguous_context, + kv_seq_lens[seq], zero); + } + ggml_backend_tensor_set( + padding_mask, mask_data.data(), 0, + mask_data.size() * sizeof(mask_data[0])); + } + ggml_backend_synchronize(backend); + + if (ggml_backend_graph_compute(backend, paged_graph) != + GGML_STATUS_SUCCESS || + ggml_backend_graph_compute(backend, contiguous_graph) != + GGML_STATUS_SUCCESS || + !compare_outputs( + paged_output, contiguous_output, n_seq, result)) { + std::fprintf( + stderr, + "layout parity failed at n_seq=%d: max_abs=%.6g " + "max_tolerance_ratio=%.6g rmse=%.6g\n", + n_seq, result.max_abs_error, + result.max_tolerance_ratio, result.rmse); + return false; + } + + int paged_iterations = options.iterations; + int contiguous_iterations = options.iterations; + if (options.iterations == 0) { + double paged_seconds_per_step = 0.0; + double contiguous_seconds_per_step = 0.0; + if (!warm_graph(backend, paged_graph, options.warmup) || + !probe_graph( + backend, paged_graph, paged_seconds_per_step) || + !warm_graph(backend, contiguous_graph, options.warmup) || + !probe_graph( + backend, contiguous_graph, + contiguous_seconds_per_step)) { + return false; + } + paged_iterations = std::clamp( + static_cast( + std::ceil( + TARGET_SAMPLE_SECONDS / + paged_seconds_per_step)), + 1, 100000); + contiguous_iterations = std::clamp( + static_cast( + std::ceil( + TARGET_SAMPLE_SECONDS / + contiguous_seconds_per_step)), + 1, 100000); + } + + // Keep each layout's sample windows together so ggml's CUDA-graph cache + // remains stable. Alternate the order across cases to reduce clock and + // thermal bias without repeatedly invalidating graph capture. + bool timing_ok = false; + if (contiguous_first) { + timing_ok = + collect_timing( + backend, contiguous_graph, options, + contiguous_iterations, + n_seq, result.contiguous) && + collect_timing( + backend, paged_graph, options, paged_iterations, + n_seq, result.paged); + } else { + timing_ok = + collect_timing( + backend, paged_graph, options, paged_iterations, + n_seq, result.paged) && + collect_timing( + backend, contiguous_graph, options, + contiguous_iterations, + n_seq, result.contiguous); + } + if (!timing_ok) return false; + + result.n_seq = n_seq; + result.padded_context = contiguous_context; + result.live_tokens = 0; + for (int context : kv_seq_lens) { + result.live_tokens += context; + } + result.padded_tokens = + static_cast(contiguous_context) * n_seq; + result.paged_pool_tokens = pool_tokens; + result.contiguous_kv_bytes = + ggml_nbytes(k_contiguous) + ggml_nbytes(v_contiguous); + result.paged_kv_bytes = + ggml_nbytes(k_paged) + ggml_nbytes(v_paged); + result.paged_metadata_bytes = + ggml_nbytes(table) + ggml_nbytes(kv_seq_lens_tensor); + result.contiguous_kv_mib = + static_cast(result.contiguous_kv_bytes) / + (1024.0 * 1024.0); + result.paged_kv_mib = + static_cast(result.paged_kv_bytes) / + (1024.0 * 1024.0); + result.paged_metadata_kib = + static_cast(result.paged_metadata_bytes) / 1024.0; + return true; +} + +bool padded_context_for( + const std::vector & contexts, + int & padded_context) { + if (contexts.empty()) return false; + const int max_context = + *std::max_element(contexts.begin(), contexts.end()); + const int64_t aligned = + ((static_cast(max_context) + + CONTIGUOUS_CONTEXT_ALIGNMENT - 1) / + CONTIGUOUS_CONTEXT_ALIGNMENT) * + CONTIGUOUS_CONTEXT_ALIGNMENT; + if (aligned <= 0 || aligned > INT32_MAX) return false; + padded_context = static_cast(aligned); + return true; +} + +std::string join_contexts(const std::vector & contexts) { + std::string joined; + for (size_t i = 0; i < contexts.size(); ++i) { + if (i != 0) joined.push_back(','); + joined += std::to_string(contexts[i]); + } + return joined; +} + +void print_ragged_result( + const Options & options, + const BenchResult & result) { + const bool cuda_graphs_enabled = + std::getenv("GGML_CUDA_DISABLE_GRAPHS") == nullptr; + const std::string contexts = + join_contexts(options.ragged_contexts); + const double step_time_ratio = + result.paged.median_us / result.contiguous.median_us; + const double throughput_ratio = + result.paged.aggregate_queries_s / + result.contiguous.aggregate_queries_s; + const double memory_ratio = + static_cast(result.contiguous_kv_bytes) / + static_cast(result.paged_kv_bytes); + const double memory_saving = + (1.0 - + static_cast(result.paged_kv_bytes) / + static_cast(result.contiguous_kv_bytes)) * + 100.0; + const uint64_t projected_contiguous_bytes = + result.contiguous_kv_bytes * QWEN_FULL_ATTENTION_LAYERS; + const uint64_t projected_paged_bytes = + result.paged_kv_bytes * QWEN_FULL_ATTENTION_LAYERS; + + std::printf( + "# padded per-sequence contiguous vs paged GPU attention; " + "ragged decode, not whole-model throughput\n" + "# D=%d,Hq=%d,Hkv=%d,block=%d,contexts=%s," + "contiguous_context_pad=%d,k=%s,v=%s,warmup=%d,samples=%d," + "cuda_graphs_env=%s\n" + "# the Qwen projection multiplies K/V storage by %d full-attention " + "layers; recurrent state and all other model memory are excluded\n" + "# paged_step_time_overhead_pct > 0 means paged is slower; " + "paged_over_padded_contiguous_query_throughput > 1 means paged " + "is faster\n", + D, N_HEAD, N_HEAD_KV, BLOCK_SIZE, contexts.c_str(), + result.padded_context, + ggml_type_name(options.k_type), + ggml_type_name(options.v_type), + options.warmup, options.samples, + cuda_graphs_enabled ? "allowed" : "disabled", + QWEN_FULL_ATTENTION_LAYERS); + std::printf( + "case,n_seq,contexts,live_tokens," + "padded_contiguous_tokens,paged_pool_tokens," + "padded_contiguous_iterations,paged_iterations," + "padded_contiguous_median_window_mean_us_per_step," + "paged_median_window_mean_us_per_step," + "padded_contiguous_p95_window_mean_us_per_step," + "paged_p95_window_mean_us_per_step," + "paged_over_padded_contiguous_step_time," + "paged_step_time_overhead_pct," + "padded_contiguous_aggregate_attention_queries_s," + "paged_aggregate_attention_queries_s," + "paged_over_padded_contiguous_query_throughput," + "padded_contiguous_kv_bytes_one_layer," + "paged_kv_bytes_one_layer," + "padded_contiguous_kv_mib_one_layer," + "paged_kv_mib_one_layer," + "padded_contiguous_kv_bytes_qwen16," + "paged_kv_bytes_qwen16," + "padded_contiguous_kv_gib_qwen16," + "paged_kv_gib_qwen16," + "padded_contiguous_over_paged_kv_ratio," + "paged_kv_saving_pct,paged_metadata_bytes," + "max_abs_error,max_tolerance_ratio,rmse\n"); + std::printf( + "ragged,%d,\"%s\",%lld,%lld,%lld,%d,%d," + "%.3f,%.3f,%.3f,%.3f,%.6f,%.3f,%.2f,%.2f,%.6f," + "%llu,%llu,%.2f,%.2f,%llu,%llu,%.3f,%.3f,%.6f,%.3f,%llu," + "%.8g,%.8g,%.8g\n", + result.n_seq, contexts.c_str(), + static_cast(result.live_tokens), + static_cast(result.padded_tokens), + static_cast(result.paged_pool_tokens), + result.contiguous.iterations, result.paged.iterations, + result.contiguous.median_us, result.paged.median_us, + result.contiguous.p95_us, result.paged.p95_us, + step_time_ratio, (step_time_ratio - 1.0) * 100.0, + result.contiguous.aggregate_queries_s, + result.paged.aggregate_queries_s, throughput_ratio, + static_cast(result.contiguous_kv_bytes), + static_cast(result.paged_kv_bytes), + result.contiguous_kv_mib, result.paged_kv_mib, + static_cast(projected_contiguous_bytes), + static_cast(projected_paged_bytes), + static_cast(projected_contiguous_bytes) / + (1024.0 * 1024.0 * 1024.0), + static_cast(projected_paged_bytes) / + (1024.0 * 1024.0 * 1024.0), + memory_ratio, memory_saving, + static_cast(result.paged_metadata_bytes), + result.max_abs_error, result.max_tolerance_ratio, result.rmse); +} + +} // namespace + +int main(int argc, char ** argv) { + Options options; + if (!parse_options(argc, argv, options)) { + print_usage(argv[0]); + return 2; + } + + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "GPU backend unavailable\n"); + return 1; + } + + std::vector results; + bool ok = true; + try { + if (!options.ragged_contexts.empty()) { + int padded_context = 0; + BenchResult result; + if (!padded_context_for( + options.ragged_contexts, padded_context) || + !run_case( + backend, options, options.ragged_contexts, + padded_context, true, true, result)) { + std::fprintf( + stderr, "ragged attention benchmark failed\n"); + ok = false; + } else { + results.push_back(result); + } + } else { + const std::array sequence_counts = {1, 2, 4, 8}; + results.reserve(sequence_counts.size()); + for (size_t case_index = 0; + case_index < sequence_counts.size(); + ++case_index) { + const int n_seq = sequence_counts[case_index]; + const std::vector contexts( + n_seq, options.context); + BenchResult result; + if (!run_case( + backend, options, contexts, options.context, + false, case_index % 2 == 0, result)) { + std::fprintf( + stderr, + "attention benchmark failed at n_seq=%d\n", + n_seq); + ok = false; + break; + } + results.push_back(result); + } + } + } catch (const std::exception & error) { + std::fprintf(stderr, "benchmark setup failed: %s\n", error.what()); + ok = false; + } + + if (ok) { + if (!options.ragged_contexts.empty()) { + print_ragged_result(options, results.front()); + ggml_backend_free(backend); + return 0; + } + + const bool cuda_graphs_enabled = + std::getenv("GGML_CUDA_DISABLE_GRAPHS") == nullptr; + std::printf( + "# native contiguous vs paged GPU attention; " + "not whole-model throughput\n" + "# D=%d,Hq=%d,Hkv=%d,block=%d,context=%d,k=%s,v=%s," + "warmup=%d,samples=%d,cuda_graphs_env=%s\n" + "# paged_step_time_overhead_pct > 0 means paged is slower; " + "paged_over_contiguous_query_throughput > 1 means paged " + "is faster\n", + D, N_HEAD, N_HEAD_KV, BLOCK_SIZE, options.context, + ggml_type_name(options.k_type), + ggml_type_name(options.v_type), + options.warmup, options.samples, + cuda_graphs_enabled ? "allowed" : "disabled"); + std::printf( + "n_seq,context,k_type,v_type," + "contiguous_iterations,paged_iterations," + "contiguous_median_window_mean_us_per_step," + "paged_median_window_mean_us_per_step," + "contiguous_p95_window_mean_us_per_step," + "paged_p95_window_mean_us_per_step," + "paged_over_contiguous_step_time," + "paged_step_time_overhead_pct," + "contiguous_aggregate_attention_queries_s," + "paged_aggregate_attention_queries_s," + "paged_over_contiguous_query_throughput," + "contiguous_scaling_vs_b1," + "paged_scaling_vs_b1," + "contiguous_batch_efficiency," + "paged_batch_efficiency," + "contiguous_kv_mib,paged_kv_mib," + "paged_metadata_kib," + "max_abs_error,max_tolerance_ratio,rmse\n"); + + const double contiguous_baseline = + results.front().contiguous.aggregate_queries_s; + const double paged_baseline = + results.front().paged.aggregate_queries_s; + for (const BenchResult & result : results) { + const double contiguous_scaling = + result.contiguous.aggregate_queries_s / + contiguous_baseline; + const double paged_scaling = + result.paged.aggregate_queries_s / paged_baseline; + const double step_time_ratio = + result.paged.median_us / + result.contiguous.median_us; + const double throughput_ratio = + result.paged.aggregate_queries_s / + result.contiguous.aggregate_queries_s; + std::printf( + "%d,%d,%s,%s,%d,%d," + "%.3f,%.3f,%.3f,%.3f," + "%.6f,%.3f,%.2f,%.2f,%.6f," + "%.4f,%.4f,%.4f,%.4f," + "%.2f,%.2f,%.3f," + "%.8g,%.8g,%.8g\n", + result.n_seq, options.context, + ggml_type_name(options.k_type), + ggml_type_name(options.v_type), + result.contiguous.iterations, + result.paged.iterations, + result.contiguous.median_us, + result.paged.median_us, + result.contiguous.p95_us, + result.paged.p95_us, + step_time_ratio, + (step_time_ratio - 1.0) * 100.0, + result.contiguous.aggregate_queries_s, + result.paged.aggregate_queries_s, + throughput_ratio, + contiguous_scaling, + paged_scaling, + contiguous_scaling / result.n_seq, + paged_scaling / result.n_seq, + result.contiguous_kv_mib, + result.paged_kv_mib, + result.paged_metadata_kib, + result.max_abs_error, + result.max_tolerance_ratio, + result.rmse); + } + } + + ggml_backend_free(backend); + return ok ? 0 : 1; +} diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp new file mode 100644 index 000000000..ca9749b5a --- /dev/null +++ b/server/test/test_paged_attention.cpp @@ -0,0 +1,257 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int D = 256; +constexpr int N_HEAD = 24; +constexpr int N_HEAD_KV = 4; +constexpr int BLOCK_SIZE = 16; +constexpr int MAX_BLOCKS = 3; +constexpr int N_SEQ = 6; +constexpr int PHYSICAL_BLOCKS = MAX_BLOCKS * N_SEQ; +constexpr int POOL_TOKENS = PHYSICAL_BLOCKS * BLOCK_SIZE; + +const std::array KV_SEQ_LENS = {1, 15, 16, 17, 31, 33}; + +// Rows are sequences. Every live block is private, and deliberately +// shuffled so a contiguous-cache implementation cannot pass this test. +// Two in-range entries are deliberately invalid — one negative, one past +// the physical pool — to pin down the kernel's bounds guard: it must skip +// those blocks (matching the oracle) instead of reading out of bounds. +const std::array BLOCK_TABLE = { + 11, 2, 15, + 4, 17, 7, + 9, 0, 13, + 5, -1, 1, // token 16 of seq 3 maps to a negative block + 14, 8, 3, + 10, 99, 12, // tokens 16-31 of seq 5 map past the pool +}; + +bool block_is_valid(int32_t block) { + return block >= 0 && block < PHYSICAL_BLOCKS; +} + +std::vector quantize_rows(ggml_type type, + const std::vector & values, + int64_t rows) { + std::vector result( + ggml_row_size(type, D) * static_cast(rows)); + const size_t written = ggml_quantize_chunk( + type, values.data(), result.data(), 0, rows, D, nullptr); + if (written != result.size()) { + std::fprintf(stderr, + "quantize size mismatch for %s: %zu != %zu\n", + ggml_type_name(type), written, result.size()); + return {}; + } + return result; +} + +std::vector dequantize_rows(ggml_type type, + const std::vector & values, + int64_t rows) { + const ggml_type_traits * traits = ggml_get_type_traits(type); + if (!traits || !traits->to_float) return {}; + + const size_t row_bytes = ggml_row_size(type, D); + std::vector result(static_cast(rows) * D); + for (int64_t row = 0; row < rows; ++row) { + traits->to_float(values.data() + row * row_bytes, + result.data() + row * D, D); + } + return result; +} + +std::vector reference_attention( + const std::vector & q, + const std::vector & k, + const std::vector & v) { + std::vector output(q.size(), 0.0f); + const float scale = 1.0f / std::sqrt(static_cast(D)); + const int q_per_kv = N_HEAD / N_HEAD_KV; + + for (int seq = 0; seq < N_SEQ; ++seq) { + const int kv_seq_len = KV_SEQ_LENS[seq]; + for (int head = 0; head < N_HEAD; ++head) { + const int kv_head = head / q_per_kv; + const float * q_row = + q.data() + (static_cast(head) * N_SEQ + seq) * D; + + std::vector scores(kv_seq_len); + float max_score = -INFINITY; + for (int token = 0; token < kv_seq_len; ++token) { + const int block = + BLOCK_TABLE[seq * MAX_BLOCKS + token / BLOCK_SIZE]; + if (!block_is_valid(block)) { + // Mirrors the kernel: invalid blocks contribute nothing. + scores[token] = -INFINITY; + continue; + } + const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + const float * k_row = + k.data() + + (static_cast(kv_head) * POOL_TOKENS + physical) * D; + float dot = 0.0f; + for (int d = 0; d < D; ++d) dot += q_row[d] * k_row[d]; + scores[token] = dot * scale; + max_score = std::max(max_score, scores[token]); + } + + float denominator = 0.0f; + for (float & score : scores) { + score = std::exp(score - max_score); + denominator += score; + } + + float * out_row = + output.data() + (static_cast(head) * N_SEQ + seq) * D; + for (int token = 0; token < kv_seq_len; ++token) { + const int block = + BLOCK_TABLE[seq * MAX_BLOCKS + token / BLOCK_SIZE]; + if (!block_is_valid(block)) continue; + const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + const float * v_row = + v.data() + + (static_cast(kv_head) * POOL_TOKENS + physical) * D; + const float probability = scores[token] / denominator; + for (int d = 0; d < D; ++d) { + out_row[d] += probability * v_row[d]; + } + } + } + } + return output; +} + +bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { + ggml_init_params params{}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * q = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_SEQ, N_HEAD); + ggml_tensor * k = + ggml_new_tensor_3d(ctx, k_type, D, POOL_TOKENS, N_HEAD_KV); + ggml_tensor * v = + ggml_new_tensor_3d(ctx, v_type, D, POOL_TOKENS, N_HEAD_KV); + ggml_tensor * table = + ggml_new_tensor_2d(ctx, GGML_TYPE_I32, MAX_BLOCKS, N_SEQ); + ggml_tensor * kv_seq_lens = + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N_SEQ); + for (ggml_tensor * input : {q, k, v, table, kv_seq_lens}) { + ggml_set_input(input); + } + + ggml_tensor * output = ggml_paged_attn( + ctx, q, k, v, table, kv_seq_lens, + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + + ggml_gallocr_t allocator = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(allocator, graph)) { + ggml_gallocr_free(allocator); + ggml_free(ctx); + return false; + } + + std::vector q_data(static_cast(D) * N_SEQ * N_HEAD); + std::vector k_source( + static_cast(D) * POOL_TOKENS * N_HEAD_KV); + std::vector v_source(k_source.size()); + for (size_t i = 0; i < q_data.size(); ++i) { + q_data[i] = std::sin(static_cast(i) * 0.013f) * 0.25f; + } + for (size_t i = 0; i < k_source.size(); ++i) { + k_source[i] = std::cos(static_cast(i) * 0.007f) * 0.20f; + v_source[i] = std::sin(static_cast(i) * 0.011f + 0.3f) * 0.30f; + } + + const int64_t cache_rows = static_cast(POOL_TOKENS) * N_HEAD_KV; + const std::vector k_data = + quantize_rows(k_type, k_source, cache_rows); + const std::vector v_data = + quantize_rows(v_type, v_source, cache_rows); + const std::vector k_reference = + dequantize_rows(k_type, k_data, cache_rows); + const std::vector v_reference = + dequantize_rows(v_type, v_data, cache_rows); + bool ok = !k_data.empty() && !v_data.empty() && + !k_reference.empty() && !v_reference.empty(); + + if (ok) { + ggml_backend_tensor_set(q, q_data.data(), 0, + q_data.size() * sizeof(q_data[0])); + ggml_backend_tensor_set(k, k_data.data(), 0, k_data.size()); + ggml_backend_tensor_set(v, v_data.data(), 0, v_data.size()); + ggml_backend_tensor_set(table, BLOCK_TABLE.data(), 0, + BLOCK_TABLE.size() * sizeof(BLOCK_TABLE[0])); + ggml_backend_tensor_set( + kv_seq_lens, KV_SEQ_LENS.data(), 0, + KV_SEQ_LENS.size() * sizeof(KV_SEQ_LENS[0])); + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + } + + float max_abs_error = INFINITY; + if (ok) { + std::vector actual(ggml_nelements(output)); + ggml_backend_tensor_get(output, actual.data(), 0, + actual.size() * sizeof(actual[0])); + const std::vector expected = + reference_attention(q_data, k_reference, v_reference); + max_abs_error = 0.0f; + for (size_t i = 0; i < actual.size(); ++i) { + if (!std::isfinite(actual[i])) { + ok = false; + break; + } + max_abs_error = + std::max(max_abs_error, std::fabs(actual[i] - expected[i])); + } + ok = ok && max_abs_error < 3.0e-3f; + } + + std::printf("paged attention K=%-4s V=%-4s max_abs=%.6g %s\n", + ggml_type_name(k_type), ggml_type_name(v_type), + max_abs_error, ok ? "PASS" : "FAIL"); + ggml_gallocr_free(allocator); + ggml_free(ctx); + return ok; +} + +} // namespace + +int main() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "GPU backend unavailable\n"); + return 1; + } + + const ggml_type types[] = { + GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, + }; + bool ok = true; + for (ggml_type k_type : types) { + for (ggml_type v_type : types) { + ok = run_case(backend, k_type, v_type) && ok; + } + } + + ggml_backend_free(backend); + return ok ? 0 : 1; +} diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp new file mode 100644 index 000000000..b5f027222 --- /dev/null +++ b/server/test/test_paged_kv_pool.cpp @@ -0,0 +1,291 @@ +#include "../src/common/paged_kv_pool.h" +#include "../src/common/paged_attention_config.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static int failures = 0; +static int checks = 0; + +#define CHECK(expr) do { \ + ++checks; \ + if (!(expr)) { \ + ++failures; \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #expr); \ + } \ +} while (0) + +static PagedKvSequenceHandle acquire(PagedKvPool & pool, + PagedKvRequestId request_id) { + PagedKvSequenceHandle handle; + CHECK(pool.acquire(request_id, handle) == PagedKvStatus::Ok); + return handle; +} + +static PagedKvSequenceSnapshot sequence(PagedKvPool & pool, + PagedKvSequenceHandle handle) { + PagedKvSequenceSnapshot snapshot; + CHECK(pool.sequence(handle, snapshot) == PagedKvStatus::Ok); + return snapshot; +} + +static bool equals(const std::vector & actual, + std::initializer_list expected) { + return actual == std::vector(expected); +} + +static void test_block_boundaries() { + const uint32_t lengths[] = {1, 15, 16, 17, 31, 32, 33}; + for (uint32_t length : lengths) { + PagedKvPool pool(8, 2); + const auto handle = acquire(pool, 1000 + length); + const auto append = pool.append(handle, length); + CHECK(append.status == PagedKvStatus::Ok); + CHECK(append.write_slots.size() == length); + + const auto snapshot = sequence(pool, handle); + const uint32_t expected_blocks = (length + 15) / 16; + CHECK(snapshot.kv_seq_len == length); + CHECK(snapshot.block_table.size() == expected_blocks); + CHECK(pool.free_block_count() == 8 - expected_blocks); + + for (uint32_t i = 0; i < length; ++i) { + const auto & slot = append.write_slots[i]; + CHECK(slot.logical_position == i); + CHECK(slot.physical_block == i / 16); + CHECK(slot.block_offset == i % 16); + CHECK(slot.physical_token_index == i); + } + } +} + +static void test_reserve_then_append() { + PagedKvPool pool(4, 2); + const auto handle = acquire(pool, 7); + + CHECK(pool.reserve(handle, 17) == PagedKvStatus::Ok); + auto snapshot = sequence(pool, handle); + CHECK(snapshot.kv_seq_len == 0); + CHECK(equals(snapshot.block_table, {0, 1})); + CHECK(pool.free_block_count() == 2); + + const auto append = pool.append(handle, 17); + CHECK(append.status == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 2); + snapshot = sequence(pool, handle); + CHECK(snapshot.kv_seq_len == 17); + CHECK(equals(snapshot.block_table, {0, 1})); + + CHECK(pool.reserve(handle, 16) == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, handle).block_table, {0, 1})); +} + +static void test_compact_append() { + PagedKvPool pool(8, 1); + const auto handle = acquire(pool, 123); + + const auto prompt = pool.append_compact(handle, 33); + CHECK(prompt.status == PagedKvStatus::Ok); + CHECK(prompt.token_count == 33); + CHECK(prompt.first.logical_position == 0); + CHECK(prompt.first.physical_token_index == 0); + CHECK(prompt.last.logical_position == 32); + CHECK(prompt.last.physical_block == 2); + CHECK(prompt.last.block_offset == 0); + CHECK(sequence(pool, handle).kv_seq_len == 33); + + const auto token = pool.append_compact(handle, 1); + CHECK(token.status == PagedKvStatus::Ok); + CHECK(token.token_count == 1); + CHECK(token.first.logical_position == 33); + CHECK(token.first.physical_token_index == 33); + CHECK(token.last.logical_position == token.first.logical_position); + + const auto empty = pool.append_compact(handle, 0); + CHECK(empty.status == PagedKvStatus::Ok); + CHECK(empty.token_count == 0); + CHECK(sequence(pool, handle).kv_seq_len == 34); +} + +static void test_noncontiguous_reuse_and_isolation() { + PagedKvPool pool(6, 3); + const auto first = acquire(pool, 101); + const auto second = acquire(pool, 202); + + CHECK(pool.append(first, 17)); + CHECK(pool.append(second, 17)); + CHECK(equals(sequence(pool, first).block_table, {0, 1})); + CHECK(equals(sequence(pool, second).block_table, {2, 3})); + + CHECK(pool.release(first) == PagedKvStatus::Ok); + const auto second_more = pool.append(second, 16); + CHECK(second_more.status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, second).block_table, {2, 3, 0})); + CHECK(second_more.write_slots.back().physical_block == 0); + CHECK(second_more.write_slots.back().block_offset == 0); + + const auto third = acquire(pool, 303); + CHECK(pool.append(third, 17)); + CHECK(equals(sequence(pool, third).block_table, {1, 4})); + CHECK(equals(sequence(pool, second).block_table, {2, 3, 0})); + + CHECK(pool.cancel(303) == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, second).block_table, {2, 3, 0})); + CHECK(pool.free_block_count() == 3); +} + +static void test_exhaustion_rolls_back() { + PagedKvPool pool(3, 2); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pool.append(first, 17)); + CHECK(pool.append(second, 1)); + CHECK(pool.free_block_count() == 0); + + const auto before = sequence(pool, first); + const auto failed_append = pool.append(first, 16); + CHECK(failed_append.status == PagedKvStatus::BlocksExhausted); + CHECK(failed_append.write_slots.empty()); + const auto after_append = sequence(pool, first); + CHECK(after_append.kv_seq_len == before.kv_seq_len); + CHECK(after_append.block_table == before.block_table); + CHECK(pool.free_block_count() == 0); + + CHECK(pool.reserve(first, 33) == PagedKvStatus::BlocksExhausted); + const auto after_reserve = sequence(pool, first); + CHECK(after_reserve.kv_seq_len == before.kv_seq_len); + CHECK(after_reserve.block_table == before.block_table); + CHECK(pool.free_block_count() == 0); + + const auto fits_existing_blocks = pool.append(first, 15); + CHECK(fits_existing_blocks.status == PagedKvStatus::Ok); + CHECK(sequence(pool, first).kv_seq_len == 32); + + PagedKvSequenceHandle unchanged{77, 88}; + CHECK(pool.acquire(3, unchanged) == + PagedKvStatus::SequenceSlotsExhausted); + CHECK(unchanged == (PagedKvSequenceHandle{77, 88})); + CHECK(pool.active_sequence_count() == 2); +} + +static void test_request_identity_and_stale_handles() { + PagedKvPool pool(4, 2); + const auto old_handle = acquire(pool, 9001); + const auto other_handle = acquire(pool, 42); + CHECK(old_handle.slot == 0); + CHECK(other_handle.slot == 1); + + PagedKvSequenceHandle found; + CHECK(pool.lookup(9001, found) == PagedKvStatus::Ok); + CHECK(found == old_handle); + CHECK(pool.release(old_handle) == PagedKvStatus::Ok); + CHECK(pool.lookup(9001, found) == PagedKvStatus::RequestNotFound); + + const auto replacement = acquire(pool, 123456); + CHECK(replacement.slot == old_handle.slot); + CHECK(replacement.generation != old_handle.generation); + CHECK(pool.append(old_handle, 1).status == PagedKvStatus::StaleHandle); + CHECK(pool.reserve(old_handle, 16) == PagedKvStatus::StaleHandle); + CHECK(pool.release(old_handle) == PagedKvStatus::StaleHandle); + PagedKvSequenceSnapshot stale_snapshot; + CHECK(pool.sequence(old_handle, stale_snapshot) == + PagedKvStatus::StaleHandle); + CHECK(pool.cancel(9001) == PagedKvStatus::RequestNotFound); + + CHECK(pool.append(replacement, 1)); + pool.reset(); + CHECK(pool.active_sequence_count() == 0); + CHECK(pool.free_block_count() == 4); + CHECK(pool.append(replacement, 1).status == PagedKvStatus::StaleHandle); + + const auto after_reset = acquire(pool, 555); + CHECK(after_reset.slot == 0); + CHECK(after_reset.generation != replacement.generation); + CHECK(pool.append(after_reset, 1)); + CHECK(sequence(pool, after_reset).block_table.front() == 0); +} + +static void test_metadata_snapshot() { + PagedKvPool pool(7, 4); + const auto first = acquire(pool, 11); + const auto second = acquire(pool, 22); + CHECK(pool.append(first, 17)); + CHECK(pool.append(second, 1)); + CHECK(pool.reserve(second, 33) == PagedKvStatus::Ok); + + const auto metadata = pool.metadata_snapshot(); + CHECK(metadata.block_size == 16); + CHECK(metadata.physical_block_count == 7); + CHECK(metadata.sequences.size() == 2); + CHECK(equals(metadata.block_table, {0, 1, 2, 3, 4})); + + const auto & first_meta = metadata.sequences[0]; + CHECK(first_meta.request_id == 11); + CHECK(first_meta.generation == first.generation); + CHECK(first_meta.sequence_slot == first.slot); + CHECK(first_meta.kv_seq_len == 17); + CHECK(first_meta.block_table_offset == 0); + CHECK(first_meta.block_count == 2); + + const auto & second_meta = metadata.sequences[1]; + CHECK(second_meta.request_id == 22); + CHECK(second_meta.generation == second.generation); + CHECK(second_meta.sequence_slot == second.slot); + CHECK(second_meta.kv_seq_len == 1); + CHECK(second_meta.block_table_offset == 2); + CHECK(second_meta.block_count == 3); +} + +static void test_duplicate_request_is_transactional() { + PagedKvPool pool(2, 2); + const auto first = acquire(pool, 88); + PagedKvSequenceHandle output{9, 10}; + CHECK(pool.acquire(88, output) == PagedKvStatus::DuplicateRequest); + CHECK(output == (PagedKvSequenceHandle{9, 10})); + CHECK(pool.active_sequence_count() == 1); + CHECK(pool.lookup(88, output) == PagedKvStatus::Ok); + CHECK(output == first); +} + +static void test_paged_attention_compatibility() { + PagedAttentionOptions options; + options.enabled = true; + options.architecture = "qwen35"; + options.max_ctx = 131072; + CHECK(validate_paged_attention_options(options).empty()); + + options.pflash = true; + CHECK(!validate_paged_attention_options(options).empty()); + options.pflash = false; + options.kvflash = true; + CHECK(!validate_paged_attention_options(options).empty()); + options.kvflash = false; + options.fa_window = 4096; + CHECK(!validate_paged_attention_options(options).empty()); + options.fa_window = 0; + options.max_ctx = 0; + CHECK(!validate_paged_attention_options(options).empty()); +} + +int main() { + test_block_boundaries(); + test_reserve_then_append(); + test_compact_append(); + test_noncontiguous_reuse_and_isolation(); + test_exhaustion_rolls_back(); + test_request_identity_and_stale_handles(); + test_metadata_snapshot(); + test_duplicate_request_is_transactional(); + test_paged_attention_compatibility(); + + if (failures != 0) { + std::fprintf(stderr, "%d checks, %d failures\n", checks, failures); + return 1; + } + std::printf("OK test_paged_kv_pool (%d checks)\n", checks); + return 0; +} From 8c809f579b1ee3d071fac45f5e8ea6557e79c6ce Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 23 Jul 2026 11:26:23 +0000 Subject: [PATCH 2/3] perf(qwen35): parallelize paged attention Rework the decode kernel: one warp now computes all GQA heads of a K/V group per K/V row loaded (batched q8_1 dots), scores are tiled per 32 tokens with one skip-if-unchanged accumulator rescale, and softmax uses exp2f with log2(e) folded into the Q prescale. Q is quantized once per step by a small pre-pass instead of per (warp x partition), and partition partials are stored normalized in f16, halving scratch traffic. The block table and sequence length move into the persistent target cache next to the K/V pool, so decode steps upload four bytes per fresh block instead of the whole live table every token. On an RTX 3090 (q4_0 K/V) vs the padded-contiguous fattn decode: uniform 4K-32K batches go from +30..104% overhead to 10-22% faster, 128K to 21-51% faster, ragged 128K+7x8K to 29% faster. Short ragged batches stay behind until the stream-K rework staged with continuous batching. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QwBCUqGbQV4LSm6qUkZS3H Claude-Session: https://claude.ai/code/session_01Vkkox5Yo66K2Rn1oPuVoX7 --- optimizations/paged_attention/README.md | 90 +- server/CMakeLists.txt | 4 + server/deps/llama.cpp/ggml/include/ggml.h | 5 +- .../llama.cpp/ggml/src/ggml-backend-meta.cpp | 22 +- .../ggml/src/ggml-cuda/paged-attn.cu | 877 +++++++++++++++--- server/deps/llama.cpp/ggml/src/ggml.c | 6 +- server/src/common/paged_kv_pool.cpp | 30 +- server/src/common/paged_kv_pool.h | 47 +- server/src/internal.h | 9 +- server/src/qwen35/graph_builders.cpp | 17 +- server/src/qwen35/qwen35_backend.cpp | 35 +- server/src/qwen35/qwen35_target_graph.cpp | 40 +- server/test/bench_paged_attention.cpp | 485 +++++----- server/test/test_kvflash_pool_sizing.cpp | 7 +- server/test/test_paged_attention.cpp | 187 ++-- server/test/test_paged_kv_pool.cpp | 95 +- 16 files changed, 1439 insertions(+), 517 deletions(-) diff --git a/optimizations/paged_attention/README.md b/optimizations/paged_attention/README.md index 501af73ce..4cef7a05c 100644 --- a/optimizations/paged_attention/README.md +++ b/optimizations/paged_attention/README.md @@ -88,7 +88,8 @@ these unsupported combinations explicit. `bench_paged_attention` compares one native batched contiguous-attention operation with one paged-attention operation at `n_seq=1,2,4,8`, using Qwen's -D=256, Hq=24, and Hkv=4 dimensions: +D=256, Hq=24, and Hkv=4 dimensions. Its default run also includes one +representative ragged eight-request capacity case: The paged kernel, Qwen integration, numerical test, and benchmark build on both CUDA and HIP. HIP numerical coverage has been validated on gfx1151. @@ -103,15 +104,57 @@ cmake --build "$BUILD_DIR" --target bench_paged_attention Both layouts receive the same logical queries and identical quantized K/V rows. The paged inputs use private, interleaved physical blocks for every -sequence, and an untimed output comparison must pass before timing starts. -Each sample window queues repeated graph executions and synchronizes once. - -The CSV reports median and p95 window-averaged step time, paged overhead -relative to contiguous attention, aggregate attention-query throughput, -scaling efficiency, K/V memory, block-table metadata, and output error. -Automatic calibration gives each layout an approximately equal-duration -sample window. All sequences have the same context length, so this is a direct -kernel/layout comparison. +sequence. Before timing, both GPU outputs are compared with a double-precision +CPU oracle, and neither path may exceed `5e-3` maximum absolute error. The +focused paged-attention numerical test uses a tighter `5e-4` bound. Direct +layout error remains a diagnostic because different GPU reduction trees can +diverge at long context. Each sample window queues repeated graph executions +and synchronizes once. + +CTest runs the numerical test twice: the default case covers partitioned +long-context attention plus negative/over-capacity length clamps, while +`paged_attention_direct` forces one partition to cover the scratch-free +normalization and empty-sequence path independently of device occupancy. + +The CSV reports median step time, paged overhead, aggregate attention-query +throughput, exact K/V and paging-metadata bytes, memory saving, layout error, +and both paths' CPU-oracle error. Automatic calibration gives each layout an +approximately equal-duration sample window. + +The default run emits uniform `n_seq=1,2,4,8` rows for direct kernel/layout +comparison, followed by a ragged eight-request capacity row with one +`--context` request and seven requests at `--context / 16`. At the default 4K +context this is `4096 + 7 x 256`; it has the same 16:1 long/short mix and +82.03% K/V saving as the documented `128K + 7 x 8K` profile while remaining +quick to run. + +The CUDA decode kernel groups the six Qwen query heads that share each K/V +head in one thread block. It also divides long contexts across independently +computed partitions and merges their softmax statistics with a stable +max/sum reduction. The partition count is occupancy-aware for ordinary +contexts and grows with cache capacity for long contexts, without reading a +device context length on the host or changing CUDA graph topology. + +The host-side launch bound the graph carries is padded to the same 256-token +stride the dense path uses for its flash-attention span. The ggml-cuda graph +cache compares whole tensors, op params included, so an exact per-step kv +length would invalidate the cached graph on every token; padding keeps the +node byte-identical within each 256-token window and leaves re-capture on the +same cadence as the contiguous path. Exact per-sequence extents still come +from `kv_seq_lens` on device, so the padding only over-sizes the partition +grid — partitions past the real length exit with a zero-weight sentinel. + +The following reference result was measured on an RTX 3090 with CUDA graphs, +Q4_0 K/V, a 4096-token context, 10 warmups, and 7 automatically calibrated +samples. Times are median microseconds per attention step; the ratio is +`paged / contiguous`, so values above 1 are paged-kernel overhead. + +| Sequences | Contiguous (us) | Paged (us) | Ratio | +|----------:|----------------:|-----------:|------:| +| 1 | 81.549 | 106.759 | 1.31x | +| 2 | 142.287 | 209.283 | 1.47x | +| 4 | 254.292 | 393.838 | 1.55x | +| 8 | 486.698 | 733.469 | 1.51x | To expose the allocation trade-off separately, pass a comma-separated set of live context lengths: @@ -128,23 +171,32 @@ specifically a **padded per-sequence contiguous** layout: every sequence owns `align256(max(contexts))` K/V slots and an F16 mask hides its unused suffix. The paged layout allocates only `sum(align16(context_i))` slots and interleaves the sequences' unique physical blocks in that compact pool. Both paths receive -identical logical Q/K/V values, and output parity is mandatory before either -path is timed. The CUDA contiguous kernel derives each live prefix from the -mask and skips fully masked tail tiles, so the timing comparison does not -charge it for all padded tokens even though its K/V allocation remains padded. +identical logical Q/K/V values. As in the uniform benchmark, both paths must +pass the benchmark's CPU-oracle ceiling before timing; direct parity remains +diagnostic. The CUDA contiguous kernel derives each live prefix from the mask +and skips fully masked tail tiles, so the timing comparison does not charge it +for all padded tokens even though its K/V allocation remains padded. The ragged CSV reports live, contiguous-padded, and paged-pool token counts; -kernel timings and output error; exact one-layer K/V bytes; and a K/V-only -projection across Qwen3.6's 16 full-attention layers. The projection excludes -model weights, activations, recurrent-layer state, allocator workspace, and -other runtime memory, so it is a capacity comparison rather than a claim that -one layout will or will not fit on a particular GPU. +kernel timings and output error; and exact one-layer K/V and paging-metadata +bytes. The Qwen3.6 projection below multiplies K/V storage by its 16 +full-attention layers and excludes model weights, activations, recurrent-layer +state, allocator workspace, and other runtime memory, so it is a capacity +comparison rather than a claim that one layout will or will not fit on a +particular GPU. For Q4_0 K/V, the eight-user profile above projects to 18.000 GiB for padded per-sequence contiguous K/V and 3.234 GiB for paged K/V: a 5.57x reduction, or 82.03% less K/V storage. Uniform `8 x 128K` contexts would use 18.000 GiB in both layouts, apart from small paging metadata and block rounding. +On the same RTX 3090 setup, the ragged case measured 3.122 ms for contiguous +attention and 4.095 ms for paged attention (`1.31x`), while reducing the +one-layer K/V allocation from 1152 MiB to 207 MiB. This shows both sides of +the trade-off: the current paged kernel still has compute overhead, but it +keeps the highly ragged eight-request workload close to the native kernel +while admitting it with 82.03% less K/V storage. + A second GGML design can use one unified contiguous arena plus a sequence-aware mask. It can approach sum-sized storage while compact, but each query attends over a shared high-water span, and request churn can leave diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 6d444d579..b2b46db4b 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1195,6 +1195,10 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) add_test(NAME paged_attention COMMAND test_paged_attention) + add_test(NAME paged_attention_direct + COMMAND test_paged_attention --direct) + set_tests_properties(paged_attention_direct PROPERTIES + ENVIRONMENT "GGML_CUDA_PAGED_ATTN_FORCE_PARTITIONS=1") endif() if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR DFLASH27B_GPU_BACKEND STREQUAL "hip") diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 809564704..303a4fbfa 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2465,6 +2465,8 @@ extern "C" { // k/v: [D, pool_tokens, n_head_kv] // block_table: [max_blocks, n_seq] I32 // kv_seq_lens: [n_seq] I32 (valid cached K/V tokens per sequence) + // max_kv_seq_len: host-known maximum of kv_seq_lens, used to size the + // CUDA/HIP launch without reading block-table capacity // res: [D, n_seq, n_head] F32 // // A logical token t is read from physical row @@ -2479,7 +2481,8 @@ extern "C" { struct ggml_tensor * block_table, struct ggml_tensor * kv_seq_lens, float scale, - int block_size); + int block_size, + int max_kv_seq_len); // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 494b0ed05..c9b8190cf 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -731,22 +731,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co }; auto handle_paged_attn = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - const bool mirrored = - src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED; - if (mirrored) { - return src_ss[0]; - } - - // Head-sharded Q/K/V are valid when every device preserves the same - // GQA ratio. The generic ratio validation below enforces that; routing - // metadata itself must be mirrored. - GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); + // Paged attention currently rejects layer-split targets, so the meta + // backend can only see a fully mirrored op. Add sharded semantics when + // paged layer splitting is implemented and can be exercised end to end. + GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); return src_ss[0]; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu index 3de091832..826e89d50 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -2,20 +2,166 @@ #include "fattn-common.cuh" +#include #include +#include #include -// Decode uses one logical warp per (sequence, query head). Keeping Q in the -// same register formats as fattn-vec lets the quantized K dot products and V -// dequantizers share their well-tested F16/Q4_0/Q8_0 implementations. -template +static constexpr int PAGED_ATTN_MAX_PARTITIONS = 128; +static constexpr int PAGED_ATTN_BLOCKS_PER_PARTITION = 64; + +// All scores are computed in the log2 domain: log2(e) is folded into the same +// Q prescale that already carries the 1/sqrt(D) attention scale, so every +// softmax exponential uses the fast exp2f SFU path. +static constexpr float PAGED_ATTN_LOG2E = 1.44269504088896340736f; + +// Each warp handles n_batch_heads query heads that share one GQA K/V head for +// one (sequence, context partition). Batching the heads into a single warp +// amortizes the K/V loads and dequantization across the whole group: each K +// row is read once and dotted against every batched Q, and each V row is +// dequantized once and accumulated with every batched weight. Warps covering +// the remaining heads of a K/V group (and further K/V groups) are colocated +// in the same block so the GPU caches can still dedupe overlapping reads. +// +// The token loop works on WARP_SIZE-token score tiles: lane t owns the score +// of tile token t, so the per-token dot/reduce chains are independent, the +// accumulator rescale runs once per tile instead of once per token, and it is +// skipped entirely when the tile does not raise the running maximum. +// +// Quantized Q is produced once per decode step by paged_attn_quantize_q and +// read back from global memory here, so context partitions do not repeat the +// quantization work. Long contexts are split between blocks and merged below; +// the direct specialization avoids scratch for a single partition. + +// Batched K·Q dots: one K row is loaded once per lane and dotted against all +// n_batch_heads quantized Q vectors. These mirror the fattn-common vec_dot +// implementations for the three supported K cache types. + +template +static __device__ __forceinline__ void multi_vec_dot_kq_f16( + const char * __restrict__ K_c, +#ifdef V_DOT2_F32_F16_AVAILABLE + const half2 (&Q_reg)[nbh][(D/2)/WARP_SIZE], +#else + const float2 (&Q_reg)[nbh][(D/2)/WARP_SIZE], +#endif + float (&sum)[nbh]) { + const half2 * K_h2 = (const half2 *) K_c; + constexpr int cpy_nb = ggml_cuda_get_max_cpy_bytes(); + constexpr int cpy_ne = cpy_nb / 4; + +#pragma unroll + for (int k0 = 0; k0 < D/2; k0 += WARP_SIZE*cpy_ne) { + __align__(16) half2 tmp[cpy_ne]; + ggml_cuda_memcpy_1(tmp, K_h2 + k0 + threadIdx.x*cpy_ne); +#pragma unroll + for (int k1 = 0; k1 < cpy_ne; ++k1) { +#pragma unroll + for (int h = 0; h < nbh; ++h) { +#ifdef V_DOT2_F32_F16_AVAILABLE + ggml_cuda_mad(sum[h], tmp[k1] , Q_reg[h][k0/WARP_SIZE + k1]); +#else + ggml_cuda_mad(sum[h], __half22float2(tmp[k1]), Q_reg[h][k0/WARP_SIZE + k1]); +#endif + } + } + } +} + +template +static __device__ __forceinline__ void multi_vec_dot_kq_q4_0( + const char * __restrict__ K_c, + const int (&Q_q8)[nbh][D/(sizeof(int)*WARP_SIZE)], + const float2 (&Q_ds)[nbh][D/(sizeof(int)*WARP_SIZE)], + float (&sum)[nbh]) { + const block_q4_0 * K_q4_0 = (const block_q4_0 *) K_c; + +#pragma unroll + for (int k0 = 0; k0 < int(D/sizeof(int)); k0 += WARP_SIZE) { + const int k_KQ = k0 + threadIdx.x; + const int ib = k_KQ / QI8_1; + const int iqs4 = k_KQ % QI4_0; + const int shift = k_KQ & (QI8_1/2); + + int v; + ggml_cuda_memcpy_1(&v, K_q4_0[ib].qs + sizeof(int)*iqs4); + v = (v >> shift) & 0x0F0F0F0F; + const float K_d = __half2float(K_q4_0[ib].d); + +#pragma unroll + for (int h = 0; h < nbh; ++h) { + const int sumi = ggml_cuda_dp4a(v, Q_q8[h][k0/WARP_SIZE], 0); + const float2 ds = Q_ds[h][k0/WARP_SIZE]; + sum[h] += K_d * (sumi*ds.x - (8/QI8_1)*ds.y); + } + } +} + +template +static __device__ __forceinline__ void multi_vec_dot_kq_q8_0( + const char * __restrict__ K_c, + const int (&Q_q8)[nbh][D/(sizeof(int)*WARP_SIZE)], + const float2 (&Q_ds)[nbh][D/(sizeof(int)*WARP_SIZE)], + float (&sum)[nbh]) { + const block_q8_0 * K_q8_0 = (const block_q8_0 *) K_c; + +#pragma unroll + for (int k0 = 0; k0 < int(D/sizeof(int)); k0 += WARP_SIZE) { + const int k_KQ = k0 + threadIdx.x; + const int ib = k_KQ / QI8_0; + const int iqs = k_KQ % QI8_0; + + int v; + ggml_cuda_memcpy_1(&v, K_q8_0[ib].qs + 4*iqs); + const float K_d = __half2float(K_q8_0[ib].d); + +#pragma unroll + for (int h = 0; h < nbh; ++h) { + const int sumi = ggml_cuda_dp4a(v, Q_q8[h][k0/WARP_SIZE], 0); + sum[h] += K_d * Q_ds[h][k0/WARP_SIZE].x * sumi; + } + } +} + +// Quantizes scale*log2(e)*Q to q8_1 once per decode step. The decode kernel +// used to requantize Q per (warp, partition) through shared memory; at long +// contexts that repeated the same work up to PAGED_ATTN_MAX_PARTITIONS times +// per head. One warp per (head, sequence) row. +template +static __global__ void paged_attn_quantize_q( + const char * __restrict__ q, + int * __restrict__ q_i32, + float2 * __restrict__ q_ds, + int64_t q_nb1, int64_t q_nb2, + int32_t n_seq, + float scale) { + const int head = blockIdx.x; + const int seq = blockIdx.y; + + const float * q_row = (const float *) (q + (int64_t) seq * q_nb1 + (int64_t) head * q_nb2); + const int64_t row = (int64_t) head * n_seq + seq; + int * yq32 = q_i32 + row * (D / (int) sizeof(int)); + float2 * yds = q_ds + row * (D / QK8_1); + +#pragma unroll + for (int i0 = 0; i0 < D / (int) sizeof(int); i0 += WARP_SIZE) { + quantize_q8_1_to_shared( + q_row + i0 * sizeof(int), scale, yq32 + i0, yds + i0 / QI8_1); + } +} + +template static __global__ void paged_attn_decode( const char * __restrict__ q, const char * __restrict__ k, const char * __restrict__ v, + const int * __restrict__ q_i32_glob, + const float2 * __restrict__ q_ds_glob, const char * __restrict__ block_table, const char * __restrict__ kv_seq_lens, char * __restrict__ dst, + half * __restrict__ partial_acc, + float2 * __restrict__ partial_meta, int64_t q_nb1, int64_t q_nb2, int64_t k_nb1, int64_t k_nb2, int64_t v_nb1, int64_t v_nb2, @@ -28,139 +174,254 @@ static __global__ void paged_attn_decode( int32_t pool_tokens, int32_t max_blocks, int32_t block_size, + int32_t n_partitions, + int32_t min_partitions, float scale) { constexpr int nthreads = WARP_SIZE; constexpr int values_per_load = 4; constexpr int values_per_lane = D / nthreads; static_assert(D % (nthreads * values_per_load) == 0, "unsupported head size"); - const int head = blockIdx.x; - const int seq = blockIdx.y; - const int lane = threadIdx.x; - if (seq >= n_seq || head >= n_head) { + // The launcher guarantees n_batch_heads divides the GQA ratio, the ratio + // divided by n_batch_heads divides blockDim.y, and the grid covers exact + // sequence/partition extents, so every warp maps to live heads and no + // bounds return is needed. + const int gqa_ratio = n_head / n_head_kv; + const int warps_per_group = gqa_ratio / n_batch_heads; + const int warp = threadIdx.y; + const int kv_head = + (int) blockIdx.x * ((int) blockDim.y / warps_per_group) + + warp / warps_per_group; + const int head0 = + kv_head * gqa_ratio + (warp % warps_per_group) * n_batch_heads; + const int seq = blockIdx.y; + const int partition = blockIdx.z; + const int lane = threadIdx.x; + + const int32_t kv_seq_len_raw = + *(const int32_t *) (kv_seq_lens + (int64_t) seq * ksl_nb0); + const int64_t table_capacity = + (int64_t) max_blocks * block_size; + const int32_t kv_seq_len = kv_seq_len_raw <= 0 + ? 0 + : (kv_seq_len_raw < table_capacity + ? kv_seq_len_raw + : (int32_t) table_capacity); + const int32_t n_logical_blocks = + (kv_seq_len + block_size - 1) / block_size; + const int32_t context_partitions = + (n_logical_blocks + PAGED_ATTN_BLOCKS_PER_PARTITION - 1) / + PAGED_ATTN_BLOCKS_PER_PARTITION; + const int32_t requested_partitions = + context_partitions > min_partitions + ? context_partitions + : min_partitions; + const int32_t available_partitions = + n_logical_blocks < n_partitions ? n_logical_blocks : n_partitions; + const int32_t active_partitions = + requested_partitions < available_partitions + ? requested_partitions + : available_partitions; + + if (partition >= active_partitions) { +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + const int64_t output_row = (int64_t) (head0 + h) * n_seq + seq; + if constexpr (write_partials) { + if (lane == 0) { + partial_meta[output_row * n_partitions + partition] = + make_float2(-FLT_MAX, 0.0f); + } + } else { + float * o_row = + (float *) (dst + (int64_t) seq * dst_nb1 + + (int64_t) (head0 + h) * dst_nb2); +#pragma unroll + for (int i = lane; i < D; i += nthreads) { + o_row[i] = 0.0f; + } + } + } return; } - const float * q_row = (const float *) (q + (int64_t) seq * q_nb1 + (int64_t) head * q_nb2); - float * o_row = (float *) (dst + (int64_t) seq * dst_nb1 + (int64_t) head * dst_nb2); + const int32_t logical_block_begin = + ((int64_t) n_logical_blocks * partition) / active_partitions; + const int32_t logical_block_end = + ((int64_t) n_logical_blocks * (partition + 1)) / + active_partitions; + const int32_t token_begin = logical_block_begin * block_size; + const int32_t token_end_blocks = logical_block_end * block_size; + const int32_t token_end = + kv_seq_len < token_end_blocks ? kv_seq_len : token_end_blocks; - constexpr bool quantize_q = - type_K != GGML_TYPE_F16; - constexpr int q_registers = (D / 2) / nthreads; + constexpr bool quantize_q = type_K != GGML_TYPE_F16; + constexpr int q_registers = (D / 2) / nthreads; + constexpr int q_i32_per_lane = D / (sizeof(int) * nthreads); #ifdef V_DOT2_F32_F16_AVAILABLE - half2 q_reg[q_registers]; + half2 q_reg[n_batch_heads][q_registers]; #else - float2 q_reg[q_registers]; + float2 q_reg[n_batch_heads][q_registers]; #endif - int q_i32[D / (sizeof(int) * nthreads)]; - float2 q_ds [D / (sizeof(int) * nthreads)]; - - __shared__ int q_i32_shared[D / sizeof(int)]; - __shared__ float2 q_ds_shared[D / QK8_1]; + int q_i32[n_batch_heads][q_i32_per_lane]; + float2 q_ds [n_batch_heads][q_i32_per_lane]; if constexpr (quantize_q) { - // Quantizing scale*Q matches fattn-vec and avoids an extra multiply for - // every cached K row. + // Read back the q8_1 rows produced by paged_attn_quantize_q. The rows + // are tiny and shared by every partition, so they stay L2-resident. #pragma unroll - for (int i0 = 0; i0 < D / (int) sizeof(int); i0 += nthreads) { - quantize_q8_1_to_shared( - q_row + i0 * sizeof(int), scale, - q_i32_shared + i0, q_ds_shared + i0 / QI8_1); - } - __syncthreads(); - + for (int h = 0; h < n_batch_heads; ++h) { + const int64_t row = (int64_t) (head0 + h) * n_seq + seq; + const int * yq32 = q_i32_glob + row * (D / (int) sizeof(int)); + const float2 * yds = q_ds_glob + row * (D / QK8_1); #pragma unroll - for (int i0 = 0; i0 < D / (int) sizeof(int); i0 += nthreads) { - const int i = i0 + lane; - q_i32[i0 / nthreads] = q_i32_shared[i]; - q_ds [i0 / nthreads] = q_ds_shared[i / QI8_1]; + for (int i0 = 0; i0 < D / (int) sizeof(int); i0 += nthreads) { + const int i = i0 + lane; + q_i32[h][i0 / nthreads] = yq32[i]; + q_ds [h][i0 / nthreads] = yds[i / QI8_1]; + } } } else { constexpr int cpy_nb = ggml_cuda_get_max_cpy_bytes(); constexpr int cpy_ne = cpy_nb / sizeof(float); - const float2 * q_f2 = (const float2 *) q_row; + // Fold log2(e) into the same prescale that carries the attention + // scale so the softmax below can use exp2f throughout. + const float scale_log2 = scale * PAGED_ATTN_LOG2E; #pragma unroll - for (int i0 = 0; i0 < D / 2; i0 += nthreads * cpy_ne) { - const int i = i0 + lane * cpy_ne; - __align__(16) float2 tmp[cpy_ne]; - ggml_cuda_memcpy_1( - tmp, q_f2 + i); - ggml_cuda_memcpy_1( - tmp + cpy_ne / 2, q_f2 + i + cpy_ne / 2); + for (int h = 0; h < n_batch_heads; ++h) { + const float2 * q_f2 = (const float2 *) + (q + (int64_t) seq * q_nb1 + (int64_t) (head0 + h) * q_nb2); +#pragma unroll + for (int i0 = 0; i0 < D / 2; i0 += nthreads * cpy_ne) { + const int i = i0 + lane * cpy_ne; + __align__(16) float2 tmp[cpy_ne]; + ggml_cuda_memcpy_1( + tmp, q_f2 + i); + ggml_cuda_memcpy_1( + tmp + cpy_ne / 2, q_f2 + i + cpy_ne / 2); #pragma unroll - for (int j = 0; j < cpy_ne; ++j) { + for (int j = 0; j < cpy_ne; ++j) { #ifdef V_DOT2_F32_F16_AVAILABLE - q_reg[i0 / nthreads + j] = - make_half2(tmp[j].x * scale, tmp[j].y * scale); + q_reg[h][i0 / nthreads + j] = + make_half2(tmp[j].x * scale_log2, tmp[j].y * scale_log2); #else - q_reg[i0 / nthreads + j] = - make_float2(tmp[j].x * scale, tmp[j].y * scale); + q_reg[h][i0 / nthreads + j] = + make_float2(tmp[j].x * scale_log2, tmp[j].y * scale_log2); #endif + } } } } - constexpr vec_dot_KQ_t vec_dot_kq = - get_vec_dot_KQ(); constexpr dequantize_V_t dequantize_v = get_dequantize_V(); - float acc[values_per_lane] = { 0.0f }; - float qk_max = -FLT_MAX; - float qk_sum = 0.0f; + float acc[n_batch_heads][values_per_lane] = {{0.0f}}; + float qk_max[n_batch_heads]; + float qk_sum[n_batch_heads]; +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + qk_max[h] = -FLT_MAX; + qk_sum[h] = 0.0f; + } - const int32_t kv_seq_len_raw = - *(const int32_t *) (kv_seq_lens + (int64_t) seq * ksl_nb0); - const int64_t table_capacity = - (int64_t) max_blocks * block_size; - const int32_t kv_seq_len = kv_seq_len_raw <= 0 - ? 0 - : (kv_seq_len_raw < table_capacity - ? kv_seq_len_raw - : (int32_t) table_capacity); - const int32_t n_logical_blocks = - (kv_seq_len + block_size - 1) / block_size; const int32_t n_physical_blocks = pool_tokens / block_size; - const int32_t kv_head = head / (n_head / n_head_kv); - - for (int32_t logical_block = 0; logical_block < n_logical_blocks; ++logical_block) { - const int32_t physical_block = - *(const int32_t *) (block_table + - (int64_t) logical_block * bt_nb0 + - (int64_t) seq * bt_nb1); - - // Invalid entries are never expected from the allocator, but skipping - // them prevents stale metadata from becoming an out-of-bounds read. - if (physical_block < 0 || physical_block >= n_physical_blocks) { - continue; + + for (int32_t tile_begin = token_begin; + tile_begin < token_end; + tile_begin += nthreads) { + const int32_t tile_len = + token_end - tile_begin < nthreads + ? token_end - tile_begin + : nthreads; + + // Lane t resolves tile token t once; both phases below fetch it from + // the owning lane instead of re-reading the block table per token. + // Invalid entries are never expected from the allocator, but mapping + // them to -1 prevents stale metadata from becoming an out-of-bounds + // read; their tokens contribute nothing, mirroring the CPU reference. + int32_t phys_mine = -1; + const int32_t my_token = tile_begin + lane; + if (my_token < token_end) { + const int32_t logical_block = my_token / block_size; + const int32_t physical_block = + *(const int32_t *) (block_table + + (int64_t) logical_block * bt_nb0 + + (int64_t) seq * bt_nb1); + if (physical_block >= 0 && physical_block < n_physical_blocks) { + phys_mine = + physical_block * block_size + my_token % block_size; + } } - const int32_t block_begin = logical_block * block_size; - const int32_t block_end = - kv_seq_len < block_begin + block_size - ? kv_seq_len - : block_begin + block_size; + float score_mine[n_batch_heads]; +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + score_mine[h] = -FLT_MAX; + } - for (int32_t token = block_begin; token < block_end; ++token) { - const int32_t physical_token = - physical_block * block_size + token - block_begin; + for (int j = 0; j < tile_len; ++j) { + const int32_t phys_j = __shfl_sync(0xFFFFFFFF, phys_mine, j, WARP_SIZE); + if (phys_j < 0) { + continue; + } const char * k_row = - k + (int64_t) physical_token * k_nb1 + - (int64_t) kv_head * k_nb2; - const char * v_row = - v + (int64_t) physical_token * v_nb1 + - (int64_t) kv_head * v_nb2; + k + (int64_t) phys_j * k_nb1 + (int64_t) kv_head * k_nb2; + + float sums[n_batch_heads] = {0.0f}; + if constexpr (type_K == GGML_TYPE_F16) { + multi_vec_dot_kq_f16(k_row, q_reg, sums); + } else if constexpr (type_K == GGML_TYPE_Q4_0) { + multi_vec_dot_kq_q4_0(k_row, q_i32, q_ds, sums); + } else { + static_assert(type_K == GGML_TYPE_F16 || type_K == GGML_TYPE_Q4_0 || + type_K == GGML_TYPE_Q8_0, "unsupported K type"); + multi_vec_dot_kq_q8_0(k_row, q_i32, q_ds, sums); + } +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + const float score = warp_reduce_sum(sums[h]); + score_mine[h] = lane == j ? score : score_mine[h]; + } + } - float score = vec_dot_kq(k_row, q_reg, q_i32, q_ds); - score = warp_reduce_sum(score); + // Tile softmax update: one accumulator rescale per tile, skipped + // entirely when the tile does not raise the running maximum (the + // overwhelmingly common case, where the rescale would multiply by 1). + float w_mine[n_batch_heads]; +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + const float tile_max = warp_reduce_max(score_mine[h]); + if (tile_max > qk_max[h]) { + const float old_scale = exp2f(qk_max[h] - tile_max); + qk_sum[h] *= old_scale; +#pragma unroll + for (int i = 0; i < values_per_lane; ++i) { + acc[h][i] *= old_scale; + } + qk_max[h] = tile_max; + } + w_mine[h] = score_mine[h] > -FLT_MAX/2 + ? exp2f(score_mine[h] - qk_max[h]) + : 0.0f; + qk_sum[h] += warp_reduce_sum(w_mine[h]); + } - const float qk_max_new = fmaxf(qk_max, score); - const float old_scale = expf(qk_max - qk_max_new); - const float weight = expf(score - qk_max_new); - qk_sum = qk_sum * old_scale + weight; - qk_max = qk_max_new; + for (int j = 0; j < tile_len; ++j) { + const int32_t phys_j = __shfl_sync(0xFFFFFFFF, phys_mine, j, WARP_SIZE); + if (phys_j < 0) { + continue; + } + const char * v_row = + v + (int64_t) phys_j * v_nb1 + (int64_t) kv_head * v_nb2; + + float weight[n_batch_heads]; +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + weight[h] = __shfl_sync(0xFFFFFFFF, w_mine[h], j, WARP_SIZE); + } #pragma unroll for (int segment = 0; @@ -172,30 +433,127 @@ static __global__ void paged_attn_decode( lane * values_per_load; dequantize_v(v_row, values, value0); #pragma unroll - for (int j = 0; j < values_per_load; ++j) { - const int ai = segment * values_per_load + j; - acc[ai] = acc[ai] * old_scale + weight * values[j]; + for (int i = 0; i < values_per_load; ++i) { + const int ai = segment * values_per_load + i; +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + acc[h][ai] += weight[h] * values[i]; + } } } } } - const float inv_sum = qk_sum > 0.0f ? 1.0f / qk_sum : 0.0f; #pragma unroll - for (int segment = 0; - segment < D / (nthreads * values_per_load); - ++segment) { - const int value0 = - segment * nthreads * values_per_load + - lane * values_per_load; + for (int h = 0; h < n_batch_heads; ++h) { + const int64_t output_row = (int64_t) (head0 + h) * n_seq + seq; + const float inv_sum = qk_sum[h] > 0.0f ? 1.0f / qk_sum[h] : 0.0f; + if constexpr (write_partials) { + // Partials are stored normalized by the partition's qk_sum: it + // keeps the values inside f16 range (halving scratch traffic) and + // lets the combine kernel reuse its weight*qk_sum coefficient. + const int64_t partial_row = + output_row * n_partitions + partition; #pragma unroll - for (int j = 0; j < values_per_load; ++j) { - o_row[value0 + j] = - acc[segment * values_per_load + j] * inv_sum; + for (int segment = 0; + segment < D / (nthreads * values_per_load); + ++segment) { + const int value0 = + segment * nthreads * values_per_load + + lane * values_per_load; +#pragma unroll + for (int i = 0; i < values_per_load; ++i) { + partial_acc[partial_row * D + value0 + i] = + __float2half(acc[h][segment * values_per_load + i] * inv_sum); + } + } + if (lane == 0) { + partial_meta[partial_row] = make_float2(qk_max[h], qk_sum[h]); + } + } else { + float * o_row = + (float *) (dst + (int64_t) seq * dst_nb1 + + (int64_t) (head0 + h) * dst_nb2); +#pragma unroll + for (int segment = 0; + segment < D / (nthreads * values_per_load); + ++segment) { + const int value0 = + segment * nthreads * values_per_load + + lane * values_per_load; +#pragma unroll + for (int i = 0; i < values_per_load; ++i) { + o_row[value0 + i] = + acc[h][segment * values_per_load + i] * inv_sum; + } + } } } } +template +__launch_bounds__(D, 1) +static __global__ void paged_attn_combine( + const half * __restrict__ partial_acc, + const float2 * __restrict__ partial_meta, + char * __restrict__ dst, + int64_t dst_nb1, + int64_t dst_nb2, + int32_t n_seq, + int32_t n_partitions) { + const int head = blockIdx.x; + const int seq = blockIdx.y; + const int tid = threadIdx.x; + // combine_grid is exactly [n_head, n_seq], so no bounds branch is needed + // above the block-wide barrier below. + + const int64_t output_row = (int64_t) head * n_seq + seq; + const int64_t partial_row = output_row * n_partitions; + + __shared__ float partition_scale[PAGED_ATTN_MAX_PARTITIONS]; + __shared__ float reduction[D / WARP_SIZE]; + + // Load each partition's metadata once, then use all warps for the stable + // max/sum reduction. Threads beyond n_partitions contribute sentinels. + // meta.x is a log2-domain maximum, hence exp2f below. + const float2 meta = + tid < n_partitions + ? partial_meta[partial_row + tid] + : make_float2(-FLT_MAX, 0.0f); + const float local_max = meta.y > 0.0f ? meta.x : -FLT_MAX; + const float global_max = + block_reduce(local_max, reduction); + __syncthreads(); + + // The stored partials are normalized by their partition's qk_sum, so a + // partition's combine coefficient is weight*qk_sum — the same product + // that forms the denominator. + const float coefficient = + meta.y > 0.0f ? exp2f(meta.x - global_max) * meta.y : 0.0f; + if (tid < n_partitions) { + partition_scale[tid] = coefficient; + } + const float denominator = + block_reduce(coefficient, reduction); + __syncthreads(); + + float numerator = 0.0f; + for (int partition = 0; partition < n_partitions; ++partition) { + const float w = partition_scale[partition]; + if (w > 0.0f) { + numerator += + w * + __half2float( + partial_acc[(partial_row + partition) * D + tid]); + } + } + + float * o_row = + (float *) (dst + (int64_t) seq * dst_nb1 + + (int64_t) head * dst_nb2); + o_row[tid] = denominator > 0.0f ? numerator / denominator : 0.0f; +} + static bool paged_attn_type_supported(ggml_type type) { return type == GGML_TYPE_F16 || type == GGML_TYPE_Q4_0 || @@ -214,6 +572,7 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { } const int32_t block_size = ggml_get_op_params_i32(dst, 1); + const int32_t max_kv_seq_len = ggml_get_op_params_i32(dst, 2); return dst->type == GGML_TYPE_F32 && q->type == GGML_TYPE_F32 && paged_attn_type_supported(k->type) && @@ -249,15 +608,22 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { kv_seq_lens->ne[2] == 1 && kv_seq_lens->ne[3] == 1 && block_size > 0 && + max_kv_seq_len > 0 && + max_kv_seq_len <= k->ne[1] && k->ne[1] % block_size == 0; } -template -static void launch_paged_attn( +// Attempts a launch with n_batch_heads query heads per warp. Returns false +// when this width cannot reach a viable occupancy on the device (register +// pressure grows with the head batch), so the caller can fall back to a +// narrower instantiation. +template +static bool try_launch_paged_attn( ggml_backend_cuda_context & ctx, ggml_tensor * dst, float scale, - int32_t block_size) { + int32_t block_size, + int32_t max_kv_seq_len) { constexpr int D = 256; const ggml_tensor * q = dst->src[0]; const ggml_tensor * k = dst->src[1]; @@ -265,15 +631,217 @@ static void launch_paged_attn( const ggml_tensor * block_table = dst->src[3]; const ggml_tensor * kv_seq_lens = dst->src[4]; - const dim3 grid((unsigned int) q->ne[2], (unsigned int) q->ne[1], 1); - const dim3 block(WARP_SIZE, 1, 1); - paged_attn_decode<<>>( + const int32_t n_head = (int32_t) q->ne[2]; + const int32_t n_head_kv = (int32_t) k->ne[2]; + const int32_t gqa_ratio = n_head / n_head_kv; + GGML_ASSERT(gqa_ratio % n_batch_heads == 0); + const int32_t warps_per_group = gqa_ratio / n_batch_heads; + + // Colocate the warps of one K/V group, then pack further K/V groups into + // the same block up to ~8 warps so small groups still fill a block. + int32_t kv_heads_per_block = 1; + for (int32_t candidate = 2; candidate <= n_head_kv; ++candidate) { + if (n_head_kv % candidate == 0 && + warps_per_group * candidate <= 8) { + kv_heads_per_block = candidate; + } + } + + int max_blocks_per_sm = 0; + // Deliberately query the write_partials instantiation: occupancy only + // steers min_partitions, which is a partition count for that variant. + // The direct variant launches solely when the count collapses to one, + // where occupancy no longer influences the topology. + while (true) { + const int warps_per_block = warps_per_group * kv_heads_per_block; + // 0 means "not queried yet"; UNLAUNCHABLE records a block size this + // device rejects, so the fallback below is decided once rather than + // on every decode step. + constexpr int UNLAUNCHABLE = -1; + static std::atomic + occupancy[GGML_CUDA_MAX_DEVICES][WARP_SIZE + 1] = {}; + std::atomic & cached = + occupancy[ctx.device][warps_per_block]; + int probe = cached.load(std::memory_order_acquire); + if (probe == 0) { + // A block size above the kernel's maxThreadsPerBlock makes CUDA + // return cudaErrorInvalidValue rather than a zero occupancy, so + // the status has to drive the fallback below instead of aborting. + const cudaError_t err = + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &probe, + paged_attn_decode, + WARP_SIZE * warps_per_block, 0); + if (err != cudaSuccess) { + probe = UNLAUNCHABLE; + // Keep the rejected query out of the context's error state so + // the next unrelated CUDA_CHECK does not inherit it. + (void) cudaGetLastError(); + } + cached.store(probe > 0 ? probe : UNLAUNCHABLE, + std::memory_order_release); + } + max_blocks_per_sm = probe > 0 ? probe : 0; + if (max_blocks_per_sm > 0 || kv_heads_per_block == 1) { + break; + } + do { + kv_heads_per_block /= 2; + } while (kv_heads_per_block > 1 && + n_head_kv % kv_heads_per_block != 0); + } + if (max_blocks_per_sm == 0) { + return false; + } + + GGML_ASSERT(n_head_kv % kv_heads_per_block == 0); + const int32_t head_groups = n_head_kv / kv_heads_per_block; + const dim3 block( + WARP_SIZE, + (unsigned int) (warps_per_group * kv_heads_per_block), + 1); + + const int64_t output_rows = q->ne[1] * q->ne[2]; + const int64_t work_groups = q->ne[1] * head_groups; + const int64_t target_blocks = + (int64_t) ggml_cuda_info().devices[ctx.device].nsm * + max_blocks_per_sm; + int32_t min_partitions = (int32_t) + ((target_blocks + work_groups - 1) / work_groups); + if (min_partitions < 1) { + min_partitions = 1; + } + // Small batches need more context partitions each to expose enough work; + // large batches already fill the device, where extra partitions mostly + // repeat the per-partition fixed costs. Scale the cap so the total + // partition count stays roughly constant across batch sizes. + int32_t partition_limit = + PAGED_ATTN_MAX_PARTITIONS / (int32_t) q->ne[1]; + if (partition_limit < 32) { + partition_limit = 32; + } + if (min_partitions > partition_limit) { + min_partitions = partition_limit; + } + if (min_partitions > block_table->ne[0]) { + min_partitions = (int32_t) block_table->ne[0]; + } + + // Size the launch from the live maximum sequence length carried in the + // graph op, not the block-table capacity. Ragged sequences still clamp + // their own active partition count from kv_seq_lens on device. + const int32_t live_blocks = + (max_kv_seq_len + block_size - 1) / block_size; + const int32_t context_partitions = + (live_blocks + + PAGED_ATTN_BLOCKS_PER_PARTITION - 1) / + PAGED_ATTN_BLOCKS_PER_PARTITION; + int32_t n_partitions = + context_partitions > min_partitions + ? context_partitions + : min_partitions; + if (n_partitions > PAGED_ATTN_MAX_PARTITIONS) { + n_partitions = PAGED_ATTN_MAX_PARTITIONS; + } + if (n_partitions > live_blocks) { + n_partitions = live_blocks; + } + + // Test/debug override used to exercise both the direct and partials paths + // independently of device-specific occupancy. Values outside the valid + // grid range are ignored. Read once: every full-attention layer launches + // this on every decode token, so an environment scan per launch would sit + // in the decode hot path. + static const int forced_partitions = []() { + const char * env = + std::getenv("GGML_CUDA_PAGED_ATTN_FORCE_PARTITIONS"); + return env ? std::atoi(env) : 0; + }(); + if (forced_partitions >= 1 && + forced_partitions <= PAGED_ATTN_MAX_PARTITIONS && + forced_partitions <= block_table->ne[0]) { + min_partitions = forced_partitions; + n_partitions = forced_partitions; + } + + constexpr bool quantize_q = type_K != GGML_TYPE_F16; + int * q_i32_glob = nullptr; + float2 * q_ds_glob = nullptr; + ggml_cuda_pool_alloc q_i32_alloc(ctx.pool()); + ggml_cuda_pool_alloc q_ds_alloc(ctx.pool()); + if (quantize_q) { + q_i32_glob = q_i32_alloc.alloc(output_rows * (D / sizeof(int))); + q_ds_glob = q_ds_alloc.alloc(output_rows * (D / QK8_1)); + const dim3 quantize_grid( + (unsigned int) q->ne[2], (unsigned int) q->ne[1], 1); + paged_attn_quantize_q + <<>>( + (const char *) q->data, + q_i32_glob, + q_ds_glob, + q->nb[1], q->nb[2], + (int32_t) q->ne[1], + scale * PAGED_ATTN_LOG2E); + } + + const dim3 grid( + (unsigned int) head_groups, + (unsigned int) q->ne[1], + (unsigned int) n_partitions); + + if (n_partitions == 1) { + paged_attn_decode + <<>>( + (const char *) q->data, + (const char *) k->data, + (const char *) v->data, + q_i32_glob, + q_ds_glob, + (const char *) block_table->data, + (const char *) kv_seq_lens->data, + (char *) dst->data, + nullptr, + nullptr, + q->nb[1], q->nb[2], + k->nb[1], k->nb[2], + v->nb[1], v->nb[2], + block_table->nb[0], block_table->nb[1], + kv_seq_lens->nb[0], + dst->nb[1], dst->nb[2], + (int32_t) q->ne[1], + n_head, + n_head_kv, + (int32_t) k->ne[1], + (int32_t) block_table->ne[0], + block_size, + n_partitions, + min_partitions, + scale); + return true; + } + + const size_t partial_rows = + (size_t) output_rows * n_partitions; + const size_t partial_values = partial_rows * D; + ggml_cuda_pool_alloc acc_scratch( + ctx.pool(), partial_values); + ggml_cuda_pool_alloc meta_scratch( + ctx.pool(), partial_rows); + half * partial_acc = acc_scratch.ptr; + float2 * partial_meta = meta_scratch.ptr; + + paged_attn_decode + <<>>( (const char *) q->data, (const char *) k->data, (const char *) v->data, + q_i32_glob, + q_ds_glob, (const char *) block_table->data, (const char *) kv_seq_lens->data, (char *) dst->data, + partial_acc, + partial_meta, q->nb[1], q->nb[2], k->nb[1], k->nb[2], v->nb[1], v->nb[2], @@ -281,12 +849,69 @@ static void launch_paged_attn( kv_seq_lens->nb[0], dst->nb[1], dst->nb[2], (int32_t) q->ne[1], - (int32_t) q->ne[2], - (int32_t) k->ne[2], + n_head, + n_head_kv, (int32_t) k->ne[1], (int32_t) block_table->ne[0], block_size, + n_partitions, + min_partitions, scale); + + const dim3 combine_grid( + (unsigned int) q->ne[2], + (unsigned int) q->ne[1], + 1); + paged_attn_combine + <<>>( + partial_acc, + partial_meta, + (char *) dst->data, + dst->nb[1], + dst->nb[2], + (int32_t) q->ne[1], + n_partitions); + return true; +} + +template +static void launch_paged_attn( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + float scale, + int32_t block_size, + int32_t max_kv_seq_len) { + const int32_t gqa_ratio = + (int32_t) (dst->src[0]->ne[2] / dst->src[1]->ne[2]); + + // Debug/tuning override to pin the per-warp head batch (1, 3, or 6) so + // bench_paged_attention can A/B the widths directly. Read once — this + // sits in the per-layer decode hot path. + static const int forced_batch_heads = []() { + const char * env = + std::getenv("GGML_CUDA_PAGED_ATTN_BATCH_HEADS"); + return env ? std::atoi(env) : 0; + }(); + const int batch_heads_limit = + forced_batch_heads >= 1 ? forced_batch_heads : 6; + + // Widest head batch first: register pressure can make a wide batch + // unlaunchable on some devices, so each width falls back to the next. + if (gqa_ratio % 6 == 0 && batch_heads_limit >= 6 && + try_launch_paged_attn( + ctx, dst, scale, block_size, max_kv_seq_len)) { + return; + } + if (gqa_ratio % 3 == 0 && batch_heads_limit >= 3 && + try_launch_paged_attn( + ctx, dst, scale, block_size, max_kv_seq_len)) { + return; + } + if (try_launch_paged_attn( + ctx, dst, scale, block_size, max_kv_seq_len)) { + return; + } + GGML_ABORT("paged attention kernel has zero occupancy"); } template @@ -294,16 +919,20 @@ static void launch_paged_attn_v( ggml_backend_cuda_context & ctx, ggml_tensor * dst, float scale, - int32_t block_size) { + int32_t block_size, + int32_t max_kv_seq_len) { switch (dst->src[2]->type) { case GGML_TYPE_F16: - launch_paged_attn(ctx, dst, scale, block_size); + launch_paged_attn( + ctx, dst, scale, block_size, max_kv_seq_len); break; case GGML_TYPE_Q4_0: - launch_paged_attn(ctx, dst, scale, block_size); + launch_paged_attn( + ctx, dst, scale, block_size, max_kv_seq_len); break; case GGML_TYPE_Q8_0: - launch_paged_attn(ctx, dst, scale, block_size); + launch_paged_attn( + ctx, dst, scale, block_size, max_kv_seq_len); break; default: GGML_ABORT("unsupported paged-attention V type: %s", @@ -319,16 +948,20 @@ void ggml_cuda_paged_attn( float scale; memcpy(&scale, dst->op_params, sizeof(scale)); const int32_t block_size = ggml_get_op_params_i32(dst, 1); + const int32_t max_kv_seq_len = ggml_get_op_params_i32(dst, 2); switch (dst->src[1]->type) { case GGML_TYPE_F16: - launch_paged_attn_v(ctx, dst, scale, block_size); + launch_paged_attn_v( + ctx, dst, scale, block_size, max_kv_seq_len); break; case GGML_TYPE_Q4_0: - launch_paged_attn_v(ctx, dst, scale, block_size); + launch_paged_attn_v( + ctx, dst, scale, block_size, max_kv_seq_len); break; case GGML_TYPE_Q8_0: - launch_paged_attn_v(ctx, dst, scale, block_size); + launch_paged_attn_v( + ctx, dst, scale, block_size, max_kv_seq_len); break; default: GGML_ABORT("unsupported paged-attention K type: %s", diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index c1a16f29e..695046978 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5618,7 +5618,8 @@ struct ggml_tensor * ggml_paged_attn( struct ggml_tensor * block_table, struct ggml_tensor * kv_seq_lens, float scale, - int block_size) { + int block_size, + int max_kv_seq_len) { GGML_ASSERT(q->type == GGML_TYPE_F32); GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_Q8_0); GGML_ASSERT(v->type == GGML_TYPE_F16 || v->type == GGML_TYPE_Q4_0 || v->type == GGML_TYPE_Q8_0); @@ -5640,11 +5641,14 @@ struct ggml_tensor * ggml_paged_attn( GGML_ASSERT(block_size > 0); GGML_ASSERT(k->ne[1] % block_size == 0); + GGML_ASSERT(max_kv_seq_len > 0); + GGML_ASSERT(max_kv_seq_len <= k->ne[1]); struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, q->ne); ggml_set_op_params_f32(result, 0, scale); ggml_set_op_params_i32(result, 1, block_size); + ggml_set_op_params_i32(result, 2, max_kv_seq_len); result->op = GGML_OP_PAGED_ATTN; result->src[0] = q; diff --git a/server/src/common/paged_kv_pool.cpp b/server/src/common/paged_kv_pool.cpp index 8a437f412..53b4ae120 100644 --- a/server/src/common/paged_kv_pool.cpp +++ b/server/src/common/paged_kv_pool.cpp @@ -81,12 +81,11 @@ PagedKvStatus PagedKvPool::acquire(PagedKvRequestId request_id, uint64_t generation = sequence.generation + 1; if (generation == 0) generation = 1; - // A free slot owns no blocks and has no live request mapping. + // release_sequence(), reset(), and construction leave every free slot + // empty; acquire only installs its new identity. sequence.request_id = request_id; sequence.generation = generation; - sequence.kv_seq_len = 0; sequence.active = true; - sequence.block_table.clear(); request_to_slot_.emplace(request_id, slot); take_lowest(free_sequence_slots_); @@ -139,22 +138,20 @@ PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, return result; } - std::vector allocated_blocks; - allocated_blocks.reserve(additional_blocks); - auto free_it = free_blocks_.begin(); - for (uint32_t i = 0; i < additional_blocks; ++i, ++free_it) { - allocated_blocks.push_back(*free_it); - } - + // Reserve before mutating the pool. Once capacity is available, + // PagedKvWriteSlot insertion cannot allocate or throw. result.write_slots.reserve(token_count); + + // extend_block_table() grows the block table before consuming free blocks, + // so a failed vector allocation leaves the pool unchanged. + result.status = extend_block_table(*sequence, required_blocks); + if (result.status != PagedKvStatus::Ok) return result; + for (uint32_t i = 0; i < token_count; ++i) { const uint32_t logical_position = old_kv_seq_len + i; const uint32_t logical_block = logical_position / block_size_; const uint32_t block_offset = logical_position % block_size_; - const uint32_t physical_block = - logical_block < current_blocks - ? sequence->block_table[logical_block] - : allocated_blocks[logical_block - current_blocks]; + const uint32_t physical_block = sequence->block_table[logical_block]; result.write_slots.push_back({ logical_position, physical_block, @@ -163,11 +160,6 @@ PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, }); } - // Allocation state changes only after every write slot has been built. - for (uint32_t block : allocated_blocks) free_blocks_.erase(block); - sequence->block_table.insert(sequence->block_table.end(), - allocated_blocks.begin(), - allocated_blocks.end()); sequence->kv_seq_len = new_kv_seq_len; result.status = PagedKvStatus::Ok; return result; diff --git a/server/src/common/paged_kv_pool.h b/server/src/common/paged_kv_pool.h index 38e9cc875..fa0bbe95b 100644 --- a/server/src/common/paged_kv_pool.h +++ b/server/src/common/paged_kv_pool.h @@ -4,8 +4,12 @@ // from one shared physical pool, so sequences grow in block granularity // instead of reserving max-context capacity up front. The pool tracks // indices only — it owns no K/V storage. Callers translate the returned -// block/slot indices into offsets within their own pooled K/V tensors and -// upload metadata_snapshot() to drive the GPU paged-attention kernels. +// block/slot indices into offsets within their own pooled K/V tensors. +// +// The current single-request backend consumes sequence() directly. The +// request lookup, reservation, cancellation, capacity, and flattened +// metadata APIs below are intentionally retained for the continuous-batching +// scheduler, where several live request tables will be uploaded together. // // Not thread-safe; callers must serialize access. @@ -45,6 +49,8 @@ struct PagedKvSequenceHandle { uint64_t generation = 0; }; +// Kept as a complete comparison pair for continuous-batching scheduler state, +// where handles will be compared across admission and cancellation events. constexpr bool operator==(const PagedKvSequenceHandle & lhs, const PagedKvSequenceHandle & rhs) { return lhs.slot == rhs.slot && lhs.generation == rhs.generation; @@ -86,6 +92,8 @@ struct PagedKvAppendSpan { }; // Copy of one sequence's bookkeeping state, as returned by sequence(). +// The identity echo lets a future batching scheduler reject a snapshot that +// raced with cancellation and slot reuse before it submits GPU work. struct PagedKvSequenceSnapshot { PagedKvRequestId request_id = 0; PagedKvSequenceHandle handle; @@ -93,9 +101,11 @@ struct PagedKvSequenceSnapshot { std::vector block_table; }; -// Per-sequence entry of a metadata snapshot. Each entry addresses its own -// range [block_table_offset, block_table_offset + block_count) in -// PagedKvMetadataSnapshot::block_table. +// Per-sequence entry of a future continuous-batch metadata upload. Each entry +// addresses its own range [block_table_offset, block_table_offset + +// block_count) in PagedKvMetadataSnapshot::block_table. request_id, +// generation, and sequence_slot are CPU-side scheduler identity; only the +// lengths, offsets, counts, and flattened table are GPU routing material. struct PagedKvSequenceMetadata { PagedKvRequestId request_id = 0; uint64_t generation = 0; @@ -105,9 +115,10 @@ struct PagedKvSequenceMetadata { uint32_t block_count = 0; }; -// Flattened view of every active sequence, ordered by slot index and laid -// out for upload to the GPU paged-attention kernels: `block_table` is the -// concatenation of all per-sequence block tables. +// Flattened view of every active sequence, ordered by slot index. This is +// reserved for the continuous-batching integration; the current backend does +// not upload it yet. `block_table` is the concatenation of all per-sequence +// block tables. struct PagedKvMetadataSnapshot { uint32_t block_size = 0; uint32_t physical_block_count = 0; @@ -127,7 +138,10 @@ class PagedKvPool { uint32_t block_size() const { return block_size_; } uint32_t physical_block_count() const { return physical_block_count_; } - uint32_t max_sequences() const { return static_cast(sequences_.size()); } + // Admission capacity for the future continuous-batching scheduler. + uint32_t max_sequences() const { + return static_cast(sequences_.size()); + } uint32_t active_sequence_count() const { return static_cast(request_to_slot_.size()); } @@ -140,14 +154,16 @@ class PagedKvPool { PagedKvStatus acquire(PagedKvRequestId request_id, PagedKvSequenceHandle & out_handle); - // Fetch the live handle of an already-acquired request. + // Fetch the live handle of an already-acquired request. Continuous + // batching will use this when scheduler events carry request ids. PagedKvStatus lookup(PagedKvRequestId request_id, PagedKvSequenceHandle & out_handle) const; // Grow the sequence's block table until it can hold `token_capacity` // tokens in total (never shrinks). This does not advance kv_seq_len; a - // later append() fills the pre-allocated blocks first. Lets callers fail - // fast before prefill instead of mid-append. + // later append() fills the pre-allocated blocks first. The single-request + // backend does not need this, but batched admission can use it to fail an + // entire scheduling round before launching any prefills. PagedKvStatus reserve(PagedKvSequenceHandle handle, uint32_t token_capacity); @@ -168,8 +184,8 @@ class PagedKvPool { // any copy of it) becomes stale. PagedKvStatus release(PagedKvSequenceHandle handle); - // release() by request id, for callers that no longer hold the handle - // (e.g. client disconnect). + // release() by request id, for continuous-batching cancellation paths + // that no longer hold the handle (e.g. client disconnect). PagedKvStatus cancel(PagedKvRequestId request_id); // Drop every sequence and reclaim all blocks. Every outstanding handle @@ -180,7 +196,8 @@ class PagedKvPool { PagedKvStatus sequence(PagedKvSequenceHandle handle, PagedKvSequenceSnapshot & out_sequence) const; - // Flattened copy of all active sequences, ready for device upload. + // Flattened copy of all active sequences for the planned batched device + // metadata upload. PagedKvMetadataSnapshot metadata_snapshot() const; private: diff --git a/server/src/internal.h b/server/src/internal.h index b474bc836..1026cdd57 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -418,6 +418,14 @@ struct TargetCache { // kv_k_rotated) query per full-attention layer, written by the graph // when QwenGraphInputs::q_capture is set. F32 [head_dim, n_head, n_fa]. ggml_tensor * q_cap = nullptr; + + // Paged-attention metadata, resident next to the K/V pool (only when the + // cache was created with paged_attention). Living here instead of as + // gallocr-managed graph inputs lets decode steps update them append-only + // — one table entry per new 16-token block plus a 4-byte length per step + // — instead of re-uploading the whole live table before every compute. + ggml_tensor * paged_block_table = nullptr; // I32 [max_blocks, 1] + ggml_tensor * paged_kv_seq_lens = nullptr; // I32 [1] }; // Snapshot the current SSM+conv state into TargetCache::*_snap tensors. @@ -614,7 +622,6 @@ struct QwenGraphInputs { ggml_tensor * paged_block_table = nullptr; // [max_blocks,n_seqs] i32 // [n_seqs] i32; valid cached K/V tokens per sequence. ggml_tensor * paged_kv_seq_lens = nullptr; - int paged_block_size = 0; // Capture the LAST token's post-RoPE/post-rotation Q per full-attention // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 0c919620f..e581927f1 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -326,15 +326,13 @@ bool build_target_step( if (paged_attention) { if (n_tokens != 1 || with_mask || fa_window != 0) return false; - const int max_blocks = paged_block_count(cache.max_ctx); - sg.paged_block_table = ggml_new_tensor_2d( - sg.ctx, GGML_TYPE_I32, max_blocks, 1); - ggml_set_name(sg.paged_block_table, "paged_block_table"); - ggml_set_input(sg.paged_block_table); - sg.paged_kv_seq_lens = - ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); - ggml_set_name(sg.paged_kv_seq_lens, "paged_kv_seq_lens"); - ggml_set_input(sg.paged_kv_seq_lens); + // The paging metadata lives in the persistent target cache (next to + // the K/V pool), not as gallocr graph inputs: contents survive graph + // execution and rebuilds, so the backend uploads only what changed + // between decode steps. + if (!cache.paged_block_table || !cache.paged_kv_seq_lens) return false; + sg.paged_block_table = cache.paged_block_table; + sg.paged_kv_seq_lens = cache.paged_kv_seq_lens; } sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); @@ -375,7 +373,6 @@ bool build_target_step( gi.kv_write_rows = sg.kv_write_rows; gi.paged_block_table = sg.paged_block_table; gi.paged_kv_seq_lens = sg.paged_kv_seq_lens; - gi.paged_block_size = paged_attention ? PAGED_BLOCK_SIZE : 0; gi.q_capture = capture_qk; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 88cfa5129..d743f9148 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -534,10 +534,18 @@ bool Qwen35Backend::prepare_paged_decode_step(uint32_t logical_position) { return false; } - // Graph inputs are not persistent allocations: gallocr may reuse this - // tensor's buffer after its last attention consumer. Keep the stable table - // on the host, but upload every live entry before every compute. - if (paged_block_table_size_ == 0 || slot.block_offset == 0) { + // The block table is a persistent cache allocation (created next to the + // K/V pool), so device contents survive graph execution and rebuilds. A + // fresh sequence uploads the whole live table once — covering the prefill + // blocks — and afterwards each step uploads at most one new entry, when + // the sequence grows into a fresh block. Steps inside a block upload + // nothing here. + const bool new_block = slot.block_offset == 0; + if (paged_block_table_size_ == 0 || + (new_block && logical_block != paged_block_table_size_)) { + // The second condition covers a non-append allocation, which the + // pool never produces today; re-uploading everything keeps the + // device table correct if that ever changes. PagedKvSequenceSnapshot sequence; const PagedKvStatus status = paged_kv_pool_->sequence(*paged_sequence_, sequence); @@ -568,10 +576,17 @@ bool Qwen35Backend::prepare_paged_decode_step(uint32_t logical_position) { (int32_t)sequence.block_table[i]; } paged_block_table_size_ = sequence.block_table.size(); + ggml_backend_tensor_set( + sg_.paged_block_table, paged_block_table_upload_.data(), 0, + paged_block_table_size_ * sizeof(paged_block_table_upload_[0])); + } else if (new_block) { + const int32_t entry = (int32_t)slot.physical_block; + paged_block_table_upload_[logical_block] = entry; + ggml_backend_tensor_set( + sg_.paged_block_table, &entry, + logical_block * sizeof(entry), sizeof(entry)); + paged_block_table_size_ = logical_block + 1; } - ggml_backend_tensor_set( - sg_.paged_block_table, paged_block_table_upload_.data(), 0, - paged_block_table_size_ * sizeof(paged_block_table_upload_[0])); const int32_t kv_seq_len = (int32_t)logical_position + 1; ggml_backend_tensor_set(sg_.paged_kv_seq_lens, &kv_seq_len, 0, @@ -1861,8 +1876,10 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, if (IS_EOS_TOK(first_tok, w_)) return true; // Preserve the existing AR position semantics: the first sampled // token advances cur_pos without a cache write. The page table must - // include that zero-initialized logical row so subsequent attention - // sees exactly the same context as the dense path. + // include that logical row so subsequent attention sees exactly the + // same context as the dense path. Prefill does not refresh this row, + // so after the first request it may contain residual K/V from the + // previous request. if (!append_paged_gap((uint32_t)committed)) return false; committed++; cache_.cur_pos = committed; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index edcb705a2..19b256071 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -217,6 +217,22 @@ bool create_target_cache_partial(const TargetWeights & w, head_dim, w.n_head, n_full_attn); ggml_set_name(out.q_cap, "q_cap"); + // Paged-attention metadata is persistent like the K/V pool itself so + // decode steps can update it append-only; a gallocr graph input would + // need every live entry re-uploaded before every compute because its + // buffer region may be recycled between attention consumers. + if (paged_attention) { + out.paged_block_table = ggml_new_tensor_2d( + out.base_ctx, GGML_TYPE_I32, paged_block_count(max_ctx), 1); + ggml_set_name(out.paged_block_table, "paged_block_table"); + out.paged_kv_seq_lens = + ggml_new_tensor_1d(out.base_ctx, GGML_TYPE_I32, 1); + ggml_set_name(out.paged_kv_seq_lens, "paged_kv_seq_lens"); + } else { + out.paged_block_table = nullptr; + out.paged_kv_seq_lens = nullptr; + } + out.base_buf = ggml_backend_alloc_ctx_tensors(out.base_ctx, backend); if (!out.base_buf) { set_last_error("ggml_backend_alloc_ctx_tensors failed for base cache"); @@ -578,8 +594,7 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * kv_write_rows = nullptr, ggml_tensor ** q_fa_out = nullptr, // post-RoPE/post-rotation Q [head_dim, n_tokens, n_head] ggml_tensor * paged_block_table = nullptr, - ggml_tensor * paged_kv_seq_lens = nullptr, - int paged_block_size = 0 + ggml_tensor * paged_kv_seq_lens = nullptr ) { const int head_dim = w.n_embd_head_k; const int n_head = w.n_head; @@ -721,10 +736,24 @@ static ggml_tensor * build_full_attn_block( const float kq_scale = 1.0f / std::sqrt((float)head_dim); ggml_tensor * attn = nullptr; if (paged_block_table) { - GGML_ASSERT(paged_kv_seq_lens && paged_block_size > 0); + GGML_ASSERT(paged_kv_seq_lens); + // The launch bound lands in op_params, and the ggml-cuda graph cache + // memcmps the whole ggml_tensor: a live kv_len here would differ on + // every decode step and force a re-capture per token. Pad it on the + // same 256-token stride the dense path uses for win_len_padded, so the + // paged node's properties are stable within each window. Exact + // per-sequence extents still come from kv_seq_lens on device; a larger + // bound only over-sizes the partition grid, and partitions past the + // real length exit with a zero-weight sentinel. + int paged_launch_len = + ((kv_len + fattn_stride - 1) / fattn_stride) * fattn_stride; + // Never claim more than the physical pool (max_ctx may not be + // 256-aligned); ggml_paged_attn asserts max_kv_seq_len <= k->ne[1]. + paged_launch_len = std::min(paged_launch_len, (int)cache_k->ne[1]); attn = ggml_paged_attn(ctx, Qfa, cache_k, cache_v, paged_block_table, paged_kv_seq_lens, - kq_scale, paged_block_size); + kq_scale, PAGED_BLOCK_SIZE, + paged_launch_len); } else { // fa_window > 0: attend only to the last fa_window positions (cuts FA // cost during spec-decode verify at long contexts). Paged attention @@ -1250,8 +1279,7 @@ QwenGraphOutputs build_qwen35_graph( in.kv_write_rows, want_q_cap ? &q_fa : nullptr, in.paged_block_table, - in.paged_kv_seq_lens, - in.paged_block_size); + in.paged_kv_seq_lens); if (want_q_cap && q_fa) { // Last token's Q, all heads: src [head_dim, 1, n_head] view of // [head_dim, n_tokens, n_head]; dst = q_cap plane fa_idx diff --git a/server/test/bench_paged_attention.cpp b/server/test/bench_paged_attention.cpp index d4a9cfe50..df2723e39 100644 --- a/server/test/bench_paged_attention.cpp +++ b/server/test/bench_paged_attention.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace { @@ -21,11 +22,17 @@ constexpr int D = 256; constexpr int N_HEAD = 24; constexpr int N_HEAD_KV = 4; constexpr int BLOCK_SIZE = 16; -constexpr int PROTOTYPE_ROWS = 256; +// Keep the synthetic cache period independent of both 16-token pages and +// 256-token contiguous-attention tiles so adjacent tiles are not identical. +constexpr int PROTOTYPE_ROW_PERIOD = 521; constexpr int CONTIGUOUS_CONTEXT_ALIGNMENT = 256; -constexpr int QWEN_FULL_ATTENTION_LAYERS = 16; constexpr size_t MAX_RAGGED_SEQUENCES = 64; constexpr double TARGET_SAMPLE_SECONDS = 0.20; +// The focused numerical test enforces a tighter 5e-4 paged-kernel bound. +// This benchmark also accepts GGML's native long-context HIP reduction, +// whose observed max error is below 5e-3, but refuses to time either path +// when it exceeds that gross-correctness ceiling. +constexpr double MAX_ORACLE_ERROR = 5.0e-3; struct Options { int context = 4096; @@ -40,7 +47,6 @@ struct Options { struct TimingStats { int iterations = 0; double median_us = 0.0; - double p95_us = 0.0; double aggregate_queries_s = 0.0; }; @@ -55,12 +61,15 @@ struct BenchResult { uint64_t contiguous_kv_bytes = 0; uint64_t paged_kv_bytes = 0; uint64_t paged_metadata_bytes = 0; - double contiguous_kv_mib = 0.0; - double paged_kv_mib = 0.0; - double paged_metadata_kib = 0.0; - double max_abs_error = 0.0; - double max_tolerance_ratio = 0.0; - double rmse = 0.0; + double layout_max_abs_error = 0.0; + double contiguous_oracle_max_abs = 0.0; + double paged_oracle_max_abs = 0.0; +}; + +struct CaseResult { + std::string name; + std::vector contexts; + BenchResult metrics; }; struct CachePair { @@ -88,6 +97,8 @@ void print_usage(const char * program) { " [--iterations N] [--k-type TYPE] [--v-type TYPE]\n" " [--ragged-contexts N1,N2,...]\n" " TYPE: f16, q4_0, or q8_0\n" + " The default run measures uniform batches 1,2,4,8 plus an\n" + " 8-sequence profile with one context N and seven at N/16.\n" " --ragged-contexts runs one ragged case against a per-sequence\n" " contiguous cache padded to a 256-token maximum context.\n" " Between 2 and %zu sequence lengths are accepted.\n" @@ -135,15 +146,15 @@ bool parse_context_list( if (end == cursor || parsed < 1 || parsed > INT32_MAX || (*end != ',' && *end != '\0') || contexts.size() >= MAX_RAGGED_SEQUENCES) { - contexts.clear(); - return false; + break; } context = static_cast(parsed); contexts.push_back(context); if (*end == '\0') { - if (contexts.size() >= 2) return true; - contexts.clear(); - return false; + if (contexts.size() >= 2) { + return true; + } + break; } cursor = end + 1; } @@ -210,19 +221,118 @@ std::vector make_prototype_rows( if (row_bytes == 0) return {}; std::vector source( - static_cast(PROTOTYPE_ROWS) * D); + static_cast(PROTOTYPE_ROW_PERIOD) * D); for (size_t i = 0; i < source.size(); ++i) { source[i] = std::sin(static_cast(i) * 0.011f + phase) * 0.25f; } std::vector result( - static_cast(PROTOTYPE_ROWS) * row_bytes); + static_cast(PROTOTYPE_ROW_PERIOD) * row_bytes); const size_t written = ggml_quantize_chunk( - type, source.data(), result.data(), 0, PROTOTYPE_ROWS, D, nullptr); + type, source.data(), result.data(), 0, + PROTOTYPE_ROW_PERIOD, D, nullptr); return written == result.size() ? result : std::vector{}; } +std::vector make_prototype_rows_f32( + ggml_type type, + float phase) { + const std::vector quantized = + make_prototype_rows(type, phase); + const ggml_type_traits * traits = ggml_get_type_traits(type); + if (quantized.empty() || !traits || !traits->to_float) return {}; + + const size_t row_bytes = ggml_row_size(type, D); + std::vector result( + static_cast(PROTOTYPE_ROW_PERIOD) * D); + for (int row = 0; row < PROTOTYPE_ROW_PERIOD; ++row) { + traits->to_float( + quantized.data() + static_cast(row) * row_bytes, + result.data() + static_cast(row) * D, D); + } + return result; +} + +std::vector make_reference_output( + ggml_type k_type, + ggml_type v_type, + const std::vector & q, + const std::vector & kv_seq_lens) { + const int n_seq = static_cast(kv_seq_lens.size()); + const std::vector k = + make_prototype_rows_f32(k_type, 0.1f); + const std::vector v = + make_prototype_rows_f32(v_type, 0.7f); + if (k.empty() || v.empty()) return {}; + + std::vector output( + static_cast(n_seq) * N_HEAD * D); + const double scale = 1.0 / std::sqrt(static_cast(D)); + const int q_per_kv = N_HEAD / N_HEAD_KV; + + // The synthetic cache repeats every 521 source rows. Counting each row's + // multiplicity makes a double-precision 128K oracle inexpensive while + // preserving the exact logical token sequence used by both GPU layouts. + for (int seq = 0; seq < n_seq; ++seq) { + for (int head = 0; head < N_HEAD; ++head) { + const int kv_head = head / q_per_kv; + std::array counts{}; + for (int token = 0; token < kv_seq_lens[seq]; ++token) { + const int prototype_row = + (static_cast(token) * 17 + + (seq * N_HEAD_KV + kv_head) * 7) % + PROTOTYPE_ROW_PERIOD; + ++counts[prototype_row]; + } + + const float * q_row = + q.data() + + (static_cast(seq) * N_HEAD + head) * D; + std::array scores{}; + double max_score = -INFINITY; + for (int row = 0; row < PROTOTYPE_ROW_PERIOD; ++row) { + if (counts[row] == 0) continue; + const float * k_row = + k.data() + static_cast(row) * D; + double dot = 0.0; + for (int d = 0; d < D; ++d) { + dot += + static_cast(q_row[d]) * k_row[d]; + } + scores[row] = dot * scale; + max_score = std::max(max_score, scores[row]); + } + + std::array weights{}; + double denominator = 0.0; + for (int row = 0; row < PROTOTYPE_ROW_PERIOD; ++row) { + if (counts[row] == 0) continue; + weights[row] = + static_cast(counts[row]) * + std::exp(scores[row] - max_score); + denominator += weights[row]; + } + + float * output_row = + output.data() + + (static_cast(seq) * N_HEAD + head) * D; + for (int d = 0; d < D; ++d) { + double numerator = 0.0; + for (int row = 0; row < PROTOTYPE_ROW_PERIOD; ++row) { + if (counts[row] == 0) continue; + numerator += + weights[row] * + v[static_cast(row) * D + d]; + } + output_row[d] = + static_cast(numerator / denominator); + } + } + } + return output; +} + bool cache_size_fits(ggml_type type, int64_t rows) { const size_t row_bytes = ggml_row_size(type, D); return rows > 0 && row_bytes > 0 && @@ -277,13 +387,13 @@ CachePair make_cache_pair( static_cast(kv_head) * pool_tokens + physical_token; // Keep sequence and head identity in the source pattern. - // Using logical_row % PROTOTYPE_ROWS would collapse both at - // common contexts such as 4096. + // Mixing the dimensions also avoids making the prototype + // depend on either cache layout's row numbering. const size_t prototype_row = static_cast( (static_cast(token) * 17 + (seq * N_HEAD_KV + kv_head) * 7) % - PROTOTYPE_ROWS); + PROTOTYPE_ROW_PERIOD); std::memcpy( result.contiguous.data() + @@ -413,10 +523,6 @@ bool collect_timing( } std::sort(sample_us.begin(), sample_us.end()); - const size_t p95_index = std::min( - sample_us.size() - 1, - static_cast( - std::ceil(sample_us.size() * 0.95) - 1.0)); result.iterations = iterations; const size_t upper_middle = sample_us.size() / 2; result.median_us = @@ -425,7 +531,6 @@ bool collect_timing( : (sample_us[upper_middle - 1] + sample_us[upper_middle]) / 2.0; - result.p95_us = sample_us[p95_index]; result.aggregate_queries_s = static_cast(n_seq) * 1.0e6 / result.median_us; return true; @@ -434,6 +539,7 @@ bool collect_timing( bool compare_outputs( ggml_tensor * paged_output, ggml_tensor * contiguous_output, + const std::vector & reference, int n_seq, BenchResult & result) { std::vector paged(ggml_nelements(paged_output)); @@ -445,8 +551,8 @@ bool compare_outputs( contiguous_output, contiguous.data(), 0, contiguous.size() * sizeof(contiguous[0])); - double squared_error = 0.0; - bool ok = true; + if (reference.size() != contiguous.size()) return false; + for (int seq = 0; seq < n_seq; ++seq) { for (int head = 0; head < N_HEAD; ++head) { for (int d = 0; d < D; ++d) { @@ -457,8 +563,11 @@ bool compare_outputs( const float paged_value = paged[paged_index]; const float contiguous_value = contiguous[contiguous_index]; + const float reference_value = + reference[contiguous_index]; if (!std::isfinite(paged_value) || - !std::isfinite(contiguous_value)) { + !std::isfinite(contiguous_value) || + !std::isfinite(reference_value)) { return false; } @@ -466,28 +575,30 @@ bool compare_outputs( std::fabs( static_cast(paged_value) - contiguous_value); - result.max_abs_error = - std::max(result.max_abs_error, abs_error); - squared_error += abs_error * abs_error; - - const double tolerance = - 1.0e-4 + - 1.0e-3 * - std::fabs(static_cast(contiguous_value)); - const double tolerance_ratio = abs_error / tolerance; - result.max_tolerance_ratio = + result.layout_max_abs_error = + std::max(result.layout_max_abs_error, abs_error); + + const double contiguous_oracle_error = + std::fabs( + static_cast(contiguous_value) - + reference_value); + const double paged_oracle_error = + std::fabs( + static_cast(paged_value) - + reference_value); + result.contiguous_oracle_max_abs = + std::max( + result.contiguous_oracle_max_abs, + contiguous_oracle_error); + result.paged_oracle_max_abs = std::max( - result.max_tolerance_ratio, - tolerance_ratio); - ok = ok && tolerance_ratio <= 1.0; + result.paged_oracle_max_abs, + paged_oracle_error); } } } - result.rmse = - std::sqrt( - squared_error / - static_cast(D * N_HEAD * n_seq)); - return ok; + return result.contiguous_oracle_max_abs <= MAX_ORACLE_ERROR && + result.paged_oracle_max_abs <= MAX_ORACLE_ERROR; } bool run_case( @@ -557,6 +668,10 @@ bool run_case( } } } + const std::vector reference_output = + make_reference_output( + options.k_type, options.v_type, contiguous_q, kv_seq_lens); + if (reference_output.empty()) return false; ggml_init_params params{}; params.mem_size = 8 * 1024 * 1024; @@ -606,7 +721,7 @@ bool run_case( ggml_tensor * paged_output = ggml_paged_attn( ctx.value, q_paged, k_paged, v_paged, table, kv_seq_lens_tensor, - 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE); + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, max_context); ggml_tensor * contiguous_output = ggml_flash_attn_ext( ctx.value, q_contiguous, k_contiguous, v_contiguous, padding_mask, 1.0f / std::sqrt(static_cast(D)), @@ -683,15 +798,22 @@ bool run_case( if (ggml_backend_graph_compute(backend, paged_graph) != GGML_STATUS_SUCCESS || ggml_backend_graph_compute(backend, contiguous_graph) != - GGML_STATUS_SUCCESS || - !compare_outputs( - paged_output, contiguous_output, n_seq, result)) { + GGML_STATUS_SUCCESS) { + std::fprintf( + stderr, + "attention graph execution failed at n_seq=%d\n", + n_seq); + return false; + } + if (!compare_outputs( + paged_output, contiguous_output, reference_output, + n_seq, result)) { std::fprintf( stderr, - "layout parity failed at n_seq=%d: max_abs=%.6g " - "max_tolerance_ratio=%.6g rmse=%.6g\n", - n_seq, result.max_abs_error, - result.max_tolerance_ratio, result.rmse); + "attention CPU-oracle check failed at n_seq=%d: " + "contiguous_max_abs=%.6g paged_max_abs=%.6g limit=%.6g\n", + n_seq, result.contiguous_oracle_max_abs, + result.paged_oracle_max_abs, MAX_ORACLE_ERROR); return false; } @@ -763,14 +885,6 @@ bool run_case( ggml_nbytes(k_paged) + ggml_nbytes(v_paged); result.paged_metadata_bytes = ggml_nbytes(table) + ggml_nbytes(kv_seq_lens_tensor); - result.contiguous_kv_mib = - static_cast(result.contiguous_kv_bytes) / - (1024.0 * 1024.0); - result.paged_kv_mib = - static_cast(result.paged_kv_bytes) / - (1024.0 * 1024.0); - result.paged_metadata_kib = - static_cast(result.paged_metadata_bytes) / 1024.0; return true; } @@ -799,100 +913,69 @@ std::string join_contexts(const std::vector & contexts) { return joined; } -void print_ragged_result( +void print_results( const Options & options, - const BenchResult & result) { + const std::vector & results) { const bool cuda_graphs_enabled = std::getenv("GGML_CUDA_DISABLE_GRAPHS") == nullptr; - const std::string contexts = - join_contexts(options.ragged_contexts); - const double step_time_ratio = - result.paged.median_us / result.contiguous.median_us; - const double throughput_ratio = - result.paged.aggregate_queries_s / - result.contiguous.aggregate_queries_s; - const double memory_ratio = - static_cast(result.contiguous_kv_bytes) / - static_cast(result.paged_kv_bytes); - const double memory_saving = - (1.0 - - static_cast(result.paged_kv_bytes) / - static_cast(result.contiguous_kv_bytes)) * - 100.0; - const uint64_t projected_contiguous_bytes = - result.contiguous_kv_bytes * QWEN_FULL_ATTENTION_LAYERS; - const uint64_t projected_paged_bytes = - result.paged_kv_bytes * QWEN_FULL_ATTENTION_LAYERS; - std::printf( - "# padded per-sequence contiguous vs paged GPU attention; " - "ragged decode, not whole-model throughput\n" - "# D=%d,Hq=%d,Hkv=%d,block=%d,contexts=%s," - "contiguous_context_pad=%d,k=%s,v=%s,warmup=%d,samples=%d," - "cuda_graphs_env=%s\n" - "# the Qwen projection multiplies K/V storage by %d full-attention " - "layers; recurrent state and all other model memory are excluded\n" + "# native padded-contiguous vs paged GPU attention; " + "not whole-model throughput or peak device memory\n" + "# D=%d,Hq=%d,Hkv=%d,block=%d,k=%s,v=%s,warmup=%d," + "samples=%d,cuda_graphs_env=%s,oracle_max_abs_limit=%.6g\n" "# paged_step_time_overhead_pct > 0 means paged is slower; " - "paged_over_padded_contiguous_query_throughput > 1 means paged " - "is faster\n", - D, N_HEAD, N_HEAD_KV, BLOCK_SIZE, contexts.c_str(), - result.padded_context, + "paged_kv_saving_pct > 0 means paged uses less K/V storage\n", + D, N_HEAD, N_HEAD_KV, BLOCK_SIZE, ggml_type_name(options.k_type), ggml_type_name(options.v_type), options.warmup, options.samples, cuda_graphs_enabled ? "allowed" : "disabled", - QWEN_FULL_ATTENTION_LAYERS); + MAX_ORACLE_ERROR); std::printf( "case,n_seq,contexts,live_tokens," "padded_contiguous_tokens,paged_pool_tokens," "padded_contiguous_iterations,paged_iterations," - "padded_contiguous_median_window_mean_us_per_step," - "paged_median_window_mean_us_per_step," - "padded_contiguous_p95_window_mean_us_per_step," - "paged_p95_window_mean_us_per_step," - "paged_over_padded_contiguous_step_time," + "padded_contiguous_median_us_per_step," + "paged_median_us_per_step," "paged_step_time_overhead_pct," "padded_contiguous_aggregate_attention_queries_s," "paged_aggregate_attention_queries_s," - "paged_over_padded_contiguous_query_throughput," "padded_contiguous_kv_bytes_one_layer," "paged_kv_bytes_one_layer," - "padded_contiguous_kv_mib_one_layer," - "paged_kv_mib_one_layer," - "padded_contiguous_kv_bytes_qwen16," - "paged_kv_bytes_qwen16," - "padded_contiguous_kv_gib_qwen16," - "paged_kv_gib_qwen16," - "padded_contiguous_over_paged_kv_ratio," "paged_kv_saving_pct,paged_metadata_bytes," - "max_abs_error,max_tolerance_ratio,rmse\n"); - std::printf( - "ragged,%d,\"%s\",%lld,%lld,%lld,%d,%d," - "%.3f,%.3f,%.3f,%.3f,%.6f,%.3f,%.2f,%.2f,%.6f," - "%llu,%llu,%.2f,%.2f,%llu,%llu,%.3f,%.3f,%.6f,%.3f,%llu," - "%.8g,%.8g,%.8g\n", - result.n_seq, contexts.c_str(), - static_cast(result.live_tokens), - static_cast(result.padded_tokens), - static_cast(result.paged_pool_tokens), - result.contiguous.iterations, result.paged.iterations, - result.contiguous.median_us, result.paged.median_us, - result.contiguous.p95_us, result.paged.p95_us, - step_time_ratio, (step_time_ratio - 1.0) * 100.0, - result.contiguous.aggregate_queries_s, - result.paged.aggregate_queries_s, throughput_ratio, - static_cast(result.contiguous_kv_bytes), - static_cast(result.paged_kv_bytes), - result.contiguous_kv_mib, result.paged_kv_mib, - static_cast(projected_contiguous_bytes), - static_cast(projected_paged_bytes), - static_cast(projected_contiguous_bytes) / - (1024.0 * 1024.0 * 1024.0), - static_cast(projected_paged_bytes) / - (1024.0 * 1024.0 * 1024.0), - memory_ratio, memory_saving, - static_cast(result.paged_metadata_bytes), - result.max_abs_error, result.max_tolerance_ratio, result.rmse); + "layout_max_abs_error,contiguous_oracle_max_abs_error," + "paged_oracle_max_abs_error\n"); + for (const CaseResult & item : results) { + const BenchResult & result = item.metrics; + const double step_time_ratio = + result.paged.median_us / result.contiguous.median_us; + const double memory_saving = + (1.0 - + static_cast(result.paged_kv_bytes) / + static_cast(result.contiguous_kv_bytes)) * + 100.0; + std::printf( + "%s,%d,\"%s\",%lld,%lld,%lld,%d,%d," + "%.3f,%.3f,%.3f,%.2f,%.2f,%llu,%llu,%.3f,%llu," + "%.8g,%.8g,%.8g\n", + item.name.c_str(), result.n_seq, + join_contexts(item.contexts).c_str(), + static_cast(result.live_tokens), + static_cast(result.padded_tokens), + static_cast(result.paged_pool_tokens), + result.contiguous.iterations, result.paged.iterations, + result.contiguous.median_us, result.paged.median_us, + (step_time_ratio - 1.0) * 100.0, + result.contiguous.aggregate_queries_s, + result.paged.aggregate_queries_s, + static_cast(result.contiguous_kv_bytes), + static_cast(result.paged_kv_bytes), + memory_saving, + static_cast(result.paged_metadata_bytes), + result.layout_max_abs_error, + result.contiguous_oracle_max_abs, + result.paged_oracle_max_abs); + } } } // namespace @@ -910,36 +993,36 @@ int main(int argc, char ** argv) { return 1; } - std::vector results; + std::vector results; bool ok = true; try { if (!options.ragged_contexts.empty()) { int padded_context = 0; - BenchResult result; + CaseResult result{"ragged", options.ragged_contexts, {}}; if (!padded_context_for( options.ragged_contexts, padded_context) || !run_case( backend, options, options.ragged_contexts, - padded_context, true, true, result)) { + padded_context, true, true, result.metrics)) { std::fprintf( stderr, "ragged attention benchmark failed\n"); ok = false; } else { - results.push_back(result); + results.push_back(std::move(result)); } } else { const std::array sequence_counts = {1, 2, 4, 8}; - results.reserve(sequence_counts.size()); + results.reserve(sequence_counts.size() + 1); for (size_t case_index = 0; case_index < sequence_counts.size(); ++case_index) { const int n_seq = sequence_counts[case_index]; const std::vector contexts( n_seq, options.context); - BenchResult result; + CaseResult result{"uniform", contexts, {}}; if (!run_case( backend, options, contexts, options.context, - false, case_index % 2 == 0, result)) { + false, case_index % 2 == 0, result.metrics)) { std::fprintf( stderr, "attention benchmark failed at n_seq=%d\n", @@ -947,7 +1030,29 @@ int main(int argc, char ** argv) { ok = false; break; } - results.push_back(result); + results.push_back(std::move(result)); + } + + // The default run also includes the same 16:1 long/short request + // mix used by the documented 128K + 7 x 8K capacity example, + // scaled by --context so ordinary benchmark runs remain quick. + if (ok) { + std::vector contexts( + 8, std::max(1, options.context / 16)); + contexts.front() = options.context; + int padded_context = 0; + CaseResult result{"ragged", contexts, {}}; + if (!padded_context_for(contexts, padded_context) || + !run_case( + backend, options, contexts, padded_context, + true, true, result.metrics)) { + std::fprintf( + stderr, + "default ragged attention benchmark failed\n"); + ok = false; + } else { + results.push_back(std::move(result)); + } } } } catch (const std::exception & error) { @@ -956,95 +1061,7 @@ int main(int argc, char ** argv) { } if (ok) { - if (!options.ragged_contexts.empty()) { - print_ragged_result(options, results.front()); - ggml_backend_free(backend); - return 0; - } - - const bool cuda_graphs_enabled = - std::getenv("GGML_CUDA_DISABLE_GRAPHS") == nullptr; - std::printf( - "# native contiguous vs paged GPU attention; " - "not whole-model throughput\n" - "# D=%d,Hq=%d,Hkv=%d,block=%d,context=%d,k=%s,v=%s," - "warmup=%d,samples=%d,cuda_graphs_env=%s\n" - "# paged_step_time_overhead_pct > 0 means paged is slower; " - "paged_over_contiguous_query_throughput > 1 means paged " - "is faster\n", - D, N_HEAD, N_HEAD_KV, BLOCK_SIZE, options.context, - ggml_type_name(options.k_type), - ggml_type_name(options.v_type), - options.warmup, options.samples, - cuda_graphs_enabled ? "allowed" : "disabled"); - std::printf( - "n_seq,context,k_type,v_type," - "contiguous_iterations,paged_iterations," - "contiguous_median_window_mean_us_per_step," - "paged_median_window_mean_us_per_step," - "contiguous_p95_window_mean_us_per_step," - "paged_p95_window_mean_us_per_step," - "paged_over_contiguous_step_time," - "paged_step_time_overhead_pct," - "contiguous_aggregate_attention_queries_s," - "paged_aggregate_attention_queries_s," - "paged_over_contiguous_query_throughput," - "contiguous_scaling_vs_b1," - "paged_scaling_vs_b1," - "contiguous_batch_efficiency," - "paged_batch_efficiency," - "contiguous_kv_mib,paged_kv_mib," - "paged_metadata_kib," - "max_abs_error,max_tolerance_ratio,rmse\n"); - - const double contiguous_baseline = - results.front().contiguous.aggregate_queries_s; - const double paged_baseline = - results.front().paged.aggregate_queries_s; - for (const BenchResult & result : results) { - const double contiguous_scaling = - result.contiguous.aggregate_queries_s / - contiguous_baseline; - const double paged_scaling = - result.paged.aggregate_queries_s / paged_baseline; - const double step_time_ratio = - result.paged.median_us / - result.contiguous.median_us; - const double throughput_ratio = - result.paged.aggregate_queries_s / - result.contiguous.aggregate_queries_s; - std::printf( - "%d,%d,%s,%s,%d,%d," - "%.3f,%.3f,%.3f,%.3f," - "%.6f,%.3f,%.2f,%.2f,%.6f," - "%.4f,%.4f,%.4f,%.4f," - "%.2f,%.2f,%.3f," - "%.8g,%.8g,%.8g\n", - result.n_seq, options.context, - ggml_type_name(options.k_type), - ggml_type_name(options.v_type), - result.contiguous.iterations, - result.paged.iterations, - result.contiguous.median_us, - result.paged.median_us, - result.contiguous.p95_us, - result.paged.p95_us, - step_time_ratio, - (step_time_ratio - 1.0) * 100.0, - result.contiguous.aggregate_queries_s, - result.paged.aggregate_queries_s, - throughput_ratio, - contiguous_scaling, - paged_scaling, - contiguous_scaling / result.n_seq, - paged_scaling / result.n_seq, - result.contiguous_kv_mib, - result.paged_kv_mib, - result.paged_metadata_kib, - result.max_abs_error, - result.max_tolerance_ratio, - result.rmse); - } + print_results(options, results); } ggml_backend_free(backend); diff --git a/server/test/test_kvflash_pool_sizing.cpp b/server/test/test_kvflash_pool_sizing.cpp index fc4dd64b3..c4d2d70c9 100644 --- a/server/test/test_kvflash_pool_sizing.cpp +++ b/server/test/test_kvflash_pool_sizing.cpp @@ -18,9 +18,14 @@ static void expect(bool cond, const char * msg) { } int main() { + const int max_ctx = 131072; + + setenv("DFLASH_KVFLASH", "0", 1); + expect(kvflash_pool_from_env(max_ctx) == 0, + "explicit zero should disable KVFlash"); + setenv("DFLASH_KVFLASH", "auto", 1); unsetenv("DFLASH_KVFLASH_MAX_POOL"); - const int max_ctx = 131072; // No budget supplied -> fallback fraction of max_ctx (the buggy placement // path). scorer_expected toggles 1/2 vs 1/4. diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index ca9749b5a..d4efeaf72 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -4,10 +4,10 @@ #include "ggml-cuda.h" #include -#include #include #include #include +#include #include namespace { @@ -16,29 +16,66 @@ constexpr int D = 256; constexpr int N_HEAD = 24; constexpr int N_HEAD_KV = 4; constexpr int BLOCK_SIZE = 16; -constexpr int MAX_BLOCKS = 3; -constexpr int N_SEQ = 6; -constexpr int PHYSICAL_BLOCKS = MAX_BLOCKS * N_SEQ; -constexpr int POOL_TOKENS = PHYSICAL_BLOCKS * BLOCK_SIZE; - -const std::array KV_SEQ_LENS = {1, 15, 16, 17, 31, 33}; - -// Rows are sequences. Every live block is private, and deliberately -// shuffled so a contiguous-cache implementation cannot pass this test. -// Two in-range entries are deliberately invalid — one negative, one past -// the physical pool — to pin down the kernel's bounds guard: it must skip -// those blocks (matching the oracle) instead of reading out of bounds. -const std::array BLOCK_TABLE = { - 11, 2, 15, - 4, 17, 7, - 9, 0, 13, - 5, -1, 1, // token 16 of seq 3 maps to a negative block - 14, 8, 3, - 10, 99, 12, // tokens 16-31 of seq 5 map past the pool +constexpr float MAX_ABS_ERROR = 5.0e-4f; + +struct TestCase { + const char * name; + int max_blocks; + std::vector kv_seq_lens; + bool corrupt_blocks; }; -bool block_is_valid(int32_t block) { - return block >= 0 && block < PHYSICAL_BLOCKS; +int clamped_seq_len(const TestCase & test_case, int seq) { + return std::max( + 0, std::min( + test_case.kv_seq_lens[seq], + test_case.max_blocks * BLOCK_SIZE)); +} + +int count_physical_blocks(const TestCase & test_case) { + int result = 0; + for (size_t seq = 0; seq < test_case.kv_seq_lens.size(); ++seq) { + const int kv_seq_len = + clamped_seq_len(test_case, static_cast(seq)); + result += (kv_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE; + } + return result; +} + +bool block_is_valid(int32_t block, int physical_blocks) { + return block >= 0 && block < physical_blocks; +} + +std::vector make_block_table( + const TestCase & test_case, + int physical_blocks) { + const int n_seq = static_cast(test_case.kv_seq_lens.size()); + std::vector result( + static_cast(test_case.max_blocks) * n_seq, -1); + + // 37 is coprime with both cases' live-block counts, so this maps every + // logical page to a unique shuffled physical page. + int ordinal = 0; + for (int seq = 0; seq < n_seq; ++seq) { + const int kv_seq_len = clamped_seq_len(test_case, seq); + const int n_blocks = + (kv_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE; + for (int logical = 0; logical < n_blocks; ++logical) { + result[seq * test_case.max_blocks + logical] = + (ordinal * 37 + 11) % physical_blocks; + ++ordinal; + } + } + + if (test_case.corrupt_blocks) { + // Two in-range entries are deliberately invalid — one negative, one + // past the physical pool — to pin down the physical-block bounds + // guard. Sequence 9 spans 65 blocks, so its corrupt page also + // exercises the partitioned long-context variant. + result[4 * test_case.max_blocks + 1] = -1; + result[9 * test_case.max_blocks + 7] = physical_blocks + 9; + } + return result; } std::vector quantize_rows(ggml_type type, @@ -73,26 +110,32 @@ std::vector dequantize_rows(ggml_type type, } std::vector reference_attention( + const TestCase & test_case, + const std::vector & block_table, + int pool_tokens, + int physical_blocks, const std::vector & q, const std::vector & k, const std::vector & v) { std::vector output(q.size(), 0.0f); const float scale = 1.0f / std::sqrt(static_cast(D)); const int q_per_kv = N_HEAD / N_HEAD_KV; + const int n_seq = static_cast(test_case.kv_seq_lens.size()); - for (int seq = 0; seq < N_SEQ; ++seq) { - const int kv_seq_len = KV_SEQ_LENS[seq]; + for (int seq = 0; seq < n_seq; ++seq) { + const int kv_seq_len = clamped_seq_len(test_case, seq); for (int head = 0; head < N_HEAD; ++head) { const int kv_head = head / q_per_kv; const float * q_row = - q.data() + (static_cast(head) * N_SEQ + seq) * D; + q.data() + (static_cast(head) * n_seq + seq) * D; std::vector scores(kv_seq_len); float max_score = -INFINITY; for (int token = 0; token < kv_seq_len; ++token) { const int block = - BLOCK_TABLE[seq * MAX_BLOCKS + token / BLOCK_SIZE]; - if (!block_is_valid(block)) { + block_table[ + seq * test_case.max_blocks + token / BLOCK_SIZE]; + if (!block_is_valid(block, physical_blocks)) { // Mirrors the kernel: invalid blocks contribute nothing. scores[token] = -INFINITY; continue; @@ -100,7 +143,7 @@ std::vector reference_attention( const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; const float * k_row = k.data() + - (static_cast(kv_head) * POOL_TOKENS + physical) * D; + (static_cast(kv_head) * pool_tokens + physical) * D; float dot = 0.0f; for (int d = 0; d < D; ++d) dot += q_row[d] * k_row[d]; scores[token] = dot * scale; @@ -114,15 +157,16 @@ std::vector reference_attention( } float * out_row = - output.data() + (static_cast(head) * N_SEQ + seq) * D; + output.data() + (static_cast(head) * n_seq + seq) * D; for (int token = 0; token < kv_seq_len; ++token) { const int block = - BLOCK_TABLE[seq * MAX_BLOCKS + token / BLOCK_SIZE]; - if (!block_is_valid(block)) continue; + block_table[ + seq * test_case.max_blocks + token / BLOCK_SIZE]; + if (!block_is_valid(block, physical_blocks)) continue; const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; const float * v_row = v.data() + - (static_cast(kv_head) * POOL_TOKENS + physical) * D; + (static_cast(kv_head) * pool_tokens + physical) * D; const float probability = scores[token] / denominator; for (int d = 0; d < D; ++d) { out_row[d] += probability * v_row[d]; @@ -133,7 +177,16 @@ std::vector reference_attention( return output; } -bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { +bool run_case(ggml_backend_t backend, + const TestCase & test_case, + ggml_type k_type, + ggml_type v_type) { + const int n_seq = static_cast(test_case.kv_seq_lens.size()); + const int physical_blocks = count_physical_blocks(test_case); + const int pool_tokens = physical_blocks * BLOCK_SIZE; + const std::vector block_table = + make_block_table(test_case, physical_blocks); + ggml_init_params params{}; params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; @@ -141,22 +194,25 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { if (!ctx) return false; ggml_tensor * q = - ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_SEQ, N_HEAD); + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, n_seq, N_HEAD); ggml_tensor * k = - ggml_new_tensor_3d(ctx, k_type, D, POOL_TOKENS, N_HEAD_KV); + ggml_new_tensor_3d(ctx, k_type, D, pool_tokens, N_HEAD_KV); ggml_tensor * v = - ggml_new_tensor_3d(ctx, v_type, D, POOL_TOKENS, N_HEAD_KV); + ggml_new_tensor_3d(ctx, v_type, D, pool_tokens, N_HEAD_KV); ggml_tensor * table = - ggml_new_tensor_2d(ctx, GGML_TYPE_I32, MAX_BLOCKS, N_SEQ); + ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, test_case.max_blocks, n_seq); ggml_tensor * kv_seq_lens = - ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N_SEQ); + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seq); for (ggml_tensor * input : {q, k, v, table, kv_seq_lens}) { ggml_set_input(input); } ggml_tensor * output = ggml_paged_attn( ctx, q, k, v, table, kv_seq_lens, - 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE); + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, + *std::max_element(test_case.kv_seq_lens.begin(), + test_case.kv_seq_lens.end())); ggml_set_output(output); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); @@ -169,9 +225,9 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { return false; } - std::vector q_data(static_cast(D) * N_SEQ * N_HEAD); + std::vector q_data(static_cast(D) * n_seq * N_HEAD); std::vector k_source( - static_cast(D) * POOL_TOKENS * N_HEAD_KV); + static_cast(D) * pool_tokens * N_HEAD_KV); std::vector v_source(k_source.size()); for (size_t i = 0; i < q_data.size(); ++i) { q_data[i] = std::sin(static_cast(i) * 0.013f) * 0.25f; @@ -181,7 +237,7 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { v_source[i] = std::sin(static_cast(i) * 0.011f + 0.3f) * 0.30f; } - const int64_t cache_rows = static_cast(POOL_TOKENS) * N_HEAD_KV; + const int64_t cache_rows = static_cast(pool_tokens) * N_HEAD_KV; const std::vector k_data = quantize_rows(k_type, k_source, cache_rows); const std::vector v_data = @@ -198,11 +254,13 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { q_data.size() * sizeof(q_data[0])); ggml_backend_tensor_set(k, k_data.data(), 0, k_data.size()); ggml_backend_tensor_set(v, v_data.data(), 0, v_data.size()); - ggml_backend_tensor_set(table, BLOCK_TABLE.data(), 0, - BLOCK_TABLE.size() * sizeof(BLOCK_TABLE[0])); ggml_backend_tensor_set( - kv_seq_lens, KV_SEQ_LENS.data(), 0, - KV_SEQ_LENS.size() * sizeof(KV_SEQ_LENS[0])); + table, block_table.data(), 0, + block_table.size() * sizeof(block_table[0])); + ggml_backend_tensor_set( + kv_seq_lens, test_case.kv_seq_lens.data(), 0, + test_case.kv_seq_lens.size() * + sizeof(test_case.kv_seq_lens[0])); ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; } @@ -212,7 +270,9 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(actual[0])); const std::vector expected = - reference_attention(q_data, k_reference, v_reference); + reference_attention( + test_case, block_table, pool_tokens, physical_blocks, + q_data, k_reference, v_reference); max_abs_error = 0.0f; for (size_t i = 0; i < actual.size(); ++i) { if (!std::isfinite(actual[i])) { @@ -222,11 +282,11 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { max_abs_error = std::max(max_abs_error, std::fabs(actual[i] - expected[i])); } - ok = ok && max_abs_error < 3.0e-3f; + ok = ok && max_abs_error < MAX_ABS_ERROR; } - std::printf("paged attention K=%-4s V=%-4s max_abs=%.6g %s\n", - ggml_type_name(k_type), ggml_type_name(v_type), + std::printf("paged attention %-11s K=%-4s V=%-4s max_abs=%.6g %s\n", + test_case.name, ggml_type_name(k_type), ggml_type_name(v_type), max_abs_error, ok ? "PASS" : "FAIL"); ggml_gallocr_free(allocator); ggml_free(ctx); @@ -235,20 +295,43 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type) { } // namespace -int main() { +int main(int argc, char ** argv) { ggml_backend_t backend = ggml_backend_cuda_init(0); if (!backend) { std::fprintf(stderr, "GPU backend unavailable\n"); return 1; } + const TestCase partitioned_case{ + "partitioned", + 65, + // Retains page boundaries and >64 blocks, while pinning both context + // clamps: negative becomes empty and over-capacity becomes 65 blocks. + {-7, 1, 15, 16, 17, 31, 33, 257, 511, 1025, 2000}, + true, + }; + const TestCase direct_case{ + "direct", + 64, + {0, 1, 15, 16, 17, 257, 511}, + false, + }; + const bool direct = + argc == 2 && std::strcmp(argv[1], "--direct") == 0; + if (argc > 2 || (argc == 2 && !direct)) { + std::fprintf(stderr, "usage: %s [--direct]\n", argv[0]); + ggml_backend_free(backend); + return 2; + } + const TestCase & test_case = direct ? direct_case : partitioned_case; + const ggml_type types[] = { GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, }; bool ok = true; for (ggml_type k_type : types) { for (ggml_type v_type : types) { - ok = run_case(backend, k_type, v_type) && ok; + ok = run_case(backend, test_case, k_type, v_type) && ok; } } diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp index b5f027222..2b09058a2 100644 --- a/server/test/test_paged_kv_pool.cpp +++ b/server/test/test_paged_kv_pool.cpp @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include using namespace dflash::common; @@ -63,6 +65,20 @@ static void test_block_boundaries() { } } +static void test_nondefault_block_size() { + PagedKvPool pool(5, 3, 7); + CHECK(pool.block_size() == 7); + CHECK(pool.max_sequences() == 3); + + const auto handle = acquire(pool, 77); + const auto append = pool.append(handle, 15); + CHECK(append.status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, handle).block_table, {0, 1, 2})); + CHECK(append.write_slots[6].block_offset == 6); + CHECK(append.write_slots[7].physical_block == 1); + CHECK(append.write_slots[14].physical_token_index == 14); +} + static void test_reserve_then_append() { PagedKvPool pool(4, 2); const auto handle = acquire(pool, 7); @@ -82,6 +98,14 @@ static void test_reserve_then_append() { CHECK(pool.reserve(handle, 16) == PagedKvStatus::Ok); CHECK(equals(sequence(pool, handle).block_table, {0, 1})); + + const auto before_zero = sequence(pool, handle); + const auto append_zero = pool.append(handle, 0); + CHECK(append_zero.status == PagedKvStatus::Ok); + CHECK(append_zero.write_slots.empty()); + const auto after_zero = sequence(pool, handle); + CHECK(after_zero.kv_seq_len == before_zero.kv_seq_len); + CHECK(after_zero.block_table == before_zero.block_table); } static void test_compact_append() { @@ -207,21 +231,34 @@ static void test_request_identity_and_stale_handles() { CHECK(after_reset.generation != replacement.generation); CHECK(pool.append(after_reset, 1)); CHECK(sequence(pool, after_reset).block_table.front() == 0); + + const PagedKvSequenceHandle out_of_range{ + std::numeric_limits::max(), 1}; + CHECK(pool.append(out_of_range, 1).status == + PagedKvStatus::StaleHandle); + CHECK(pool.reserve(out_of_range, 1) == PagedKvStatus::StaleHandle); + CHECK(pool.release(out_of_range) == PagedKvStatus::StaleHandle); + PagedKvSequenceSnapshot invalid_snapshot; + CHECK(pool.sequence(out_of_range, invalid_snapshot) == + PagedKvStatus::StaleHandle); } static void test_metadata_snapshot() { - PagedKvPool pool(7, 4); + PagedKvPool pool(8, 4); const auto first = acquire(pool, 11); const auto second = acquire(pool, 22); + const auto third = acquire(pool, 33); CHECK(pool.append(first, 17)); CHECK(pool.append(second, 1)); - CHECK(pool.reserve(second, 33) == PagedKvStatus::Ok); + CHECK(pool.append(third, 17)); + CHECK(pool.release(second) == PagedKvStatus::Ok); + CHECK(pool.reserve(third, 33) == PagedKvStatus::Ok); const auto metadata = pool.metadata_snapshot(); CHECK(metadata.block_size == 16); - CHECK(metadata.physical_block_count == 7); + CHECK(metadata.physical_block_count == 8); CHECK(metadata.sequences.size() == 2); - CHECK(equals(metadata.block_table, {0, 1, 2, 3, 4})); + CHECK(equals(metadata.block_table, {0, 1, 3, 4, 2})); const auto & first_meta = metadata.sequences[0]; CHECK(first_meta.request_id == 11); @@ -231,13 +268,13 @@ static void test_metadata_snapshot() { CHECK(first_meta.block_table_offset == 0); CHECK(first_meta.block_count == 2); - const auto & second_meta = metadata.sequences[1]; - CHECK(second_meta.request_id == 22); - CHECK(second_meta.generation == second.generation); - CHECK(second_meta.sequence_slot == second.slot); - CHECK(second_meta.kv_seq_len == 1); - CHECK(second_meta.block_table_offset == 2); - CHECK(second_meta.block_count == 3); + const auto & third_meta = metadata.sequences[1]; + CHECK(third_meta.request_id == 33); + CHECK(third_meta.generation == third.generation); + CHECK(third_meta.sequence_slot == third.slot); + CHECK(third_meta.kv_seq_len == 17); + CHECK(third_meta.block_table_offset == 2); + CHECK(third_meta.block_count == 3); } static void test_duplicate_request_is_transactional() { @@ -251,6 +288,40 @@ static void test_duplicate_request_is_transactional() { CHECK(output == first); } +static bool constructor_rejects(uint32_t physical_blocks, + uint32_t max_sequences, + uint32_t block_size) { + try { + PagedKvPool pool(physical_blocks, max_sequences, block_size); + } catch (const std::invalid_argument &) { + return true; + } + return false; +} + +static void test_invalid_arguments() { + CHECK(constructor_rejects(0, 1, 16)); + CHECK(constructor_rejects( + 0, std::numeric_limits::max(), 16)); + CHECK(constructor_rejects(1, 0, 16)); + CHECK(constructor_rejects(1, 1, 0)); + CHECK(constructor_rejects( + 2, 1, std::numeric_limits::max())); + + PagedKvPool pool( + 1, 1, std::numeric_limits::max()); + const auto handle = acquire(pool, 99); + CHECK(pool.append(handle, 1).status == PagedKvStatus::Ok); + const auto before = sequence(pool, handle); + const auto overflow = + pool.append(handle, std::numeric_limits::max()); + CHECK(overflow.status == PagedKvStatus::InvalidArgument); + CHECK(overflow.write_slots.empty()); + const auto after = sequence(pool, handle); + CHECK(after.kv_seq_len == before.kv_seq_len); + CHECK(after.block_table == before.block_table); +} + static void test_paged_attention_compatibility() { PagedAttentionOptions options; options.enabled = true; @@ -273,6 +344,7 @@ static void test_paged_attention_compatibility() { int main() { test_block_boundaries(); + test_nondefault_block_size(); test_reserve_then_append(); test_compact_append(); test_noncontiguous_reuse_and_isolation(); @@ -280,6 +352,7 @@ int main() { test_request_identity_and_stale_handles(); test_metadata_snapshot(); test_duplicate_request_is_transactional(); + test_invalid_arguments(); test_paged_attention_compatibility(); if (failures != 0) { From e28f8d311cf8000cab5210d63bdbfbc15c938905 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Fri, 24 Jul 2026 19:54:02 +0000 Subject: [PATCH 3/3] fix(qwen35): write first AR token at committed position --- server/src/common/paged_kv_pool.cpp | 9 ++++- server/src/qwen35/qwen35_backend.cpp | 60 +++++++++++----------------- server/src/qwen35/qwen35_backend.h | 1 - 3 files changed, 32 insertions(+), 38 deletions(-) diff --git a/server/src/common/paged_kv_pool.cpp b/server/src/common/paged_kv_pool.cpp index 53b4ae120..5f080632a 100644 --- a/server/src/common/paged_kv_pool.cpp +++ b/server/src/common/paged_kv_pool.cpp @@ -81,12 +81,19 @@ PagedKvStatus PagedKvPool::acquire(PagedKvRequestId request_id, uint64_t generation = sequence.generation + 1; if (generation == 0) generation = 1; + // Insert before mutating anything: emplace() is the only step here that can + // throw (rehash allocation), and a slot marked active while its request is + // missing from the map would desync active_sequence_count() (map size) from + // metadata_snapshot() (which walks sequences_ by the active flag). Past this + // point every step is non-allocating, so the claim completes or the pool is + // untouched. + request_to_slot_.emplace(request_id, slot); + // release_sequence(), reset(), and construction leave every free slot // empty; acquire only installs its new identity. sequence.request_id = request_id; sequence.generation = generation; sequence.active = true; - request_to_slot_.emplace(request_id, slot); take_lowest(free_sequence_slots_); out_handle = {slot, generation}; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index d743f9148..cf78dbaa1 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -458,22 +458,6 @@ bool Qwen35Backend::begin_paged_sequence(uint32_t prompt_tokens) { return true; } -bool Qwen35Backend::append_paged_gap(uint32_t logical_position) { - if (!cfg_.paged_attention) return true; - if (!paged_kv_pool_ || !paged_sequence_) return false; - PagedKvAppendSpan append = - paged_kv_pool_->append_compact(*paged_sequence_, /*token_count=*/1); - if (!append || append.token_count != 1 || - append.last.logical_position != logical_position) { - std::fprintf(stderr, - "[paged-attention] logical gap allocation failed at %u (%s)\n", - logical_position, paged_kv_status_string(append.status)); - set_last_error("paged attention logical gap allocation failed"); - return false; - } - return true; -} - bool Qwen35Backend::prepare_paged_decode_step(uint32_t logical_position) { if (!cfg_.paged_attention) return true; if (!paged_kv_pool_ || !paged_sequence_ || !sg_.kv_write_rows || @@ -1130,7 +1114,12 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, // positions past max_ctx) instead of finishing cleanly. The HTTP // admission gate normally guarantees prompt+n_gen <= max_ctx, but // daemon-command callers are not required to. - const int max_ar_n_gen = cfg_.device.max_ctx - committed; + // + // AR decode emits its first token from the prefill logits without + // writing a K/V row, so the last position it forwards is + // committed+n_gen-2: one token more than the remaining context + // still fits. + const int max_ar_n_gen = cfg_.device.max_ctx - committed + 1; if (max_ar_n_gen <= 0) { // do_ar_decode would return true on a non-positive count and // report an empty completion as success. @@ -1758,15 +1747,20 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, // normal sampling resumes (model writes visible answer). bool budget_close_started = false; int close_inject_pos = 0; - // Capture entry KV position so the budget check is in the + // Capture the entry emit count so the budget check is in the // "generated since entry" frame, not the absolute KV frame. // n_gen is the gen-only count (or the remaining-budget remap done by - // spec-decode tail-off); subtracting committed_now (absolute KV = - // prompt_len + tokens generated this call) directly would treat - // prompt-length tokens as if they were generated output, firing - // force-close prompt_len tokens early on prompted requests and - // potentially going negative after spec-decode tail-off. - const int committed_at_entry = committed; + // spec-decode tail-off); measuring against the absolute KV position + // (prompt_len + tokens generated this call) would treat prompt-length + // tokens as if they were generated output, firing force-close + // prompt_len tokens early on prompted requests and potentially going + // negative after spec-decode tail-off. + // + // Count emitted tokens rather than KV positions: the first AR token is + // sampled from the prefill logits and stays pending until the first loop + // iteration forwards it, so `committed` lags the emit count by one for + // the whole loop and is not a usable "tokens generated" proxy. + const size_t out_tokens_at_entry = out_tokens.size(); auto maybe_force_close = [&](int32_t & tok, int committed_now) { if (budget_hook.close_token_ids.empty()) return; @@ -1789,11 +1783,12 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, if (budget_close_started) return; // Check if budget has tightened to the force-close trigger. - // generated = tokens produced in THIS do_ar_decode call; + // generated = tokens already emitted by THIS do_ar_decode call + // (`tok` is the candidate for the next one, not yet pushed); // remaining = budget headroom, measured against n_gen (the // requested gen count or tail-off remap, never against the // absolute KV position which would mis-count the prompt). - const int generated = committed_now - committed_at_entry; + const int generated = (int)(out_tokens.size() - out_tokens_at_entry); int remaining = n_gen - generated; if (remaining <= budget_hook.hard_limit_remaining) { // Don't trigger if the model already sampled the first close @@ -1830,7 +1825,6 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, if (n_gen <= 0) return true; auto t_dec0_ar = std::chrono::steady_clock::now(); - const size_t out_tokens_at_entry = out_tokens.size(); const int _min_floor = dflash_min_tokens_floor(); static const int _repeat_guard = []{ const int explicit_guard = @@ -1874,15 +1868,9 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, io.emit(first_tok); if (kvflash_active()) kvflash_history_.push_back(first_tok); if (IS_EOS_TOK(first_tok, w_)) return true; - // Preserve the existing AR position semantics: the first sampled - // token advances cur_pos without a cache write. The page table must - // include that logical row so subsequent attention sees exactly the - // same context as the dense path. Prefill does not refresh this row, - // so after the first request it may contain residual K/V from the - // previous request. - if (!append_paged_gap((uint32_t)committed)) return false; - committed++; - cache_.cur_pos = committed; + // The first token is pending: the prefill logits produced it, but its + // K/V row has not been written yet. The first loop iteration below + // forwards it at `committed`; only that compute may advance the cache. } // AR decode loop for remaining tokens diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 58bf06456..a408438df 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -331,7 +331,6 @@ class Qwen35Backend : public ModelBackend { bool * degenerate_close_out = nullptr); bool begin_paged_sequence(uint32_t prompt_tokens); - bool append_paged_gap(uint32_t logical_position); bool prepare_paged_decode_step(uint32_t logical_position); void end_paged_sequence();