Qwen 3.8 long context: paged KV with SSD tier, Quest sparse decode, MTP composition

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 `<think>` block.
+  also opens a live `<think>` 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

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<Float>.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<UInt32>.size, index: 6)
+        p1.setBytes(&nq,  length: MemoryLayout<UInt32>.size, index: 7)
+        p1.setBytes(&nkv, length: MemoryLayout<UInt32>.size, index: 8)
+        p1.setBytes(&st,  length: MemoryLayout<UInt32>.size, index: 9)
+        p1.setBuffer(pageTable, offset: pageTableOffset, index: 10)
+        p1.setBytes(&cl,  length: MemoryLayout<UInt32>.size, index: 11)
+        p1.setBytes(&nc,  length: MemoryLayout<UInt32>.size, index: 12)
+        p1.setBytes(&sc2, length: MemoryLayout<Float>.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<UInt32>.size, index: 4)
+        p2.setBytes(&nc2, length: MemoryLayout<UInt32>.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,

Sources/Mference/Kernels/Attention/KVPageKernels.swift

@@ -0,0 +1,180 @@
+import Foundation
+import Metal
+
+/// Wrappers for the paged-KV maintenance kernels: page seal summaries
+/// (`kv_page_minmax`) and Quest page criticality scores
+/// (`attention_page_scores`). Both ride the token command buffer.
+final class KVPageKernels {
+    private let ctx: MetalContext
+    private let psoMinMax: MTLComputePipelineState
+    private let psoScores: MTLComputePipelineState
+    private let psoFlashInit: MTLComputePipelineState
+    private let psoFlashUpdate: MTLComputePipelineState
+    private let psoFlashFinalize: MTLComputePipelineState
+
+    private static let threadsPerGroup = 256
+    private static let flashSimdgroupsPerTG = 8
+
+    init(context: MetalContext) throws {
+        self.ctx = context
+        self.psoMinMax = try context.pipeline("kv_page_minmax")
+        self.psoScores = try context.pipeline("attention_page_scores")
+        self.psoFlashInit = try context.pipeline("attention_prefill_flash_init")
+        self.psoFlashUpdate = try context.pipeline("attention_prefill_flash_update")
+        self.psoFlashFinalize = try context.pipeline("attention_prefill_flash_finalize")
+    }
+
+    // MARK: - Blocked prefill attention
+
+    /// Reset the running online-softmax state for a chunk's queries.
+    func encodeFlashInit(commandBuffer: MTLCommandBuffer,
+                         mState: MTLBuffer, dState: MTLBuffer, oState: MTLBuffer,
+                         rows: UInt32, headDim: UInt32) {
+        guard let enc = commandBuffer.makeComputeCommandEncoder() else { return }
+        enc.setComputePipelineState(psoFlashInit)
+        enc.setBuffer(mState, offset: 0, index: 0)
+        enc.setBuffer(dState, offset: 0, index: 1)
+        enc.setBuffer(oState, offset: 0, index: 2)
+        var r = rows, hd = headDim
+        enc.setBytes(&r,  length: MemoryLayout<UInt32>.size, index: 3)
+        enc.setBytes(&hd, length: MemoryLayout<UInt32>.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<UInt32>.size, index: 7)
+        enc.setBytes(&qs,  length: MemoryLayout<UInt32>.size, index: 8)
+        enc.setBytes(&hd,  length: MemoryLayout<UInt32>.size, index: 9)
+        enc.setBytes(&nq,  length: MemoryLayout<UInt32>.size, index: 10)
+        enc.setBytes(&nkv, length: MemoryLayout<UInt32>.size, index: 11)
+        enc.setBytes(&ws,  length: MemoryLayout<UInt32>.size, index: 12)
+        enc.setBytes(&wt,  length: MemoryLayout<UInt32>.size, index: 13)
+        enc.setBytes(&qst, length: MemoryLayout<UInt32>.size, index: 14)
+        enc.setBytes(&sc,  length: MemoryLayout<Float>.size,  index: 15)
+        enc.setBytes(&cz,  length: MemoryLayout<UInt32>.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<UInt32>.size, index: 4)
+        enc.setBytes(&hd, length: MemoryLayout<UInt32>.size, index: 5)
+        enc.setBytes(&nq, length: MemoryLayout<UInt32>.size, index: 6)
+        enc.setBytes(&os, length: MemoryLayout<UInt32>.size, index: 7)
+        let total = Int(queryCount * numQHeads * headDim)
+        let width = min(Self.threadsPerGroup, psoFlashFinalize.maxTotalThreadsPerThreadgroup)
+        enc.dispatchThreads(MTLSize(width: total, height: 1, depth: 1),
+                            threadsPerThreadgroup: MTLSize(width: width, height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    /// Reduce a page's K rows to element-wise min/max vectors, written to
+    /// the page's slot in the metadata buffer.
+    func encodePageMinMax(commandBuffer: MTLCommandBuffer,
+                          kPool: MTLBuffer,
+                          slot: UInt32,
+                          validTokens: UInt32,
+                          metadata: MTLBuffer,
+                          metadataOffset: Int,
+                          numKVHeads: UInt32,
+                          headDim: UInt32) {
+        guard let enc = commandBuffer.makeComputeCommandEncoder() else { return }
+        enc.setComputePipelineState(psoMinMax)
+        enc.setBuffer(kPool, offset: 0, index: 0)
+        enc.setBuffer(metadata, offset: metadataOffset, index: 1)
+        var s = slot, vt = validTokens, nkv = numKVHeads, hd = headDim
+        enc.setBytes(&s,   length: MemoryLayout<UInt32>.size, index: 2)
+        enc.setBytes(&vt,  length: MemoryLayout<UInt32>.size, index: 3)
+        enc.setBytes(&nkv, length: MemoryLayout<UInt32>.size, index: 4)
+        enc.setBytes(&hd,  length: MemoryLayout<UInt32>.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<UInt32>.size, index: 3)
+        enc.setBytes(&hd,  length: MemoryLayout<UInt32>.size, index: 4)
+        enc.setBytes(&nq,  length: MemoryLayout<UInt32>.size, index: 5)
+        enc.setBytes(&nkv, length: MemoryLayout<UInt32>.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()
+    }
+}