Add packed paged KV formats with bounded workspace admission

Libraries/MLXLMCommon/ContinuousBatchingV2/AdmissionOpportunisticWorkspace.swift

@@ -0,0 +1,10 @@
+extension AdmissionV2 {
+    /// Optional acceleration must not consume credit promised to later direct
+    /// attention calls. Charge the complete candidate as additional nonbackend
+    /// memory, atomically, before creating an arena or writing any KV. Its lease
+    /// remains independent of request cancellation, capacity changes and other
+    /// in-flight steps. Returning it after GPU retirement refunds exactly once.
+    func reserveOpportunisticWorkspace(bytes: Int) throws -> CBv2CheckpointReservation {
+        try reserveTransient(bytes: bytes)
+    }
+}

Libraries/MLXLMCommon/ContinuousBatchingV2/AdmissionWorkspaceFloor.swift

@@ -0,0 +1,85 @@
+/// Immutable upper bound for one request's overlapping attention workspaces.
+/// It must cover every supported chunk/decode shape at the reserved length,
+/// including two submitted steps. It is physical scratch, not KV storage.
+public struct CBv2RequestWorkspaceProjection: Sendable {
+    private let project: @Sendable (Int) -> Int?
+    private let projectAggregate: (@Sendable (Int, Int) -> Int?)?
+
+    public init(bytesForTokens: @escaping @Sendable (Int) -> Int?) {
+        project = bytesForTokens
+        projectAggregate = nil
+    }
+
+    /// The aggregate must bound the sum of individual allowances for every
+    /// split of totalTokens into at most maximumRequests positive lengths.
+    /// Existing and retired owners are charged separately by admission.
+    public init(
+        bytesForTokens: @escaping @Sendable (Int) -> Int?,
+        bytesForAggregate: @escaping @Sendable (Int, Int) -> Int?
+    ) {
+        project = bytesForTokens
+        projectAggregate = bytesForAggregate
+    }
+
+    func bytes(forTokens tokens: Int) -> Int? {
+        guard tokens >= 0 else { return nil }
+        if tokens == 0 { return 0 }
+        guard let value = project(tokens), value >= 0 else { return nil }
+        return value
+    }
+
+    func bytes(totalTokens: Int, maximumRequests: Int) -> Int? {
+        guard totalTokens >= 0, maximumRequests >= 0 else { return nil }
+        if totalTokens == 0 { return 0 }
+        guard maximumRequests > 0 else { return nil }
+        let count = min(totalTokens, maximumRequests)
+        if let projectAggregate {
+            guard let value = projectAggregate(totalTokens, count), value >= 0 else { return nil }
+            return value
+        }
+        // Generic projections have no algebraic contract beyond monotonicity.
+        // Each positive request is at most the aggregate length; retain the
+        // conservative fallback until a projection supplies a tighter proof.
+        guard let single = bytes(forTokens: totalTokens) else { return nil }
+        let (result, overflow) = single.multipliedReportingOverflow(by: count)
+        return overflow ? nil : result
+    }
+}
+
+/// Prepaid request allowances and actual leased workspaces are independent
+/// owners. Live leases consume existing credit once; retiring credit stays
+/// charged until GPU work releases it, so a new request cannot reuse it early.
+struct CBv2AdmissionWorkspaceFloor {
+    var prepaidBytes = 0
+    var leasedBytes = 0
+    private(set) var retiredBytes = 0
+
+    /// Credits retired while any workspace is live cannot be reused for a
+    /// newly admitted request's future workspace. Conservatively retain that
+    /// part of the old credit until outstanding leased bytes fall below it.
+    mutating func replacePrepaid(_ bytes: Int) {
+        precondition(bytes >= 0)
+        if bytes < prepaidBytes {
+            let released = prepaidBytes - bytes
+            retiredBytes += min(released, max(0, leasedBytes - retiredBytes))
+        }
+        prepaidBytes = bytes
+    }
+
+    func replacingPrepaid(_ bytes: Int) -> Self {
+        var result = self
+        result.replacePrepaid(bytes)
+        return result
+    }
+
+    mutating func releaseLease(_ bytes: Int) {
+        precondition(bytes >= 0 && bytes <= leasedBytes)
+        leasedBytes -= bytes
+        retiredBytes = min(retiredBytes, leasedBytes)
+    }
+
+    func overhead(prepaid: Int? = nil, leased: Int? = nil) -> Int {
+        let state = prepaid.map(replacingPrepaid) ?? self
+        return max(state.retiredBytes, (leased ?? state.leasedBytes) - state.prepaidBytes)
+    }
+}

