diff --git a/README.md b/README.md index 398d12f..e8b7afc 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,10 @@ Mference currently runs five pinned instruction checkpoints: multimodal checkpoint, fully resident in ~15 GB (24 GB Macs). Ships MTP speculative decoding with byte-identical greedy output — 15.0 tok/s decode on a 24 GB M5 (2.35× mlx-vlm on the same checkpoint). Its chat template - also opens a live `` block. + also opens a live `` block. Long contexts (past RAM, up to the + model's 262k) run a paged KV cache with an SSD spill tier and query-aware + sparse decode at full FP16 — see + [docs/QWEN38_LONG_CONTEXT.md](docs/QWEN38_LONG_CONTEXT.md). The runtime, streaming installer, CLI, native Mac app, and loopback OpenAI-compatible server are written in Swift and Metal. Mference is diff --git a/Sources/Mference/Kernels/Attention/Attention.swift b/Sources/Mference/Kernels/Attention/Attention.swift index 80b2dec..9e78f72 100644 --- a/Sources/Mference/Kernels/Attention/Attention.swift +++ b/Sources/Mference/Kernels/Attention/Attention.swift @@ -45,6 +45,9 @@ final class Attention { private let psoCombineQwenFullChunks16: MTLComputePipelineState private let psoCombineQwen38FullChunks16: MTLComputePipelineState private let psoCombineFullChunks16: MTLComputePipelineState + private let psoPagedPartial: MTLComputePipelineState + private let psoPagedPartialQwen38: MTLComputePipelineState + private let psoPagedPartialQwen38Chunks16: MTLComputePipelineState /// Mirrors `kAttnThreads` in `attention.metal`. The kernel was authored /// with a hardcoded 256-thread group so its threadgroup-memory scratch @@ -167,6 +170,18 @@ final class Attention { numQHeads: 16, numKVHeads: 2, numChunks: 16) + self.psoPagedPartial = try context.pipeline("attention_decode_paged_partial") + self.psoPagedPartialQwen38 = try Self.specializedPipeline(context, + "attention_decode_paged_partial", + headDim: 256, + numQHeads: 24, + numKVHeads: 4) + self.psoPagedPartialQwen38Chunks16 = try Self.specializedPipeline(context, + "attention_decode_paged_partial", + headDim: 256, + numQHeads: 24, + numKVHeads: 4, + numChunks: 16) let md = Self.maxQHeads * Self.maxChunks guard let m = context.device.makeBuffer(length: md * MemoryLayout.size, options: .storageModeShared), @@ -273,6 +288,85 @@ final class Attention { } + /// Paged full attention over a page-table selection (KVPageStore pools). + /// `pageTable` holds one `uint32` pool slot per selected 64-token page in + /// ascending logical order; `selTokens` counts selected tokens (the last + /// page may be partial). Split geometry and the combine pass match + /// `encodeFull`, so with a full identity selection the result is + /// bit-identical to the contiguous kernel. + func encodeFullPaged(commandBuffer: MTLCommandBuffer, + q: MTLBuffer, qOffset: Int = 0, + kPool: MTLBuffer, + vPool: MTLBuffer, + pageTable: MTLBuffer, pageTableOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + headDim: UInt32, + numQHeads: UInt32, + numKVHeads: UInt32, + selTokens: UInt32, + scale: Float? = nil) { + precondition(numQHeads % numKVHeads == 0, + "numQHeads must be a multiple of numKVHeads for GQA") + precondition(headDim <= 512, + "head_dim must be <= 512 (kernel scratch is sized for the full-attn case)") + precondition(selTokens > 0, "paged attention requires at least one selected token") + precondition(Int(numQHeads) <= Self.maxQHeads, + "numQHeads \(numQHeads) exceeds split-KV scratch (max \(Self.maxQHeads))") + let sc = scale ?? Self.defaultScale(headDim: headDim) + let geometry = Self.splitGeometry(numQHeads: numQHeads, + numKVHeads: numKVHeads, + seqLen: selTokens, + kvStart: 0, + preferGQASWA: false) + let nChunks = geometry.numChunks + let partialPSO = pagedPartialPipeline(headDim: headDim, + numQHeads: numQHeads, + numKVHeads: numKVHeads, + numChunks: nChunks) + let tgWidth = min(Self.threadsPerGroup, Int(partialPSO.maxTotalThreadsPerThreadgroup)) + + guard let p1 = commandBuffer.makeComputeCommandEncoder() else { return } + p1.setComputePipelineState(partialPSO) + p1.setBuffer(q, offset: qOffset, index: 0) + p1.setBuffer(kPool, offset: 0, index: 1) + p1.setBuffer(vPool, offset: 0, index: 2) + p1.setBuffer(mPartial, offset: 0, index: 3) + p1.setBuffer(dPartial, offset: 0, index: 4) + p1.setBuffer(oPartial, offset: 0, index: 5) + var hd = headDim, nq = numQHeads, nkv = numKVHeads, st = selTokens + var cl = UInt32(geometry.chunkLength), nc = UInt32(nChunks), sc2 = sc + p1.setBytes(&hd, length: MemoryLayout.size, index: 6) + p1.setBytes(&nq, length: MemoryLayout.size, index: 7) + p1.setBytes(&nkv, length: MemoryLayout.size, index: 8) + p1.setBytes(&st, length: MemoryLayout.size, index: 9) + p1.setBuffer(pageTable, offset: pageTableOffset, index: 10) + p1.setBytes(&cl, length: MemoryLayout.size, index: 11) + p1.setBytes(&nc, length: MemoryLayout.size, index: 12) + p1.setBytes(&sc2, length: MemoryLayout.size, index: 13) + p1.dispatchThreadgroups(MTLSize(width: geometry.partialThreadgroups, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: tgWidth, height: 1, depth: 1)) + p1.endEncoding() + + guard let p2 = commandBuffer.makeComputeCommandEncoder() else { return } + let combinePSO = combinePipeline(headDim: headDim, + numQHeads: numQHeads, + numKVHeads: numKVHeads, + numChunks: nChunks) + p2.setComputePipelineState(combinePSO) + p2.setBuffer(mPartial, offset: 0, index: 0) + p2.setBuffer(dPartial, offset: 0, index: 1) + p2.setBuffer(oPartial, offset: 0, index: 2) + p2.setBuffer(out, offset: outOffset, index: 3) + var hd2 = headDim, nc2 = UInt32(nChunks) + p2.setBytes(&hd2, length: MemoryLayout.size, index: 4) + p2.setBytes(&nc2, length: MemoryLayout.size, index: 5) + let combineTGWidth = min(Self.threadsPerGroup, + Int(combinePSO.maxTotalThreadsPerThreadgroup)) + p2.dispatchThreadgroups(MTLSize(width: Int(numQHeads), height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: combineTGWidth, height: 1, depth: 1)) + p2.endEncoding() + } + /// Two-pass split-KV (Flash-Decoding) dispatch shared by SWA and full /// attention — they differ only by `kvStart`. Pass 1 fans the head's /// `[kvStart, seqLen)` range across `chunkCount` threadgroups per head; @@ -426,6 +520,16 @@ final class Attention { return useGQAPartial ? psoGQAPartial : psoPartial } + private func pagedPartialPipeline(headDim: UInt32, + numQHeads: UInt32, + numKVHeads: UInt32, + numChunks: Int) -> MTLComputePipelineState { + if headDim == 256 && numQHeads == 24 && numKVHeads == 4 { + return numChunks == 16 ? psoPagedPartialQwen38Chunks16 : psoPagedPartialQwen38 + } + return psoPagedPartial + } + private func combinePipeline(headDim: UInt32, numQHeads: UInt32, numKVHeads: UInt32, diff --git a/Sources/Mference/Kernels/Attention/KVPageKernels.swift b/Sources/Mference/Kernels/Attention/KVPageKernels.swift new file mode 100644 index 0000000..3de486a --- /dev/null +++ b/Sources/Mference/Kernels/Attention/KVPageKernels.swift @@ -0,0 +1,180 @@ +import Foundation +import Metal + +/// Wrappers for the paged-KV maintenance kernels: page seal summaries +/// (`kv_page_minmax`) and Quest page criticality scores +/// (`attention_page_scores`). Both ride the token command buffer. +final class KVPageKernels { + private let ctx: MetalContext + private let psoMinMax: MTLComputePipelineState + private let psoScores: MTLComputePipelineState + private let psoFlashInit: MTLComputePipelineState + private let psoFlashUpdate: MTLComputePipelineState + private let psoFlashFinalize: MTLComputePipelineState + + private static let threadsPerGroup = 256 + private static let flashSimdgroupsPerTG = 8 + + init(context: MetalContext) throws { + self.ctx = context + self.psoMinMax = try context.pipeline("kv_page_minmax") + self.psoScores = try context.pipeline("attention_page_scores") + self.psoFlashInit = try context.pipeline("attention_prefill_flash_init") + self.psoFlashUpdate = try context.pipeline("attention_prefill_flash_update") + self.psoFlashFinalize = try context.pipeline("attention_prefill_flash_finalize") + } + + // MARK: - Blocked prefill attention + + /// Reset the running online-softmax state for a chunk's queries. + func encodeFlashInit(commandBuffer: MTLCommandBuffer, + mState: MTLBuffer, dState: MTLBuffer, oState: MTLBuffer, + rows: UInt32, headDim: UInt32) { + guard let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(psoFlashInit) + enc.setBuffer(mState, offset: 0, index: 0) + enc.setBuffer(dState, offset: 0, index: 1) + enc.setBuffer(oState, offset: 0, index: 2) + var r = rows, hd = headDim + enc.setBytes(&r, length: MemoryLayout.size, index: 3) + enc.setBytes(&hd, length: MemoryLayout.size, index: 4) + let width = min(Self.threadsPerGroup, psoFlashInit.maxTotalThreadsPerThreadgroup) + enc.dispatchThreads(MTLSize(width: Int(rows), height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: width, height: 1, depth: 1)) + enc.endEncoding() + } + + /// Fold one KV window into the running state. `pageTable` maps the + /// window's 64-token pages to slots in `kPool`/`vPool` (staging passes a + /// stride-2 identity for the interleaved [K|V] spill layout). `causal` + /// applies `p <= q_pos` for the chunk's own pages. + func encodeFlashUpdate(commandBuffer: MTLCommandBuffer, + q: MTLBuffer, + kPool: MTLBuffer, kPoolOffset: Int = 0, + vPool: MTLBuffer, vPoolOffset: Int = 0, + pageTable: MTLBuffer, pageTableOffset: Int = 0, + mState: MTLBuffer, dState: MTLBuffer, oState: MTLBuffer, + queryCount: UInt32, + qStartPosition: UInt32, + headDim: UInt32, + numQHeads: UInt32, + numKVHeads: UInt32, + windowStartPosition: UInt32, + windowTokens: UInt32, + qStrideElements: UInt32, + scale: Float, + causal: Bool) { + guard windowTokens > 0, let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(psoFlashUpdate) + enc.setBuffer(q, offset: 0, index: 0) + enc.setBuffer(kPool, offset: kPoolOffset, index: 1) + enc.setBuffer(vPool, offset: vPoolOffset, index: 2) + enc.setBuffer(pageTable, offset: pageTableOffset, index: 3) + enc.setBuffer(mState, offset: 0, index: 4) + enc.setBuffer(dState, offset: 0, index: 5) + enc.setBuffer(oState, offset: 0, index: 6) + var qc = queryCount, qs = qStartPosition, hd = headDim + var nq = numQHeads, nkv = numKVHeads + var ws = windowStartPosition, wt = windowTokens, qst = qStrideElements + var sc = scale + var cz: UInt32 = causal ? 1 : 0 + enc.setBytes(&qc, length: MemoryLayout.size, index: 7) + enc.setBytes(&qs, length: MemoryLayout.size, index: 8) + enc.setBytes(&hd, length: MemoryLayout.size, index: 9) + enc.setBytes(&nq, length: MemoryLayout.size, index: 10) + enc.setBytes(&nkv, length: MemoryLayout.size, index: 11) + enc.setBytes(&ws, length: MemoryLayout.size, index: 12) + enc.setBytes(&wt, length: MemoryLayout.size, index: 13) + enc.setBytes(&qst, length: MemoryLayout.size, index: 14) + enc.setBytes(&sc, length: MemoryLayout.size, index: 15) + enc.setBytes(&cz, length: MemoryLayout.size, index: 16) + let rows = Int(queryCount * numQHeads) + let groups = (rows + Self.flashSimdgroupsPerTG - 1) / Self.flashSimdgroupsPerTG + enc.dispatchThreadgroups( + MTLSize(width: groups, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: Self.flashSimdgroupsPerTG * 32, + height: 1, depth: 1)) + enc.endEncoding() + } + + /// Normalize the running state into the chunk's FP16 attention output. + func encodeFlashFinalize(commandBuffer: MTLCommandBuffer, + mState: MTLBuffer, dState: MTLBuffer, oState: MTLBuffer, + out: MTLBuffer, + queryCount: UInt32, + headDim: UInt32, + numQHeads: UInt32, + oStrideElements: UInt32) { + guard let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(psoFlashFinalize) + enc.setBuffer(mState, offset: 0, index: 0) + enc.setBuffer(dState, offset: 0, index: 1) + enc.setBuffer(oState, offset: 0, index: 2) + enc.setBuffer(out, offset: 0, index: 3) + var qc = queryCount, hd = headDim, nq = numQHeads, os = oStrideElements + enc.setBytes(&qc, length: MemoryLayout.size, index: 4) + enc.setBytes(&hd, length: MemoryLayout.size, index: 5) + enc.setBytes(&nq, length: MemoryLayout.size, index: 6) + enc.setBytes(&os, length: MemoryLayout.size, index: 7) + let total = Int(queryCount * numQHeads * headDim) + let width = min(Self.threadsPerGroup, psoFlashFinalize.maxTotalThreadsPerThreadgroup) + enc.dispatchThreads(MTLSize(width: total, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: width, height: 1, depth: 1)) + enc.endEncoding() + } + + /// Reduce a page's K rows to element-wise min/max vectors, written to + /// the page's slot in the metadata buffer. + func encodePageMinMax(commandBuffer: MTLCommandBuffer, + kPool: MTLBuffer, + slot: UInt32, + validTokens: UInt32, + metadata: MTLBuffer, + metadataOffset: Int, + numKVHeads: UInt32, + headDim: UInt32) { + guard let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(psoMinMax) + enc.setBuffer(kPool, offset: 0, index: 0) + enc.setBuffer(metadata, offset: metadataOffset, index: 1) + var s = slot, vt = validTokens, nkv = numKVHeads, hd = headDim + enc.setBytes(&s, length: MemoryLayout.size, index: 2) + enc.setBytes(&vt, length: MemoryLayout.size, index: 3) + enc.setBytes(&nkv, length: MemoryLayout.size, index: 4) + enc.setBytes(&hd, length: MemoryLayout.size, index: 5) + let width = min(Self.threadsPerGroup, psoMinMax.maxTotalThreadsPerThreadgroup) + enc.dispatchThreadgroups(MTLSize(width: 1, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: width, height: 1, depth: 1)) + enc.endEncoding() + } + + /// Score every sealed page of one layer against the current query. + /// `metadataOffset` addresses the layer's metadata base; `scores` receives + /// one float per page (read back by the CPU after the token completes — + /// the lag-one selection input for the next token). + func encodePageScores(commandBuffer: MTLCommandBuffer, + q: MTLBuffer, qOffset: Int = 0, + metadata: MTLBuffer, + metadataOffset: Int, + scores: MTLBuffer, + scoresOffset: Int, + numPages: UInt32, + headDim: UInt32, + numQHeads: UInt32, + numKVHeads: UInt32) { + guard numPages > 0, let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(psoScores) + enc.setBuffer(q, offset: qOffset, index: 0) + enc.setBuffer(metadata, offset: metadataOffset, index: 1) + enc.setBuffer(scores, offset: scoresOffset, index: 2) + var np = numPages, hd = headDim, nq = numQHeads, nkv = numKVHeads + enc.setBytes(&np, length: MemoryLayout.size, index: 3) + enc.setBytes(&hd, length: MemoryLayout.size, index: 4) + enc.setBytes(&nq, length: MemoryLayout.size, index: 5) + enc.setBytes(&nkv, length: MemoryLayout.size, index: 6) + let width = min(Self.threadsPerGroup, psoScores.maxTotalThreadsPerThreadgroup) + enc.dispatchThreadgroups(MTLSize(width: Int(numPages), height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: width, height: 1, depth: 1)) + enc.endEncoding() + } +} diff --git a/Sources/Mference/Metal/Attention/attention.metal b/Sources/Mference/Metal/Attention/attention.metal index 9e21502..d5c803c 100644 --- a/Sources/Mference/Metal/Attention/attention.metal +++ b/Sources/Mference/Metal/Attention/attention.metal @@ -223,6 +223,111 @@ void attention_decode_partial( } } +// ============================================================================ +// Paged decode attention — split-KV partial over a page-table selection. +// +// The KV cache lives in per-layer page pools (KVPageStore): fixed 64-token +// pages at arbitrary pool slots. `page_table[i]` is the pool slot of the +// i-th *selected* page, in ascending logical-position order; `sel_tokens` +// counts the selected logical tokens (the last listed page may be a partial +// tail). Softmax runs over exactly the selected subset — the sparse +// (Quest-style) decode path. With an identity table and a full selection the +// accumulation order matches `attention_decode_partial` bit for bit. +// Combine pass is shared (`attention_decode_combine`). +// ============================================================================ + +constant constexpr uint kAttnPageTokens = 64; + +[[kernel, max_total_threads_per_threadgroup(kAttnThreads)]] +void attention_decode_paged_partial( + device const half* Q [[buffer(0)]], + device const half* K_pool [[buffer(1)]], + device const half* V_pool [[buffer(2)]], + device float* m_out [[buffer(3)]], // [num_q_heads * num_chunks] + device float* d_out [[buffer(4)]], // [num_q_heads * num_chunks] + device float* o_out [[buffer(5)]], // [num_q_heads * num_chunks * head_dim] + constant uint& head_dim [[buffer(6)]], + constant uint& num_q_heads [[buffer(7)]], + constant uint& num_kv_heads [[buffer(8)]], + constant uint& sel_tokens [[buffer(9)]], + device const uint* page_table [[buffer(10)]], + constant uint& chunk_len [[buffer(11)]], + constant uint& num_chunks [[buffer(12)]], + constant float& scale [[buffer(13)]], + uint tg_id [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simdgroups [[simdgroups_per_threadgroup]] +) { + threadgroup float q_smem[kAttnMaxHeadDim]; + threadgroup float reduce_scratch[kAttnMaxSimdGroups]; + threadgroup float bcast; + const uint HD = attn_fc_head_dim(head_dim); + const uint NQ = attn_fc_num_q_heads(num_q_heads); + const uint NKV = attn_fc_num_kv_heads(num_kv_heads); + const uint NC = attn_fc_num_chunks(num_chunks); + + const uint q_head = tg_id / NC; + const uint chunk = tg_id % NC; + const uint l_start = chunk * chunk_len; + uint l_end = l_start + chunk_len; + if (l_end > sel_tokens) { l_end = sel_tokens; } + + const uint kv_head = q_head / (NQ / NKV); + + device const half* Q_row = Q + uint(q_head) * HD; + for (uint i = lid; i < HD; i += lsize) { + q_smem[i] = float(Q_row[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + constexpr uint kPerThread = (kAttnMaxHeadDim + kAttnThreads - 1) / kAttnThreads; + float o_local[kPerThread]; + for (uint k = 0; k < kPerThread; ++k) { o_local[k] = 0.0f; } + + float m_run = -INFINITY; + float d_run = 0.0f; + + for (uint l = l_start; l < l_end; ++l) { + const uint slot = page_table[l / kAttnPageTokens]; + const uint phys_p = slot * kAttnPageTokens + (l % kAttnPageTokens); + device const half* K_row = K_pool + (phys_p * NKV + kv_head) * HD; + device const half* V_row = V_pool + (phys_p * NKV + kv_head) * HD; + + float partial = 0.0f; + for (uint i = lid; i < HD; i += lsize) { + partial = fma(q_smem[i], float(K_row[i]), partial); + } + float s = block_reduce_sum(partial, + simd_lane_id, simd_group_id, simdgroups, + reduce_scratch, &bcast); + s *= attn_fc_scale(scale); + + const float m_new = max(m_run, s); + const float alpha = attn_softmax_exp(m_run - m_new); + const float p_exp = attn_softmax_exp(s - m_new); + d_run = d_run * alpha + p_exp; + + uint slot_i = 0; + for (uint i = lid; i < HD; i += lsize) { + o_local[slot_i] = o_local[slot_i] * alpha + p_exp * float(V_row[i]); + slot_i += 1; + } + m_run = m_new; + } + + const uint base = uint(q_head) * NC + chunk; + if (lid == 0) { m_out[base] = m_run; d_out[base] = d_run; } + device float* o_row = o_out + base * HD; + uint slot_i = 0; + for (uint i = lid; i < HD; i += lsize) { + o_row[i] = o_local[slot_i]; + slot_i += 1; + } +} + [[kernel, max_total_threads_per_threadgroup(kAttnThreads)]] void attention_decode_gqa_swa_partial( device const half* Q [[buffer(0)]], @@ -339,6 +444,243 @@ void attention_decode_gqa_swa_partial( } } +// ============================================================================ +// Blocked (streamed) prefill attention — beyond-RAM contexts. +// +// A prefill chunk's queries attend the whole past, which no longer fits the +// pool. The past streams through staging windows; each window dispatch folds +// its positions into per-(query, q-head) running online-softmax state +// (m, d, o[head_dim], FP32 in device memory), and a finalize pass writes +// out = o/d in FP16. The chunk's own pages run as a final window with the +// causal predicate p <= q_pos. +// +// Geometry: one simdgroup per (query token, q head) — 32 lanes split the +// head_dim, positions run sequentially with simd-level reduction only (no +// threadgroup barriers). K/V rows resolve through a page table like the +// paged decode kernel; staging passes an identity table. +// ============================================================================ + +constant constexpr uint kFlashSimdgroupsPerTG = 8; + +[[kernel]] +void attention_prefill_flash_init( + device float* m_state [[buffer(0)]], // [query_count * num_q_heads] + device float* d_state [[buffer(1)]], + device float* o_state [[buffer(2)]], // [query_count * num_q_heads * head_dim] + constant uint& rows [[buffer(3)]], // query_count * num_q_heads + constant uint& head_dim [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid < rows) { + m_state[gid] = -INFINITY; + d_state[gid] = 0.0f; + } + const uint total = rows * head_dim; + for (uint i = gid; i < total; i += rows) { + // rows threads stride the o_state clear; grid is sized to `rows`. + o_state[i] = 0.0f; + } +} + +[[kernel, max_total_threads_per_threadgroup(kFlashSimdgroupsPerTG * 32)]] +void attention_prefill_flash_update( + device const half* Q [[buffer(0)]], // [query_count, q_stride] FP16 + device const half* K_pool [[buffer(1)]], + device const half* V_pool [[buffer(2)]], + device const uint* page_table [[buffer(3)]], // window pages -> pool/staging slots + device float* m_state [[buffer(4)]], + device float* d_state [[buffer(5)]], + device float* o_state [[buffer(6)]], + constant uint& query_count [[buffer(7)]], + constant uint& q_start [[buffer(8)]], // global position of query row 0 + constant uint& head_dim [[buffer(9)]], + constant uint& num_q_heads [[buffer(10)]], + constant uint& num_kv_heads [[buffer(11)]], + constant uint& win_start [[buffer(12)]], // global position of window token 0 (page-aligned) + constant uint& win_tokens [[buffer(13)]], + constant uint& q_stride [[buffer(14)]], // elements per query row + constant float& scale [[buffer(15)]], + constant uint& causal [[buffer(16)]], // 1: apply p <= q_pos within the window + uint tg_id [[threadgroup_position_in_grid]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]] +) { + const uint row = tg_id * kFlashSimdgroupsPerTG + simd_group_id; + const uint total_rows = query_count * num_q_heads; + if (row >= total_rows) { return; } + const uint t = row / num_q_heads; + const uint q_head = row % num_q_heads; + const uint kv_head = q_head / (num_q_heads / num_kv_heads); + const uint q_pos = q_start + t; + + // Effective window span for this query under the causal predicate. + uint span = win_tokens; + if (causal != 0u) { + if (q_pos + 1 <= win_start) { return; } + span = min(span, q_pos + 1 - win_start); + } + if (span == 0u) { return; } + + device const half* Q_row = Q + t * q_stride + q_head * head_dim; + + // Lane-strided registers: head_dim <= 512 -> at most 16 elems per lane. + constexpr uint kMaxPerLane = kAttnMaxHeadDim / 32; + float q_reg[kMaxPerLane]; + const uint per_lane = (head_dim + 31) / 32; + for (uint k = 0; k < per_lane; ++k) { + const uint i = simd_lane_id + k * 32; + q_reg[k] = i < head_dim ? float(Q_row[i]) : 0.0f; + } + + float m_run = m_state[row]; + float d_run = d_state[row]; + device float* o_row = o_state + row * head_dim; + float o_reg[kMaxPerLane]; + for (uint k = 0; k < per_lane; ++k) { + const uint i = simd_lane_id + k * 32; + o_reg[k] = i < head_dim ? o_row[i] : 0.0f; + } + + const uint elems = num_kv_heads * head_dim; + for (uint w = 0; w < span; ++w) { + const uint slot = page_table[w / kAttnPageTokens]; + const uint phys = slot * kAttnPageTokens + (w % kAttnPageTokens); + device const half* K_row = K_pool + phys * elems + kv_head * head_dim; + device const half* V_row = V_pool + phys * elems + kv_head * head_dim; + + float partial = 0.0f; + for (uint k = 0; k < per_lane; ++k) { + const uint i = simd_lane_id + k * 32; + if (i < head_dim) { partial = fma(q_reg[k], float(K_row[i]), partial); } + } + const float s = simd_sum(partial) * scale; + + const float m_new = max(m_run, s); + const float alpha = attn_softmax_exp(m_run - m_new); + const float p_exp = attn_softmax_exp(s - m_new); + d_run = d_run * alpha + p_exp; + for (uint k = 0; k < per_lane; ++k) { + const uint i = simd_lane_id + k * 32; + if (i < head_dim) { + o_reg[k] = o_reg[k] * alpha + p_exp * float(V_row[i]); + } + } + m_run = m_new; + } + + if (simd_lane_id == 0) { m_state[row] = m_run; d_state[row] = d_run; } + for (uint k = 0; k < per_lane; ++k) { + const uint i = simd_lane_id + k * 32; + if (i < head_dim) { o_row[i] = o_reg[k]; } + } +} + +[[kernel]] +void attention_prefill_flash_finalize( + device const float* m_state [[buffer(0)]], + device const float* d_state [[buffer(1)]], + device const float* o_state [[buffer(2)]], + device half* out [[buffer(3)]], // [query_count, o_stride] + constant uint& query_count [[buffer(4)]], + constant uint& head_dim [[buffer(5)]], + constant uint& num_q_heads [[buffer(6)]], + constant uint& o_stride [[buffer(7)]], + uint gid [[thread_position_in_grid]] +) { + const uint total = query_count * num_q_heads * head_dim; + if (gid >= total) { return; } + const uint row = gid / head_dim; + const uint i = gid % head_dim; + const uint t = row / num_q_heads; + const uint q_head = row % num_q_heads; + const float d = d_state[row]; + const float v = d > 0.0f ? o_state[row * head_dim + i] / d : 0.0f; + out[t * o_stride + q_head * head_dim + i] = half(v); +} + +// ============================================================================ +// KV page maintenance kernels (paged long-context path). +// ============================================================================ + +// Element-wise min/max over a sealed page's K rows — the Quest (arXiv +// 2406.10774) page summary used to estimate a page's attention criticality +// without reading it. One threadgroup; runs once per page seal. +// `metadata` points at the page's metadata slot: min[NKV*HD] then max[NKV*HD]. +[[kernel, max_total_threads_per_threadgroup(kAttnThreads)]] +void kv_page_minmax( + device const half* K_pool [[buffer(0)]], + device half* metadata [[buffer(1)]], + constant uint& slot [[buffer(2)]], + constant uint& valid_tokens [[buffer(3)]], + constant uint& num_kv_heads [[buffer(4)]], + constant uint& head_dim [[buffer(5)]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]] +) { + const uint elems = num_kv_heads * head_dim; + const uint base = slot * kAttnPageTokens * elems; + for (uint e = lid; e < elems; e += lsize) { + float mn = INFINITY; + float mx = -INFINITY; + for (uint t = 0; t < valid_tokens; ++t) { + const float v = float(K_pool[base + t * elems + e]); + mn = min(mn, v); + mx = max(mx, v); + } + metadata[e] = half(mn); + metadata[elems + e] = half(mx); + } +} + +// Quest page criticality: for each page, per q-head upper bound of q·k over +// the page — sum_d max(q_d·min_d, q_d·max_d) against the head's kv-head +// min/max summary — reduced with max over q heads. One threadgroup per page; +// `metadata` points at the layer's metadata base. +[[kernel, max_total_threads_per_threadgroup(kAttnThreads)]] +void attention_page_scores( + device const half* Q [[buffer(0)]], // [num_q_heads, head_dim] + device const half* metadata [[buffer(1)]], + device float* scores [[buffer(2)]], // [num_pages] + constant uint& num_pages [[buffer(3)]], + constant uint& head_dim [[buffer(4)]], + constant uint& num_q_heads [[buffer(5)]], + constant uint& num_kv_heads [[buffer(6)]], + uint tg_id [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simdgroups [[simdgroups_per_threadgroup]] +) { + threadgroup float reduce_scratch[kAttnMaxSimdGroups]; + threadgroup float bcast; + const uint page = tg_id; + if (page >= num_pages) { return; } + + const uint elems = num_kv_heads * head_dim; + device const half* mn = metadata + page * 2 * elems; + device const half* mx = mn + elems; + + const uint q_per_kv = num_q_heads / num_kv_heads; + float best = -INFINITY; + for (uint qh = 0; qh < num_q_heads; ++qh) { + const uint kvh = qh / q_per_kv; + device const half* Q_row = Q + qh * head_dim; + device const half* lo = mn + kvh * head_dim; + device const half* hi = mx + kvh * head_dim; + float partial = 0.0f; + for (uint i = lid; i < head_dim; i += lsize) { + const float q = float(Q_row[i]); + partial += max(q * float(lo[i]), q * float(hi[i])); + } + const float s = block_reduce_sum(partial, + simd_lane_id, simd_group_id, simdgroups, + reduce_scratch, &bcast); + best = max(best, s); + } + if (lid == 0) { scores[page] = best; } +} + [[kernel, max_total_threads_per_threadgroup(kAttnThreads)]] void attention_decode_combine( device const float* m_in [[buffer(0)]], // [num_q_heads * num_chunks] diff --git a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift index d0fd278..598784e 100644 --- a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift +++ b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift @@ -21,6 +21,14 @@ public enum RuntimeExpertCachePolicy: String, Codable, Sendable { case lru } +/// Paged KV cache for long contexts (Qwen 3.8 full-attention layers): fixed +/// 64-token pages with an SSD spill tier and Quest-style query-aware sparse +/// decode. `.off` keeps the linear FP16 cache and the dense decode path. +public enum RuntimeKVPagedPolicy: String, Codable, Sendable { + case off + case on +} + public struct RuntimeConfiguration: Sendable, Equatable { /// 96 and 128 are the near-resident rungs: large wired LFU sets for hosts /// with RAM to spare but not enough to cache the whole expert pool. @@ -37,6 +45,14 @@ public struct RuntimeConfiguration: Sendable, Equatable { /// Opt-in approximate singleton-decode head for Maple checkpoints that /// retain FlashHead tensors. Prefill continues to use the exact head. public let useMapleFlashHead: Bool + public let kvPagedPolicy: RuntimeKVPagedPolicy + /// Sparse decode selection budget, in 64-token pages. + public let kvTopKPages: Int + public let kvSinkPages: Int + public let kvRecentPages: Int + /// Pool residency per full-attention layer, in pages. `nil` sizes the + /// pool to the full context (everything resident; SSD tier idle). + public let kvPoolPagesPerLayer: Int? public init(expertCacheSlots: Int = 16, expertCachePolicy: RuntimeExpertCachePolicy = .lfu, @@ -45,7 +61,12 @@ public struct RuntimeConfiguration: Sendable, Equatable { prefillChunkTokens: Int = 128, prefillAttentionPath: RuntimePrefillAttentionPath = .fullTensorOps2DPreferred, forceLogitsHead: Bool = false, - useMapleFlashHead: Bool = false) { + useMapleFlashHead: Bool = false, + kvPagedPolicy: RuntimeKVPagedPolicy = .off, + kvTopKPages: Int = 60, + kvSinkPages: Int = 2, + kvRecentPages: Int = 4, + kvPoolPagesPerLayer: Int? = nil) { precondition(Self.allowedExpertCacheSlots.contains(expertCacheSlots), "unsupported expert-cache slot count") precondition(Self.allowedPrefillChunkTokens.contains(prefillChunkTokens), @@ -58,6 +79,13 @@ public struct RuntimeConfiguration: Sendable, Equatable { self.prefillAttentionPath = prefillAttentionPath self.headPath = forceLogitsHead ? .logits : .fusedRows self.useMapleFlashHead = useMapleFlashHead + precondition(kvTopKPages >= 0 && kvSinkPages >= 0 && kvRecentPages >= 1, + "invalid paged-KV selection parameters") + self.kvPagedPolicy = kvPagedPolicy + self.kvTopKPages = kvTopKPages + self.kvSinkPages = kvSinkPages + self.kvRecentPages = kvRecentPages + self.kvPoolPagesPerLayer = kvPoolPagesPerLayer } public static var production: RuntimeConfiguration { @@ -108,6 +136,31 @@ public struct RuntimeConfiguration: Sendable, Equatable { physicalMemoryBytes: physicalMemoryBytes)) } + /// Auto pool sizing for the paged KV cache: everything resident when it + /// fits, otherwise whatever RAM remains after weights and headroom + /// (~physical − 22 GiB on the 24 GiB M5 → ~2 GiB of pool ≈ 32k resident + /// tokens per full-attention layer), never below 1 GiB. Measured: a + /// 4 GiB pool beside 14 GiB of weights pushed the host into compression + /// and cost ~2× decode; 2 GiB keeps full speed and the SSD tier absorbs + /// the rest. + public static func defaultKVPoolPagesPerLayer( + config: ArchConfig, + maxContext: Int, + physicalMemoryBytes: UInt64 = ProcessInfo.processInfo.physicalMemory + ) -> Int { + let pageTokens = 64 + let pagesPerLayer = (maxContext + pageTokens - 1) / pageTokens + let numFull = config.fullAttentionLayerMask.lazy.filter { $0 == 1 }.count + guard numFull > 0 else { return pagesPerLayer } + let pagePairBytes = 2 * pageTokens * config.numFullKVHeads * config.fullHeadDim * 2 + let gib = UInt64(1) << 30 + let headroom = UInt64(22) * gib + let budget = max(gib, physicalMemoryBytes > headroom + ? physicalMemoryBytes - headroom : gib) + let budgetPages = Int(budget) / (numFull * pagePairBytes) + return max(1, min(pagesPerLayer, budgetPages)) + } + public var fp16RingEnabled: Bool { true } public var rdadviseEnabled: Bool { rdadvisePolicy != .off } public var prefillConfig: PrefillRuntimeConfig { diff --git a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift index df39d61..2cddec0 100644 --- a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift +++ b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift @@ -151,6 +151,83 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR private let kv: KVCacheManager private let gdnState: GDNStateManager + /// Paged long-context state (kvPagedPolicy == .on), shared with the + /// MTP speculator so round and plain tokens drive one cursor. + private var pagedKV: Qwen38PagedKVRuntime? + + /// Scratch for the blocked (streamed) prefill attention: FP32 running + /// online-softmax state per (query, q-head) plus two staging buffers the + /// past KV windows stream through (ring of two so a window's pread can + /// overlap the previous window's GPU pass). + private struct BlockedPrefillScratch { + static let windowPages = 128 // 8k tokens, 32 MiB per stage + let chunkTokens: Int + let mState: MTLBuffer + let dState: MTLBuffer + let oState: MTLBuffer + let stages: [MTLBuffer] + /// Stride-2 identity table addressing the interleaved [K|V] staging + /// layout (K page i at slot 2i, its V page at +1 K-page offset). + let stagingTable: MTLBuffer + let tailTable: MTLBuffer + + init(device: MTLDevice, config: ArchConfig, chunkTokens: Int, + pagesPerLayer: Int, kPageBytes: Int) throws { + self.chunkTokens = chunkTokens + let rows = chunkTokens * config.numHeads + let headDim = config.fullHeadDim + guard let m = device.makeBuffer(length: rows * 4, options: .storageModeShared), + let d = device.makeBuffer(length: rows * 4, options: .storageModeShared), + let o = device.makeBuffer(length: rows * headDim * 4, + options: .storageModeShared) else { + throw KVPageStoreError.allocationFailed("blocked prefill state") + } + m.label = "kvpage.flash.m"; d.label = "kvpage.flash.d"; o.label = "kvpage.flash.o" + self.mState = m; self.dState = d; self.oState = o + + var stages: [MTLBuffer] = [] + for i in 0..<2 { + guard let s = device.makeBuffer(length: Self.windowPages * 2 * kPageBytes, + options: .storageModeShared) else { + throw KVPageStoreError.allocationFailed("blocked prefill staging") + } + s.label = "kvpage.flash.stage\(i)" + stages.append(s) + } + self.stages = stages + + var identity = (0.. BlockedPrefillScratch { + if let scratch = blockedPrefillScratch, scratch.chunkTokens >= chunkTokens { + return scratch + } + let scratch = try BlockedPrefillScratch(device: ctx.device, + config: cfg, + chunkTokens: chunkTokens, + pagesPerLayer: paged.store.geometry.pagesPerLayer, + kPageBytes: paged.store.geometry.kPageBytes) + blockedPrefillScratch = scratch + return scratch + } + // Kernels private let embedInt4: EmbedLookupInt4 private let rms: RMSNorm @@ -253,12 +330,20 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR self.cfg = cfg self.maxContext = maxContext self.useFusedGreedyHead = runtimeConfiguration.headPath == .fusedRows + let paged = runtimeConfiguration.kvPagedPolicy == .on self.kv = try KVCacheManager(device: context.device, config: cfg, maxContext: maxContext, fp16RingEnabled: runtimeConfiguration.fp16RingEnabled, slidingWindow: cfg.slidingWindow, - maxPrefillChunkTokens: runtimeConfiguration.prefillConfig.chunkTokens) + maxPrefillChunkTokens: runtimeConfiguration.prefillConfig.chunkTokens, + pagedFullAttention: paged) + if paged { + self.pagedKV = try Qwen38PagedKVRuntime(context: context, + config: cfg, + maxContext: maxContext, + runtimeConfiguration: runtimeConfiguration) + } self.gdnState = try GDNStateManager(device: context.device, config: cfg) self.embedInt4 = try EmbedLookupInt4(context: context) @@ -375,7 +460,8 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR gdnState: gdnState, layers: layers, maxContext: maxContext, - mlpWeightBits: mlpWeightBits) + mlpWeightBits: mlpWeightBits, + paged: pagedKV) } } @@ -404,6 +490,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR kv.reset() gdnState.reset() mtp?.reset() + pagedKV?.resetState() } public var continuationPosition: Int { kv.position } @@ -554,7 +641,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR } prefillChunkState.markDirty(startPosition: startPosition, tokenCount: t) - let cb = try commandBuffer() + var cb = try commandBuffer() let emb = model.embedding prefillEmbed.encode(commandBuffer: cb, @@ -578,11 +665,11 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR encodeLinearAttentionPrefill(cb, layer: layer, layerIndex: index, scratch: scratch, tokenCount: t) } else { - try encodeGatedFullAttentionPrefill(cb, layer: layer, - layerIndex: index, - scratch: scratch, - tokenCount: t, - startPosition: startPosition) + cb = try encodeGatedFullAttentionPrefill(cb, layer: layer, + layerIndex: index, + scratch: scratch, + tokenCount: t, + startPosition: startPosition) } elementwise.encodeResidualAdd(commandBuffer: cb, hidden: scratch.hidden, @@ -657,6 +744,14 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR if writeFinalHead, emitGreedyHead { lastGreedyToken = greedyTokenBuf.contents().load(as: UInt32.self) } + if let paged = pagedKV { + // Page summaries for chunk-sealed pages were encoded in the chunk + // command buffer itself, so nothing is pending here. The chunk + // moved the frontier; any scores from an earlier token are stale + // and the next decode token takes the warmup selection. + paged.store.advance(by: t) + for i in 0.. MTLCommandBuffer { guard let q = layer.q, let k = layer.k, let v = layer.v, let o = layer.o, let qNormW = layer.qNorm, let kNormW = layer.kNorm else { preconditionFailure("full-attention layer without attention tensors") @@ -814,42 +912,252 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR theta: Float(cfg.fullRopeTheta), rotaryDim: rotaryDim, eps: Self.epsilon) - try copyStagedKVToCache(cb, layer: layerIndex, - startPosition: startPosition, - tokenCount: t, - keySource: scratch.kStage, - valueSource: scratch.vStage, - bytesPerToken: kvDim * MemoryLayout.stride) - let params = PrefillAttentionParams( - startPosition: UInt32(startPosition), - queryCount: UInt32(t), - headDim: UInt32(headDim), - numQHeads: UInt32(cfg.numHeads), - numKVHeads: UInt32(numKV), - kvValidCount: UInt32(startPosition + t), - slidingWindow: UInt32(startPosition + t), - kvTokenStrideElements: UInt32(kvDim), - qTokenStrideElements: UInt32(qDim), - oTokenStrideElements: UInt32(qDim), - scale: Float(cfg.attentionScale)) - prefillAttention.encodeCausal( - commandBuffer: cb, - q: scratch.attnQ, - k: kv.keyBuffer(layer: layerIndex, validTokenCount: startPosition + t), - v: kv.valueBuffer(layer: layerIndex, validTokenCount: startPosition + t), - out: scratch.attnOut, - params: params) - elementwise.encodeSigmoidGateMul(commandBuffer: cb, + var activeCB = cb + let bytesPerToken = kvDim * MemoryLayout.stride + let endPages = (startPosition + t + KVPageGeometry.tokensPerPage - 1) + / KVPageGeometry.tokensPerPage + if let paged = pagedKV, + !(paged.store.identityMappingIntact && endPages <= paged.poolPagesPerLayer) { + // Blocked (streamed) path: the chunk's KV scatters into whatever + // pool slots are free while the sealed past streams from the + // spill file through staging windows, folded into FP32 running + // softmax state. Exact — only the summation order differs from + // the tensor-ops path. + try encodePagedChunkKVScatter(cb, paged: paged, layerIndex: layerIndex, + startPosition: startPosition, tokenCount: t, + keySource: scratch.kStage, + valueSource: scratch.vStage, + bytesPerToken: bytesPerToken) + try encodePagedChunkMinMax(cb, paged: paged, layerIndex: layerIndex, + startPosition: startPosition, tokenCount: t) + try finish(cb) + let blocked = try ensureBlockedPrefillScratch(chunkTokens: scratch.chunkTokens, + paged: paged) + try runBlockedAttention(paged: paged, blocked: blocked, + layerIndex: layerIndex, + queryCount: t, startPosition: startPosition, + scratch: scratch) + activeCB = try commandBuffer() + } else { + try copyStagedKVToCache(cb, layer: layerIndex, + startPosition: startPosition, + tokenCount: t, + keySource: scratch.kStage, + valueSource: scratch.vStage, + bytesPerToken: bytesPerToken) + if let paged = pagedKV { + try encodePagedChunkMinMax(cb, paged: paged, layerIndex: layerIndex, + startPosition: startPosition, tokenCount: t) + } + let params = PrefillAttentionParams( + startPosition: UInt32(startPosition), + queryCount: UInt32(t), + headDim: UInt32(headDim), + numQHeads: UInt32(cfg.numHeads), + numKVHeads: UInt32(numKV), + kvValidCount: UInt32(startPosition + t), + slidingWindow: UInt32(startPosition + t), + kvTokenStrideElements: UInt32(kvDim), + qTokenStrideElements: UInt32(qDim), + oTokenStrideElements: UInt32(qDim), + scale: Float(cfg.attentionScale)) + prefillAttention.encodeCausal( + commandBuffer: cb, + q: scratch.attnQ, + k: pagedKV?.store.kPoolBuffer(layer: layerIndex) + ?? kv.keyBuffer(layer: layerIndex, validTokenCount: startPosition + t), + v: pagedKV?.store.vPoolBuffer(layer: layerIndex) + ?? kv.valueBuffer(layer: layerIndex, validTokenCount: startPosition + t), + out: scratch.attnOut, + params: params) + } + elementwise.encodeSigmoidGateMul(commandBuffer: activeCB, out: scratch.attnOut, gate: scratch.attnGate, count: t * qDim) - encodePrefillInt4Projection(cb, + encodePrefillInt4Projection(activeCB, weights: o, x: scratch.attnOut, y: scratch.h1, rows: D, columns: qDim, tokenCount: t, xStrideElements: qDim, yStrideElements: D) + return activeCB + } + + /// Scatter a chunk's staged K/V rows into their (possibly non-identity) + /// pool page slots — one blit per touched page per stream. + private func encodePagedChunkKVScatter(_ cb: MTLCommandBuffer, + paged: Qwen38PagedKVRuntime, + layerIndex: Int, + startPosition: Int, + tokenCount: Int, + keySource: MTLBuffer, + valueSource: MTLBuffer, + bytesPerToken: Int) throws { + let pageTokens = KVPageGeometry.tokensPerPage + let firstPage = startPosition / pageTokens + let lastPage = (startPosition + tokenCount - 1) / pageTokens + guard let blit = cb.makeBlitCommandEncoder() else { + throw Qwen38ForwardRunnerError.commandFailed( + "unable to create Qwen 3.8 paged KV scatter encoder") + } + for page in firstPage...lastPage { + let writeStart = max(page * pageTokens, startPosition) + let writeEnd = min((page + 1) * pageTokens, startPosition + tokenCount) + let count = writeEnd - writeStart + guard count > 0 else { continue } + let kDst = try paged.store.kSlot(layer: layerIndex, position: writeStart) + let vDst = try paged.store.vSlot(layer: layerIndex, position: writeStart) + let srcOffset = (writeStart - startPosition) * bytesPerToken + blit.copy(from: keySource, sourceOffset: srcOffset, + to: kDst.buffer, destinationOffset: kDst.offset, + size: count * bytesPerToken) + blit.copy(from: valueSource, sourceOffset: srcOffset, + to: vDst.buffer, destinationOffset: vDst.offset, + size: count * bytesPerToken) + } + blit.endEncoding() + } + + /// Quest min/max summaries for every page this chunk finishes filling — + /// encoded in the chunk's own command buffer while the rows are + /// guaranteed resident, so a long prefill never triggers a refetch storm + /// at first decode. + private func encodePagedChunkMinMax(_ cb: MTLCommandBuffer, + paged: Qwen38PagedKVRuntime, + layerIndex: Int, + startPosition: Int, + tokenCount: Int) throws { + let pageTokens = KVPageGeometry.tokensPerPage + let firstSealed = startPosition / pageTokens + let sealedEnd = (startPosition + tokenCount) / pageTokens + guard sealedEnd > firstSealed, + let ordinal = paged.store.fullLayerOrdinal(forLayer: layerIndex) else { return } + let g = paged.store.geometry + for page in firstSealed...stride, + out: attnOut, + headDim: UInt32(headDim), + numQHeads: UInt32(cfg.numHeads), + numKVHeads: UInt32(numKV), + selTokens: UInt32(selection.selTokens), + scale: Float(cfg.attentionScale)) + // Score every sealed page against this token's query — the + // selection input for the *next* token (lag-one policy). + paged.encodeScores(commandBuffer: cb, ordinal: ordinal, + q: qScratch, qOffset: 0, + sealedPages: position / KVPageGeometry.tokensPerPage) + } else { + attention.encodeFull(commandBuffer: cb, + q: qScratch, + k: kSlot.buffer, kOffset: 0, + v: vSlot.buffer, vOffset: 0, + out: attnOut, + headDim: UInt32(headDim), + numQHeads: UInt32(cfg.numHeads), + numKVHeads: UInt32(numKV), + seqLen: seqLen, + scale: Float(cfg.attentionScale)) + } elementwise.encodeSigmoidGateMul(commandBuffer: cb, out: attnOut, gate: attnGateScratch, diff --git a/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift b/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift index c405595..f108f3c 100644 --- a/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift +++ b/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift @@ -66,6 +66,10 @@ final class Qwen38MTPSpeculator { private let ctx: MetalContext private let cfg: ArchConfig private let kv: KVCacheManager + /// Paged long-context state shared with the runner; nil in linear mode. + /// The verify pass writes draft rows through the page store and runs + /// per-position paged sparse attention over the round's pinned selection. + private let paged: Qwen38PagedKVRuntime? private let gdnState: GDNStateManager private let layers: [Qwen38ForwardRunner.LayerTensors] private let maxContext: Int @@ -209,12 +213,14 @@ final class Qwen38MTPSpeculator { gdnState: GDNStateManager, layers: [Qwen38ForwardRunner.LayerTensors], maxContext: Int, - mlpWeightBits: Int) throws -> Qwen38MTPSpeculator? { + mlpWeightBits: Int, + paged: Qwen38PagedKVRuntime? = nil) throws -> Qwen38MTPSpeculator? { guard (try? model.resident(name: "mtp.fc.weight")) != nil else { return nil } return try Qwen38MTPSpeculator(model: model, context: context, config: config, kv: kv, gdnState: gdnState, layers: layers, maxContext: maxContext, - mlpWeightBits: mlpWeightBits) + mlpWeightBits: mlpWeightBits, + paged: paged) } private init(model: Model, @@ -224,11 +230,13 @@ final class Qwen38MTPSpeculator { gdnState: GDNStateManager, layers: [Qwen38ForwardRunner.LayerTensors], maxContext: Int, - mlpWeightBits: Int) throws { + mlpWeightBits: Int, + paged: Qwen38PagedKVRuntime?) throws { self.model = model self.ctx = context self.cfg = config self.kv = kv + self.paged = paged self.gdnState = gdnState self.layers = layers self.maxContext = maxContext @@ -456,6 +464,7 @@ final class Qwen38MTPSpeculator { "Qwen 3.8 MTP continuation restore did not complete") } kv.rewind(to: expectedPosition) + paged?.store.rewind(to: expectedPosition) arenaRoundBase = -1 arenaCaptured = 0 } @@ -479,7 +488,23 @@ final class Qwen38MTPSpeculator { /// A round needs the cursor caught up and room for at least one draft. func canRunRound(position: Int) -> Bool { - position == kv.position && position + 2 <= maxContext + // A round needs a prior decoded position: the drafter seeds from the + // previous token's final-norm hidden, and `basePair = position - 1` + // anchors its RoPE — position 0 decodes plainly. + guard position >= 1, position == kv.position, + position + 2 <= maxContext else { return false } + guard let paged else { return true } + // Sparse selections break the byte-identity contract: a round reuses + // one page table across its verify rows where plain decode reselects + // per token, and accepted rows are emitted without their own score + // pass. Rounds therefore run only while the selection through the + // round's span — plus one position of margin, so the token feeding + // the first sparse selection's lag-one scores is always + // plain-decoded — covers the entire context. Past that point decode + // falls back to plain paged tokens. + let kRound = min(draftCount, maxContext - position - 1) + let horizon = min(position + kRound + 1, maxContext - 1) + return paged.selectionIsExhaustive(at: horizon, maxSpanTokens: kRound + 1) } /// Run one draft/verify/accept round for `produce(bonus, position)`. @@ -537,8 +562,15 @@ final class Qwen38MTPSpeculator { // 2. Verify. let verifyStart = clock_gettime_nsec_np(CLOCK_UPTIME_RAW) let verifyTokens = [bonus] + drafts + if let paged { + // Round selection: lag-one scores as in plain decode, with the + // table extended over the pages the verify span writes into. + try paged.prepareSelections(position: P) + try paged.appendVerifySpan(position: P, count: verifyTokens.count) + } try runVerifyPass(tokens: verifyTokens, startPosition: P) kv.advance(by: verifyTokens.count) + paged?.store.advance(by: verifyTokens.count) arenaRoundBase = P arenaCaptured = verifyTokens.count - 1 @@ -560,6 +592,14 @@ final class Qwen38MTPSpeculator { if rolledBack { stats.rollbacks += 1 kv.rewind(to: P + accepted + 1) + paged?.store.rewind(to: P + accepted + 1) + } + if let paged { + // Net-sealed pages need their summaries next command buffer, and + // the round's score pass (bonus-position query) feeds the next + // selection. + paged.noteAdvance(from: P, to: P + accepted + 1) + paged.readBackScores(sealedPages: P / KVPageGeometry.tokensPerPage) } try runAcceptPass(bonus: bonus, emitted: emitted, position: P, restoreSlot: rolledBack ? accepted : nil) @@ -584,6 +624,7 @@ final class Qwen38MTPSpeculator { for (i, token) in tokens.enumerated() { ids[i] = UInt32(bitPattern: token) } let cb = try commandBuffer() + try paged?.encodePendingMetadata(commandBuffer: cb) let emb = model.embedding prefillEmbed.encode(commandBuffer: cb, table: emb.buffer, tableOffset: Int(emb.offset), @@ -674,23 +715,46 @@ final class Qwen38MTPSpeculator { packed: vQPacked, q: vQ, gate: vGate, heads: cfg.numHeads, dim: headDim, rows: t) - // Stage -> cache (contiguous slots for a linear full-attention layer). - let kRange = kv.kRange(layer: layerIndex, start: startPosition, count: t) - let vRange = kv.vRange(layer: layerIndex, start: startPosition, count: t) + // Stage -> cache: contiguous slots in linear mode, per-page scatter + // through the page store in paged mode (the span may cross a page + // boundary into a freshly allocated slot). guard let blit = cb.makeBlitCommandEncoder() else { throw Qwen38ForwardRunnerError.commandFailed( "unable to create Qwen 3.8 MTP verify KV blit encoder") } - blit.copy(from: vKStage, sourceOffset: 0, - to: kRange.buffer, destinationOffset: kRange.offset, - size: t * kvDim * h) - blit.copy(from: vVStage, sourceOffset: 0, - to: vRange.buffer, destinationOffset: vRange.offset, - size: t * kvDim * h) + if let paged { + let pageTokens = KVPageGeometry.tokensPerPage + let firstPage = startPosition / pageTokens + let lastPage = (startPosition + t - 1) / pageTokens + for page in firstPage...lastPage { + let writeStart = max(page * pageTokens, startPosition) + let writeEnd = min((page + 1) * pageTokens, startPosition + t) + let kDst = try paged.store.kSlot(layer: layerIndex, position: writeStart) + let vDst = try paged.store.vSlot(layer: layerIndex, position: writeStart) + let srcOffset = (writeStart - startPosition) * kvDim * h + blit.copy(from: vKStage, sourceOffset: srcOffset, + to: kDst.buffer, destinationOffset: kDst.offset, + size: (writeEnd - writeStart) * kvDim * h) + blit.copy(from: vVStage, sourceOffset: srcOffset, + to: vDst.buffer, destinationOffset: vDst.offset, + size: (writeEnd - writeStart) * kvDim * h) + } + } else { + let kRange = kv.kRange(layer: layerIndex, start: startPosition, count: t) + let vRange = kv.vRange(layer: layerIndex, start: startPosition, count: t) + blit.copy(from: vKStage, sourceOffset: 0, + to: kRange.buffer, destinationOffset: kRange.offset, + size: t * kvDim * h) + blit.copy(from: vVStage, sourceOffset: 0, + to: vRange.buffer, destinationOffset: vRange.offset, + size: t * kvDim * h) + } blit.endEncoding() for i in 0...stride, + out: vAttnOut, outOffset: i * qDim * h, + headDim: UInt32(headDim), + numQHeads: UInt32(cfg.numHeads), + numKVHeads: UInt32(numKV), + selTokens: UInt32(base + i + 1), + scale: Float(cfg.attentionScale)) + } + // Bonus-position query scores every sealed page for the next + // selection — the bonus token is always committed, so its query + // is always the right one to rank against. + paged.encodeScores(commandBuffer: cb, ordinal: ordinal, + q: vQ, qOffset: 0, + sealedPages: startPosition / KVPageGeometry.tokensPerPage) + } else { + for i in 0.. 0, "maxContext must be positive") precondition(maxPrefillChunkTokens > 0, "maxPrefillChunkTokens must be positive") self.config = config @@ -109,7 +114,7 @@ public final class KVCacheManager { for layer in 0..= 0 && recentPages >= 1 && topKPages >= 0, + "invalid selector parameters (recent must cover the tail)") + self.sinkPages = sinkPages + self.recentPages = recentPages + self.topKPages = topKPages + } + + public struct Selection: Equatable, Sendable { + /// Ascending page indices. + public let pages: [Int] + /// Selected logical token count; the tail page contributes only its + /// valid rows. + public let selTokens: Int + } + + /// - Parameters: + /// - scores: Quest score per *sealed* page, from the previous token's + /// query. Empty on the first decode token after a prefill (warmup: + /// sinks + recent only, plus trailing fill up to the top-k budget). + /// May cover fewer pages than `sealedPages` — a page that sealed on + /// the previous token has no score yet; unscored trailing pages are + /// recent by construction and the recent window covers them. + /// - sealedPages: pages fully written (64 valid tokens each). + /// - tailValidTokens: valid rows in the unsealed tail page (up to a + /// full 64 for a page written but not yet sealed by `advance`); 0 + /// when the position sits exactly on a page boundary. + public func select(scores: [Float], + sealedPages: Int, + tailValidTokens: Int) -> Selection { + precondition(tailValidTokens >= 0 && tailValidTokens <= 64, + "tailValidTokens must be 0...64") + let totalPages = sealedPages + (tailValidTokens > 0 ? 1 : 0) + guard totalPages > 0 else { return Selection(pages: [], selTokens: 0) } + + var picked = Set() + for page in 0.. 0 { + if scores.isEmpty { + // Warmup: no query yet — fill the budget with the most recent + // sealed pages, the best context-free prior. + var page = totalPages - 1 + var budget = topKPages + while budget > 0 && page >= 0 { + if picked.insert(page).inserted { budget -= 1 } + page -= 1 + } + } else { + // Deterministic top-k: score desc, index asc on ties. + let scoredPages = min(scores.count, sealedPages) + let candidates = (0.. scores[$1] : $0 < $1 } + for page in candidates.prefix(topKPages) { picked.insert(page) } + } + } + + let pages = picked.sorted() + let tailIndex = tailValidTokens > 0 ? totalPages - 1 : -1 + let selTokens = pages.reduce(0) { acc, page in + acc + (page == tailIndex ? tailValidTokens : 64) + } + return Selection(pages: pages, selTokens: selTokens) + } + + /// True when `select` over `totalPages` pages provably picks every page + /// no matter what the scores contain, provided at most + /// `maxUnscoredSealedPages` trailing sealed pages lack Quest scores + /// (plain decode lags one page; speculative rounds can lag by their + /// span). Gap pages — neither sink nor recent — are picked by top-k + /// only when scored, so any unscored page must sit inside the recent + /// window for coverage to be unconditional. + public func coversEntireContext(totalPages: Int, + maxUnscoredSealedPages: Int) -> Bool { + guard totalPages > 0 else { return true } + let sinks = min(sinkPages, totalPages) + let recentStart = max(0, totalPages - recentPages) + let unionCount = sinks + (totalPages - recentStart) + - max(0, sinks - recentStart) + let gap = totalPages - unionCount + if gap == 0 { return true } + return gap <= topKPages && recentPages >= 1 + maxUnscoredSealedPages + } +} diff --git a/Sources/Mference/Runtime/KVCache/KVPageStore.swift b/Sources/Mference/Runtime/KVCache/KVPageStore.swift new file mode 100644 index 0000000..f20ed6c --- /dev/null +++ b/Sources/Mference/Runtime/KVCache/KVPageStore.swift @@ -0,0 +1,588 @@ +import Foundation +import Darwin +import Metal + +/// Page geometry for the paged full-attention KV store. Pages are fixed at +/// 64 tokens; per layer a page is one K block and one V block of +/// `tokensPerPage * tokenStrideBytes` each. +public struct KVPageGeometry: Sendable, Equatable { + /// Fixed page size. 64 tokens * 2048 B/token = 128 KiB per K (or V) page + /// for Qwen 3.8 — the measured NVMe random-read sweet spot when K and V + /// are fetched together (256 KiB at ~2.3 GiB/s). + public static let tokensPerPage = 64 + + /// Absolute layer indices with full attention (mask value 1), in order. + public let fullLayerOrdinals: [Int] + /// Bytes per token per layer for one of K or V: + /// `numFullKVHeads * fullHeadDim * sizeof(FP16)`. + public let tokenStrideBytes: Int + public let maxContext: Int + public let numKVHeads: Int + public let headDim: Int + + public var pagesPerLayer: Int { + (maxContext + Self.tokensPerPage - 1) / Self.tokensPerPage + } + /// Bytes of one K (or V) page for one layer. + public var kPageBytes: Int { Self.tokensPerPage * tokenStrideBytes } + + /// Spill-file offset of a page's K block. The V block follows at + /// `+kPageBytes`. Layer-major so a layer's pages stream sequentially. + public func fileOffset(layerOrdinal: Int, pageIndex: Int) -> Int { + (layerOrdinal * pagesPerLayer + pageIndex) * 2 * kPageBytes + } + + /// Quest metadata per page: element-wise min and max over the page's K + /// rows, per kv-head — 2 * numKVHeads * headDim FP16 values. + public var metadataBytesPerPage: Int { 2 * numKVHeads * headDim * 2 } + + public func metadataOffset(layerOrdinal: Int, pageIndex: Int) -> Int { + (layerOrdinal * pagesPerLayer + pageIndex) * metadataBytesPerPage + } +} + +public enum KVPageStoreError: Error, Equatable { + case allocationFailed(String) + case poolExhausted(layer: Int, pageIndex: Int) + case spillFileFailed(String) + case ioFailed(operation: String, errno: Int32) + case notAFullAttentionLayer(Int) + case pageNotSealed(layer: Int, pageIndex: Int) +} + +/// Paged FP16 K/V storage for full-attention layers with an SSD spill tier. +/// +/// Each full-attention layer owns a K pool and a V pool of +/// `poolPagesPerLayer` page slots. Pages seal on 64-token boundaries as the +/// position cursor advances; sealed pages are written behind to a sparse +/// layer-major spill file and become evictable under LRU pressure. Unsealed +/// (tail) pages and explicitly pinned pages never evict. Fetching a spilled +/// page is a single 2x128 KiB `pread` pair into a free or victim slot. +/// +/// Mirrors `KVCacheManager`'s `kSlot`/`vSlot`/`advance`/`reset` surface for +/// the layers it owns; linear/GDN layers remain the runner's concern. +public final class KVPageStore { + public let geometry: KVPageGeometry + public let spillFileName = "kvpages.spill" + /// Per-page Quest min/max metadata, shared storage so score kernels read + /// it directly and seal kernels write it in the token command buffer. + public let metadataBuffer: MTLBuffer + + public private(set) var position: Int = 0 + public private(set) var sealedPageCount: Int = 0 + /// True while every page has only ever landed at slot == pageIndex and + /// nothing has been evicted — the pool region [0, position) is then one + /// linear cache and the fast (tensor-ops) prefill path applies. + public private(set) var identityMappingIntact = true + + private let poolPagesPerLayer: Int + private let kPools: [MTLBuffer] // [fullLayerOrdinal] + private let vPools: [MTLBuffer] + private let ordinalByLayer: [Int: Int] + + private enum PageState: UInt8 { case untouched, unsealed, sealed } + + private struct PageKey: Hashable { + let ordinal: Int + let pageIndex: Int + } + + // All parallel arrays are indexed [fullLayerOrdinal][...]. + private var pageSlot: [[Int32]] // [ordinal][pageIndex] -> slot or -1 + private var slotPage: [[Int32]] // [ordinal][slot] -> pageIndex or -1 + private var pageState: [[PageState]] + /// (ordinal, pageIndex) pairs queued for write-behind but possibly not on + /// disk yet. Caller-thread-only; the spill queue never touches it — a + /// `flushSpills()` barrier is the one synchronization point, after which + /// every queued write has landed and the set empties. + private var pendingSpills = Set() + private var pagePinned: [[Bool]] + private var slotLastUse: [[UInt64]] + private var useTick: UInt64 = 0 + + private let spillFD: Int32 + private let spillQueue = DispatchQueue(label: "mference.kvpage.spill") + private let spillDirectory: URL + /// First write-behind failure (ENOSPC, short write, I/O error), recorded + /// on the spill queue and surfaced by the next spill-file read — pages + /// past a failed spill are unreliable on disk, so the run must fail + /// cleanly instead of fetching garbage. A separate box because spill + /// closures must not retain the store: its deinit synchronizes on the + /// spill queue, and releasing the last reference on that queue would + /// deadlock. + private final class SpillErrorBox { + private let lock = NSLock() + private var error: KVPageStoreError? + func record(_ newError: KVPageStoreError) { + lock.lock() + defer { lock.unlock() } + if error == nil { error = newError } + } + var value: KVPageStoreError? { + lock.lock() + defer { lock.unlock() } + return error + } + func clear() { + lock.lock() + defer { lock.unlock() } + error = nil + } + } + private let spillError = SpillErrorBox() + + public init(device: MTLDevice, + config: ArchConfig, + maxContext: Int, + poolPagesPerLayer: Int, + spillDirectory: URL) throws { + precondition(maxContext > 0, "maxContext must be positive") + precondition(poolPagesPerLayer > 0, "poolPagesPerLayer must be positive") + + let ordinals = config.fullAttentionLayerMask.enumerated() + .filter { $0.element == 1 }.map(\.offset) + precondition(!ordinals.isEmpty, "config has no full-attention layers") + + let stride = config.numFullKVHeads * config.fullHeadDim * 2 + self.geometry = KVPageGeometry(fullLayerOrdinals: ordinals, + tokenStrideBytes: stride, + maxContext: maxContext, + numKVHeads: config.numFullKVHeads, + headDim: config.fullHeadDim) + self.poolPagesPerLayer = poolPagesPerLayer + self.ordinalByLayer = Dictionary(uniqueKeysWithValues: + ordinals.enumerated().map { ($0.element, $0.offset) }) + + let poolBytes = poolPagesPerLayer * geometry.kPageBytes + var ks: [MTLBuffer] = [] + var vs: [MTLBuffer] = [] + ks.reserveCapacity(ordinals.count) + vs.reserveCapacity(ordinals.count) + for layer in ordinals { + guard let k = device.makeBuffer(length: poolBytes, options: .storageModeShared), + let v = device.makeBuffer(length: poolBytes, options: .storageModeShared) else { + throw KVPageStoreError.allocationFailed("kv page pool layer \(layer)") + } + k.label = "kvpage.K.layer\(layer)" + v.label = "kvpage.V.layer\(layer)" + ks.append(k) + vs.append(v) + } + self.kPools = ks + self.vPools = vs + + let metadataLength = ordinals.count * geometry.pagesPerLayer + * geometry.metadataBytesPerPage + guard let metadata = device.makeBuffer(length: metadataLength, + options: .storageModeShared) else { + throw KVPageStoreError.allocationFailed("kv page metadata") + } + metadata.label = "kvpage.metadata" + self.metadataBuffer = metadata + + let n = ordinals.count + let pages = geometry.pagesPerLayer + self.pageSlot = Array(repeating: Array(repeating: -1, count: pages), count: n) + self.slotPage = Array(repeating: Array(repeating: -1, count: poolPagesPerLayer), count: n) + self.pageState = Array(repeating: Array(repeating: .untouched, count: pages), count: n) + self.pagePinned = Array(repeating: Array(repeating: false, count: pages), count: n) + self.slotLastUse = Array(repeating: Array(repeating: 0, count: poolPagesPerLayer), count: n) + + self.spillDirectory = spillDirectory + let path = spillDirectory.appendingPathComponent(spillFileName).path + let fd = open(path, O_RDWR | O_CREAT, 0o600) + guard fd >= 0 else { + throw KVPageStoreError.spillFileFailed("open(\(path)) errno \(errno)") + } + // Sparse-extend to full capacity; APFS allocates blocks only on write. + let fullSize = off_t(ordinals.count * geometry.pagesPerLayer * 2 * geometry.kPageBytes) + guard ftruncate(fd, fullSize) == 0 else { + close(fd) + throw KVPageStoreError.spillFileFailed("ftruncate errno \(errno)") + } + // The pool is the cache; keep spill I/O out of the unified page cache. + _ = fcntl(fd, F_NOCACHE, 1) + self.spillFD = fd + } + + deinit { + spillQueue.sync {} + close(spillFD) + try? FileManager.default.removeItem( + at: spillDirectory.appendingPathComponent(spillFileName)) + } + + // MARK: - Layer mapping + + public func fullLayerOrdinal(forLayer layer: Int) -> Int? { ordinalByLayer[layer] } + + public func kPoolBuffer(layer: Int) -> MTLBuffer { kPools[requireOrdinal(layer)] } + public func vPoolBuffer(layer: Int) -> MTLBuffer { vPools[requireOrdinal(layer)] } + + // MARK: - Unsealed writes (decode/prefill hot path) + + /// Write target for `layer`'s K projection at `position`. Allocates the + /// page slot on first touch (possibly evicting an LRU sealed page). + public func kSlot(layer: Int, position: Int) throws -> (buffer: MTLBuffer, offset: Int) { + let (ordinal, slot, within) = try writableSlot(layer: layer, position: position) + return (kPools[ordinal], slot * geometry.kPageBytes + within * geometry.tokenStrideBytes) + } + + /// Write target for `layer`'s V projection at `position`. Same slot index + /// as K, in the distinct V pool. + public func vSlot(layer: Int, position: Int) throws -> (buffer: MTLBuffer, offset: Int) { + let (ordinal, slot, within) = try writableSlot(layer: layer, position: position) + return (vPools[ordinal], slot * geometry.kPageBytes + within * geometry.tokenStrideBytes) + } + + private func writableSlot(layer: Int, + position: Int) throws -> (ordinal: Int, slot: Int, within: Int) { + precondition(position >= 0 && position < geometry.maxContext, + "position \(position) out of range 0..<\(geometry.maxContext)") + let ordinal = requireOrdinal(layer) + let pageIndex = position / KVPageGeometry.tokensPerPage + precondition(pageState[ordinal][pageIndex] != .sealed, + "write to sealed page \(pageIndex) of layer \(layer)") + let slot = try residentSlot(ordinal: ordinal, layer: layer, pageIndex: pageIndex, + allocateAs: .unsealed) + return (ordinal, slot, position % KVPageGeometry.tokensPerPage) + } + + /// Contiguous K range for a prefill chunk write, valid under the + /// identity slot mapping (pool sized to the full context, no evictions): + /// pages are allocated sequentially so page `i` sits at slot `i` and a + /// multi-page span is one linear region of the pool. The beyond-RAM + /// prefill path (blocked attention) replaces this. + public func contiguousKRange(layer: Int, start: Int, + count: Int) throws -> (buffer: MTLBuffer, offset: Int, stride: Int) { + let ordinal = try contiguousRangeOrdinal(layer: layer, start: start, count: count) + return (kPools[ordinal], start * geometry.tokenStrideBytes, geometry.tokenStrideBytes) + } + + public func contiguousVRange(layer: Int, start: Int, + count: Int) throws -> (buffer: MTLBuffer, offset: Int, stride: Int) { + let ordinal = try contiguousRangeOrdinal(layer: layer, start: start, count: count) + return (vPools[ordinal], start * geometry.tokenStrideBytes, geometry.tokenStrideBytes) + } + + private func contiguousRangeOrdinal(layer: Int, start: Int, count: Int) throws -> Int { + precondition(count > 0 && start >= 0 && start + count <= geometry.maxContext, + "range \(start)..<\(start + count) exceeds maxContext") + let ordinal = requireOrdinal(layer) + let firstPage = start / KVPageGeometry.tokensPerPage + let lastPage = (start + count - 1) / KVPageGeometry.tokensPerPage + for page in firstPage...lastPage { + let slot = try residentSlot(ordinal: ordinal, layer: layer, pageIndex: page, + allocateAs: .unsealed) + precondition(slot == page, + "contiguous KV range requires the identity slot mapping " + + "(page \(page) at slot \(slot)); use the blocked prefill path") + } + return ordinal + } + + // MARK: - Position / sealing + + /// Advance the position cursor. Every page fully crossed by the cursor is + /// sealed across all full-attention layers and queued for write-behind. + public func advance(by count: Int) { + precondition(count >= 0, "advance count must be non-negative") + precondition(position + count <= geometry.maxContext, + "advance would exceed maxContext") + let sealedBefore = position / KVPageGeometry.tokensPerPage + position += count + let sealedAfter = position / KVPageGeometry.tokensPerPage + for pageIndex in sealedBefore..= 0 && newPosition <= position, + "rewind target \(newPosition) outside 0...\(position)") + let keepSealed = newPosition / KVPageGeometry.tokensPerPage + let currentSealed = position / KVPageGeometry.tokensPerPage + position = newPosition + guard currentSealed > keepSealed else { return } + flushSpills() + for pageIndex in keepSealed..= 0, + "rewind found an evicted in-span page") + pageState[ordinal][pageIndex] = .unsealed + unsealedAny = true + } + if unsealedAny { sealedPageCount -= 1 } + } + } + + private func seal(pageIndex: Int) { + var sealedAny = false + for ordinal in 0.. Bool { + pageSlot[requireOrdinal(layer)][pageIndex] >= 0 + } + + public func residentPageCount(layer: Int) -> Int { + slotPage[requireOrdinal(layer)].lazy.filter { $0 >= 0 }.count + } + + /// Fetch `pageIndex` into the pool if spilled; returns its slot. Sealed + /// pages only — unsealed pages are always resident. + @discardableResult + public func ensureResident(layer: Int, pageIndex: Int) throws -> Int { + let ordinal = requireOrdinal(layer) + guard pageState[ordinal][pageIndex] != .untouched else { + throw KVPageStoreError.pageNotSealed(layer: layer, pageIndex: pageIndex) + } + return try residentSlot(ordinal: ordinal, layer: layer, pageIndex: pageIndex, + allocateAs: nil) + } + + public func pin(layer: Int, pageIndex: Int) { + pagePinned[requireOrdinal(layer)][pageIndex] = true + } + + public func unpin(layer: Int, pageIndex: Int) { + pagePinned[requireOrdinal(layer)][pageIndex] = false + } + + /// Pool-slot table for the paged attention kernel: one `uint32` slot per + /// selected page, in the caller's (ascending-position) order. Fetches any + /// spilled selection member. The selection must fit the pool. + public func pageTable(layer: Int, selectedPages: [Int]) throws -> [UInt32] { + var table = [UInt32]() + table.reserveCapacity(selectedPages.count) + for page in selectedPages { + table.append(UInt32(try ensureResident(layer: layer, pageIndex: page))) + } + return table + } + + // MARK: - Streamed window reads (blocked prefill) + + /// One sequential `pread` of `pageCount` interleaved [K page | V page] + /// pairs starting at `firstPage` into `staging` — the blocked-prefill + /// streaming path. Pages must be sealed and spilled (`flushSpills()` + /// first). The flash kernel addresses the interleaved layout with a + /// stride-2 page table and a V base offset of one K page. + public func readSpilledSpan(layer: Int, firstPage: Int, pageCount: Int, + into staging: MTLBuffer) throws { + precondition(pageCount > 0, "empty span read") + try throwIfSpillFailed() + let ordinal = requireOrdinal(layer) + let bytes = pageCount * 2 * geometry.kPageBytes + precondition(staging.length >= bytes, "staging buffer too small for span") + precondition(firstPage + pageCount <= geometry.pagesPerLayer, "span out of range") + let offset = off_t(geometry.fileOffset(layerOrdinal: ordinal, pageIndex: firstPage)) + var done = 0 + while done < bytes { + let n = pread(spillFD, staging.contents() + done, bytes - done, offset + off_t(done)) + guard n > 0 else { + throw KVPageStoreError.ioFailed(operation: "pread span", errno: errno) + } + done += n + } + } + + /// Current pool slot of a resident page (unsealed or fetched); the + /// blocked-prefill tail window builds its page table from these. + public func residentSlot(layer: Int, pageIndex: Int) -> Int? { + let slot = pageSlot[requireOrdinal(layer)][pageIndex] + return slot >= 0 ? Int(slot) : nil + } + + // MARK: - Reset + + /// Drop all pages, rewind the cursor, and return pool pages to the OS. + public func reset() { + flushSpills() + // A fresh run depends on no failed write: the file is re-punched + // below and every page rewritten before it can be read again. + spillError.clear() + position = 0 + sealedPageCount = 0 + identityMappingIntact = true + let pages = geometry.pagesPerLayer + for ordinal in 0.. 0 { _ = posix_madvise(buffer.contents(), len, POSIX_MADV_DONTNEED) } + } + // Punch the file back to sparse. + let fullSize = off_t(kPools.count * geometry.pagesPerLayer * 2 * geometry.kPageBytes) + _ = ftruncate(spillFD, 0) + _ = ftruncate(spillFD, fullSize) + } + + // MARK: - Internals + + private func requireOrdinal(_ layer: Int) -> Int { + guard let ordinal = ordinalByLayer[layer] else { + preconditionFailure("layer \(layer) is not a full-attention layer") + } + return ordinal + } + + /// Resolve (and if needed allocate or fetch) the slot for a page. + /// `allocateAs == .unsealed` permits first-touch allocation of a fresh + /// page; `nil` requires the page to exist already (fetch path). + private func residentSlot(ordinal: Int, layer: Int, pageIndex: Int, + allocateAs: PageState?) throws -> Int { + useTick += 1 + if pageSlot[ordinal][pageIndex] >= 0 { + let slot = Int(pageSlot[ordinal][pageIndex]) + slotLastUse[ordinal][slot] = useTick + return slot + } + let slot = try claimSlot(ordinal: ordinal, layer: layer, pageIndex: pageIndex) + switch pageState[ordinal][pageIndex] { + case .untouched: + guard allocateAs == .unsealed else { + throw KVPageStoreError.pageNotSealed(layer: layer, pageIndex: pageIndex) + } + pageState[ordinal][pageIndex] = .unsealed + case .sealed: + try fetch(ordinal: ordinal, pageIndex: pageIndex, slot: slot) + case .unsealed: + preconditionFailure("unsealed page \(pageIndex) lost residency") + } + pageSlot[ordinal][pageIndex] = Int32(slot) + slotPage[ordinal][slot] = Int32(pageIndex) + slotLastUse[ordinal][slot] = useTick + return slot + } + + private func claimSlot(ordinal: Int, layer: Int, pageIndex: Int) throws -> Int { + if let free = slotPage[ordinal].firstIndex(of: -1) { + if free != pageIndex { identityMappingIntact = false } + return free + } + identityMappingIntact = false + + var victim = -1 + var victimUse = UInt64.max + for slot in 0..= 0 else { + throw KVPageStoreError.poolExhausted(layer: layer, pageIndex: pageIndex) + } + let victimPage = Int(slotPage[ordinal][victim]) + if pendingSpills.contains(PageKey(ordinal: ordinal, pageIndex: victimPage)) { + // Write-behind has not landed yet; barrier so the eviction cannot + // outrun its own spill. + flushSpills() + } + pageSlot[ordinal][victimPage] = -1 + slotPage[ordinal][victim] = -1 + return victim + } + + private func enqueueSpill(ordinal: Int, pageIndex: Int) { + let slot = Int(pageSlot[ordinal][pageIndex]) + precondition(slot >= 0, "sealing a non-resident page") + let pageBytes = geometry.kPageBytes + let kSrc = kPools[ordinal].contents() + slot * pageBytes + let vSrc = vPools[ordinal].contents() + slot * pageBytes + let offset = off_t(geometry.fileOffset(layerOrdinal: ordinal, pageIndex: pageIndex)) + let fd = spillFD + spillQueue.async { [box = spillError] in + // The slot cannot be reused while its spill is pending (eviction + // barriers on this queue first), so the pointers stay valid. + do { + try Self.writeFully(fd: fd, from: kSrc, count: pageBytes, + offset: offset) + try Self.writeFully(fd: fd, from: vSrc, count: pageBytes, + offset: offset + off_t(pageBytes)) + } catch let error as KVPageStoreError { + box.record(error) + } catch { + box.record(.ioFailed(operation: "pwrite spill", errno: EIO)) + } + } + } + + /// `pwrite` until `count` bytes land, resuming short writes and EINTR. + static func writeFully(fd: Int32, from source: UnsafeRawPointer, + count: Int, offset: off_t) throws { + var done = 0 + while done < count { + let n = pwrite(fd, source + done, count - done, offset + off_t(done)) + guard n > 0 else { + if n < 0 && errno == EINTR { continue } + throw KVPageStoreError.ioFailed(operation: "pwrite spill", + errno: n < 0 ? errno : ENOSPC) + } + done += n + } + } + + func recordSpillError(_ error: KVPageStoreError) { + spillError.record(error) + } + + /// The first recorded write-behind failure, if any. + public var spillFailure: KVPageStoreError? { spillError.value } + + private func throwIfSpillFailed() throws { + if let error = spillFailure { throw error } + } + + private func fetch(ordinal: Int, pageIndex: Int, slot: Int) throws { + try throwIfSpillFailed() + let pageBytes = geometry.kPageBytes + let kDst = kPools[ordinal].contents() + slot * pageBytes + let vDst = vPools[ordinal].contents() + slot * pageBytes + let offset = off_t(geometry.fileOffset(layerOrdinal: ordinal, pageIndex: pageIndex)) + let readK = pread(spillFD, kDst, pageBytes, offset) + guard readK == pageBytes else { + throw KVPageStoreError.ioFailed(operation: "pread K", errno: errno) + } + let readV = pread(spillFD, vDst, pageBytes, offset + off_t(pageBytes)) + guard readV == pageBytes else { + throw KVPageStoreError.ioFailed(operation: "pread V", errno: errno) + } + } +} diff --git a/Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift b/Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift new file mode 100644 index 0000000..d352975 --- /dev/null +++ b/Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift @@ -0,0 +1,251 @@ +import Foundation +import Metal + +/// Shared paged long-context state for Qwen 3.8 (kvPagedPolicy == .on): the +/// page store owning full-attention KV, the selection policy and its pinned +/// working set, the per-token page tables and Quest score buffers, and the +/// metadata bookkeeping. Both the plain decode path (`Qwen38ForwardRunner`) +/// and the MTP speculative verify path (`Qwen38MTPSpeculator`) drive one +/// instance, so cursors, pins, and scores stay coherent across round and +/// plain tokens. +final class Qwen38PagedKVRuntime { + let store: KVPageStore + let kernels: KVPageKernels + let selector: KVPageSelector + /// [numFull][pagesPerLayer] float — Quest scores written per token, + /// read back after the command buffer completes (lag-one selection). + let scoresBuf: MTLBuffer + /// [numFull][pagesPerLayer] uint32 — CPU-built page tables bound by the + /// paged attention kernel. + let tablesBuf: MTLBuffer + var lastScores: [[Float]] + var pendingMetadata: [Int] = [] + var selections: [KVPageSelector.Selection] + /// Pages pinned for the in-flight token or verify round, per ordinal — + /// the selection must survive its own fetches under a tight pool, where + /// LRU alone could evict an earlier selection member to admit a later + /// one. + var pinnedSelections: [[Int]] + let poolPagesPerLayer: Int + + private let numQHeads: Int + private let numKVHeads: Int + private let headDim: Int + private let spillDir: URL + + init(context: MetalContext, config: ArchConfig, maxContext: Int, + runtimeConfiguration: RuntimeConfiguration) throws { + let device = context.device + self.numQHeads = config.numHeads + self.numKVHeads = config.numFullKVHeads + self.headDim = config.fullHeadDim + let pagesPerLayer = (maxContext + KVPageGeometry.tokensPerPage - 1) + / KVPageGeometry.tokensPerPage + let poolPages = min( + runtimeConfiguration.kvPoolPagesPerLayer + ?? RuntimeConfiguration.defaultKVPoolPagesPerLayer(config: config, + maxContext: maxContext), + pagesPerLayer) + self.poolPagesPerLayer = poolPages + let dir = FileManager.default.temporaryDirectory + .appendingPathComponent("mference-kvpages-\(UUID().uuidString)", + isDirectory: true) + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + self.spillDir = dir + self.store = try KVPageStore(device: device, + config: config, + maxContext: maxContext, + poolPagesPerLayer: poolPages, + spillDirectory: dir) + self.kernels = try KVPageKernels(context: context) + self.selector = KVPageSelector(sinkPages: runtimeConfiguration.kvSinkPages, + recentPages: runtimeConfiguration.kvRecentPages, + topKPages: runtimeConfiguration.kvTopKPages) + let numFull = store.geometry.fullLayerOrdinals.count + let tableEntries = numFull * store.geometry.pagesPerLayer + guard let scores = device.makeBuffer(length: tableEntries * 4, + options: .storageModeShared), + let tables = device.makeBuffer(length: tableEntries * 4, + options: .storageModeShared) else { + throw KVPageStoreError.allocationFailed("paged KV selection buffers") + } + scores.label = "kvpage.scores" + tables.label = "kvpage.tables" + self.scoresBuf = scores + self.tablesBuf = tables + self.lastScores = Array(repeating: [], count: numFull) + self.selections = Array(repeating: .init(pages: [], selTokens: 0), + count: numFull) + self.pinnedSelections = Array(repeating: [], count: numFull) + // The pinned per-token selection (plus the unsealed tail and one + // slot of eviction slack) must fit the pool. + let worstSelection = runtimeConfiguration.kvSinkPages + + runtimeConfiguration.kvRecentPages + + runtimeConfiguration.kvTopKPages + 2 + guard poolPages >= min(pagesPerLayer, worstSelection) else { + throw KVPageStoreError.allocationFailed( + "kv pool (\(poolPages) pages/layer) smaller than the selection budget") + } + } + + deinit { try? FileManager.default.removeItem(at: spillDir) } + + func resetState() { + store.reset() + for i in 0.. Bool { + let totalPages = position / KVPageGeometry.tokensPerPage + 1 + let lagPages = (maxSpanTokens + KVPageGeometry.tokensPerPage - 1) + / KVPageGeometry.tokensPerPage + return selector.coversEntireContext(totalPages: totalPages, + maxUnscoredSealedPages: max(1, lagPages)) + } + + /// Extend each layer's table past the selection tail with the unsealed + /// pages a speculative verify span [position, position + count) writes + /// into, so per-position paged attention can address the whole span. + /// Call after `prepareSelections(position:)`. + func appendVerifySpan(position: Int, count: Int) throws { + let g = store.geometry + let tailPage = position / KVPageGeometry.tokensPerPage + let lastPage = (position + count - 1) / KVPageGeometry.tokensPerPage + guard lastPage > tailPage else { return } + let tables = tablesBuf.contents() + .bindMemory(to: UInt32.self, + capacity: g.fullLayerOrdinals.count * g.pagesPerLayer) + for (ordinal, layerIndex) in g.fullLayerOrdinals.enumerated() { + let base = ordinal * g.pagesPerLayer + var entry = selections[ordinal].pages.count + for page in (tailPage + 1)...lastPage { + let slot = try store.kSlot(layer: layerIndex, + position: page * KVPageGeometry.tokensPerPage) + let slotIndex = slot.offset / g.kPageBytes + store.pin(layer: layerIndex, pageIndex: page) + pinnedSelections[ordinal].append(page) + tables[base + entry] = UInt32(slotIndex) + entry += 1 + } + } + } + + /// Selected logical tokens strictly before the tail page's first row for + /// a verify round at `position` — verify position `i` attends + /// `verifyBaseTokens + (position % 64) + i + 1` logical tokens. + func verifyBaseTokens(ordinal: Int) -> Int { + 64 * max(0, selections[ordinal].pages.count - 1) + } + + /// Quest min/max summaries for pages sealed by earlier tokens, encoded + /// before this command buffer's score pass reads them. + func encodePendingMetadata(commandBuffer cb: MTLCommandBuffer) throws { + guard !pendingMetadata.isEmpty else { return } + let g = store.geometry + for pageIndex in pendingMetadata { + for (ordinal, layerIndex) in g.fullLayerOrdinals.enumerated() { + let slot = try store.ensureResident(layer: layerIndex, + pageIndex: pageIndex) + kernels.encodePageMinMax( + commandBuffer: cb, + kPool: store.kPoolBuffer(layer: layerIndex), + slot: UInt32(slot), + validTokens: UInt32(KVPageGeometry.tokensPerPage), + metadata: store.metadataBuffer, + metadataOffset: g.metadataOffset(layerOrdinal: ordinal, + pageIndex: pageIndex), + numKVHeads: UInt32(numKVHeads), + headDim: UInt32(headDim)) + } + } + pendingMetadata.removeAll(keepingCapacity: true) + } + + /// Quest criticality of every sealed page against `q` — the selection + /// input for the next token. + func encodeScores(commandBuffer cb: MTLCommandBuffer, + ordinal: Int, q: MTLBuffer, qOffset: Int, + sealedPages: Int) { + guard sealedPages > 0 else { return } + let g = store.geometry + kernels.encodePageScores( + commandBuffer: cb, + q: q, qOffset: qOffset, + metadata: store.metadataBuffer, + metadataOffset: g.metadataOffset(layerOrdinal: ordinal, pageIndex: 0), + scores: scoresBuf, + scoresOffset: ordinal * g.pagesPerLayer * MemoryLayout.stride, + numPages: UInt32(sealedPages), + headDim: UInt32(headDim), + numQHeads: UInt32(numQHeads), + numKVHeads: UInt32(numKVHeads)) + } + + func readBackScores(sealedPages: Int) { + guard sealedPages > 0 else { return } + let g = store.geometry + let numFull = g.fullLayerOrdinals.count + let ptr = scoresBuf.contents() + .bindMemory(to: Float.self, capacity: numFull * g.pagesPerLayer) + for ordinal in 0.. before { + pendingMetadata.append(contentsOf: before.. System message for --chat (repeatable). --max-new Generated-token limit (default 1024). --max-context Context limit in tokens (default 4096). + --kv-paged Paged KV cache with SSD spill + Quest sparse + decode (Qwen 3.8; default auto: on above 32k + context). Exact when everything fits RAM. + --kv-topk Sparse decode budget in 64-token pages + (default 60 ≈ 3.8k attended tokens/layer). + --kv-pool-pages Resident pool per full-attention layer in + pages (default auto: sized from RAM). --temperature Sampling temperature (default 0.2; 0 = greedy). --top-k Top-k truncation, 1...256 (default 64; 0 = off). --top-p Nucleus truncation (default 0.95). @@ -183,6 +203,9 @@ extension Args { var prefillChunk = PrefillChunkChoice.auto var flashHead = false var verification = ModelIntegrityPolicy.fullSha256 + var kvPaged = "auto" + var kvTopKPages = 60 + var kvPoolPages: Int? = nil var index = 0 while index < argv.count { @@ -220,6 +243,28 @@ extension Args { throw ArgsError.invalidValue(flag: flag, value: value) } maxContext = parsed + case "--kv-paged": + let value = try takeValue(argv, &index, flag: flag) + guard ["on", "off", "auto"].contains(value) else { + throw ArgsError.invalidValue(flag: flag, value: value) + } + kvPaged = value + case "--kv-topk": + let value = try takeValue(argv, &index, flag: flag) + guard let parsed = Int(value), parsed >= 0 else { + throw ArgsError.invalidValue(flag: flag, value: value) + } + kvTopKPages = parsed + case "--kv-pool-pages": + let value = try takeValue(argv, &index, flag: flag) + if value == "auto" { + kvPoolPages = nil + } else { + guard let parsed = Int(value), parsed > 0 else { + throw ArgsError.invalidValue(flag: flag, value: value) + } + kvPoolPages = parsed + } case "--temperature": let value = try takeValue(argv, &index, flag: flag) guard let parsed = Float(value), parsed >= 0 else { @@ -330,7 +375,10 @@ extension Args { rdadvise: rdadvise, prefillChunk: prefillChunk, flashHead: flashHead, - verification: verification) + verification: verification, + kvPaged: kvPaged, + kvTopKPages: kvTopKPages, + kvPoolPages: kvPoolPages) } private static func takeValue(_ argv: [String], diff --git a/Sources/MferenceCLI/Run.swift b/Sources/MferenceCLI/Run.swift index d22a779..a9b440b 100644 --- a/Sources/MferenceCLI/Run.swift +++ b/Sources/MferenceCLI/Run.swift @@ -75,7 +75,10 @@ public func run(args: Args, rdadvisePolicy: RDAdvicePolicyMode.parse(args.rdadvise), prefillChunkTokens: prefillChunkTokens, forceLogitsHead: !config.isPureGreedy, - useMapleFlashHead: args.flashHead) + useMapleFlashHead: args.flashHead, + kvPagedPolicy: kvPagedPolicy(for: args), + kvTopKPages: args.kvTopKPages, + kvPoolPagesPerLayer: args.kvPoolPages) guard MTLCreateSystemDefaultDevice() != nil else { return errored(stderr, "no Metal device", 1) @@ -225,6 +228,17 @@ private func errored(_ stderr: FileHandle, _ message: String, _ code: Int32) -> return RunResult(exitCode: code) } +/// "auto" enables the paged KV cache above 32k context — the point where the +/// linear FP16 full-attention cache (2 GiB there, growing 64 KiB/token) +/// stops being the sensible default on consumer RAM. +private func kvPagedPolicy(for args: Args) -> RuntimeKVPagedPolicy { + switch args.kvPaged { + case "on": return .on + case "off": return .off + default: return args.maxContext > 32_768 ? .on : .off + } +} + private func structuredEvents(_ decoder: StructuredAssistantDecoder?, tokenID: Int32, text: String) throws -> [StructuredAssistantEvent] { @@ -283,7 +297,10 @@ private func runChat(args: Args, rdadvisePolicy: RDAdvicePolicyMode.parse(args.rdadvise), prefillChunkTokens: prefillChunkTokens, forceLogitsHead: !baseConfig.isPureGreedy, - useMapleFlashHead: args.flashHead) + useMapleFlashHead: args.flashHead, + kvPagedPolicy: kvPagedPolicy(for: args), + kvTopKPages: args.kvTopKPages, + kvPoolPagesPerLayer: args.kvPoolPages) guard MTLCreateSystemDefaultDevice() != nil else { return errored(stderr, "no Metal device", 1) diff --git a/Tests/Mference/Core/CLI/CLIArgumentsTests.swift b/Tests/Mference/Core/CLI/CLIArgumentsTests.swift index 27fb1d2..6ccbe43 100644 --- a/Tests/Mference/Core/CLI/CLIArgumentsTests.swift +++ b/Tests/Mference/Core/CLI/CLIArgumentsTests.swift @@ -72,6 +72,7 @@ import Mference "--temperature", "--top-k", "--top-p", "--repetition-penalty", "--seed", "--stop", "--prefill-chunk", "--quiet", "--help", "--rdadvise", "--expert-cache-slots", "--flash-head", "--verify", + "--kv-paged", "--kv-topk", "--kv-pool-pages", ] let words = Args.usage.split { $0.isWhitespace || $0 == "(" || $0 == ")" } let options = Set(words.map(String.init).filter { $0.hasPrefix("--") }) diff --git a/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift b/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift new file mode 100644 index 0000000..1727747 --- /dev/null +++ b/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift @@ -0,0 +1,212 @@ +import Testing +import Foundation +import Metal +@testable import Mference +import MferenceValidationSupport + +/// `kv_page_minmax` and `attention_page_scores` against CPU references, plus +/// `KVPageSelector` policy behavior. +@Suite struct KVPageKernelsTests { + + private static let pageTokens = 64 + + // MARK: kv_page_minmax + + @Test func pageMinMax_matchesCPUReference() throws { + let numKVHeads = 4, headDim = 256 + let elems = numKVHeads * headDim + let poolSlots = 3, slot = 2 // non-zero slot exercises addressing + var rng = SeedTree(0x3117).key("minmax") + let pool = (0.. (linear: [Float16], paged: [Float16]) { + let cb = ctx.queue.makeCommandBuffer()! + kernel.encodeFull(commandBuffer: cb, + q: q, k: kLinear, v: vLinear, out: outLinear, + headDim: UInt32(headDim), + numQHeads: UInt32(numQHeads), + numKVHeads: UInt32(numKVHeads), + seqLen: UInt32(seqLen)) + cb.commit(); cb.waitUntilCompleted() + + let cb2 = ctx.queue.makeCommandBuffer()! + kernel.encodeFullPaged(commandBuffer: cb2, + q: q, kPool: kPool, vPool: vPool, + pageTable: pageTable, + out: outPaged, + headDim: UInt32(headDim), + numQHeads: UInt32(numQHeads), + numKVHeads: UInt32(numKVHeads), + selTokens: UInt32(seqLen)) + cb2.commit(); cb2.waitUntilCompleted() + + let n = numQHeads * headDim + let lp = outLinear.contents().assumingMemoryBound(to: Float16.self) + let pp = outPaged.contents().assumingMemoryBound(to: Float16.self) + return ((0.. (MetalContext, KVPageStore, URL) { + let ctx = try MetalContext() + let dir = FileManager.default.temporaryDirectory + .appendingPathComponent("kvpage-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + let store = try KVPageStore(device: ctx.device, + config: config, + maxContext: maxContext, + poolPagesPerLayer: poolPagesPerLayer, + spillDirectory: dir) + return (ctx, store, dir) + } + + // MARK: geometry + + @Test func geometry_matchesQwen38Config() throws { + let (_, store, _) = try makeStore(maxContext: 512, poolPagesPerLayer: 8) + #expect(store.geometry.fullLayerOrdinals.count == 16) + // mask 1 at layers 3, 7, 11, ... + #expect(store.geometry.fullLayerOrdinals.first == 3) + #expect(store.geometry.fullLayerOrdinals.last == 63) + #expect(store.geometry.tokenStrideBytes == Self.tokenStride) + #expect(store.geometry.pagesPerLayer == 512 / Self.pageTokens) + #expect(store.geometry.kPageBytes == Self.pageTokens * Self.tokenStride) + } + + @Test func geometry_rejectsNonFullAttentionLayer() throws { + let (_, store, _) = try makeStore() + // Layer 0 is a linear-attention layer in Qwen 3.8. + #expect(store.fullLayerOrdinal(forLayer: 0) == nil) + #expect(store.fullLayerOrdinal(forLayer: 3) == 0) + #expect(store.fullLayerOrdinal(forLayer: 7) == 1) + } + + // MARK: unsealed writes + + @Test func kSlotAndVSlot_addressWithinUnsealedPage() throws { + let (_, store, _) = try makeStore() + // Position 0 lands in page 0 at offset 0. + let k0 = try store.kSlot(layer: 3, position: 0) + #expect(k0.offset % Self.tokenStride == 0) + // Position 65 lands in page 1 at within-page row 1. + for p in 0...65 { _ = try store.kSlot(layer: 3, position: p) } + let k65 = try store.kSlot(layer: 3, position: 65) + let v65 = try store.vSlot(layer: 3, position: 65) + #expect(k65.offset % Self.tokenStride == 0) + // K and V use the same slot index in distinct pool buffers. + #expect(k65.offset == v65.offset) + #expect(k65.buffer !== v65.buffer) + } + + @Test func advance_sealsCrossedPagesAcrossAllLayers() throws { + let (_, store, _) = try makeStore() + for p in 0..() + for ord in 0.. (URL, MetalContext, Qwen38ForwardRunner) { + let dir = try Qwen38ToySynthetic.write() + let ctx = try MetalContext() + let model = try Model.load(directoryURL: dir, + device: ctx.device, + expecting: .qwen38Toy()) + let config = RuntimeConfiguration(prefillEnabled: true, + kvPagedPolicy: paged ? .on : .off, + kvTopKPages: topK, + kvSinkPages: sink, + kvRecentPages: recent, + kvPoolPagesPerLayer: poolPages) + let runner = try Qwen38ForwardRunner(model: model, + context: ctx, + maxContext: maxContext, + runtimeConfiguration: config) + return (dir, ctx, runner) + } + + private func makeLogits(_ ctx: MetalContext) throws -> MTLBuffer { + guard let buf = ctx.device.makeBuffer( + length: Self.vocab * MemoryLayout.stride, + options: .storageModeShared) else { + throw ModelError.residentBufferWrapFailed + } + return buf + } + + private static func prompt(_ count: Int, seed: Int = 0) -> [Int32] { + (0.. [UInt32] { + var stream: [UInt32] = [] + try await prefill(runner, tokens: turn1, start: 0, logits: logits) + stream.append(runner.lastGreedyToken) + var position = 200 + var token: Int32 = 9 + for _ in 0..<10 { + try await runner.produce(token: token, position: position, into: logits) + stream.append(runner.lastGreedyToken) + token = Int32(runner.lastGreedyToken % UInt32(Self.vocab)) + position += 1 + } + try runner.prepareForContinuation(expectedPosition: position) + try await prefill(runner, tokens: turn2, start: position, logits: logits) + stream.append(runner.lastGreedyToken) + position += turn2.count + for _ in 0..<5 { + try await runner.produce(token: token, position: position, into: logits) + stream.append(runner.lastGreedyToken) + token = Int32(runner.lastGreedyToken % UInt32(Self.vocab)) + position += 1 + } + #expect(runner.continuationPosition == position) + return stream + } + + let first = try await run() + runner.reset() + let second = try await run() + #expect(first == second) + #expect(first.allSatisfy { $0 < UInt32(Self.vocab) }) + } +} diff --git a/Tests/Mference/Core/Runtime/Qwen38PagedKVParityTests.swift b/Tests/Mference/Core/Runtime/Qwen38PagedKVParityTests.swift new file mode 100644 index 0000000..d33a9c4 --- /dev/null +++ b/Tests/Mference/Core/Runtime/Qwen38PagedKVParityTests.swift @@ -0,0 +1,151 @@ +import Foundation +import Metal +import Testing +@testable import Mference + +/// E2E parity for the paged long-context mode against the qwen38 toy: with a +/// selection budget that covers every page, paged decode must reproduce the +/// dense runner's exact greedy stream (the paged kernel is bit-identical +/// under full selection), through pure decode, reset, and prefill + decode. +/// A tight-budget sparse run must still decode without faulting. +@Suite struct Qwen38PagedKVParityTests { + private static let vocab = 1024 + + private func makeRunner(maxContext: Int, + paged: Bool, + topK: Int = 1024, + sink: Int = 2, + recent: Int = 4) throws -> (URL, MetalContext, Qwen38ForwardRunner) { + let dir = try Qwen38ToySynthetic.write() + let ctx = try MetalContext() + let model = try Model.load(directoryURL: dir, + device: ctx.device, + expecting: .qwen38Toy()) + let config = RuntimeConfiguration(prefillEnabled: true, + kvPagedPolicy: paged ? .on : .off, + kvTopKPages: topK, + kvSinkPages: sink, + kvRecentPages: recent) + let runner = try Qwen38ForwardRunner(model: model, + context: ctx, + maxContext: maxContext, + runtimeConfiguration: config) + return (dir, ctx, runner) + } + + private func makeLogits(_ ctx: MetalContext) throws -> MTLBuffer { + guard let buf = ctx.device.makeBuffer( + length: Self.vocab * MemoryLayout.stride, + options: .storageModeShared) else { + throw ModelError.residentBufferWrapFailed + } + return buf + } + + private static func prompt(_ count: Int) -> [Int32] { + (0.. [UInt32] { + var out: [UInt32] = [] + var token: Int32 = 3 + for position in 0..<130 { + try await runner.produce(token: token, position: position, into: logits) + out.append(runner.lastGreedyToken) + token = Int32(runner.lastGreedyToken % UInt32(Self.vocab)) + } + return out + } + let first = try await run() + runner.reset() + let second = try await run() + #expect(first == second) + } + + /// A genuinely sparse budget (fewer pages than exist) must stay stable + /// and produce valid tokens — the quality gate for real sparsity runs on + /// the real model, not the toy. + @Test func sparseBudget_decodesWithoutFaulting() async throws { + let (dir, ctx, runner) = try makeRunner(maxContext: 512, paged: true, + topK: 1, sink: 1, recent: 2) + defer { try? FileManager.default.removeItem(at: dir) } + let logits = try makeLogits(ctx) + + var token: Int32 = 11 + for position in 0..<400 { // 6+ pages; selection covers at most 4 + try await runner.produce(token: token, position: position, into: logits) + #expect(runner.lastGreedyToken < UInt32(Self.vocab)) + token = Int32(runner.lastGreedyToken % UInt32(Self.vocab)) + } + #expect(runner.continuationPosition == 400) + } +} diff --git a/Tests/Mference/Core/Runtime/Qwen38PagedMTPTests.swift b/Tests/Mference/Core/Runtime/Qwen38PagedMTPTests.swift new file mode 100644 index 0000000..9c8d3ec --- /dev/null +++ b/Tests/Mference/Core/Runtime/Qwen38PagedMTPTests.swift @@ -0,0 +1,174 @@ +import Foundation +import Metal +import Testing +@testable import Mference +@testable import MferenceRepackCore + +/// MTP speculative decoding composed with the paged KV mode: with a +/// covering selection budget, spec rounds run per-position paged attention +/// over the same tables as plain paged decode, so the greedy stream must be +/// byte-identical to plain paged decode — through page-boundary crossings, +/// rollbacks (cursor rewind un-seals pages), and reset. A tight pool run +/// exercises rounds whose selections fetch from the spill file. +@Suite(.serialized) struct Qwen38PagedMTPTests { + private static let vocab = 1024 + + private static func makeAttachedDirectory() throws -> URL { + let dir = try Qwen38ToySynthetic.write() + let shard = try Qwen38ToySynthetic.writeMTPShard() + defer { try? FileManager.default.removeItem(at: shard) } + _ = try MTPAttachTool.run(gturboDirectory: dir.path, shardPath: shard.path) + return dir + } + + private func makeRunner(_ dir: URL, + maxContext: Int, + paged: Bool, + poolPages: Int? = nil, + topK: Int = 1024, + sink: Int = 2, + recent: Int = 4) throws -> (MetalContext, Qwen38ForwardRunner) { + let ctx = try MetalContext() + let model = try Model.load(directoryURL: dir, + device: ctx.device, + expecting: .qwen38Toy()) + let config = RuntimeConfiguration(prefillEnabled: true, + kvPagedPolicy: paged ? .on : .off, + kvTopKPages: topK, + kvSinkPages: sink, + kvRecentPages: recent, + kvPoolPagesPerLayer: poolPages) + let runner = try Qwen38ForwardRunner(model: model, + context: ctx, + maxContext: maxContext, + runtimeConfiguration: config) + return (ctx, runner) + } + + private func makeLogits(_ ctx: MetalContext) throws -> MTLBuffer { + guard let buf = ctx.device.makeBuffer( + length: Self.vocab * MemoryLayout.stride, + options: .storageModeShared) else { + throw ModelError.residentBufferWrapFailed + } + return buf + } + + /// Greedy self-drive for `steps` tokens from a fixed seed. + private func drive(_ runner: Qwen38ForwardRunner, + _ logits: MTLBuffer, + steps: Int, + seed: Int32 = 7) async throws -> [UInt32] { + var out: [UInt32] = [] + var token = seed + for position in 0.. 0) + #expect(mtp.stats.rounds == roundsAtCrossover) + } + + /// Tight pool: while the selection is exhaustive, rounds run under pool + /// pressure; past the coverage gate, plain paged decode fetches sparse + /// selections from the spill file. Both phases must replay + /// deterministically. + @Test func pagedSpecDecode_tightPool_replaysDeterministically() async throws { + let dir = try Self.makeAttachedDirectory() + defer { try? FileManager.default.removeItem(at: dir) } + let (ctx, runner) = try makeRunner(dir, maxContext: 512, paged: true, + poolPages: 6, topK: 1, + sink: 1, recent: 2) + let logits = try makeLogits(ctx) + + let first = try await drive(runner, logits, steps: 300) + runner.reset() + let second = try await drive(runner, logits, steps: 300) + #expect(first == second) + #expect(first.allSatisfy { $0 < UInt32(Self.vocab) }) + } +} diff --git a/docs/QWEN38_LONG_CONTEXT.md b/docs/QWEN38_LONG_CONTEXT.md new file mode 100644 index 0000000..dc1454c --- /dev/null +++ b/docs/QWEN38_LONG_CONTEXT.md @@ -0,0 +1,121 @@ +# Qwen 3.8 long context: paged KV with an SSD tier + +Qwen 3.8-27B advertises a 262,144-token context, but its full-attention KV +cache costs 64 KiB/token — 16 GiB at full context, which cannot sit beside +14 GiB of weights on a 24 GB machine. And even if it could, dense attention +would read all 16 GiB per decoded token (~181 ms at the M5's ~95 GB/s): +long context is a *sparsity* problem before it is a *capacity* problem. + +The paged KV mode solves both at full FP16 precision (no KV quantization): + +- **Capacity**: full-attention KV lives in fixed 64-token pages. A bounded + RAM pool (auto-sized from physical memory, ~4 GiB on a 24 GB host) holds + the working set; sealed pages write behind to a sparse, layer-major spill + file on SSD and evict under LRU pressure. The other 48 of 64 layers are + Gated-DeltaNet linear attention whose 144 MiB recurrent state is constant + and always exact — the architecture already carries most long-range signal + outside the cache that grows. +- **Decode sparsity**: each token attends sink pages + the recent window + + the top-k pages ranked by Quest criticality (arXiv 2406.10774): per page, + element-wise min/max of its K rows summarize the page; `Σ_d max(q·min, + q·max)` bounds the page's attention mass for the current query. Scores + compute on-GPU each token and select for the *next* token (lag-one), + which hides page-fetch latency in the inter-token gap. Selected-but- + spilled pages return via one 256 KiB `pread` each (~2.3 GiB/s measured). +- **Exact prefill at any depth**: a chat turn appended beyond the pool runs + the blocked streamed path — the sealed past flows through two staging + buffers (one sequential `pread` per 8k-token window, overlapped with the + previous window's GPU pass) and folds into FP32 running-softmax state; + the chunk's own pages fold causally from the pool. Same math as the + resident path, so prefill stays exact; only decode is sparse. + +## Controls + +``` +--kv-paged default auto: on above 32k --max-context +--kv-topk decode budget beyond sinks+recent + (default 60 pages ≈ 3.8k attended tokens/layer) +--kv-pool-pages resident pool per full-attn layer + (default auto: sized from RAM; 64 KiB tokens + resident per layer on a 24 GB host) +``` + +`RuntimeConfiguration`: `kvPagedPolicy`, `kvTopKPages`, `kvSinkPages`, +`kvRecentPages`, `kvPoolPagesPerLayer`. MTP speculative decode composes +with paged mode: verify rounds write draft rows through the page store and +run per-position paged attention over the round's pinned selection, with +cursor rewinds un-sealing pages on rejected drafts — byte-identical to +plain paged decode (`Qwen38PagedMTPTests`). + +Speculative rounds run only while the page selection is **exhaustive** +(every context page fits the sinks + recent + top-k budget, ≈ 4.2k tokens +at the defaults). A round reuses one page table across its verify rows +where plain decode reselects per token, so under a sparse selection the +speculative stream could drift from the plain paged stream; the gate +(`Qwen38MTPSpeculator.canRunRound`) hands decode off to plain paged tokens +just before the selection turns sparse, preserving byte-identity across +the crossover. Raise `--kv-topk` to extend the exact-MTP window. + +## Correctness + +- The paged decode kernel is **bit-identical** to the contiguous kernel + under a full selection, for any page→slot scattering + (`PagedAttentionParityTests`). +- E2E: paged mode with a covering budget reproduces the dense runner's + greedy stream exactly through decode, prefill + continuation, and reset + (`Qwen38PagedKVParityTests`), on the toy and on the real checkpoint. +- Blocked streamed prefill agrees with the dense head (argmax) with sealed + pages spilled and streamed back, including mid-page chunk boundaries; + growing-chat flows replay deterministically under eviction, fetch, and + selection pinning (`Qwen38BlockedPrefillTests`). +- Live selections are pinned for their token so their own fetches cannot + evict them; page summaries are computed in the same command buffer that + seals a page, so long prefills never trigger a metadata refetch storm. + +## Measured (M5 MacBook Pro, 24 GB, real 27B checkpoint, 2026-08-15) + +Plain decode (no MTP attach): + +| run | context | pool | prefill | decode | notes | +|---|---|---|---|---|---| +| dense baseline | 4k | — | — | 7.8 tok/s | `--kv-paged off` | +| paged, all resident | 4k | auto | — | 8.0 tok/s | output identical to dense | +| paged + SSD, needle @30% | 5.4k prompt | 72 pages (4.6k tok) | 39 tok/s | 5.9 tok/s | passkey retrieved exactly | +| paged + SSD, needle @45% | 12.2k prompt | 128 pages (8k tok) | 27 tok/s | 6.3 tok/s | passkey retrieved exactly through spill | +| capacity smoke | 262,144 max-context | auto (2 GiB) | — | 7.9 tok/s | full-context settings, no decode regression | + +With MTP speculative decode attached (byte-identical greedy): + +| run | context | decode | notes | +|---|---|---|---| +| dense + MTP | 4k | 16.7 tok/s | reference | +| paged + MTP | 4k | **16.8 tok/s** | identical output to dense+MTP, zero paging overhead | +| paged + SSD needle + MTP | 5.4k prompt, 4.6k-token pool | 10.5 tok/s † | passkey retrieved exactly | +| 262k settings + MTP | 262,144 max-context | **14.0 tok/s** | | + +† Measured before the exactness gate. At the default budget the needle's +5.4k context exceeds the exhaustive-selection window, so MTP now hands +those decodes to plain paged tokens (≈ the 5.9 tok/s plain rate); raise +`--kv-topk` past the context length to keep speculative rounds running +exactly. + +Notes: + +- The 64 KiB/token page geometry puts a K+V page pair at 256 KiB — the + measured sweet spot of this NVMe's random-read curve (2.3 GiB/s; 4.96 + GiB/s sequential for the blocked-prefill streams). +- Reaching 262k *by prefill* costs ≈ an hour of compute (quadratic + attention term at ~60 tok/s prefill); the paged mode's target workload is + the growing chat, which pays that incrementally per turn. + +## Files + +- `Sources/Mference/Runtime/KVCache/KVPageStore.swift` — pools, spill file, + LRU, pinning, metadata layout, streamed span reads. +- `Sources/Mference/Runtime/KVCache/KVPageSelector.swift` — sinks + recent + + top-k policy. +- `Sources/Mference/Kernels/Attention/KVPageKernels.swift` + + `Metal/Attention/attention.metal` — paged decode partial, page scores, + page min/max, blocked-prefill flash init/update/finalize. +- `Qwen38ForwardRunner` — paged decode/prefill integration. +- Design spec: `docs/superpowers/specs/2026-08-15-qwen38-longctx-paged-kv-ssd-design.md`. diff --git a/docs/superpowers/specs/2026-08-15-qwen38-longctx-paged-kv-ssd-design.md b/docs/superpowers/specs/2026-08-15-qwen38-longctx-paged-kv-ssd-design.md new file mode 100644 index 0000000..43f4220 --- /dev/null +++ b/docs/superpowers/specs/2026-08-15-qwen38-longctx-paged-kv-ssd-design.md @@ -0,0 +1,75 @@ +# Qwen 3.8 long context: paged KV cache with SSD tier and query-aware sparse decode + +**Date:** 2026-08-15 +**Target:** 128k–262k context for the Qwen 3.8-27B-4bit growing-chat workload on the 24 GB M5, full-fp16 KV (no KV quantization), decode ≥ ~5 tok/s at 262k. +**Branch:** `claude/qwen38-longctx-ssd-kv` (off main 579382f). + +## Why this shape + +Measured constraints (this machine, this model): + +- Qwen 3.8 is hybrid: only **16 of 64 layers** are full attention (4 KV heads × 256 head_dim → **64 KiB/token** fp16 K+V). The other 48 are GDN linear-attention layers whose recurrent state is **144 MiB constant** — exact long-range signal that never grows. Full-attn KV at 262k = **16 GiB**; does not fit beside 14 GiB of weights in 24 GB. +- Dense attention at 262k is dead regardless of storage tier: 16 GiB/token read ≈ 181 ms/token from RAM (95 GB/s), 3.1 s/token from SSD. +- Measured NVMe (O_NOCACHE): 4.96 GiB/s sequential; **2.3 GiB/s at 256 KiB random** (0.11 ms/read); 0.76 GiB/s at 64 KiB. +- Therefore: **sparsity, not bandwidth**. Quest-style (arXiv 2406.10774) query-aware top-k page selection reads ~2–5% of pages per token at full precision. KVSwap (arXiv 2511.11907) validates the disk-resident + in-RAM-metadata shape for unified-memory devices. + +Napkin at 262k: ~66 selected pages/layer × 64 tok/page ≈ 4k attended tokens/layer → 16 layers × 16 MiB = 256 MiB/token read from RAM pool ≈ 2.8 ms; SSD traffic only on selection misses (temporal locality keeps steady-state misses at a few pages/token ≈ ≤1 ms amortized). + +## Architecture + +Feature-flagged paged mode for the Qwen 3.8 runner only. Dense path untouched when off. + +### 1. KVPageStore (`Sources/Mference/Runtime/KVCache/KVPageStore.swift`) + +Owns full-attn-layer KV in **pages of 64 tokens** (per layer: 128 KiB K + 128 KiB V). + +- **RAM pool:** per full-attn layer, one K pool buffer + one V pool buffer, `poolSlots` slots each. Auto-sized from free memory (default cap ~4 GiB total → ~65k tokens fully resident; beyond that LRU). Unsealed (tail) pages and pinned pages never evict. +- **SSD spill file:** sparse file in the model dir (`kvspill-` v1), layer-major layout: `offset(layer, page) = layerOrdinal · 1 GiB + page · 256 KiB` (K then V halves). Write-behind on seal via dedicated IO queue (pread/pwrite, F_NOCACHE like PreadExpertStreamer). Fetch = 256 KiB pread into a free/evicted slot. +- **Page metadata (Quest):** at seal, per (layer, page): element-wise min and max of the 64 post-RoPE K rows, per kv-head → 2 × 4 × 256 fp16 = 4 KiB. RAM-resident always (256 MiB at 262k). Computed on GPU (`kv_page_minmax` appended to the sealing CB). +- **States:** unsealed → sealed-resident(+clean-on-disk) → spilled. LRU over unpinned sealed pages. + +### 2. Kernels (extend `attention.metal`) + +- `attention_decode_paged_partial` (+ combine reuse): identical online-softmax to `attention_decode_partial`, but iterates logical selected slots and resolves K/V rows through a **page table** (`uint32` pool slot per selected page; last entry may be the partial tail page with `tail_valid` tokens). Grid/threadgroup geometry unchanged; existing `attention_decode_combine` merges partials as-is. +- `attention_page_scores`: per page, Quest criticality `score = max over q-heads of Σ_d max(q_d·minK_d, q_d·maxK_d)` against that head's kv-head metadata. One TG per page; output `float` per page per layer into a shared-storage scores buffer (read back after the token's CB completes). +- `kv_page_minmax`: reduce a sealed page's K rows to min/max vectors. +- (M5) `attention_prefill_blocked`: prefill attention with carry-in/out running state (m, d, o per query row) so past KV can stream through a bounded window of pool slots, block by block — exact, sequential-read-friendly. + +### 3. Selection policy (`KVPageSelector`, CPU) + +Per token, per layer: **sinks** (first 2 pages, StreamingLLM-style) ∪ **recent** (last 4 pages incl. unsealed tail) ∪ **top-k by score** (default ~60 pages), sorted ascending. Selection uses the *previous* token's scores (lag-one; standard, hides fetch latency behind the inter-CB gap): after CB t−1 completes → read scores → top-k → issue async fetches for misses in layer order → build page tables → encode CB t (per-layer wait on that layer's fetch just before encoding its attention). First decode token after a prefill uses sinks+recent+trailing-k (one-token warmup), corrected from the second token. + +### 4. Runner integration (paged mode) + +- Decode: QKV GEMV writes K/V into the unsealed page slot (same `kSlot`/`vSlot` shape via KVPageStore); RoPE in place; paged attention with the token's page table; `attention_page_scores` appended per layer (for the next token). Page seal on 64-token boundary appends `kv_page_minmax` and enqueues write-behind. +- Prefill (M4, in-RAM): pool slots allocated sequentially so each layer's pool region is contiguous → existing chunked-prefill kernels write/read it unchanged (guarded). Beyond-RAM growing-chat appends (M5) switch to `attention_prefill_blocked` streaming past pages from SSD (~2.5 s per 2k-token chunk at 200k context — acceptable turn latency; reads are sequential per layer). +- MTP speculative decode: auto-disabled in paged mode v1 (logged); re-enable later. +- GDN layers: untouched. + +### 5. Config & CLI + +`RuntimeConfiguration`: `kvPagedMode` (off/on/auto — auto = on when `maxContext > 32768`), `kvTopKPages`, `kvPoolBytes`, fixed `pageTokens = 64`. CLI `--kv-paged`, `--kv-topk`, `--max-context` up to 262144 in paged mode. Env mirrors (`MFERENCE_KV_PAGED`, …) per existing conventions. + +## Correctness & quality gates + +1. **Unit:** KVPageStore seal/spill/fetch/LRU/slot-map/file-offset tests; metadata vs CPU reference; selector determinism (Swift Testing, no GPU needed for store logic). +2. **Kernel parity:** paged attention with all pages selected + identity slot map ≡ contiguous kernel (same split geometry) within existing parity tolerances; score kernel vs CPU reference. +3. **E2E exact parity:** paged mode with k = ∞ (everything selected, all resident) must produce the identical greedy token stream as dense mode at 4k on the real model. +4. **Sparse quality:** needle-in-haystack + long-doc QA harness at 32k/64k (RAM-only) and 128k+ (SSD tier); acceptance = needle retrieval unimpaired at default k. +5. **Perf:** decode tok/s at 4k (no regression when off; bounded regression when on), 64k, 128k, 262k; page-fetch stall histogram; prefill throughput for streamed appends. + +## Milestones + +- **M1** KVPageStore + file format + LRU + metadata layout (CPU, TDD). +- **M2** Paged decode kernel + parity tests. +- **M3** Score kernel + selector + lag-one plumbing. +- **M4** Runner paged mode, in-RAM (≤ ~64k), E2E parity + needle at 32k. +- **M5** SSD tier live: spill/fetch/LRU + blocked streamed prefill → 128k–262k. +- **M6** Bench sweep, docs (`docs/QWEN38_LONG_CONTEXT.md`), memory entry. + +## Risks + +- Lag-one selection quality: mitigated by sinks+recent pinning and per-token correction; measured by needle harness before SSD work starts. +- Pool fragmentation vs prefill contiguity: M4 guards on contiguity, M5 removes the assumption via blocked prefill. +- 256 MiB metadata at 262k: acceptable v1; fp8 metadata is a known follow-up. +- One-CB-per-token invariant is preserved; all new GPU work rides the existing token CB.