perf: streaming I/O overlap — preadv coalescing, Inkling shadow prefetch + pipelined prefill

Sources/Mference/Infrastructure/Streaming/PreadExpertStreamer.swift

@@ -275,15 +275,23 @@ public final class PreadExpertStreamer: @unchecked Sendable {
         precondition(plan.assignedSlots.count == plan.experts.count,
                      "expert cache plan slot count mismatch")
 
+        let missFileOffsets = try plan.misses.map { index in
+            try fileOffsetForExpert(plan.experts[index])
+        }
+        let runs = Self.coalescedReadRuns(offsets: missFileOffsets,
+                                          stride: layout.expertStride)
         let errorLock = NSLock()
         nonisolated(unsafe) var firstError: Error?
-        DispatchQueue.concurrentPerform(iterations: plan.misses.count) { missOffset in
-            let index = plan.misses[missOffset]
+        DispatchQueue.concurrentPerform(iterations: runs.count) { runIndex in
+            let run = runs[runIndex]
+            let destinations = run.map {
+                self.slotPointers[plan.assignedSlots[plan.misses[$0]]]
+            }
             do {
-                _ = try self.loadExpert(
-                    layer: 0,
-                    expert: plan.experts[index],
-                    slot: plan.assignedSlots[index])
+                try self.readScattered(
+                    into: destinations,
+                    fileOffset: missFileOffsets[run[0]],
+                    strideBytes: Int(self.layout.expertStride))
             } catch {
                 errorLock.lock()
                 if firstError == nil { firstError = error }
@@ -378,15 +386,23 @@ public final class PreadExpertStreamer: @unchecked Sendable {
         _ reservation: [(expert: Int, slot: Int)]
     ) -> UInt64 {
         guard !reservation.isEmpty else { return 0 }
+        let fileOffsets = reservation.map { entry in
+            (try? fileOffsetForExpert(entry.expert)) ?? UInt64.max
+        }
+        let readable = reservation.indices.filter { fileOffsets[$0] != UInt64.max }
+        let runs = Self.coalescedReadRuns(offsets: readable.map { fileOffsets[$0] },
+                                          stride: layout.expertStride)
         let loadedLock = NSLock()
         nonisolated(unsafe) var loaded: [Int] = []
-        DispatchQueue.concurrentPerform(iterations: reservation.count) { index in
-            let entry = reservation[index]
-            guard (try? self.loadExpert(layer: 0,
-                                        expert: entry.expert,
-                                        slot: entry.slot)) != nil else { return }
+        DispatchQueue.concurrentPerform(iterations: runs.count) { runIndex in
+            let entries = runs[runIndex].map { readable[$0] }
+            let destinations = entries.map { self.slotPointers[reservation[$0].slot] }
+            guard (try? self.readScattered(
+                into: destinations,
+                fileOffset: fileOffsets[entries[0]],
+                strideBytes: Int(self.layout.expertStride))) != nil else { return }
             loadedLock.lock()
-            loaded.append(index)
+            loaded.append(contentsOf: entries)
             loadedLock.unlock()
         }
 
@@ -558,6 +574,70 @@ public final class PreadExpertStreamer: @unchecked Sendable {
             maxCallNanos: maxCallNanos)
     }
 
+    private func fileOffsetForExpert(_ expert: Int) throws -> UInt64 {
+        let regionOffset = layout.expertOffset(layer: 0, expert: expert)
+        guard regionOffset + layout.expertStride <= layout.streamSize else {
+            throw StreamerError.offsetOutOfRange(regionOffset)
+        }
+        return layout.streamOffset + regionOffset
+    }
+
+    /// Groups reads of `stride` bytes at `offsets` into runs that are exactly
+    /// contiguous on disk, so each run can be fetched with one scattered
+    /// `preadv` instead of one random `pread` per expert. Returns runs of
+    /// indices into `offsets`, each run in ascending disk order. Duplicate
+    /// offsets never share a run: their reads would overlap.
+    static func coalescedReadRuns(offsets: [UInt64], stride: UInt64) -> [[Int]] {
+        let sorted = offsets.indices.sorted { offsets[$0] < offsets[$1] }
+        var runs: [[Int]] = []
+        for index in sorted {
+            if let last = runs.last?.last, offsets[index] == offsets[last] &+ stride {
+                runs[runs.count - 1].append(index)
+            } else {
+                runs.append([index])
+            }
+        }
+        return runs
+    }
+
+    /// Reads `destinations.count * strideBytes` contiguous file bytes starting
+    /// at `fileOffset`, scattering `strideBytes` into each destination in
+    /// order. Single-destination runs use the plain `pread` path.
+    private func readScattered(into destinations: [UnsafeMutableRawPointer],
+                               fileOffset: UInt64,
+                               strideBytes: Int) throws {
+        guard destinations.count > 1 else {
+            return try readFull(into: destinations[0],
+                                fileOffset: fileOffset,
+                                count: strideBytes)
+        }
+        let total = destinations.count * strideBytes
+        var filled = 0
+        while filled < total {
+            let startIndex = filled / strideBytes
+            let within = filled % strideBytes
+            var vectors = [iovec(
+                iov_base: destinations[startIndex].advanced(by: within),
+                iov_len: strideBytes - within)]
+            for index in (startIndex + 1)..<destinations.count {
+                vectors.append(iovec(iov_base: destinations[index],
+                                     iov_len: strideBytes))
+            }
+            let readCount = vectors.withUnsafeBufferPointer { buffer in
+                preadv(fd, buffer.baseAddress, Int32(buffer.count),
+                       off_t(fileOffset) + off_t(filled))
+            }
+            if readCount < 0 {
+                throw StreamerError.preadFailed(errno: errno)
+            }
+            if readCount == 0 {
+                throw StreamerError.sizeMismatch(expected: UInt64(total),
+                                                 actual: UInt64(filled))
+            }
+            filled += readCount
+        }
+    }
+
     private func readFull(into destination: UnsafeMutableRawPointer,
                           fileOffset: UInt64,
                           count: Int) throws {

Sources/Mference/Kernels/MoE/SpeculativeRouterInkling.swift

@@ -0,0 +1,110 @@
+import Foundation
+import Metal
+
+/// PILOT router-lookahead for Inkling-Small: layer L+1's sigmoid router
+/// applied to layer L's post-attention state, encoded into layer L's command
+/// buffer so the CPU learns a *guess* at the next layer's expert set at the
+/// same wake as the real router readback — no extra command buffer, no extra
+/// sync. The Inkling twin of `SpeculativeRouterDSV4`.
+///
+/// Deliberately a separate object from `InklingKernels` rather than an extra
+/// method on it: it must not share that class's `routerLogits` singleton (the
+/// real router's logits are live in the same command buffer), and keeping it
+/// in its own file means the speculative path cannot perturb the validated
+/// decode router. Same kernels (`router_gemv_bf16_r4` +
+/// `inkling_router_select`, both built without function constants exactly as
+/// `InklingKernels` builds them), so the ranking — sigmoid scoring with the
+/// per-expert gate bias — is bit-identical to what layer L+1 will compute for
+/// real. Recall depends entirely on that fidelity; only the *input
+/// activation* is an approximation.
+final class SpeculativeRouterInkling {
+
+    private let gemvPSO: MTLComputePipelineState
+    private let selectPSO: MTLComputePipelineState
+    private let logits: MTLBuffer
+    /// Predicted expert ids, `topK` x UInt32. Read by the CPU on the router
+    /// signal.
+    let predictedIndices: MTLBuffer
+    /// The select kernel writes routing weights and shared-expert gammas
+    /// alongside the indices; the prediction never uses them, but the kernel
+    /// contract needs the buffers.
+    private let discardedWeights: MTLBuffer
+    private let discardedGammas: MTLBuffer
+
+    init(context: MetalContext,
+         numRouted: Int,
+         numShared: Int,
+         topK: Int) throws {
+        self.gemvPSO = try context.pipeline("router_gemv_bf16_r4")
+        self.selectPSO = try context.pipeline("inkling_router_select")
+
+        guard let logits = context.device.makeBuffer(
+            length: (numRouted + numShared) * MemoryLayout<Float>.stride,
+            options: .storageModeShared),
+            let indices = context.device.makeBuffer(
+                length: topK * MemoryLayout<UInt32>.stride,
+                options: .storageModeShared),
+            let weights = context.device.makeBuffer(
+                length: topK * MemoryLayout<Float16>.stride,
+                options: .storageModeShared),
+            let gammas = context.device.makeBuffer(
+                length: max(numShared, 1) * MemoryLayout<Float>.stride,
+                options: .storageModeShared) else {
+            throw MetalError.noDevice
+        }
+        logits.label = "inkling.pilot_router_logits"
+        indices.label = "inkling.pilot_router_indices"
+        self.logits = logits
+        self.predictedIndices = indices
+        self.discardedWeights = weights
+        self.discardedGammas = gammas
+    }
+
+    /// Encodes the lookahead GEMV + sigmoid top-k. Must be encoded *before*
+    /// the router signal so the CPU sees the prediction at the same wake as
+    /// the real indices. Dispatch geometry mirrors
+    /// `InklingKernels.encodeRouter` exactly.
+    func encodePrediction(commandBuffer cb: MTLCommandBuffer,
+                          weights: MTLBuffer, weightsOffset: Int,
+                          hidden: MTLBuffer,
+                          onesScale: MTLBuffer,
+                          gateBias: MTLBuffer, gateBiasOffset: Int,
+                          globalScale: MTLBuffer, globalScaleOffset: Int,
+                          numRouted: UInt32, numShared: UInt32,
+                          topK: UInt32, routeScale: Float,
+                          d: UInt32) {
+        var total = numRouted + numShared
+        var dim = d
+        if let enc = cb.makeComputeCommandEncoder() {
+            enc.setComputePipelineState(gemvPSO)
+            enc.setBuffer(weights, offset: weightsOffset, index: 0)
+            enc.setBuffer(hidden, offset: 0, index: 1)
+            enc.setBuffer(onesScale, offset: 0, index: 2)
+            enc.setBuffer(logits, offset: 0, index: 3)
+            enc.setBytes(&total, length: 4, index: 4)
+            enc.setBytes(&dim, length: 4, index: 5)
+            enc.dispatchThreadgroups(
+                MTLSize(width: (Int(total) + 3) / 4, height: 1, depth: 1),
+                threadsPerThreadgroup: MTLSize(width: 128, height: 1, depth: 1))
+            enc.endEncoding()
+        }
+
+        guard let enc = cb.makeComputeCommandEncoder() else { return }
+        var nr = numRouted, ns = numShared, tk = topK, rs = routeScale
+        enc.setComputePipelineState(selectPSO)
+        enc.setBuffer(logits, offset: 0, index: 0)
+        enc.setBuffer(gateBias, offset: gateBiasOffset, index: 1)
+        enc.setBuffer(globalScale, offset: globalScaleOffset, index: 2)
+        enc.setBuffer(predictedIndices, offset: 0, index: 3)
+        enc.setBuffer(discardedWeights, offset: 0, index: 4)
+        enc.setBuffer(discardedGammas, offset: 0, index: 5)
+        enc.setBytes(&nr, length: 4, index: 6)
+        enc.setBytes(&ns, length: 4, index: 7)
+        enc.setBytes(&tk, length: 4, index: 8)
+        enc.setBytes(&rs, length: 4, index: 9)
+        enc.dispatchThreadgroups(
+            MTLSize(width: 1, height: 1, depth: 1),
+            threadsPerThreadgroup: MTLSize(width: 32, height: 1, depth: 1))
+        enc.endEncoding()
+    }
+}