From 99f00a6f74d934f6bc78253f808b875dea3bf871 Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Sat, 15 Aug 2026 21:44:05 -0400 Subject: [PATCH 01/12] docs: spec for Qwen 3.8 paged KV + SSD tier + query-aware sparse decode --- ...8-15-qwen38-longctx-paged-kv-ssd-design.md | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 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. From 75a2feebef6a68c0d8623a1e9e7eced75c596a10 Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Sat, 15 Aug 2026 21:49:29 -0400 Subject: [PATCH 02/12] =?UTF-8?q?feat:=20KVPageStore=20=E2=80=94=20paged?= =?UTF-8?q?=20full-attn=20KV=20with=20SSD=20spill=20tier,=20LRU,=20Quest?= =?UTF-8?q?=20metadata=20layout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../Runtime/KVCache/KVPageStore.swift | 422 ++++++++++++++++++ .../Runtime/KVCache/KVPageStoreTests.swift | 238 ++++++++++ 2 files changed, 660 insertions(+) create mode 100644 Sources/Mference/Runtime/KVCache/KVPageStore.swift create mode 100644 Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift diff --git a/Sources/Mference/Runtime/KVCache/KVPageStore.swift b/Sources/Mference/Runtime/KVCache/KVPageStore.swift new file mode 100644 index 0000000..2793170 --- /dev/null +++ b/Sources/Mference/Runtime/KVCache/KVPageStore.swift @@ -0,0 +1,422 @@ +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 + + 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 + + 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) + } + + // 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.. 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: - Reset + + /// Drop all pages, rewind the cursor, and return pool pages to the OS. + public func reset() { + flushSpills() + position = 0 + sealedPageCount = 0 + 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) { return free } + + 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 { + // The slot cannot be reused while its spill is pending (eviction + // barriers on this queue first), so the pointers stay valid. + let wroteK = pwrite(fd, kSrc, pageBytes, offset) + let wroteV = pwrite(fd, vSrc, pageBytes, offset + off_t(pageBytes)) + precondition(wroteK == pageBytes && wroteV == pageBytes, + "kv spill pwrite failed: errno \(errno)") + } + } + + private func fetch(ordinal: Int, pageIndex: Int, slot: Int) throws { + 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/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift b/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift new file mode 100644 index 0000000..abad74b --- /dev/null +++ b/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift @@ -0,0 +1,238 @@ +import Testing +import Foundation +import Darwin +import Metal +@testable import Mference + +/// Tests `KVPageStore` page geometry, unsealed-slot addressing, seal + +/// write-behind spill, LRU eviction, fetch round-trips, pinning, and the +/// layer-major spill-file layout against the Qwen 3.8 config. +@Suite struct KVPageStoreTests { + + private let config = ArchConfig.qwen38_27B + + /// Qwen 3.8 full-attn stride: 4 kv-heads * 256 head_dim * FP16 = 2048 B. + private static let tokenStride = 4 * 256 * 2 + private static let pageTokens = KVPageGeometry.tokensPerPage + + private func makeStore(maxContext: Int = 512, + poolPagesPerLayer: Int = 8) throws -> (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.. Date: Sat, 15 Aug 2026 21:52:11 -0400 Subject: [PATCH 03/12] feat: paged split-KV decode attention kernel, bit-identical to contiguous path --- .../Kernels/Attention/Attention.swift | 104 +++++++++ .../Mference/Metal/Attention/attention.metal | 105 +++++++++ .../Attention/PagedAttentionParityTests.swift | 221 ++++++++++++++++++ 3 files changed, 430 insertions(+) create mode 100644 Tests/Mference/Core/Kernels/Attention/PagedAttentionParityTests.swift 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/Metal/Attention/attention.metal b/Sources/Mference/Metal/Attention/attention.metal index 9e21502..de2f2ba 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)]], diff --git a/Tests/Mference/Core/Kernels/Attention/PagedAttentionParityTests.swift b/Tests/Mference/Core/Kernels/Attention/PagedAttentionParityTests.swift new file mode 100644 index 0000000..71671b0 --- /dev/null +++ b/Tests/Mference/Core/Kernels/Attention/PagedAttentionParityTests.swift @@ -0,0 +1,221 @@ +import Testing +import Foundation +import Metal +@testable import Mference +import MferenceValidationSupport + +/// Parity for `attention_decode_paged_partial`: with the full selection the +/// paged kernel must reproduce the contiguous `encodeFull` result exactly — +/// identical split geometry and accumulation order, so bit-identical FP16 +/// output — regardless of how pages are scattered across pool slots. +@Suite struct PagedAttentionParityTests { + + private static let pageTokens = 64 + + private struct Harness { + let ctx: MetalContext + let kernel: Attention + let q: MTLBuffer + let kLinear: MTLBuffer // [seqLen, numKVHeads, headDim] logical order + let vLinear: MTLBuffer + let kPool: MTLBuffer // page-scattered copy + let vPool: MTLBuffer + let pageTable: MTLBuffer + let outLinear: MTLBuffer + let outPaged: MTLBuffer + let headDim: Int + let numQHeads: Int + let numKVHeads: Int + let seqLen: Int + + /// `slotOf[pageIndex]` scatters logical pages across pool slots. + init(seqLen: Int, headDim: Int, numQHeads: Int, numKVHeads: Int, + slotOf: [Int], seed: UInt64) throws { + self.headDim = headDim + self.numQHeads = numQHeads + self.numKVHeads = numKVHeads + self.seqLen = seqLen + let pages = (seqLen + pageTokens - 1) / pageTokens + precondition(slotOf.count == pages) + + var rng = SeedTree(seed).key("paged-attn") + let qVals = (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.. Date: Sat, 15 Aug 2026 21:55:15 -0400 Subject: [PATCH 04/12] feat: Quest page scores + seal min/max kernels; KVPageSelector policy --- .../Kernels/Attention/KVPageKernels.swift | 74 +++++++ .../Mference/Metal/Attention/attention.metal | 83 ++++++++ .../Runtime/KVCache/KVPageSelector.swift | 75 ++++++++ .../Attention/KVPageKernelsTests.swift | 181 ++++++++++++++++++ 4 files changed, 413 insertions(+) create mode 100644 Sources/Mference/Kernels/Attention/KVPageKernels.swift create mode 100644 Sources/Mference/Runtime/KVCache/KVPageSelector.swift create mode 100644 Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift diff --git a/Sources/Mference/Kernels/Attention/KVPageKernels.swift b/Sources/Mference/Kernels/Attention/KVPageKernels.swift new file mode 100644 index 0000000..86bec85 --- /dev/null +++ b/Sources/Mference/Kernels/Attention/KVPageKernels.swift @@ -0,0 +1,74 @@ +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 static let threadsPerGroup = 256 + + init(context: MetalContext) throws { + self.ctx = context + self.psoMinMax = try context.pipeline("kv_page_minmax") + self.psoScores = try context.pipeline("attention_page_scores") + } + + /// 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 de2f2ba..7e1adfc 100644 --- a/Sources/Mference/Metal/Attention/attention.metal +++ b/Sources/Mference/Metal/Attention/attention.metal @@ -444,6 +444,89 @@ void attention_decode_gqa_swa_partial( } } +// ============================================================================ +// 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/KVCache/KVPageSelector.swift b/Sources/Mference/Runtime/KVCache/KVPageSelector.swift new file mode 100644 index 0000000..0468e03 --- /dev/null +++ b/Sources/Mference/Runtime/KVCache/KVPageSelector.swift @@ -0,0 +1,75 @@ +import Foundation + +/// Which pages a decode token attends: StreamingLLM-style sink pages, the +/// recent window (including the unsealed tail), and the top-k sealed pages by +/// Quest criticality score. Pure policy — no storage, no GPU. +public struct KVPageSelector: Sendable, Equatable { + public let sinkPages: Int + public let recentPages: Int + public let topKPages: Int + + public init(sinkPages: Int = 2, recentPages: Int = 4, topKPages: Int = 60) { + precondition(sinkPages >= 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 (`sealedPages` entries), + /// 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). + /// - sealedPages: pages fully written (64 valid tokens each). + /// - tailValidTokens: valid rows in the unsealed tail page; 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") + precondition(scores.isEmpty || scores.count >= sealedPages, + "scores must cover every sealed page") + 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 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) + } +} diff --git a/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift b/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift new file mode 100644 index 0000000..4d8288f --- /dev/null +++ b/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift @@ -0,0 +1,181 @@ +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.. Date: Sat, 15 Aug 2026 22:06:12 -0400 Subject: [PATCH 05/12] =?UTF-8?q?feat:=20Qwen=203.8=20paged=20KV=20decode?= =?UTF-8?q?=20mode=20=E2=80=94=20sparse=20Quest=20selection,=20exact=20und?= =?UTF-8?q?er=20full=20budget?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit kvPagedPolicy=on hands full-attn KV to KVPageStore; decode runs the paged split-KV kernel over sinks+recent+top-k pages with lag-one Quest scores; page seals compute min/max metadata on the token CB. Prefill writes stream into the identity-mapped pool, existing kernels unchanged. MTP defers to plain decode in paged mode (v1). E2E parity: paged==dense greedy streams. --- .../Configuration/RuntimeConfiguration.swift | 30 +- .../Inference/Qwen38ForwardRunner.swift | 275 ++++++++++++++++-- .../Runtime/KVCache/KVCacheManager.swift | 19 +- .../Runtime/KVCache/KVPageSelector.swift | 24 +- .../Runtime/KVCache/KVPageStore.swift | 33 +++ .../Runtime/Qwen38PagedKVParityTests.swift | 151 ++++++++++ 6 files changed, 499 insertions(+), 33 deletions(-) create mode 100644 Tests/Mference/Core/Runtime/Qwen38PagedKVParityTests.swift diff --git a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift index d0fd278..f1a959c 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 { diff --git a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift index df39d61..5d89825 100644 --- a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift +++ b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift @@ -151,6 +151,82 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR private let kv: KVCacheManager private let gdnState: GDNStateManager + /// Paged long-context state (kvPagedPolicy == .on): the page store owns + /// full-attention KV, decode runs Quest-selected sparse attention, and + /// sealed pages carry min/max metadata for the next token's selection. + private final class PagedKV { + 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] + private let spillDir: URL + + init(context: MetalContext, config: ArchConfig, maxContext: Int, + runtimeConfiguration: RuntimeConfiguration) throws { + let device = context.device + let pagesPerLayer = (maxContext + KVPageGeometry.tokensPerPage - 1) + / KVPageGeometry.tokensPerPage + let poolPages = runtimeConfiguration.kvPoolPagesPerLayer ?? pagesPerLayer + 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) + // The per-token selection must fit the pool alongside pinned + // sinks/recents and the unsealed tail. + let worstSelection = runtimeConfiguration.kvSinkPages + + runtimeConfiguration.kvRecentPages + + runtimeConfiguration.kvTopKPages + 1 + guard poolPages >= min(pagesPerLayer, worstSelection) else { + throw KVPageStoreError.allocationFailed( + "kv pool (\(poolPages) pages/layer) smaller than the selection budget") + } + } + + func resetState() { + store.reset() + for i in 0.. sealedBefore { + paged.pendingMetadata.append(contentsOf: sealedBefore.. 0 else { return } + let g = paged.store.geometry + let numFull = g.fullLayerOrdinals.count + let ptr = paged.scoresBuf.contents() + .bindMemory(to: Float.self, capacity: numFull * g.pagesPerLayer) + for ordinal in 0...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). + let sealedPages = position / KVPageGeometry.tokensPerPage + if sealedPages > 0 { + paged.kernels.encodePageScores( + commandBuffer: cb, + q: qScratch, + metadata: paged.store.metadataBuffer, + metadataOffset: g.metadataOffset(layerOrdinal: ordinal, pageIndex: 0), + scores: paged.scoresBuf, + scoresOffset: ordinal * g.pagesPerLayer * MemoryLayout.stride, + numPages: UInt32(sealedPages), + headDim: UInt32(headDim), + numQHeads: UInt32(cfg.numHeads), + numKVHeads: UInt32(numKV)) + } + } 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/KVCache/KVCacheManager.swift b/Sources/Mference/Runtime/KVCache/KVCacheManager.swift index fa05920..4847447 100644 --- a/Sources/Mference/Runtime/KVCache/KVCacheManager.swift +++ b/Sources/Mference/Runtime/KVCache/KVCacheManager.swift @@ -69,13 +69,18 @@ public final class KVCacheManager { private static let fp16Size = 2 + /// `pagedFullAttention` hands full-attention KV storage to a + /// `KVPageStore`: those layers get the shared placeholder here (no linear + /// allocation — 16 GiB at 262k context) while this manager keeps serving + /// the position cursor and the linear-layer placeholders. public init(device: MTLDevice, config: ArchConfig, maxContext: Int, fp16RingEnabled: Bool = false, slidingWindow: Int? = nil, maxPrefillChunkTokens: Int = 128, - fp16RingCapacityOverride: Int? = nil) throws { + fp16RingCapacityOverride: Int? = nil, + pagedFullAttention: Bool = false) throws { precondition(maxContext > 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.. Selection { - precondition(tailValidTokens >= 0 && tailValidTokens < 64, - "tailValidTokens must be 0..<64") - precondition(scores.isEmpty || scores.count >= sealedPages, - "scores must cover every sealed page") + 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) } @@ -58,7 +59,8 @@ public struct KVPageSelector: Sendable, Equatable { } } else { // Deterministic top-k: score desc, index asc on ties. - let candidates = (0.. scores[$1] : $0 < $1 } for page in candidates.prefix(topKPages) { picked.insert(page) } diff --git a/Sources/Mference/Runtime/KVCache/KVPageStore.swift b/Sources/Mference/Runtime/KVCache/KVPageStore.swift index 2793170..39e9f5c 100644 --- a/Sources/Mference/Runtime/KVCache/KVPageStore.swift +++ b/Sources/Mference/Runtime/KVCache/KVPageStore.swift @@ -217,6 +217,39 @@ public final class KVPageStore { 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 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) + } +} From f0dd14f91febf2bc3af555fc3cb9ede0f021ab3f Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Sat, 15 Aug 2026 22:23:11 -0400 Subject: [PATCH 06/12] =?UTF-8?q?feat:=20SSD=20tier=20live=20=E2=80=94=20b?= =?UTF-8?q?locked=20streamed=20prefill,=20selection=20pinning,=20CLI=20fla?= =?UTF-8?q?gs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pool smaller than context: sealed pages spill (write-behind, layer-major sparse file) and evict under LRU; decode fetches selection misses with the live selection pinned so its own fetches cannot evict it. Beyond-RAM chat appends run the blocked prefill: chunk KV scatters into free slots, the sealed past streams via one sequential pread per window through a 2-stage ring overlapped with flash-update dispatches (FP32 carry state), tail pages fold causally from the pool. In-chunk min/max metadata avoids a refetch storm at first decode. CLI: --kv-paged on/off/auto, --kv-topk, --kv-pool-pages; auto pool sizing from RAM. --- .../Kernels/Attention/KVPageKernels.swift | 106 +++++ .../Mference/Metal/Attention/attention.metal | 154 +++++++ .../Configuration/RuntimeConfiguration.swift | 22 + .../Inference/Qwen38ForwardRunner.swift | 404 +++++++++++++++--- .../Runtime/KVCache/KVPageStore.swift | 43 +- Sources/MferenceCLI/Args.swift | 52 ++- Sources/MferenceCLI/Run.swift | 21 +- .../Mference/Core/CLI/CLIArgumentsTests.swift | 1 + .../Runtime/Qwen38BlockedPrefillTests.swift | 143 +++++++ 9 files changed, 891 insertions(+), 55 deletions(-) create mode 100644 Tests/Mference/Core/Runtime/Qwen38BlockedPrefillTests.swift diff --git a/Sources/Mference/Kernels/Attention/KVPageKernels.swift b/Sources/Mference/Kernels/Attention/KVPageKernels.swift index 86bec85..3de486a 100644 --- a/Sources/Mference/Kernels/Attention/KVPageKernels.swift +++ b/Sources/Mference/Kernels/Attention/KVPageKernels.swift @@ -8,13 +8,119 @@ 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 diff --git a/Sources/Mference/Metal/Attention/attention.metal b/Sources/Mference/Metal/Attention/attention.metal index 7e1adfc..d5c803c 100644 --- a/Sources/Mference/Metal/Attention/attention.metal +++ b/Sources/Mference/Metal/Attention/attention.metal @@ -444,6 +444,160 @@ 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). // ============================================================================ diff --git a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift index f1a959c..3d7d66a 100644 --- a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift +++ b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift @@ -136,6 +136,28 @@ 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 − 20 GiB on the 24 GiB M5 → ~4 GiB of pool ≈ 65k resident + /// tokens per full-attention layer), never below 1 GiB. + 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(20) * 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 5d89825..b7213e8 100644 --- a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift +++ b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift @@ -167,6 +167,11 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR var lastScores: [[Float]] var pendingMetadata: [Int] = [] var selections: [KVPageSelector.Selection] + /// Pages pinned for the in-flight token, 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 spillDir: URL init(context: MetalContext, config: ArchConfig, maxContext: Int, @@ -174,7 +179,12 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR let device = context.device let pagesPerLayer = (maxContext + KVPageGeometry.tokensPerPage - 1) / KVPageGeometry.tokensPerPage - let poolPages = runtimeConfiguration.kvPoolPagesPerLayer ?? pagesPerLayer + 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) @@ -204,11 +214,12 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR self.lastScores = Array(repeating: [], count: numFull) self.selections = Array(repeating: .init(pages: [], selTokens: 0), count: numFull) - // The per-token selection must fit the pool alongside pinned - // sinks/recents and the unsealed tail. + 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 + 1 + + runtimeConfiguration.kvTopKPages + 2 guard poolPages >= min(pagesPerLayer, worstSelection) else { throw KVPageStoreError.allocationFailed( "kv pool (\(poolPages) pages/layer) smaller than the selection budget") @@ -220,6 +231,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR for i in 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 @@ -642,7 +727,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, @@ -666,11 +751,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, @@ -746,14 +831,11 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR lastGreedyToken = greedyTokenBuf.contents().load(as: UInt32.self) } if let paged = pagedKV { - let sealedBefore = paged.store.position / KVPageGeometry.tokensPerPage + // 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) - let sealedAfter = paged.store.position / KVPageGeometry.tokensPerPage - if sealedAfter > sealedBefore { - paged.pendingMetadata.append(contentsOf: sealedBefore.. 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") @@ -913,44 +998,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: 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: 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: PagedKV, + 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: PagedKV, + 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.. 0, "empty span read") + 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. @@ -335,6 +371,7 @@ public final class KVPageStore { flushSpills() position = 0 sealedPageCount = 0 + identityMappingIntact = true let pages = geometry.pagesPerLayer for ordinal in 0.. Int { - if let free = slotPage[ordinal].firstIndex(of: -1) { return free } + if let free = slotPage[ordinal].firstIndex(of: -1) { + if free != pageIndex { identityMappingIntact = false } + return free + } + identityMappingIntact = false var victim = -1 var victimUse = UInt64.max diff --git a/Sources/MferenceCLI/Args.swift b/Sources/MferenceCLI/Args.swift index 71a884d..dd7ffaa 100644 --- a/Sources/MferenceCLI/Args.swift +++ b/Sources/MferenceCLI/Args.swift @@ -46,6 +46,13 @@ public struct Args: Equatable, Sendable { /// written at install time instead. Mirrors the Mac app's existing /// verification control. public var verification: ModelIntegrityPolicy + /// Paged KV cache with SSD spill + sparse decode (Qwen 3.8): + /// "on" / "off" / "auto" (auto enables it above 32k context). + public var kvPaged: String + /// Sparse decode selection budget in 64-token pages. + public var kvTopKPages: Int + /// Resident pool per full-attention layer in pages; nil = auto by RAM. + public var kvPoolPages: Int? public init(model: String, prompt: String? = nil, @@ -65,7 +72,10 @@ public struct Args: Equatable, Sendable { rdadvise: String = "off", prefillChunk: PrefillChunkChoice = .auto, flashHead: Bool = false, - verification: ModelIntegrityPolicy = .fullSha256) { + verification: ModelIntegrityPolicy = .fullSha256, + kvPaged: String = "auto", + kvTopKPages: Int = 60, + kvPoolPages: Int? = nil) { self.model = model self.prompt = prompt self.messagesFile = messagesFile @@ -82,6 +92,9 @@ public struct Args: Equatable, Sendable { self.prefillChunk = prefillChunk self.flashHead = flashHead self.verification = verification + self.kvPaged = kvPaged + self.kvTopKPages = kvTopKPages + self.kvPoolPages = kvPoolPages self.seed = seed self.stops = stops self.quiet = quiet @@ -128,6 +141,13 @@ extension Args { --system 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/Runtime/Qwen38BlockedPrefillTests.swift b/Tests/Mference/Core/Runtime/Qwen38BlockedPrefillTests.swift new file mode 100644 index 0000000..64f9a9c --- /dev/null +++ b/Tests/Mference/Core/Runtime/Qwen38BlockedPrefillTests.swift @@ -0,0 +1,143 @@ +import Foundation +import Metal +import Testing +@testable import Mference + +/// The SSD tier live: a pool smaller than the context forces sealed pages to +/// spill and the blocked (streamed) prefill path to run. Blocked prefill is +/// exact — same math as the resident path, different summation order — so +/// the prefill head must agree with the dense runner. Growing-chat flows +/// (prefill → decode → prefill continuation) must be deterministic under +/// eviction, fetch, and pinning. +@Suite struct Qwen38BlockedPrefillTests { + private static let vocab = 1024 + + private func makeRunner(maxContext: Int, + paged: Bool, + poolPages: Int? = nil, + topK: Int = 0, + sink: Int = 1, + recent: Int = 2) 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, + 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) }) + } +} From a3e5f7c565b32d8d936a5bb9315c92c4c2b644d7 Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Sat, 15 Aug 2026 22:38:35 -0400 Subject: [PATCH 07/12] docs: Qwen 3.8 long-context page (bench table pending final numbers) --- docs/QWEN38_LONG_CONTEXT.md | 92 +++++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 docs/QWEN38_LONG_CONTEXT.md diff --git a/docs/QWEN38_LONG_CONTEXT.md b/docs/QWEN38_LONG_CONTEXT.md new file mode 100644 index 0000000..15b2bd4 --- /dev/null +++ b/docs/QWEN38_LONG_CONTEXT.md @@ -0,0 +1,92 @@ +# 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 defers to +plain decode while paged mode is on (v1). + +## 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, 2026-08-15) + +| 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% | TBD | 128 pages (8k tok) | TBD | TBD | TBD | +| capacity smoke | 262,144 max-context | auto | TBD | TBD | TBD | + +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`. From a1a12d727379d2239a2f632ab0b4789a089141ec Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Sat, 15 Aug 2026 22:39:15 -0400 Subject: [PATCH 08/12] docs: README pointer to Qwen 3.8 long-context mode --- README.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 From e04d8f2de32841f622d94fe85206d61ad4e27c95 Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Sat, 15 Aug 2026 23:16:47 -0400 Subject: [PATCH 09/12] feat: MTP speculative decode composes with paged KV mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shared Qwen38PagedKVRuntime (extracted from the runner) drives one cursor, pin set, and score state for both plain tokens and spec rounds. The verify pass scatters draft rows through the page store, runs per-position paged attention over the round's pinned selection (table extended across span page-crossings), and scores sealed pages with the always-committed bonus query. Cursor rewind on rejected drafts un-seals pages (KVPageStore.rewind; stale spills unreachable). canRunRound now requires position >= 1 — a round at 0 has no prior hidden to seed the drafter and anchored RoPE at -1. Auto pool headroom 20->22 GiB: measured 2x decode loss from a 4 GiB pool squeezing weights on the 24 GiB host. Byte-identity of spec vs plain decode preserved under paging (Qwen38PagedMTPTests); full suite 1078 green. --- .../Configuration/RuntimeConfiguration.swift | 9 +- .../Inference/Qwen38ForwardRunner.swift | 203 ++------------- .../Inference/Qwen38MTPSpeculator.swift | 139 ++++++++-- .../Runtime/KVCache/KVPageStore.swift | 28 +++ .../KVCache/Qwen38PagedKVRuntime.swift | 237 ++++++++++++++++++ .../Core/Runtime/Qwen38PagedMTPTests.swift | 122 +++++++++ 6 files changed, 523 insertions(+), 215 deletions(-) create mode 100644 Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift create mode 100644 Tests/Mference/Core/Runtime/Qwen38PagedMTPTests.swift diff --git a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift index 3d7d66a..598784e 100644 --- a/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift +++ b/Sources/Mference/Runtime/Configuration/RuntimeConfiguration.swift @@ -138,8 +138,11 @@ public struct RuntimeConfiguration: Sendable, Equatable { /// Auto pool sizing for the paged KV cache: everything resident when it /// fits, otherwise whatever RAM remains after weights and headroom - /// (~physical − 20 GiB on the 24 GiB M5 → ~4 GiB of pool ≈ 65k resident - /// tokens per full-attention layer), never below 1 GiB. + /// (~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, @@ -151,7 +154,7 @@ public struct RuntimeConfiguration: Sendable, Equatable { guard numFull > 0 else { return pagesPerLayer } let pagePairBytes = 2 * pageTokens * config.numFullKVHeads * config.fullHeadDim * 2 let gib = UInt64(1) << 30 - let headroom = UInt64(20) * gib + let headroom = UInt64(22) * gib let budget = max(gib, physicalMemoryBytes > headroom ? physicalMemoryBytes - headroom : gib) let budgetPages = Int(budget) / (numFull * pagePairBytes) diff --git a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift index b7213e8..2cddec0 100644 --- a/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift +++ b/Sources/Mference/Runtime/Inference/Qwen38ForwardRunner.swift @@ -151,93 +151,9 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR private let kv: KVCacheManager private let gdnState: GDNStateManager - /// Paged long-context state (kvPagedPolicy == .on): the page store owns - /// full-attention KV, decode runs Quest-selected sparse attention, and - /// sealed pages carry min/max metadata for the next token's selection. - private final class PagedKV { - 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, 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 spillDir: URL - - init(context: MetalContext, config: ArchConfig, maxContext: Int, - runtimeConfiguration: RuntimeConfiguration) throws { - let device = context.device - 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") - } - } - - func resetState() { - store.reset() - for i in 0.. BlockedPrefillScratch { + paged: Qwen38PagedKVRuntime) throws -> BlockedPrefillScratch { if let scratch = blockedPrefillScratch, scratch.chunkTokens >= chunkTokens { return scratch } @@ -423,7 +339,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR maxPrefillChunkTokens: runtimeConfiguration.prefillConfig.chunkTokens, pagedFullAttention: paged) if paged { - self.pagedKV = try PagedKV(context: context, + self.pagedKV = try Qwen38PagedKVRuntime(context: context, config: cfg, maxContext: maxContext, runtimeConfiguration: runtimeConfiguration) @@ -536,10 +452,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR isLinear: isLinear) } - // MTP speculative decode verifies against the linear KV cache; the - // paged mode owns full-attention storage, so it decodes plainly (v1). - if pagedKV == nil, - ProcessInfo.processInfo.environment["MFERENCE_MTP"] != "0" { + if ProcessInfo.processInfo.environment["MFERENCE_MTP"] != "0" { self.mtp = try Qwen38MTPSpeculator.probe(model: model, context: context, config: cfg, @@ -547,7 +460,8 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR gdnState: gdnState, layers: layers, maxContext: maxContext, - mlpWeightBits: mlpWeightBits) + mlpWeightBits: mlpWeightBits, + paged: pagedKV) } } @@ -1074,7 +988,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR /// 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: PagedKV, + paged: Qwen38PagedKVRuntime, layerIndex: Int, startPosition: Int, tokenCount: Int, @@ -1111,7 +1025,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR /// guaranteed resident, so a long prefill never triggers a refetch storm /// at first decode. private func encodePagedChunkMinMax(_ cb: MTLCommandBuffer, - paged: PagedKV, + paged: Qwen38PagedKVRuntime, layerIndex: Int, startPosition: Int, tokenCount: Int) throws { @@ -1141,7 +1055,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR /// resident tail (the chunk's own pages plus any unsealed prefix) with /// the causal predicate, into the chunk's attention output. Synchronous: /// window N+1's pread overlaps window N's GPU pass via the stage ring. - private func runBlockedAttention(paged: PagedKV, + private func runBlockedAttention(paged: Qwen38PagedKVRuntime, blocked: BlockedPrefillScratch, layerIndex: Int, queryCount t: Int, @@ -1466,7 +1380,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR // scores from the previous token), fetch any spilled members, and // build the page tables — all before the command buffer encodes. if let paged = pagedKV { - try preparePagedSelections(paged, position: position) + try paged.prepareSelections(position: position) } let D = UInt32(cfg.hiddenSize) @@ -1475,7 +1389,7 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR // Sealed pages from the previous token need their Quest min/max // summaries before this token's score pass reads them. if let paged = pagedKV { - try encodePendingPageMetadata(paged, commandBuffer: cb) + try paged.encodePendingMetadata(commandBuffer: cb) } let emb = model.embedding @@ -1577,85 +1491,15 @@ public final class Qwen38ForwardRunner: ContinuableLogitProducer, ContextWindowR lastGreedyToken = greedyTokenBuf.contents().load(as: UInt32.self) } if let paged = pagedKV { - readBackPagedScores(paged, position: position) + paged.readBackScores(sealedPages: position / KVPageGeometry.tokensPerPage) paged.store.advance() - if (position + 1) % KVPageGeometry.tokensPerPage == 0 { - paged.pendingMetadata.append(position / KVPageGeometry.tokensPerPage) - } + paged.noteAdvance(from: position, to: position + 1) } kv.advance() } // MARK: - Paged decode support - private func preparePagedSelections(_ paged: PagedKV, position: Int) throws { - let g = paged.store.geometry - let sealedPages = position / KVPageGeometry.tokensPerPage - let tailValid = position % KVPageGeometry.tokensPerPage + 1 - let tables = paged.tablesBuf.contents() - .bindMemory(to: UInt32.self, - capacity: g.fullLayerOrdinals.count * g.pagesPerLayer) - for (ordinal, layerIndex) in g.fullLayerOrdinals.enumerated() { - // Touch the tail page so it has a slot before selection maps it. - _ = try paged.store.kSlot(layer: layerIndex, position: position) - let selection = paged.selector.select(scores: paged.lastScores[ordinal], - sealedPages: sealedPages, - tailValidTokens: tailValid) - // Swap pins to the new selection before fetching: members fetched - // early must survive fetches of later members under LRU pressure. - for page in paged.pinnedSelections[ordinal] { - paged.store.unpin(layer: layerIndex, pageIndex: page) - } - var pinned: [Int] = [] - pinned.reserveCapacity(selection.pages.count) - let base = ordinal * g.pagesPerLayer - for (i, page) in selection.pages.enumerated() { - let slot = try paged.store.ensureResident(layer: layerIndex, pageIndex: page) - paged.store.pin(layer: layerIndex, pageIndex: page) - pinned.append(page) - tables[base + i] = UInt32(slot) - } - paged.pinnedSelections[ordinal] = pinned - paged.selections[ordinal] = selection - } - } - - private func encodePendingPageMetadata(_ paged: PagedKV, - commandBuffer cb: MTLCommandBuffer) throws { - guard !paged.pendingMetadata.isEmpty else { return } - let g = paged.store.geometry - for pageIndex in paged.pendingMetadata { - for (ordinal, layerIndex) in g.fullLayerOrdinals.enumerated() { - let slot = try paged.store.ensureResident(layer: layerIndex, - pageIndex: pageIndex) - paged.kernels.encodePageMinMax( - commandBuffer: cb, - kPool: paged.store.kPoolBuffer(layer: layerIndex), - slot: UInt32(slot), - validTokens: UInt32(KVPageGeometry.tokensPerPage), - metadata: paged.store.metadataBuffer, - metadataOffset: g.metadataOffset(layerOrdinal: ordinal, - pageIndex: pageIndex), - numKVHeads: UInt32(cfg.numFullKVHeads), - headDim: UInt32(cfg.fullHeadDim)) - } - } - paged.pendingMetadata.removeAll(keepingCapacity: true) - } - - private func readBackPagedScores(_ paged: PagedKV, position: Int) { - let sealedPages = position / KVPageGeometry.tokensPerPage - guard sealedPages > 0 else { return } - let g = paged.store.geometry - let numFull = g.fullLayerOrdinals.count - let ptr = paged.scoresBuf.contents() - .bindMemory(to: Float.self, capacity: numFull * g.pagesPerLayer) - for ordinal in 0.. 0 { - paged.kernels.encodePageScores( - commandBuffer: cb, - q: qScratch, - metadata: paged.store.metadataBuffer, - metadataOffset: g.metadataOffset(layerOrdinal: ordinal, pageIndex: 0), - scores: paged.scoresBuf, - scoresOffset: ordinal * g.pagesPerLayer * MemoryLayout.stride, - numPages: UInt32(sealedPages), - headDim: UInt32(headDim), - numQHeads: UInt32(cfg.numHeads), - numKVHeads: UInt32(numKV)) - } + paged.encodeScores(commandBuffer: cb, ordinal: ordinal, + q: qScratch, qOffset: 0, + sealedPages: position / KVPageGeometry.tokensPerPage) } else { attention.encodeFull(commandBuffer: cb, q: qScratch, diff --git a/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift b/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift index c405595..2cd1e11 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,10 @@ 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. + position >= 1 && position == kv.position && position + 2 <= maxContext } /// Run one draft/verify/accept round for `produce(bonus, position)`. @@ -537,8 +549,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 +579,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 +611,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 +702,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 && 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..= 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.. 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.. 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.. Date: Sat, 15 Aug 2026 23:25:28 -0400 Subject: [PATCH 10/12] docs: final long-context bench numbers incl. MTP composition --- docs/QWEN38_LONG_CONTEXT.md | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/docs/QWEN38_LONG_CONTEXT.md b/docs/QWEN38_LONG_CONTEXT.md index 15b2bd4..bbe2564 100644 --- a/docs/QWEN38_LONG_CONTEXT.md +++ b/docs/QWEN38_LONG_CONTEXT.md @@ -41,8 +41,11 @@ The paged KV mode solves both at full FP16 precision (no KV quantization): ``` `RuntimeConfiguration`: `kvPagedPolicy`, `kvTopKPages`, `kvSinkPages`, -`kvRecentPages`, `kvPoolPagesPerLayer`. MTP speculative decode defers to -plain decode while paged mode is on (v1). +`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`). ## Correctness @@ -60,15 +63,26 @@ plain decode while paged mode is on (v1). 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, 2026-08-15) +## 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% | TBD | 128 pages (8k tok) | TBD | TBD | TBD | -| capacity smoke | 262,144 max-context | auto | TBD | TBD | TBD | +| 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** | | Notes: From 9b4538d791e0504e952a7e97873835da1b026754 Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Wed, 26 Aug 2026 10:32:20 -0700 Subject: [PATCH 11/12] fix: gate MTP rounds on an exhaustive page selection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A speculative round reuses one page table across its verify rows where plain paged decode reselects per token, and accepted rows are emitted without their own Quest score pass — under a sparse selection the speculative stream could drift from the plain paged stream (Codex P1 on PR #20). Rounds now run only while the selection through the round's span, plus one position of margin, provably covers the entire context; past that point decode falls back to plain paged tokens. The margin guarantees the token feeding the first sparse selection's lag-one scores is always plain-decoded, so byte-identity holds across the crossover. KVPageSelector.coversEntireContext is conservative about score staleness: gap pages are picked by top-k only when scored, so any unscored just-sealed page must sit inside the recent window. Tests: selector coverage boundaries (agrees with select()), a rounds-stop-at-the-boundary gate test (fails without the gate: rounds kept running 83 -> 100 in the sparse tail), and a sparse-budget e2e byte-identity run across the crossover. The suite's MTP-off reference runners now disable MTP per-instance instead of setenv, which raced runners constructed concurrently by other suites. --- .../Inference/Qwen38MTPSpeculator.swift | 15 ++++- .../Runtime/KVCache/KVPageSelector.swift | 19 ++++++ .../KVCache/Qwen38PagedKVRuntime.swift | 14 +++++ .../Attention/KVPageKernelsTests.swift | 31 ++++++++++ .../Core/Runtime/Qwen38PagedMTPTests.swift | 60 +++++++++++++++++-- docs/QWEN38_LONG_CONTEXT.md | 17 +++++- 6 files changed, 150 insertions(+), 6 deletions(-) diff --git a/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift b/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift index 2cd1e11..f108f3c 100644 --- a/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift +++ b/Sources/Mference/Runtime/Inference/Qwen38MTPSpeculator.swift @@ -491,7 +491,20 @@ final class Qwen38MTPSpeculator { // 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. - position >= 1 && position == kv.position && position + 2 <= maxContext + 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)`. diff --git a/Sources/Mference/Runtime/KVCache/KVPageSelector.swift b/Sources/Mference/Runtime/KVCache/KVPageSelector.swift index 2b88e8c..7ab288e 100644 --- a/Sources/Mference/Runtime/KVCache/KVPageSelector.swift +++ b/Sources/Mference/Runtime/KVCache/KVPageSelector.swift @@ -74,4 +74,23 @@ public struct KVPageSelector: Sendable, Equatable { } 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/Qwen38PagedKVRuntime.swift b/Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift index 133e9ca..d352975 100644 --- a/Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift +++ b/Sources/Mference/Runtime/KVCache/Qwen38PagedKVRuntime.swift @@ -133,6 +133,20 @@ final class Qwen38PagedKVRuntime { } } + /// True when the selection at `position` provably includes every context + /// page regardless of score staleness — the regime where speculative + /// rounds are byte-identical to plain paged decode. `maxSpanTokens` + /// bounds how many tokens can commit between Quest score refreshes (a + /// verify span's accepted rows are emitted without their own score + /// pass), which bounds how many trailing sealed pages may be unscored. + func selectionIsExhaustive(at position: Int, maxSpanTokens: Int) -> 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. diff --git a/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift b/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift index 4d8288f..1727747 100644 --- a/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift +++ b/Tests/Mference/Core/Kernels/Attention/KVPageKernelsTests.swift @@ -172,6 +172,37 @@ import MferenceValidationSupport #expect(empty.pages.isEmpty && empty.selTokens == 0) } + @Test func selector_coversEntireContext_flipsAtTheBudgetBoundary() { + let selector = KVPageSelector(sinkPages: 2, recentPages: 4, topKPages: 60) + // 66 pages = sinks(2) + recent(4) + topk(60) exactly. + #expect(selector.coversEntireContext(totalPages: 66, maxUnscoredSealedPages: 1)) + #expect(!selector.coversEntireContext(totalPages: 67, maxUnscoredSealedPages: 1)) + #expect(selector.coversEntireContext(totalPages: 1, maxUnscoredSealedPages: 1)) + } + + @Test func selector_coversEntireContext_agreesWithSelect() { + let selector = KVPageSelector(sinkPages: 2, recentPages: 4, topKPages: 60) + for totalPages in [1, 6, 40, 65, 66, 67, 80] { + let sealed = totalPages - 1 + let sel = selector.select(scores: (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) } diff --git a/docs/QWEN38_LONG_CONTEXT.md b/docs/QWEN38_LONG_CONTEXT.md index bbe2564..dc1454c 100644 --- a/docs/QWEN38_LONG_CONTEXT.md +++ b/docs/QWEN38_LONG_CONTEXT.md @@ -47,6 +47,15 @@ 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 @@ -81,9 +90,15 @@ With MTP speculative decode attached (byte-identical greedy): |---|---|---|---| | 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 | +| 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 From cda21c40dd72098f10a59e11da68ef70ed9a3a4a Mon Sep 17 00:00:00 2001 From: NeelM0906 Date: Wed, 26 Aug 2026 10:32:30 -0700 Subject: [PATCH 12/12] fix: surface spill write failures instead of trapping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The write-behind pwrite pair asserted success from the spill queue, so ENOSPC, a short write, or an I/O error terminated the whole process (Codex P2 on PR #20). Writes now run through a complete-write loop (resuming short writes and EINTR); the first failure is recorded and thrown by the next spill-file read — fetch or readSpilledSpan — so the foreground operation fails cleanly instead of reading garbage. reset() clears the recorded failure along with the rest of the state, since every page is rewritten before it can be read again. The error box is deliberately separate from the store: spill closures must not retain the store itself, because its deinit synchronizes on the spill queue and dropping the last reference there would deadlock. --- .../Runtime/KVCache/KVPageStore.swift | 74 +++++++++++++++++-- .../Runtime/KVCache/KVPageStoreTests.swift | 40 ++++++++++ 2 files changed, 109 insertions(+), 5 deletions(-) diff --git a/Sources/Mference/Runtime/KVCache/KVPageStore.swift b/Sources/Mference/Runtime/KVCache/KVPageStore.swift index e1111f6..f20ed6c 100644 --- a/Sources/Mference/Runtime/KVCache/KVPageStore.swift +++ b/Sources/Mference/Runtime/KVCache/KVPageStore.swift @@ -103,6 +103,33 @@ public final class KVPageStore { 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, @@ -370,6 +397,7 @@ public final class KVPageStore { 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") @@ -397,6 +425,9 @@ public final class KVPageStore { /// 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 @@ -497,17 +528,50 @@ public final class KVPageStore { let vSrc = vPools[ordinal].contents() + slot * pageBytes let offset = off_t(geometry.fileOffset(layerOrdinal: ordinal, pageIndex: pageIndex)) let fd = spillFD - spillQueue.async { + 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. - let wroteK = pwrite(fd, kSrc, pageBytes, offset) - let wroteV = pwrite(fd, vSrc, pageBytes, offset + off_t(pageBytes)) - precondition(wroteK == pageBytes && wroteV == pageBytes, - "kv spill pwrite failed: errno \(errno)") + 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 diff --git a/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift b/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift index abad74b..31aaaef 100644 --- a/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift +++ b/Tests/Mference/Core/Runtime/KVCache/KVPageStoreTests.swift @@ -217,6 +217,46 @@ import Metal #expect(store.residentPageCount(layer: 3) == 0) } + // MARK: spill write failures + + @Test func writeFully_surfacesDescriptorErrors() { + let bytes: [UInt8] = [1, 2, 3, 4] + #expect(throws: KVPageStoreError.self) { + try bytes.withUnsafeBytes { buf in + try KVPageStore.writeFully(fd: -1, from: buf.baseAddress!, + count: buf.count, offset: 0) + } + } + } + + @Test func recordedSpillFailure_failsSpillReadsUntilReset() throws { + let (ctx, store, dir) = try makeStore() + defer { try? FileManager.default.removeItem(at: dir) } + _ = try store.kSlot(layer: 3, position: 0) + store.advance(by: Self.pageTokens) + store.flushSpills() + + store.recordSpillError(.ioFailed(operation: "pwrite spill", errno: ENOSPC)) + #expect(store.spillFailure != nil) + let staging = try #require(ctx.device.makeBuffer( + length: 2 * store.geometry.kPageBytes, options: .storageModeShared)) + #expect(throws: KVPageStoreError.ioFailed(operation: "pwrite spill", + errno: ENOSPC)) { + try store.readSpilledSpan(layer: 3, firstPage: 0, pageCount: 1, + into: staging) + } + + // Reset rewrites every page before it can be read again, so the + // recorded failure clears with the rest of the state. + store.reset() + #expect(store.spillFailure == nil) + _ = try store.kSlot(layer: 3, position: 0) + store.advance(by: Self.pageTokens) + store.flushSpills() + try store.readSpilledSpan(layer: 3, firstPage: 0, pageCount: 1, + into: staging) + } + // MARK: metadata layout @Test func metadataOffsets_areDistinctPerLayerAndPage() throws {