diff --git a/Sources/Mference/Infrastructure/Streaming/PreadExpertStreamer.swift b/Sources/Mference/Infrastructure/Streaming/PreadExpertStreamer.swift index f3224b0..24d46ac 100644 --- a/Sources/Mference/Infrastructure/Streaming/PreadExpertStreamer.swift +++ b/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)...stride, + options: .storageModeShared), + let indices = context.device.makeBuffer( + length: topK * MemoryLayout.stride, + options: .storageModeShared), + let weights = context.device.makeBuffer( + length: topK * MemoryLayout.stride, + options: .storageModeShared), + let gammas = context.device.makeBuffer( + length: max(numShared, 1) * MemoryLayout.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() + } +} diff --git a/Sources/Mference/Runtime/Inference/RealForwardRunner.swift b/Sources/Mference/Runtime/Inference/RealForwardRunner.swift index 28111f6..a7398aa 100644 --- a/Sources/Mference/Runtime/Inference/RealForwardRunner.swift +++ b/Sources/Mference/Runtime/Inference/RealForwardRunner.swift @@ -467,6 +467,12 @@ public final class RealForwardRunner: ChunkedPrefillRunner, ContextWindowReporti /// byte-identical). `MFERENCE_SLOT_MAP=0` is the kill-switch. static let slotMapEnabledDefault = ProcessInfo.processInfo.environment["MFERENCE_SLOT_MAP"] != "0" + /// Inkling prefill expert streaming: depth-1 pipeline — pread expert e+1 + /// while expert e's GLU runs, misses placed only in slots the in-flight + /// command buffer does not touch. `MFERENCE_INKLING_PREFILL_PIPELINE=0` + /// is the kill-switch back to the serialized fetch->encode->drain loop. + static let inklingPrefillPipelineEnabled = + ProcessInfo.processInfo.environment["MFERENCE_INKLING_PREFILL_PIPELINE"] != "0" var slotMapEnabled = RealForwardRunner.slotMapEnabledDefault /// Debug discriminator: encode the whole slot-map chain but feed the /// lookup an all-empty table, so the guarded kernels always no-op and @@ -517,8 +523,12 @@ public final class RealForwardRunner: ChunkedPrefillRunner, ContextWindowReporti /// PILOT lookahead kernels, built on first use so `off` pays nothing and /// the (concurrently edited) DSV4 init block stays untouched. private var pilotRouter: SpeculativeRouterDSV4? - /// The prediction read out of `pilotRouter` (or the hash table) at the - /// current layer's router wake, consumed by `issueSpeculativePrefetch`. + /// Inkling twin of `pilotRouter`: same lazy construction, same lifecycle, + /// family-matched scoring kernels. + private var inklingPilotRouter: SpeculativeRouterInkling? + /// The prediction read out of `pilotRouter`/`inklingPilotRouter` (or the + /// hash table) at the current layer's router wake, consumed by + /// `issueSpeculativePrefetch`. private var pilotPrediction: (layer: Int, experts: [Int])? public init(model: Model, context: MetalContext, maxContext: Int, @@ -1296,6 +1306,29 @@ public final class RealForwardRunner: ChunkedPrefillRunner, ContextWindowReporti return pilotRouter } + private func ensureInklingPilotRouter() -> SpeculativeRouterInkling? { + if let inklingPilotRouter { return inklingPilotRouter } + inklingPilotRouter = try? SpeculativeRouterInkling( + context: ctx, + numRouted: cfg.numExperts, + numShared: cfg.numSharedExperts, + topK: cfg.topKExperts) + return inklingPilotRouter + } + + /// Inkling wake-side twin of `capturePilotPrediction`: no hash-routed + /// layers exist in this family, so the prediction comes only from the + /// pilot GEMV. + private func captureInklingPilotPrediction(nextLayer: Int, gemvEncoded: Bool) { + pilotPrediction = nil + guard speculativePrefetchMode == .pilot || speculativePrefetchMode == .shadow, + nextLayer < cfg.numLayers, gemvEncoded, + let buffer = inklingPilotRouter?.predictedIndices else { return } + let ptr = buffer.contents().assumingMemoryBound(to: UInt32.self) + let cap = cfg.numExperts - 1 + pilotPrediction = (nextLayer, (0.. UInt64 { breakdown ? clock_gettime_nsec_np(CLOCK_UPTIME_RAW) : 0 } - for range in routedRanges { + // Depth-1 pipeline: while the GPU runs expert e's GLU, the CPU + // plans and preads expert e+1 into slots the in-flight buffer + // does not touch (`avoidingSlots`), so the fetch overlaps GPU + // work. `acc`/`act` hazards need no explicit sync: consecutive + // buffers on the serial queue are hazard-tracked on those + // resources, so expert accumulation order — and therefore the + // output — is byte-identical to the serialized loop. When slot + // pressure leaves no avoiding plan (or the pipeline is switched + // off), the next iteration drains first and demand-fetches, + // which *is* the old serialized behavior. + let tileScheduler = PrefillRoutedTileScheduler( + config: Self.prefillRoutedTileSchedulerConfig) + var pendingExpertCB: (cb: MTLCommandBuffer, slots: [Int])? + var prefetched: (expert: Int, plan: RoutedExpertFetchPlan, + blob: TensorView)? + // Error-path backstop: any throw below (a failed prefetch pread, + // a drain surfacing a GPU error) must not leave the in-flight + // buffer running against the persistent act/acc scratch — an + // immediate reset or retry would race it. The success path + // drains and clears `pendingExpertCB`, so this no-ops there. + defer { + if let pending = pendingExpertCB { + waitForCompletion(pending.cb) + pendingExpertCB = nil + } + } + func drainPendingExpertCB() throws { + guard let pending = pendingExpertCB else { return } + pendingExpertCB = nil + let tDrain = now() + waitForCompletion(pending.cb) + if breakdown { Self.prefillDrainNanos &+= now() - tDrain } + if let error = pending.cb.error { throw error } + } + for (index, range) in routedRanges.enumerated() { let t0 = now() - let blob = try await model.fetchRoutedExperts(layer: L, - experts: [range.expert])[0] + let blob: TensorView + let blobSlots: [Int] + if let ready = prefetched, ready.expert == range.expert { + blob = ready.blob + blobSlots = ready.plan.assignedSlots + prefetched = nil + } else { + prefetched = nil + try drainPendingExpertCB() + guard let plan = try model.planRoutedExperts( + layer: L, experts: [range.expert]) else { + throw ModelError.routedExpertPlanUnavailable(layer: L) + } + blob = try await model.fetchRoutedExperts(plan: plan)[0] + blobSlots = plan.assignedSlots + } let t1 = now() let cb = ctx.queue.makeCommandBuffer()! prefillGLU.encode(commandBuffer: cb, @@ -4421,20 +4544,36 @@ public final class RealForwardRunner: ChunkedPrefillRunner, ContextWindowReporti params: expertParams(range, base: Int(blob.offset))) cb.commit() let t2 = now() - // The next fetch may evict this expert's slot, and the shared - // `act` tile is reused, so drain before moving on. This - // serialization is why fetch does not overlap GPU work; see - // PrefillRoutedTileScheduler for the pipelined alternative the - // other families use. - waitForCompletion(cb) if breakdown { - let t3 = clock_gettime_nsec_np(CLOCK_UPTIME_RAW) Self.prefillFetchNanos &+= t1 - t0 Self.prefillEncodeNanos &+= t2 - t1 - Self.prefillDrainNanos &+= t3 - t2 Self.prefillExpertCount &+= 1 } + try drainPendingExpertCB() + pendingExpertCB = (cb, blobSlots) + + if Self.inklingPrefillPipelineEnabled, + index + 1 < routedRanges.count { + let nextExpert = routedRanges[index + 1].expert + let nextPlan = try model.planRoutedExpertsIfPossible( + layer: L, + experts: [nextExpert], + avoidingSlots: Set(blobSlots)) + let decision = tileScheduler.decide( + PrefillRoutedTileSchedulerInput( + hasPendingTile: true, + pendingAssignedSlots: blobSlots, + avoidingSlotPlanAvailable: nextPlan != nil)) + if case .prefetchNext = decision, let nextPlan { + let tFetch = now() + let nextBlob = + try await model.fetchRoutedExperts(plan: nextPlan)[0] + if breakdown { Self.prefillFetchNanos &+= now() - tFetch } + prefetched = (nextExpert, nextPlan, nextBlob) + } + } } + try drainPendingExpertCB() totalIoNanos &+= clock_gettime_nsec_np(CLOCK_UPTIME_RAW) - tIoStart // Tail: one causal channel-wise dispatch replaces N narrowing, diff --git a/Tests/Mference/Core/Infrastructure/Streaming/PreadExpertStreamerTests+Coalescing.swift b/Tests/Mference/Core/Infrastructure/Streaming/PreadExpertStreamerTests+Coalescing.swift new file mode 100644 index 0000000..c6b55f8 --- /dev/null +++ b/Tests/Mference/Core/Infrastructure/Streaming/PreadExpertStreamerTests+Coalescing.swift @@ -0,0 +1,139 @@ +import Darwin +import Foundation +import Metal +import Testing + +@testable import Mference + +/// Coalesced miss reads: misses whose blobs are adjacent on disk are fetched +/// with one scattered `preadv` per contiguous run instead of one `pread` per +/// expert. Correctness must be identical for uniform layouts, permuted +/// `expertOffsets` tables, and the speculative fill path. +extension PreadExpertStreamerTests { + + @Test func coalescedReadRuns_groupsContiguousOffsets() { + let stride = UInt64(Self.expertStride) + // Offsets 0,1s,2s are one run; 4s,5s another; 7s alone. Input order is + // scrambled to prove the grouping sorts by disk position first. + let offsets: [UInt64] = [4 * stride, 0, 2 * stride, 7 * stride, stride, 5 * stride] + let runs = PreadExpertStreamer.coalescedReadRuns(offsets: offsets, stride: stride) + #expect(runs == [[1, 4, 2], [0, 5], [3]]) + } + + @Test func coalescedReadRuns_duplicateOffsetsNeverShareARun() { + let stride = UInt64(Self.expertStride) + let offsets: [UInt64] = [0, 0, stride] + let runs = PreadExpertStreamer.coalescedReadRuns(offsets: offsets, stride: stride) + // Overlapping reads must stay separate syscalls; only one duplicate can + // extend into the following contiguous blob. + #expect(runs.count == 2) + #expect(runs.flatMap { $0 }.sorted() == [0, 1, 2]) + for run in runs { + let sortedOffsets = run.map { offsets[$0] } + #expect(sortedOffsets == sortedOffsets.sorted()) + for pair in zip(sortedOffsets, sortedOffsets.dropFirst()) { + #expect(pair.1 == pair.0 + stride) + } + } + } + + @Test func cachePlanWithAdjacentMisses_readsEverySlotCorrectly() throws { + let url = try Self.writeSyntheticLayer() + defer { try? FileManager.default.removeItem(at: url) } + let device = try MetalContext().device + let streamer = try PreadExpertStreamer( + layout: Self.makeLayout(path: url.path), device: device, + slotCount: Self.numExperts) + + // Cold cache: all four experts miss and are contiguous on disk, so this + // exercises a single multi-entry scattered read. + let results = try streamer.loadExpertsCached(experts: Array(0.. disk position: id 0 lives at blob 2, id 1 at blob 3, + // id 2 at blob 0, id 3 at blob 1. Runs form on disk order (ids 2,3 then + // 0,1), not id order. + let positions = [2, 3, 0, 1] + let layout = StreamLayout( + path: url.path, + streamOffset: Self.streamOffset, + streamSize: Self.streamSize, + expertsPerLayer: Self.numExperts, + expertStride: stride, + expertOffsets: positions.map { UInt64($0) * stride }) + let streamer = try PreadExpertStreamer( + layout: layout, device: device, slotCount: Self.numExperts) + + let results = try streamer.loadExpertsCached(experts: Array(0.. UInt16 { + UInt16(truncatingIfNeeded: v.bitPattern >> 16) + } + + private static func buffer(_ device: MTLDevice, _ values: [T]) -> MTLBuffer { + values.withUnsafeBytes { raw in + device.makeBuffer(bytes: raw.baseAddress!, length: raw.count, + options: .storageModeShared)! + } + } + + @Test func predictionMatchesRealRouterBitForBit() throws { + let context = try MetalContext() + let device = context.device + let kernels = try InklingKernels(context: context, + numRouted: Self.numRouted, + numShared: Self.numShared) + let pilot = try SpeculativeRouterInkling(context: context, + numRouted: Self.numRouted, + numShared: Self.numShared, + topK: Self.topK) + + var rng = SplitMix64(seed: 0x1_2C0F) + let total = Self.numRouted + Self.numShared + let weights = Self.buffer( + device, + (0..