Libraries/MLXLMCommon/ContinuousBatchingV2/EngineLoopV2+QuantizedScratch.swift

@@ -0,0 +1,72 @@
+import MLX
+
+extension CBv2InFlightStep {
+    /// A successor can already be running when this step retires. Release
+    /// only the workspaces captured by this exact step after its readbacks.
+    func finishQuantizedScratchAfterSynchronization() {
+        for lease in quantizedScratch { lease.finishAfterSynchronization() }
+        quantizedScratch.removeAll()
+    }
+}
+
+extension EngineLoopV2 {
+    /// Provision the complete submitted step's arena shape before constructing
+    /// any layer graph. Scheduler frontiers already include this assignment;
+    /// lookahead also covers serial speculative verification columns.
+    func beginQuantizedScratch(for plan: CBv2StepPlan) throws {
+        guard let pool = (backend as? PagedKVBackend)?.pool,
+              pool.config.quantization != nil else { return }
+        var maximumFrontier = 1
+        for assignment in plan.assignments {
+            guard let record = scheduler.record(for: assignment.id) else { continue }
+            let (frontier, overflow) = record.numComputedTokens.addingReportingOverflow(
+                CBv2PagedSpeculation.maxSpeculativeSpan)
+            guard !overflow, frontier > 0 else {
+                throw CBv2KVError.backendIneligible(reason: "quantized step frontier overflow")
+            }
+            maximumFrontier = max(maximumFrontier, frontier)
+        }
+        try pool.beginQuantizedScratch(
+            maximumQueries: max(8, scheduler.config.maxConcurrentRequests),
+            maximumAttendLength: maximumFrontier)
+    }
+
+    /// Offline scoring does not construct a scheduler step. Each completed
+    /// forward is its own bounded scope, with the same retirement rules.
+    func beginQuantizedScoringScratch(state: [CBv2SequenceKV?], tokens: MLXArray) throws {
+        guard let pool = (backend as? PagedKVBackend)?.pool,
+              pool.config.quantization != nil else { return }
+        let (end, endOverflow) = Self.positionOffset(state).addingReportingOverflow(tokens.dim(1))
+        let (frontier, frontierOverflow) = end.addingReportingOverflow(CBv2PagedSpeculation.maxSpeculativeSpan)
+        guard !endOverflow, !frontierOverflow, frontier > 0 else {
+            throw CBv2KVError.backendIneligible(reason: "quantized scoring frontier overflow")
+        }
+        try pool.beginQuantizedScratch(maximumQueries: max(8, tokens.dim(0)),
+                                       maximumAttendLength: frontier)
+    }
+
+    func attachQuantizedScratch(to step: CBv2InFlightStep?) {
+        guard let pool = (backend as? PagedKVBackend)?.pool else { return }
+        let leases = pool.takePendingQuantizedScratch()
+        guard let step else {
+            // An empty plan normally constructs no scratch. Do not return a
+            // reservation while an exceptional submitted empty plan still runs.
+            if !leases.isEmpty {
+                Stream.gpu.synchronize()
+                Stream.cpu.synchronize()
+                for lease in leases { lease.finishAfterSynchronization() }
+            }
+            return
+        }
+        step.quantizedScratch.append(contentsOf: leases)
+    }
+
+    /// Failed graph builders have no in-flight step to own their leases.
+    /// Call only after draining submitted work and discarding failed fences.
+    func discardPendingQuantizedScratchAfterSynchronization() {
+        guard let pool = (backend as? PagedKVBackend)?.pool else { return }
+        for lease in pool.takePendingQuantizedScratch() {
+            lease.finishAfterSynchronization()
+        }
+    }
+}

Libraries/MLXLMCommon/ContinuousBatchingV2/EngineLoopV2.swift

