Enforce a 90% unified-memory cap + purge KV on unload

provider-swift/Sources/ProviderCore/Diagnostics/ModelFitDiagnostic.swift

@@ -25,7 +25,11 @@ public enum ModelFitDiagnostic {
     /// requirement. `estimatedMemoryGb` is the scanner's overhead-included size
     /// (the same value the runtime passes), not the raw on-disk bytes.
     public static func requiredGb(estimatedMemoryGb: Double) -> Double {
-        ModelLoadAdmission.requiredToLoadGb(weightsGb: estimatedMemoryGb)
+        // Cap-aware headroom (activation reserve + min serveable KV) so the
+        // doctor's "needs ~X GB" matches what the runtime load gate requires.
+        ModelLoadAdmission.requiredToLoadGb(
+            weightsGb: estimatedMemoryGb,
+            headroomGb: Double(UnifiedMemoryCap.loadHeadroomBytes()) / (1024.0 * 1024.0 * 1024.0))
     }
 
     /// The memory (GB) the provider would actually have free to load a model,
@@ -44,12 +48,18 @@ public enum ModelFitDiagnostic {
         gpuActiveGb: Double = 0,
         gpuCacheGb: Double = 0
     ) -> Double {
-        ModelLoadAdmission.freeForLoadGb(
-            totalBytes: bytes(totalGb),
+        let totalBytes = bytes(totalGb)
+        // Same cap-implied reserve the running gate uses (max(configReserve,
+        // physical − 90% cap)), so the doctor verdict matches what the daemon
+        // actually enforces and never reports "fits" for a model the cap refuses.
+        let reserve = UnifiedMemoryCap.loadReserveBytes(
+            physicalBytes: totalBytes, configReserveBytes: bytes(reserveGb))
+        return ModelLoadAdmission.freeForLoadGb(
+            totalBytes: totalBytes,
             systemAvailableBytes: systemAvailableGb.map(bytes) ?? .max,
             gpuActiveBytes: bytes(gpuActiveGb),
             gpuCacheBytes: bytes(gpuCacheGb),
-            reserveBytes: bytes(reserveGb),
+            reserveBytes: reserve,
             outstandingReservationBytes: 0)
     }
 

provider-swift/Sources/ProviderCore/Inference/BatchScheduler+Telemetry.swift

@@ -39,20 +39,44 @@ extension BatchScheduler {
             return staticBudget
         }
 
-        let totalMemory = Int(ProcessInfo.processInfo.physicalMemory)
-        let osReserve = 4 * 1024 * 1024 * 1024
-        let safetyMargin = totalMemory / 10
-        let globalUsed = Int(MLX.GPU.activeMemory) + Int(MLX.GPU.cacheMemory)
-        let mlxFree = max(0, totalMemory - globalUsed)
-        // Clamp the MLX-only view to real OS-free RAM (other processes' usage),
-        // same as the load gate / GlobalKVCacheBudget; else we OOM on shared boxes.
-        let osAvailable = SystemMemory.availableBytes().map { Int(min($0, UInt64(Int.max))) } ?? Int.max
-        let realFree = min(mlxFree, osAvailable)
-        let availableHeadroom = max(0, realFree - osReserve - safetyMargin)
+        // Live KV headroom under the unified 90% cap, given current MLX usage
+        // (which reflects every co-resident model's weights + KV) clamped to real
+        // OS-free RAM and net of the activation reserve. Same helper as
+        // GlobalKVCacheBudget, so the per-scheduler budget and the shared
+        // reservation gate share one ceiling instead of competing reserves.
+        let mlxUsed = UInt64(max(0, MLX.GPU.activeMemory)) + UInt64(max(0, MLX.GPU.cacheMemory))
+        let headroomBytes = UnifiedMemoryCap.liveKVHeadroomBytes(
+            mlxUsedBytes: mlxUsed,
+            systemAvailableBytes: SystemMemory.availableBytes() ?? .max)
+        let availableHeadroom = Int(min(headroomBytes, UInt64(Int.max)))
         let liveBudget = activeTokenBudgetUsed + (availableHeadroom / kvBytesPerToken)
         return max(1024, min(staticBudget, liveBudget))
     }
 
+    /// MEASURED live KV headroom in bytes right now — `liveKVHeadroomBytes`
+    /// computed from actual MLX usage (`active + cache`, reflecting this model's
+    /// just-loaded weights plus any co-resident model). Unlike ``tokenBudgetMax``
+    /// this applies NO 1024-token floor, so it reports a true zero when the cap is
+    /// already exhausted. Used by the post-load guard to reject a model that
+    /// loaded but has no room to serve (the load gate admits on an ESTIMATE, so a
+    /// model whose real residency exceeds the estimate can land here with no KV).
+    var measuredLiveKVHeadroomBytes: UInt64 {
+        let mlxUsed = UInt64(max(0, MLX.GPU.activeMemory)) + UInt64(max(0, MLX.GPU.cacheMemory))
+        return UnifiedMemoryCap.liveKVHeadroomBytes(
+            mlxUsedBytes: mlxUsed,
+            systemAvailableBytes: SystemMemory.availableBytes() ?? .max)
+    }
+
+    /// Post-load guard: true iff this freshly-loaded model has at least the
+    /// minimum serveable KV headroom under the cap. When false, the caller must
+    /// unload + clearCache + reject — keeping the model resident would just
+    /// reject every request at the KV-reservation gate (a "loaded but
+    /// unserveable" model). Catches the case where measured residency exceeds the
+    /// load gate's `estimatedMemoryGb` estimate.
+    func hasServeableKVHeadroom() -> Bool {
+        UnifiedMemoryCap.loadIsServeable(measuredLiveKVHeadroomBytes: measuredLiveKVHeadroomBytes)
+    }
+
     /// Sum of `(promptTokens + maxTokens)` across active bridges. This
     /// is the value the P1 cumulative-budget gate in `submit()` checks
     /// against `tokenBudgetMax`.

provider-swift/Sources/ProviderCore/Inference/BatchScheduler.swift

@@ -1485,12 +1485,16 @@ public actor BatchScheduler {
             architecture: snapshot.architecture,
             weightBytes: snapshot.bytes
         )
-        let totalMemory = Int(ProcessInfo.processInfo.physicalMemory)
-        let osReserve = 4 * 1024 * 1024 * 1024
-        let safetyMargin = totalMemory / 10
-        let availableForKV = totalMemory - snapshot.bytes - osReserve - safetyMargin
+        // Static upper-bound budget from the unified 90% cap minus THIS model's
+        // measured resident weights (snapshot.bytes) and the activation reserve.
+        // Only the per-model clamp; cross-model headroom (other resident models'
+        // weights/KV) is handled live by tokenBudgetMax / the shared
+        // GlobalKVCacheBudget, which read process-global MLX usage.
+        let availableForKV = UnifiedMemoryCap.kvBudgetBytes(
+            residentWeightBytes: UInt64(max(0, snapshot.bytes)))
         if availableForKV > 0 && kvBytesPerToken > 0 {
-            self.dynamicTokenBudgetMax = max(availableForKV / kvBytesPerToken, 1024)
+            let availInt = Int(min(availableForKV, UInt64(Int.max)))
+            self.dynamicTokenBudgetMax = max(availInt / kvBytesPerToken, 1024)
         } else {
             self.dynamicTokenBudgetMax = 1024
         }
@@ -1514,6 +1518,64 @@ public actor BatchScheduler {
         await stopCurrentEngine()
     }
 
+    /// Reserve unified memory for a VLM (vision-path) request against the shared
+    /// 90% cap, via the process-wide GlobalKVCacheBudget this scheduler holds. A
+    /// vision request bypasses the batched `submitTokenized` reservation entirely
+    /// — it streams through `container.generate` directly — so without this it
+    /// commits TWO kinds of memory the cap would otherwise track only reactively:
+    ///
+    /// 1. `mediaDecodeBytes` — the transient CIImage rasters + Swift `Data` pixel
+    ///    buffers from media decode. These are NOT MLXArrays, so they are
+    ///    invisible to the cap's live MLX counters (the original blind spot).
+    /// 2. The generation KV cache — `kvBytesPerToken × maxOutputTokens`. This IS
+    ///    MLXArray-backed (eventually visible to the live counters), but the
+    ///    vision path's decode loop runs in a detached task with no per-request
+    ///    reservation, so N concurrent media requests can grow KV simultaneously
+    ///    against headroom none of them reserved — a transient over-commit the
+    ///    cap would otherwise catch only on the NEXT admission. Reserving it up
+    ///    front makes the vision path share the same preemptive 90% gate the
+    ///    batched path gets from `reserveKVForRequest`.
+    ///
+    /// Both are charged to ONE reservation id and released together when the
+    /// stream ends (decode buffers are actually freed after `prepare`, so holding
+    /// them for the whole stream is conservative — never an under-reservation).
+    /// Returns true if it fits (and was reserved) or budgeting is disabled
+    /// (nil budget, legacy "always proceed"); false if it would exceed the cap,
+    /// in which case the caller surfaces a retryable 503. Pair with
+    /// `releaseVisionRequest`. Saturating; never traps.
+    public func reserveVisionRequest(
+        requestId: String, mediaDecodeBytes: UInt64, kvTokens: Int
+    ) async -> Bool {
+        guard let kvBudget else { return true }
+        // KV bytes = kvBytesPerToken × the FULL token span the cache will hold:
+        // prompt text + image/video soft tokens + generated output (the caller
+        // computes that conservative total). Reserving only the output tokens
+        // would badly under-count — a single image expands to hundreds of vision
+        // tokens, all of which occupy KV.
+        var genKVBytes: UInt64 = 0
+        if kvBytesPerToken > 0, kvTokens > 0 {
+            let (b, overflow) = UInt64(kvBytesPerToken)
+                .multipliedReportingOverflow(by: UInt64(kvTokens))
+            genKVBytes = overflow ? .max : b
+        }
+        let (total, overflow) = mediaDecodeBytes.addingReportingOverflow(genKVBytes)
+        let bytes = overflow ? UInt64.max : total
+        return await kvBudget.reserveBytes(requestID: requestId, bytes: bytes)
+    }
+
+    /// The model's configured context window (`max_position_embeddings`), or 0 if
+    /// unknown. The KV cache can never hold more than this many prompt+vision
+    /// tokens, so the vision-path reservation clamps its prompt+vision estimate to
+    /// it (output tokens are added on top, matching the batched path's
+    /// `promptTokenCount + maxTokens`).
+    public func contextLength() -> Int { maxContextLength }
+
+    /// Release a prior `reserveVisionRequest` reservation. Safe/no-op if unknown
+    /// or budgeting is disabled.
+    public func releaseVisionRequest(requestId: String) async {
+        await kvBudget?.release(requestID: requestId)
+    }
+
     // MARK: - Submit / cancel
 
     /// Submit a pre-tokenized prompt. Used by `MultiModelBatchSchedulerEngine`
@@ -1944,33 +2006,31 @@ public actor BatchScheduler {
         self.engine = nil
         modelContainer = nil
         tokenizer = nil
-        // Drain the bounded capture pipeline: the engine is stopped (no more
-        // capture hooks fire), so finish the stream and cancel the consumer.
-        // This releases any retained KV snapshots and stops an in-flight
-        // `mgr.store` from racing the index flush / accountant deregister below,
-        // and guarantees no snapshot (or consumer Task) leaks across a swap.
+        // Drain the bounded capture pipeline FIRST (#374): the engine is stopped
+        // (no more capture hooks fire), so finish the stream and cancel the
+        // consumer. This releases retained KV snapshots and stops an in-flight
+        // `mgr.store` from racing the purge below.
         capturePipeline?.shutdown()
         capturePipeline = nil
-        // Persist any coalesced index writes before dropping the manager, so
-        // checkpoints written since the last coalesced save survive restart.
+        // Then purge this model's KV from BOTH RAM and SSD on unload (#363) —
+        // restart warmth is intentionally OFF, so no KV (memory or disk) outlives
+        // the loaded model. purgeOnUnload drains in-flight writes, clears the RAM
+        // tier, deletes the kv/<modelKey> dir, and deregisters the accountant
+        // (subsumes the old flushIndexNow + deregisterFromAccountant).
         if let mgr = checkpointManager {
-            await mgr.flushIndexNow()
-            // Phase 3: deregister from the accountant before dropping the manager.
-            await mgr.deregisterFromAccountant()
+            await mgr.purgeOnUnload()
         }
         // Drop the checkpoint manager so a stale one can't serve the next
         // model (the new model's loadModel reinstalls its own, or nil).
         checkpointManager = nil
         checkpointBoundaries = []
         checkpointLayerSignatures = []
 
-        // Close the engine-tier owner FIRST (before deregister) so no disk
-        // mutation slips through between deregistration and the dir being handed
-        // to a reloaded same-modelKey owner: a stale engine step finishing after
-        // `engine.stop()` (which doesn't fence an in-flight engineQueue step) or
-        // a late accountant eviction signal will now no-op. `engine.stop()` was
-        // already awaited above, so the GPU step loop is winding down by here.
-        engineTierOwner?.close()
+        // Purge the engine-tier owner's on-disk dir too (same kv/<modelKey> dir;
+        // whichever tier ran first already removed it, so this no-ops then).
+        // purgeDir latches `closed` first so any in-flight engine-step save that
+        // resumes after `engine.stop()` no-ops at its post-write bail.
+        engineTierOwner?.purgeDir()
         // Deregister the engine-tier owner from the accountant.
         if let accountant = diskAccountant, let token = engineTierAccountantToken {
             await accountant.deregister(token)

provider-swift/Sources/ProviderCore/Inference/GlobalKVCacheBudget.swift

@@ -14,14 +14,27 @@ public actor GlobalKVCacheBudget {
         public var systemAvailable: UInt64
     }
 
-    private let safetyFactor: Double
-    private let reserveBytes: UInt64
+    /// Cap fraction and activation reserve are nil → ``UnifiedMemoryCap``
+    /// defaults (0.90 / env / 3 GiB floor). Held as overrides so tests can pin
+    /// them; production uses the defaults so this budget and the load gate share
+    /// one policy.
+    private let capFraction: Double?
+    private let activationReserveBytes: UInt64?
+    /// Operator-configured reserve (`memory_reserve_gb`, in bytes). Held back by
+    /// the live KV gate just as the load gate holds it back, so runtime KV can't
+    /// grow into memory the operator reserved. 0 = no extra reserve (cap only).
+    private let configReserveBytes: UInt64
     private let memorySnapshot: @Sendable () -> MemorySnapshot
     private var reservations: [String: UInt64] = [:]
 
-    public init(reserveBytes: UInt64 = 0, safetyFactor: Double = 0.7) {
-        self.reserveBytes = reserveBytes
-        self.safetyFactor = Self.clampedSafetyFactor(safetyFactor)
+    public init(
+        capFraction: Double? = nil,
+        activationReserveBytes: UInt64? = nil,
+        configReserveBytes: UInt64 = 0
+    ) {
+        self.capFraction = capFraction
+        self.activationReserveBytes = activationReserveBytes
+        self.configReserveBytes = configReserveBytes
         self.memorySnapshot = {
             MemorySnapshot(
                 total: ProcessInfo.processInfo.physicalMemory,
@@ -33,12 +46,14 @@ public actor GlobalKVCacheBudget {
     }
 
     init(
-        reserveBytes: UInt64 = 0,
-        safetyFactor: Double = 0.7,
+        capFraction: Double? = nil,
+        activationReserveBytes: UInt64? = nil,
+        configReserveBytes: UInt64 = 0,
         memorySnapshot: @escaping @Sendable () -> MemorySnapshot
     ) {
-        self.reserveBytes = reserveBytes
-        self.safetyFactor = Self.clampedSafetyFactor(safetyFactor)
+        self.capFraction = capFraction
+        self.activationReserveBytes = activationReserveBytes
+        self.configReserveBytes = configReserveBytes
         self.memorySnapshot = memorySnapshot
     }
 
@@ -57,6 +72,42 @@ public actor GlobalKVCacheBudget {
         reservations.removeValue(forKey: requestID)
     }
 
+    /// Reserve an arbitrary BYTE amount against the same live cap headroom KV
+    /// uses. For non-KV unified-memory consumers that the cap would otherwise be
+    /// blind to — notably VLM media decode (CIImage rasters + Swift Data pixel
+    /// buffers live in the same unified RAM as MLX arrays but are NOT counted by
+    /// MLX.GPU.active/cache). Reserving here makes those bytes share the 90% cap:
+    /// the decode is admitted only if it fits alongside resident weights + KV +
+    /// activations, and rejected (caller surfaces 429/retry) otherwise. Returns
+    /// false if it won't fit or the id is already reserved. Pair with `release`.
+    public func reserveBytes(requestID: String, bytes: UInt64) -> Bool {
+        guard bytes > 0 else { return false }
+        guard reservations[requestID] == nil else { return false }
+        if bytes > availableReservationBytes() { return false }
+        reservations[requestID] = bytes
+        return true
+    }
+
+    /// Reserve a loading model's WEIGHT footprint for the duration of its load,
+    /// unconditionally. A model's weights are not yet visible in MLX active/cache
+    /// while `loadModelContainer` is still allocating them, so a KV reservation
+    /// granted on an ALREADY-loaded model during that window would compute its
+    /// headroom blind to the incoming weights and could push total usage past the
+    /// cap — a transient OOM on the normal serve-while-load path. Reserving the
+    /// footprint here makes those in-flight weights visible to `reserve` /
+    /// `reserveBytes`, so concurrent KV can only claim `headroom − weights`.
+    ///
+    /// Unconditional (never fails): the load gate has already admitted the model,
+    /// so this is bookkeeping for the load that WILL happen, not a second gate.
+    /// It reserves only the weight estimate, so concurrent KV that still fits
+    /// underneath is admitted; only reservations that would over-commit are
+    /// rejected (caller surfaces 429/retry). Released once the weights are
+    /// resident (and thus reflected in `mlxUsed`). Pair with `release`.
+    public func reservePendingLoad(requestID: String, bytes: UInt64) {
+        guard bytes > 0 else { return }
+        reservations[requestID] = bytes
+    }
+
     /// Atomically shrink an existing reservation to a smaller byte count,
     /// freeing the difference. `reserve`/`release` cannot express a shrink:
     /// `reserve` refuses when an entry already exists for the id, and a
@@ -88,14 +139,20 @@ public actor GlobalKVCacheBudget {
     private func availableReservationBytes() -> UInt64 {
         let snap = memorySnapshot()
         let mlxUsed = Self.saturatingAdd(snap.active, snap.cache)
-        let mlxFree = snap.total > mlxUsed ? snap.total - mlxUsed : 0
-        // Clamp the MLX-only view (blind to other processes) to real OS-free RAM,
-        // mirroring the load gate (ModelLoadAdmission.freeForLoadGb). Without it
-        // the runtime admits against memory other apps hold → jetsam OOM.
-        let realFree = min(mlxFree, snap.systemAvailable)
-        let usable = realFree > reserveBytes ? realFree - reserveBytes : 0
-        let capped = Double(usable) * safetyFactor
-        let reservationCap = capped >= Double(UInt64.max) ? UInt64.max : UInt64(capped)
+        // Bytes still committable to KV under the 90% unified-memory cap, given
+        // current MLX usage (which already reflects ALL co-resident models'
+        // weights + KV), clamped to real OS-free RAM and net of the activation
+        // reserve. This replaces the old `(free − reserve) × 0.7` formula: the
+        // single cap + activation reserve are the only knobs, so this gate, the
+        // per-scheduler live token budget, and the load gate no longer apply
+        // three different, competing discounts.
+        let reservationCap = UnifiedMemoryCap.liveKVHeadroomBytes(
+            physicalBytes: snap.total,
+            mlxUsedBytes: mlxUsed,
+            systemAvailableBytes: snap.systemAvailable,
+            activationReserveBytes: activationReserveBytes,
+            configReserveBytes: configReserveBytes,
+            capFraction: capFraction)
         let reserved = reservations.values.reduce(UInt64(0)) { partial, value in
             let (sum, overflow) = partial.addingReportingOverflow(value)
             return overflow ? UInt64.max : sum
@@ -113,8 +170,4 @@ public actor GlobalKVCacheBudget {
         return total
     }
 
-    private static func clampedSafetyFactor(_ value: Double) -> Double {
-        guard value.isFinite else { return 0.7 }
-        return min(1.0, max(0.0, value))
-    }
 }