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..7a2fc4605 --- /dev/null +++ b/optimizations/paged_attention/README.md @@ -0,0 +1,133 @@ +

+ ← lucebox-hub +

+ +

Luce Paged Attention

+ +

+ Exact fixed-block K/V for long-context serving.
+ Every sequence owns a block table; decode reads reusable 16-token physical blocks straight through it.
+ Ragged 8-request profile (128K + 7x8K) on one RTX 3090: + 82% less K/V storage and 1.35x faster attention. +

+ +--- + +## How it works + +- **Blocks.** Full-attention K/V rows live in fixed 16-token physical blocks in + one pool. Nothing is evicted and nothing is approximated — this is a layout + and allocation mechanism, not a residency policy. +- **Block tables.** Each sequence owns a logical→physical map. Decode walks it, + so a sequence never has to occupy one contiguous K/V range, and the pool only + ever holds live, block-rounded tokens. +- **One decode op.** `GGML_OP_PAGED_ATTN` on CUDA and HIP: D=256 GQA, F16/Q4_0/Q8_0 + K/V. One warp covers all query heads of a K/V group per row loaded, and long + contexts split into partitions merged by a stable split-softmax. +- **Resident metadata.** The block table and sequence length live in the + persistent target cache next to the pool, so a decode step uploads 4 bytes per + newly filled block instead of the whole table every token. +- **Prefill is unchanged.** The prompt runs through the existing exact chunked + path against a freshly reset pool, so its layout is identity-mapped. Paging + begins at autoregressive decode. + +Not to be confused with [KVFlash](../kvflash/README.md): that one bounds how much +context stays resident and evicts cold chunks. Paging keeps every token. + +## Usage + +```bash +./server/build/dflash_server model.gguf --paged-attention --max-ctx 131072 +``` + +## Compatibility + +| | | +|---|---| +| Architecture | `qwen35` — dense Qwen3.5 / Qwen3.6 | +| Placement | one local CUDA or HIP device | +| Attention | full only (`--fa-window 0`) | +| K/V types | F16, Q4_0, Q8_0 | +| Decode | autoregressive, one live sequence | +| Block size | 16 tokens, fixed | + +**Rejected at startup** (exit 2, with the reason): `--draft`, `--ddtree`, +`--target-devices`, `--target-shard-ipc-bin`, a non-zero `--fa-window`, +`--prefill-compression`, and `DFLASH_KVFLASH`. The rules live with every other +launch-admission rule in `check_feature_compatibility()`; which architecture and +placement they are allowed on is one row in `model_capabilities.h`. + +**Disabled, not rejected:** prefix, prefill, and disk snapshots. Their format +assumes contiguous K/V rows, so `--paged-attention` zeroes the caps and says so. + +**Why one sequence.** Qwen3.6 is hybrid: 48 of its 64 layers hold recurrent Gated +DeltaNet state rather than an attention cache, and that state is backend-global +today. Two live requests would mix it even with perfectly isolated K/V blocks. +Making it sequence-indexed is step 2 below. + +## Numbers + +Attention ops only, Q4_0 K/V, Qwen dims (`D=256, Hq=24, Hkv=4`). Uniform batches +of 1/2/4/8 sequences, as `paged ÷ contiguous` throughput — above 1.0x is paged +winning. + +| Context | RTX 3090 (CUDA) | Radeon 8060S (HIP) | +|---|---|---| +| 4K | 1.10–1.20x | 4.28–5.60x | +| 32K | 1.12–1.23x | 4.72–5.16x | +| 128K | 1.23–1.97x | 5.00–6.42x | + +Ragged 8-request profile, `128K + 7x8K` on the RTX 3090. Paging stores only live, +block-rounded tokens (188,416) instead of padding every sequence to 128K +(1,048,576 slots): + +| | contiguous | paged | +|---|---|---| +| attention step | 3.007 ms | **2.235 ms** (1.35x) | +| K/V, one layer | 1152 MiB | **207 MiB** | +| K/V, 16 full-attn layers | 18.0 GiB | **3.23 GiB** (−82%) | + +Two caveats. The HIP ratios are inflated by a weak native HIP +`flash_attn_ext_vec` baseline — don't read them as whole-model throughput. And +small ragged batches are still behind (`4K + 7x256` measures 0.48x), because +short sequences carry dead partitions sized by the longest sequence; the stream-K +kernel in step 4 is what fixes that. + +## Roadmap + +1. **Done.** Decode-only `ggml_paged_attn`, block manager, CUDA/HIP integration, + benchmarks. +2. **Scheduler groundwork.** Sequence-indexed DeltaNet/conv state, iteration-level + scheduling, dynamic decode batches. +3. **Paged-aware chunk prefill.** Arbitrary block-table writes, multi-token + queries over preceding paged K/V, causal attention within the chunk. +4. **Continuous admission and interleaving.** Token budgets, decode + prioritization, admission control, cancellation, prefill/decode interleaving, + plus the stream-K decode kernel for ragged batches. +5. **Variable-length batched prefill.** +6. **Prefix caching / CoW.** Shared prefix blocks with reference counting and + copy-on-write. + +## Benchmarks and tests + +```bash +cmake --build "$BUILD_DIR" --target bench_paged_attention +"$BUILD_DIR/bench_paged_attention" --context 131072 --k-type q4_0 --v-type q4_0 +``` + +Each run compares one native padded-contiguous attention op against one paged op +at `n_seq=1,2,4,8`, then adds a ragged 8-request row. Both layouts get identical +logical Q/K/V and identical quantized K/V rows, and both must pass a +double-precision CPU oracle before anything is timed. The CSV reports median step +time, aggregate query throughput, exact K/V and metadata bytes, and each path's +oracle error. It measures attention ops — not the full Qwen graph, HTTP +scheduling, or continuous-batching throughput. + +`ctest -R paged` covers the host allocator plus the nine F16/Q4_0/Q8_0 K/V +combinations, in both the partitioned and the forced-single-partition kernel +paths. + +The fixed-block/block-table shape follows the +[llama.cpp paged-attention discussion][llama-discussion]. + +[llama-discussion]: https://github.com/ggml-org/llama.cpp/discussions/21961 diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 7183b5651..88fedf8a8 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 @@ -1032,6 +1033,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}) @@ -1204,6 +1214,42 @@ 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) + 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") + 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..303a4fbfa 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,30 @@ 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) + // 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 + // 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, + int max_kv_seq_len); + // 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..c9b8190cf 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,18 @@ 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 { + // 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]; + }; + 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 +960,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..826e89d50 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -0,0 +1,971 @@ +#include "paged-attn.cuh" + +#include "fattn-common.cuh" + +#include +#include +#include +#include + +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, + 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, + 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"); + + // 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 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 int q_i32_per_lane = D / (sizeof(int) * nthreads); + +#ifdef V_DOT2_F32_F16_AVAILABLE + half2 q_reg[n_batch_heads][q_registers]; +#else + float2 q_reg[n_batch_heads][q_registers]; +#endif + int q_i32[n_batch_heads][q_i32_per_lane]; + float2 q_ds [n_batch_heads][q_i32_per_lane]; + + if constexpr (quantize_q) { + // 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 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[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); + // 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 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) { +#ifdef V_DOT2_F32_F16_AVAILABLE + q_reg[h][i0 / nthreads + j] = + make_half2(tmp[j].x * scale_log2, tmp[j].y * scale_log2); +#else + q_reg[h][i0 / nthreads + j] = + make_float2(tmp[j].x * scale_log2, tmp[j].y * scale_log2); +#endif + } + } + } + } + + constexpr dequantize_V_t dequantize_v = + get_dequantize_V(); + + 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 n_physical_blocks = pool_tokens / block_size; + + 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; + } + } + + float score_mine[n_batch_heads]; +#pragma unroll + for (int h = 0; h < n_batch_heads; ++h) { + score_mine[h] = -FLT_MAX; + } + + 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) 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]; + } + } + + // 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]); + } + + 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; + 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 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]; + } + } + } + } + } + +#pragma unroll + 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 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 || + 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); + 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) && + 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 && + max_kv_seq_len > 0 && + max_kv_seq_len <= k->ne[1] && + k->ne[1] % block_size == 0; +} + +// 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 max_kv_seq_len) { + 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 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], + 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); + + 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 +static void launch_paged_attn_v( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + float scale, + 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, max_kv_seq_len); + break; + case GGML_TYPE_Q4_0: + 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, max_kv_seq_len); + 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); + 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, max_kv_seq_len); + break; + case GGML_TYPE_Q4_0: + 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, max_kv_seq_len); + 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..695046978 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,58 @@ 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, + 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); + 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); + 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; + 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_args.h b/server/src/common/backend_args.h index 1607b1495..ffd670e11 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -25,6 +25,11 @@ struct BackendFeatureConfig { // time rather than through BackendArgs. bool routing_stats_requested = false; // --freq / --collect-routing bool adaptive_experts_requested = false; // --adaptive-experts + + // KVFlash bounded KV residency, enabled through DFLASH_KVFLASH rather than + // a flag. Recorded here because the gate cannot read the environment: it + // is linked without the pool sizing code so it stays a pure function. + bool kvflash_enabled = false; }; // A superset of all per-architecture config fields. The factory reads only @@ -58,6 +63,7 @@ struct BackendArgs { // Attention and speculative-decode options. Individual backends consume // only the fields they support. 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; // 16-token paged K/V blocks for AR decode int kq_stride_pad = 32; int draft_swa_window = 0; int draft_ctx_max = 4096; diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index ab34f2a19..d7b93d7cf 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -48,6 +48,7 @@ DFLASH_ARCH_FIELD_TRAIT(has_verify_width, verify_width); DFLASH_ARCH_FIELD_TRAIT(has_draft_swa, draft_swa_window); DFLASH_ARCH_FIELD_TRAIT(has_ddtree_mode, ddtree_mode); DFLASH_ARCH_FIELD_TRAIT(has_max_verify_tokens, max_verify_tokens); +DFLASH_ARCH_FIELD_TRAIT(has_paged_attention, paged_attention); #undef DFLASH_ARCH_FIELD_TRAIT @@ -95,6 +96,22 @@ DFLASH_CHECK_ARCH("qwen3", Qwen3BackendConfig, NoLayerSplitConfig); DFLASH_CHECK_ARCH("gemma4", Gemma4BackendConfig, Gemma4LayerSplitAdapterConfig); DFLASH_CHECK_ARCH("deepseek4", DeepSeek4BackendConfig, DeepSeek4LayerSplitAdapterConfig); +// paged_attn sits outside the bundle because the field-presence trait cannot +// separate qwen35 from qwen35moe: they share Qwen35Config, so the moe row +// carries a field its backend never reads, and pairing its Never row with +// that struct would fail a check that is really about qwen35's dispatch. +DFLASH_CHECK_ARCH_OPTION("qwen35", Qwen35Config, Qwen35LayerSplitAdapterConfig, + has_paged_attention, paged_attn); + +// So the qwen35moe row gets pinned instead of derived. Qwen35MoeBackend +// inherits the init() that builds the paged pool but overrides +// run_ar_decode_path() with a pipelined path that never reads a block table, +// so the flag would allocate a pool and decode dense. Flipping this row needs +// that decode path to learn paging first. +static_assert(arch_capabilities("qwen35moe").paged_attn == FeatureSupport::Never, + "qwen35moe shares Qwen35Config but its decode path ignores " + "paged_attention; the capability row must stay Never"); + #undef DFLASH_CHECK_ARCH #undef DFLASH_CHECK_ARCH_OPTION @@ -260,6 +277,7 @@ std::unique_ptr create_backend( 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/feature_gate.cpp b/server/src/common/feature_gate.cpp index 6ab697b45..1f2f6f96a 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -3,6 +3,9 @@ #include "feature_gate.h" #include "model_capabilities.h" +#include "paged_attention_config.h" + +#include namespace dflash::common { @@ -135,6 +138,47 @@ std::string check_feature_compatibility( "' does not support PFlash compression"; } + // ── --paged-attention × architecture, placement, and decode features + // Paged decode swaps the contiguous K/V cache for a block table owned by + // the monolithic qwen35 backend, so every rule below is about reaching + // that one code path. All are errors rather than warnings: running dense + // instead would hide the memory behavior the flag was chosen for. + if (args.paged_attention) { + if (!arch_supports_paged_attention(arch, /*is_layer_split=*/false)) { + return "--paged-attention requires a Qwen3.5/Qwen3.6 dense target " + "(architecture '" + arch + "' has no paged decode path)"; + } + // No rule for "requires a CUDA or HIP build": those are the only two + // backends this binary can be configured with, and GGML_OP_PAGED_ATTN + // is compiled into both. + if (!arch_supports_paged_attention(arch, args.device.is_layer_split()) || + args.remote_target_shard.enabled()) { + return "--paged-attention requires one local target device"; + } + if (args.draft_path != nullptr || args.remote_draft.enabled() || + args.ddtree_mode) { + return "--paged-attention requires autoregressive decode without a " + "draft or DDTree"; + } + if (args.fa_window != 0) { + return "--paged-attention requires full attention (--fa-window 0)"; + } + if (features.pflash_enabled) { + return "--paged-attention cannot be combined with PFlash prefill " + "compression"; + } + if (features.kvflash_enabled) { + return "--paged-attention cannot be combined with KVFlash"; + } + // The pool rounds max_ctx up to a whole number of blocks, so the top + // of the range is what can be rounded without overflowing int. + if (args.device.max_ctx <= 0 || + args.device.max_ctx > INT_MAX - PAGED_BLOCK_SIZE + 1) { + return "--paged-attention requires --max-ctx in [1, " + + std::to_string(INT_MAX - PAGED_BLOCK_SIZE + 1) + "]"; + } + } + // ── --ds4-prefill × architecture if (args.ds4_prefill_mode_set && arch != "deepseek4") { return "--ds4-prefill is only valid for deepseek4 models (detected '" + diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index 2189fe312..cf14f414b 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -15,7 +15,13 @@ // // Placement matters as much as architecture: laguna and gemma4 forward a // draft model only when monolithic, so those entries are FeatureSupport -// rather than bool. +// rather than bool. qwen35 paged attention is Monolithic for the same reason: +// only the single-device backend owns a paged K/V pool. +// +// qwen35 and qwen35moe share Qwen35Config, so their rows can differ on a +// column the config carries — the moe backend simply never reads the field. +// paged_attn is the one such column today; see the cross-check comment in +// backend_factory.cpp for what that costs. // // Note on "qwen36": it is not a dispatchable architecture. model_card.cpp's // family fallback has a branch for it, but there is no factory case, so a @@ -55,6 +61,7 @@ struct ArchCapabilities { FeatureSupport verify_width; // --verify-width FeatureSupport fa_window; // --fa-window FeatureSupport draft_swa; // --draft-swa + FeatureSupport paged_attn; // --paged-attention }; inline constexpr FeatureSupport kNever = FeatureSupport::Never; @@ -62,13 +69,13 @@ inline constexpr FeatureSupport kMono = FeatureSupport::Monolithic; inline constexpr FeatureSupport kBoth = FeatureSupport::Both; inline constexpr ArchCapabilities kArchCapabilities[] = { -// arch split rdraft pflash offload draft ddtree vwidth fa_win dswa - {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kBoth, kBoth}, - {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kMono, kMono}, - {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever}, - {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever}, - {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever}, - {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever}, +// arch split rdraft pflash offload draft ddtree vwidth fa_win dswa paged + {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kBoth, kBoth, kMono}, + {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kMono, kMono, kNever}, + {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever, kNever}, + {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever, kNever}, + {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever, kNever}, + {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever, kNever}, }; inline constexpr std::size_t kArchCount = @@ -106,7 +113,8 @@ constexpr bool row_has_both(const ArchCapabilities & c) { c.ddtree == FeatureSupport::Both || c.verify_width == FeatureSupport::Both || c.fa_window == FeatureSupport::Both || - c.draft_swa == FeatureSupport::Both; + c.draft_swa == FeatureSupport::Both || + c.paged_attn == FeatureSupport::Both; } constexpr bool table_rows_named() { @@ -235,4 +243,9 @@ inline bool arch_supports_draft_swa(const std::string & arch, return detail::arch_has(arch, &ArchCapabilities::draft_swa, is_layer_split); } +inline bool arch_supports_paged_attention(const std::string & arch, + bool is_layer_split) { + return detail::arch_has(arch, &ArchCapabilities::paged_attn, is_layer_split); +} + } // namespace dflash::common diff --git a/server/src/common/paged_attention_config.h b/server/src/common/paged_attention_config.h new file mode 100644 index 000000000..e78e0c346 --- /dev/null +++ b/server/src/common/paged_attention_config.h @@ -0,0 +1,23 @@ +// Paged K/V block geometry. +// +// Compatibility rules for --paged-attention live in feature_gate.cpp and the +// capability table in model_capabilities.h; only the sizing arithmetic that +// the allocator, the backend, and the gate all need is here. The one rule +// that cannot move is the 256-wide K/V head check in Qwen35Backend::init(): +// it needs loaded tensor dimensions, which GgufModelInfo does not carry. + +#pragma once + +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; +} + +} // 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..5f080632a --- /dev/null +++ b/server/src/common/paged_kv_pool.cpp @@ -0,0 +1,349 @@ +#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; + + // 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; + 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; + } + + // 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 = sequence->block_table[logical_block]; + result.write_slots.push_back({ + logical_position, + physical_block, + block_offset, + static_cast(physical_block) * block_size_ + block_offset, + }); + } + + 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..fa0bbe95b --- /dev/null +++ b/server/src/common/paged_kv_pool.h @@ -0,0 +1,252 @@ +// 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. +// +// 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. + +#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; +}; + +// 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; +} + +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(). +// 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; + uint32_t kv_seq_len = 0; + std::vector 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; + 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. 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; + 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_; } + // 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()); + } + 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. 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. 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); + + // 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 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 + // 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 for the planned batched device + // metadata 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..1026cdd57 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 { @@ -417,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. @@ -527,14 +536,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 +558,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 +616,12 @@ 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; // 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..e581927f1 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,17 @@ bool build_target_step( ggml_set_input(sg.attn_mask); } + if (paged_attention) { + if (n_tokens != 1 || with_mask || fa_window != 0) return false; + // 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); // Step-invariant KV write: only when topology can't vary per step. @@ -335,10 +347,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 +371,8 @@ 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.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..00ebcec8e 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -176,7 +176,8 @@ 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); target_backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!target_backend_) { @@ -209,6 +210,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 +289,62 @@ 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; + // check_feature_compatibility() rejects this pairing, but only for callers + // that fill BackendFeatureConfig::kvflash_enabled — the gate cannot read + // DFLASH_KVFLASH itself, and create_backend(args) defaults the field to + // false. Only the server fills it today, so this is what stops any other + // caller from paging over a cache KVFlash is also evicting from. + if (cfg_.paged_attention && kvflash_active()) { + std::fprintf(stderr, + "[paged-attention] cannot be combined with KVFlash " + "(resident pool %d tokens)\n", kvflash_tokens_); + set_last_error("paged attention cannot be combined with KVFlash"); + 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 +400,184 @@ 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::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; + } + + // 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); + 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])); + } 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; + } + + 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 +671,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 +708,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 +741,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 +944,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 +969,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 +987,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 +1057,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 +1094,37 @@ 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. + // + // 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. + 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 +1145,7 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, } } if (!decode_ok) { + end_paged_sequence(); result.fail(GenerateErrorCode::DecodeFailed); return result; } @@ -844,6 +1153,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 +1165,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); @@ -1418,15 +1734,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; @@ -1449,11 +1770,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 @@ -1490,7 +1812,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 = @@ -1534,8 +1855,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; - 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 @@ -1550,6 +1872,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 +1882,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..a408438df 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,10 @@ class Qwen35Backend : public ModelBackend { bool * forced_close_out = nullptr, bool * degenerate_close_out = nullptr); + bool begin_paged_sequence(uint32_t prompt_tokens); + 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..19b256071 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; @@ -207,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"); @@ -566,7 +592,9 @@ 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 ) { const int head_dim = w.n_embd_head_k; const int n_head = w.n_head; @@ -655,7 +683,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 +705,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 +716,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 +733,58 @@ 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); + // 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, + 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 + // 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 +1277,9 @@ 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); 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 063b20f11..f680d2a5d 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -22,6 +22,7 @@ #include "common/peer_access.h" #include "placement/pflash_placement.h" #include "placement/draft_residency.h" +#include "kvflash_pager.h" #include #include @@ -103,6 +104,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" @@ -212,6 +215,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; bool adaptive_experts_set = false; // --adaptive-experts (MoE architectures only) @@ -331,12 +336,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) { @@ -570,6 +579,11 @@ int main(int argc, char ** argv) { backend_features.routing_stats_requested = sconfig.freq_tracking || !sconfig.collect_routing_path.empty(); backend_features.adaptive_experts_requested = adaptive_experts_set; + // The gate rejects --paged-attention alongside KVFlash, and KVFlash is + // requested through the environment, so resolve it here where the pool + // sizing code is linked. + backend_features.kvflash_enabled = + kvflash_pool_from_env(bargs.device.max_ctx, KvFlashConfig{}) > 0; const BackendPreparation backend_preparation = prepare_backend(bargs, backend_features); if (!backend_preparation.ok()) { @@ -588,6 +602,24 @@ int main(int argc, char ** argv) { const ResolvedBackendPlan & backend_plan = backend_preparation.plan; const std::string & arch = backend_plan.arch(); + // Paged decode owns its K/V through a block table that the snapshot format + // cannot describe yet, so the caches it would restore into are turned off. + // This rewrites ServerConfig rather than rejecting the launch, which is why + // it lives here and not in the gate. + if (bargs.paged_attention) { + 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; + } + // Sync max_ctx: if --max-ctx was not provided, use the backend's default. // This prevents the HTTP server from accepting prompts larger than the // KV cache the backend actually allocates. diff --git a/server/test/bench_paged_attention.cpp b/server/test/bench_paged_attention.cpp new file mode 100644 index 000000000..df2723e39 --- /dev/null +++ b/server/test/bench_paged_attention.cpp @@ -0,0 +1,1069 @@ +#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 +#include + +namespace { + +constexpr int D = 256; +constexpr int N_HEAD = 24; +constexpr int N_HEAD_KV = 4; +constexpr int BLOCK_SIZE = 16; +// 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 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; + 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 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 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 { + 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" + " 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" + " --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) { + break; + } + context = static_cast(parsed); + contexts.push_back(context); + if (*end == '\0') { + if (contexts.size() >= 2) { + return true; + } + break; + } + 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_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_ROW_PERIOD) * row_bytes); + const size_t written = ggml_quantize_chunk( + 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 && + 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. + // 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_ROW_PERIOD); + + 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()); + 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.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, + const std::vector & reference, + 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])); + + 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) { + 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]; + const float reference_value = + reference[contiguous_index]; + if (!std::isfinite(paged_value) || + !std::isfinite(contiguous_value) || + !std::isfinite(reference_value)) { + return false; + } + + const double abs_error = + std::fabs( + static_cast(paged_value) - + contiguous_value); + 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.paged_oracle_max_abs, + paged_oracle_error); + } + } + } + return result.contiguous_oracle_max_abs <= MAX_ORACLE_ERROR && + result.paged_oracle_max_abs <= MAX_ORACLE_ERROR; +} + +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; + } + } + } + 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; + 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, 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)), + 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) { + 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, + "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; + } + + 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); + 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_results( + const Options & options, + const std::vector & results) { + const bool cuda_graphs_enabled = + std::getenv("GGML_CUDA_DISABLE_GRAPHS") == nullptr; + std::printf( + "# 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_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", + 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_us_per_step," + "paged_median_us_per_step," + "paged_step_time_overhead_pct," + "padded_contiguous_aggregate_attention_queries_s," + "paged_aggregate_attention_queries_s," + "padded_contiguous_kv_bytes_one_layer," + "paged_kv_bytes_one_layer," + "paged_kv_saving_pct,paged_metadata_bytes," + "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 + +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; + 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.metrics)) { + std::fprintf( + stderr, "ragged attention benchmark failed\n"); + ok = false; + } else { + results.push_back(std::move(result)); + } + } else { + const std::array sequence_counts = {1, 2, 4, 8}; + 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); + CaseResult result{"uniform", contexts, {}}; + if (!run_case( + backend, options, contexts, options.context, + false, case_index % 2 == 0, result.metrics)) { + std::fprintf( + stderr, + "attention benchmark failed at n_seq=%d\n", + n_seq); + ok = false; + break; + } + 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) { + std::fprintf(stderr, "benchmark setup failed: %s\n", error.what()); + ok = false; + } + + if (ok) { + print_results(options, results); + } + + ggml_backend_free(backend); + return ok ? 0 : 1; +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index b75eee35e..7083c1af2 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -12,8 +12,10 @@ #include "common/feature_gate.h" #include "common/model_capabilities.h" +#include "common/paged_attention_config.h" #include "placement/placement_config.h" +#include #include #include #include @@ -278,6 +280,95 @@ static void test_feature_gate_layer_split_requires_supported_arch() { TEST_ASSERT(gate_result(single, "qwen3", PlacementBackend::Cuda).empty()); } +static void test_feature_gate_paged_attention_requires_qwen35_monolithic() { + BackendArgs args; + args.model_path = "/nonexistent/model.gguf"; + args.paged_attention = true; + TEST_ASSERT(gate_result(args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_result(args, "qwen35", PlacementBackend::Hip).empty()); + + // Only qwen35 has a paged decode path. qwen35moe shares Qwen35Config, so + // its rejection is this gate's job — the factory's field-presence + // cross-check cannot tell the two apart. + for (const char * arch : {"qwen35moe", "laguna", "qwen3", + "gemma4", "deepseek4"}) { + TEST_ASSERT(!gate_result(args, arch, PlacementBackend::Cuda).empty()); + } + + // Only the monolithic qwen35 backend owns a paged K/V pool. Both + // placements are supported qwen35 launches without the flag, so the + // rejection has to come from the paged rule. + BackendArgs split = args; + TEST_ASSERT(parse_placement_device_list("cuda:0,cuda:1", split.device)); + TEST_ASSERT(!gate_result(split, "qwen35", PlacementBackend::Cuda).empty()); + + BackendArgs remote_shard = args; + remote_shard.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; + TEST_ASSERT(!gate_result( + remote_shard, "qwen35", PlacementBackend::Cuda).empty()); + + for (BackendArgs * relaxed : {&split, &remote_shard}) { + relaxed->paged_attention = false; + TEST_ASSERT(gate_result( + *relaxed, "qwen35", PlacementBackend::Cuda).empty()); + } +} + +static void test_feature_gate_paged_attention_requires_plain_ar_decode() { + BackendArgs base; + base.model_path = "/nonexistent/model.gguf"; + base.paged_attention = true; + + BackendArgs draft = base; + draft.draft_path = "/nonexistent/draft.gguf"; + TEST_ASSERT(!gate_result(draft, "qwen35", PlacementBackend::Cuda).empty()); + + BackendArgs ddtree = base; + ddtree.ddtree_mode = true; + TEST_ASSERT(!gate_result(ddtree, "qwen35", PlacementBackend::Cuda).empty()); + + BackendArgs windowed = base; + windowed.fa_window = 4096; + TEST_ASSERT(!gate_result( + windowed, "qwen35", PlacementBackend::Cuda).empty()); + + BackendFeatureConfig pflash; + pflash.pflash_enabled = true; + pflash.pflash_drafter_configured = true; + TEST_ASSERT(!gate_result( + base, "qwen35", PlacementBackend::Cuda, pflash).empty()); + + BackendFeatureConfig kvflash; + kvflash.kvflash_enabled = true; + TEST_ASSERT(!gate_result( + base, "qwen35", PlacementBackend::Cuda, kvflash).empty()); + + // The pool rounds max_ctx up to whole blocks, so both ends of the range + // are rejected: nothing to allocate, and rounding that overflows int. + BackendArgs empty_ctx = base; + empty_ctx.device.max_ctx = 0; + TEST_ASSERT(!gate_result( + empty_ctx, "qwen35", PlacementBackend::Cuda).empty()); + + BackendArgs huge_ctx = base; + huge_ctx.device.max_ctx = INT_MAX; + TEST_ASSERT(!gate_result( + huge_ctx, "qwen35", PlacementBackend::Cuda).empty()); + + BackendArgs max_ctx = base; + max_ctx.device.max_ctx = INT_MAX - PAGED_BLOCK_SIZE + 1; + TEST_ASSERT(gate_result( + max_ctx, "qwen35", PlacementBackend::Cuda).empty()); + + // None of these are rules about paged attention itself: without the flag + // every one of them is a supported qwen35 launch. + for (BackendArgs * args : {&draft, &ddtree, &windowed, &empty_ctx, + &huge_ctx}) { + args->paged_attention = false; + TEST_ASSERT(gate_result(*args, "qwen35", PlacementBackend::Cuda).empty()); + } +} + // ── Inert-flag warnings ───────────────────────────────────────────────── // Warnings must never gate admission, so each case also asserts the same // configuration passes check_feature_compatibility(). @@ -403,6 +494,12 @@ static void test_model_capability_tables() { TEST_ASSERT(!arch_supports_verify_width("qwen36", false)); TEST_ASSERT(!arch_supports_fa_window("qwen36", false)); TEST_ASSERT(!arch_supports_draft_swa("qwen36", false)); + TEST_ASSERT(!arch_supports_paged_attention("qwen36", false)); + + // Paged decode lives in the monolithic qwen35 backend alone. + TEST_ASSERT(arch_supports_paged_attention("qwen35", false)); + TEST_ASSERT(!arch_supports_paged_attention("qwen35", true)); + TEST_ASSERT(!arch_supports_paged_attention("qwen35moe", false)); } int main() { @@ -419,6 +516,8 @@ int main() { RUN_TEST(test_feature_gate_ds4_decode_options_require_monolithic_hip); RUN_TEST(test_feature_gate_remote_draft_requires_supported_arch); RUN_TEST(test_feature_gate_layer_split_requires_supported_arch); + RUN_TEST(test_feature_gate_paged_attention_requires_qwen35_monolithic); + RUN_TEST(test_feature_gate_paged_attention_requires_plain_ar_decode); RUN_TEST(test_feature_warnings_silent_when_supported); RUN_TEST(test_feature_warnings_report_inert_draft); RUN_TEST(test_feature_warnings_report_inert_decode_tunables); 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 new file mode 100644 index 000000000..d4efeaf72 --- /dev/null +++ b/server/test/test_paged_attention.cpp @@ -0,0 +1,340 @@ +#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 float MAX_ABS_ERROR = 5.0e-4f; + +struct TestCase { + const char * name; + int max_blocks; + std::vector kv_seq_lens; + bool corrupt_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, + 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 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 = 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; + + 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 * 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; + } + 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 * 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; + 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, + 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; + 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, test_case.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, + *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); + + 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, 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; + } + + 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( + 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])) { + ok = false; + break; + } + max_abs_error = + std::max(max_abs_error, std::fabs(actual[i] - expected[i])); + } + ok = ok && max_abs_error < MAX_ABS_ERROR; + } + + 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); + return ok; +} + +} // namespace + +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, test_case, 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..99b0ecf5b --- /dev/null +++ b/server/test/test_paged_kv_pool.cpp @@ -0,0 +1,342 @@ +#include "../src/common/paged_kv_pool.h" + +#include +#include +#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_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); + + 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})); + + 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() { + 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); + + 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(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.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 == 8); + CHECK(metadata.sequences.size() == 2); + CHECK(equals(metadata.block_table, {0, 1, 3, 4, 2})); + + 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 & 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() { + 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 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); +} + +int main() { + test_block_boundaries(); + test_nondefault_block_size(); + 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_invalid_arguments(); + + 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; +}