@@ -408,6 +408,10 @@ final class CBv2InFlightStep {
     var attentionMetadata: CBv2AttentionMetadataForward?
     var attentionPacket: CBv2AttentionPacketForward?
 
+    /// Per-step native workspaces for packed KV attention, separately charged
+    /// from persistent packed pages and never refunded by another step.
+    var quantizedScratch: [PagedQuantizedScratchLease] = []
+
     init(
         assignments: [(id: CBv2RequestID, numTokens: Int)],
         participants: Set<CBv2RequestID>, sampledRows: [CBv2RequestID],
@@ -1493,6 +1497,32 @@ public final class EngineLoopV2: @unchecked Sendable {
         }
     }
 
+    /// An optional prefill may spend otherwise free bytes after a heartbeat
+    /// advertised them. Retire the actual owning step before forecasting a
+    /// new arrival; subtracting its permit without completing GPU work would
+    /// let the forecast promise memory that prefix adoption could use too soon.
+    private func retireOptionalPrefillBeforeDeadlineAdmission() {
+        guard running, !draining,
+            let previous = inFlight, !previous.quantizedScratch.isEmpty,
+            let pool = (backend as? PagedKVBackend)?.pool,
+            pool.quantizedPrefillStatistics.currentAdditionalWorkspaceBytes > 0
+        else { return }
+
+        // This is the normal retirement half of a step, including watchdog
+        // coverage and cancellation/lease handling, without planning more work.
+        markStepStarted()
+        defer { markStepEnded() }
+        let now = config.clock.now()
+        boundaryClockNanos = 0
+        processCancellations(now: now)
+        processLeaseExpiry(now: now)
+        inFlight = nil
+        finalize(previous, now: now)
+        publishGauges()
+        // The already queued engineStep callback remains the continuation.
+        // It observes nil and cannot finalize this submitted step a second time.
+    }
+
     /// Atomic first-token deadline admission.
     ///
     /// The queue position is load-bearing: enqueue establishes authoritative
@@ -1568,6 +1598,12 @@ public final class EngineLoopV2: @unchecked Sendable {
                         return
                     }
 
+                    // Keep the original absolute deadline. Both real GPU drain
+                    // time and this queue wait have elapsed when the existing
+                    // verdict below reads the clock; cancellation/running state
+                    // is rechecked below before any scheduler/prefix mutation.
+                    retireOptionalPrefillBeforeDeadlineAdmission()
+
                     prefixUsageByID[request.id] = CBv2PrefixUsage(
                         outcome: prefixLookup.outcome,
                         tier: nil,
@@ -2159,7 +2195,9 @@ public final class EngineLoopV2: @unchecked Sendable {
                 let boundary = (backend as? PagedKVBackend).map { CBv2PagedWriteBoundary(pool: $0.pool) }
                 let next: CBv2InFlightStep
                 do {
+                    try beginQuantizedScratch(for: plan)
                     next = try launchChainedDecode(plan, feeding: previous.sampledTokens!)
+                    attachQuantizedScratch(to: next)
                 } catch {
                     handlePagedWriteFailure(error, plan: plan, boundary: boundary, now: stepNow)
                     publishGauges()
@@ -2284,7 +2322,9 @@ public final class EngineLoopV2: @unchecked Sendable {
         let measurement = mtpMeasurement(for: plan)
         let boundary = (backend as? PagedKVBackend).map { CBv2PagedWriteBoundary(pool: $0.pool) }
         do {
+            try beginQuantizedScratch(for: plan)
             inFlight = try (mtpRoundNeeded(plan) ? executeMTPRound(plan) : executeMixed(plan))
+            attachQuantizedScratch(to: inFlight)
         } catch {
             handlePagedWriteFailure(error, plan: plan, boundary: boundary, now: stepNow)
         }
@@ -2328,6 +2368,7 @@ public final class EngineLoopV2: @unchecked Sendable {
         Stream.gpu.synchronize()
         Stream.cpu.synchronize()
         boundary?.discardFailedGraphAfterSynchronization()
+        discardPendingQuantizedScratchAfterSynchronization()
         attentionMetadata?.discardPendingForward()
         attentionPacket?.discardPendingForward()
         (cacheProvider as? CBv2CompositionInvalidating)?.releaseBoundRows()
@@ -2703,15 +2744,24 @@ public final class EngineLoopV2: @unchecked Sendable {
             // backend recycles its pages (PR#62), and force the next real
             // step to rebind its own rows.
             (cacheProvider as? CBv2CompositionInvalidating)?.releaseBoundRows()
+            if let pool = (backend as? PagedKVBackend)?.pool {
+                let pending = pool.takePendingQuantizedScratch()
+                if !pending.isEmpty {
+                    Stream.gpu.synchronize()
+                    Stream.cpu.synchronize()
+                    for lease in pending { lease.finishAfterSynchronization() }
+                }
+            }
             backend.release(state)
             state.removeAll()
             (backend as? PagedKVBackend)?.pool.writeValidation.clearAfterRetirement()
             recurrentReservation?.release()
         }
-        if (model as? any CBv2RecurrentSteppableModel)?.recurrentStateSpec != nil {
+        if (model as? any CBv2RecurrentSteppableModel)?.recurrentStateSpec != nil
+            || (backend as? PagedKVBackend)?.pool.config.quantization != nil {
             guard let admission = capacity as? AdmissionV2 else {
                 throw CBv2KVError.backendIneligible(
-                    reason: "recurrent teacher scoring requires peak admission accounting")
+                    reason: "stateful or quantized teacher scoring requires peak admission accounting")
             }
             // makeRecurrentRequestState is a one-generation allocation check;
             // ordinary admission owns the larger committed/pending peak. This
@@ -2726,6 +2776,7 @@ public final class EngineLoopV2: @unchecked Sendable {
             tokens: MLXArray, caches: [CBv2AttendingLayerCache],
             requirement: CBv2PrefillRequirement?
         ) throws -> (logits: MLXArray, innerState: [MLXArray]) {
+            try beginQuantizedScoringScratch(state: state, tokens: tokens)
             guard let recurrent else {
                 if let requirement {
                     return (try prefillOutput(tokens: tokens, inputEmbeddings: nil,
@@ -2743,6 +2794,23 @@ public final class EngineLoopV2: @unchecked Sendable {
         }
 
         func finishForward(_ arrays: [MLXArray]) throws {
+            let scratch = (backend as? PagedKVBackend)?.pool.takePendingQuantizedScratch() ?? []
+            if !scratch.isEmpty {
+                // Offline scoring does not enter CBv2InFlightStep. Bound its
+                // private workspace lifetime to each completed forward instead
+                // of retaining every continuation's scratch until final readback.
+                defer {
+                    Stream.gpu.synchronize()
+                    Stream.cpu.synchronize()
+                    for lease in scratch { lease.finishAfterSynchronization() }
+                }
+                try withError { eval(arrays + scratch.flatMap(\.evaluationTargets)) }
+                if let evaluation = recurrentEvaluation {
+                    try evaluation.commit()
+                    recurrentEvaluation = nil
+                }
+                return
+            }
             if let evaluation = recurrentEvaluation {
                 eval(arrays)
                 StreamOrDevice.default.stream.synchronize()
@@ -3488,13 +3556,27 @@ public final class EngineLoopV2: @unchecked Sendable {
         // graph pipelining stays bounded at two steps.
         let readbackStart = CBv2StepProfiler.enabled ? CFAbsoluteTimeGetCurrent() : 0
         var host: [Int32] = []
+        let scratchCompletion = step.quantizedScratch.flatMap(\.evaluationTargets)
         if let tokens = step.sampledTokens {
+            if !scratchCompletion.isEmpty { eval([tokens] + scratchCompletion) }
             host = tokens.asArray(Int32.self)
             CBv2CoreInstrumentation.recordHostSync()
-        } else if !step.evalTargets.isEmpty {
-            eval(step.evalTargets)
+        } else if !step.evalTargets.isEmpty || !scratchCompletion.isEmpty {
+            eval(step.evalTargets + scratchCompletion)
             CBv2CoreInstrumentation.recordHostSync()
         }
+        if !scratchCompletion.isEmpty {
+            // MLX event readiness precedes Metal completion-handler retirement
+            // of temporary buffer owners. A real stream drain is required before
+            // refunding their byte lease; an evaluated scalar is not that fence.
+            // This conservative path may reduce quantized decode pipelining.
+            Stream.gpu.synchronize()
+            Stream.cpu.synchronize()
+        }
+        // Completion roots belong only to this step (never the already queued
+        // successor). Retire workspace ownership before delivering any terminal
+        // event or returning a request's prepaid admission credit.
+        step.finishQuantizedScratchAfterSynchronization()
         if CBv2StepProfiler.enabled {
             CBv2StepProfiler.record(
                 "v2.readback.wait", seconds: CFAbsoluteTimeGetCurrent() - readbackStart)