Speed up prefill on pre-Apple10 Macs

Sources/TurboFieldfare/Infrastructure/Metal/MetalContext.swift

@@ -7,6 +7,7 @@ enum MetalError: Error, CustomStringConvertible {
     case missingShaderResource(String)
     case missingFunction(String)
     case libraryCompileFailed(String)
+    case commandBufferFailed(String)
 
     public var description: String {
         switch self {
@@ -15,14 +16,54 @@ enum MetalError: Error, CustomStringConvertible {
         case .missingShaderResource(let n): return "Shader resource missing: \(n)"
         case .missingFunction(let n):     return "Metal function missing in library: \(n)"
         case .libraryCompileFailed(let s):return "Metal library compile failed: \(s)"
+        case .commandBufferFailed(let s): return "Metal command buffer failed: \(s)"
         }
     }
 }
 
-func checkCommandBufferError(_ error: (any Error)?) throws {
+func metalCommandBufferStatusName(_ status: MTLCommandBufferStatus) -> String {
+    switch status {
+    case .notEnqueued: return "notEnqueued"
+    case .enqueued:    return "enqueued"
+    case .committed:   return "committed"
+    case .scheduled:   return "scheduled"
+    case .completed:   return "completed"
+    case .error:       return "error"
+    @unknown default:  return "unknown(\(status.rawValue))"
+    }
+}
+
+/// Builds the diagnostic for a command buffer that did not complete, or nil
+/// when it did. Status is checked separately because a failed buffer is not
+/// guaranteed to carry an error object.
+func metalCommandBufferFailureDetail(label: String?,
+                                     status: MTLCommandBufferStatus,
+                                     error: (any Error)?) -> String? {
+    if status == .completed && error == nil { return nil }
+
+    var parts = ["label=\(label.map { $0.isEmpty ? "<empty>" : $0 } ?? "<none>")"]
+    parts.append("status=\(metalCommandBufferStatusName(status))")
     if let error {
-        throw error
+        let nsError = error as NSError
+        parts.append("domain=\(nsError.domain)")
+        parts.append("code=\(nsError.code)")
+        parts.append("description=\(nsError.localizedDescription)")
+        if !nsError.userInfo.isEmpty {
+            parts.append("userInfoKeys=\(nsError.userInfo.keys.sorted().joined(separator: ","))")
+        }
+    } else {
+        parts.append("error=<none>")
+    }
+    return parts.joined(separator: " ")
+}
+
+func checkCommandBufferError(_ commandBuffer: MTLCommandBuffer) throws {
+    guard let detail = metalCommandBufferFailureDetail(label: commandBuffer.label,
+                                                       status: commandBuffer.status,
+                                                       error: commandBuffer.error) else {
+        return
     }
+    throw MetalError.commandBufferFailed(detail)
 }
 
 public struct MetalFunctionConstant: Hashable, Sendable {
@@ -62,8 +103,25 @@ public final class MetalContext: @unchecked Sendable {
     private var pipelineCache: [PipelineCacheKey: MTLComputePipelineState] = [:]
     private let pipelineCacheLock = NSLock()
 
+    private static func relaxInteractivityWatchdog() {
+        #if os(macOS)
+        // The AGX driver reads this once at first device creation. Long prefill
+        // dispatches can otherwise be killed as compositor-impacting on macOS
+        // 26. Overwrite 0 preserves an operator's explicit stock-behaviour
+        // override. This relaxes the deadline; it does not guarantee survival.
+        setenv("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1", 0)
+        #endif
+    }
+
+    /// Routes every production device creation through the watchdog mitigation
+    /// before the AGX driver's process-wide one-time environment read.
+    public static func makeSystemDefaultDevice() -> MTLDevice? {
+        relaxInteractivityWatchdog()
+        return MTLCreateSystemDefaultDevice()
+    }
+
     public init() throws {
-        guard let dev = MTLCreateSystemDefaultDevice() else { throw MetalError.noDevice }
+        guard let dev = Self.makeSystemDefaultDevice() else { throw MetalError.noDevice }
         guard let q   = dev.makeCommandQueue()           else { throw MetalError.noQueue }
         self.device  = dev
         self.queue   = q

Sources/TurboFieldfare/Kernels/Attention/PrefillAttention.swift

@@ -57,13 +57,33 @@ final class PrefillAttention {
     private let psoParamsSmoke: MTLComputePipelineState
     private let psoFullTensorOps2DValidityV2: MTLComputePipelineState?
 
-    init(context: MetalContext) throws {
+    var tensorOps2DValidityV2Available: Bool {
+        psoFullTensorOps2DValidityV2 != nil
+    }
+
+    convenience init(context: MetalContext) throws {
+        try self.init(context: context, simulatingMissingTensorOps: false)
+    }
+
+    /// Tests can force the production fallback even on a host where the MSL 4
+    /// TensorOps pipeline builds.
+    init(context: MetalContext, simulatingMissingTensorOps: Bool) throws {
         self.context = context
         self.psoCausalTiled = try context.pipeline("attention_prefill_causal_tiled")
         self.psoParamsSmoke = try context.pipeline("prefill_attention_params_smoke")
-        self.psoFullTensorOps2DValidityV2 = context.device.supportsFamily(.apple10)
-            ? try? context.pipeline("attention_prefill_full_tensorops_2d_validity_v2")
-            : nil
+        if simulatingMissingTensorOps {
+            self.psoFullTensorOps2DValidityV2 = nil
+        } else {
+            do {
+                self.psoFullTensorOps2DValidityV2 = try context.pipeline(
+                    "attention_prefill_full_tensorops_2d_validity_v2")
+            } catch {
+                self.psoFullTensorOps2DValidityV2 = nil
+                FileHandle.standardError.write(Data(
+                    ("PrefillAttention: TensorOps 2D pipeline unavailable; "
+                     + "using causal-tiled fallback: \(error)\n").utf8))
+            }
+        }
     }
 
     func encodeCausal(commandBuffer: MTLCommandBuffer,
@@ -87,11 +107,14 @@ final class PrefillAttention {
 
         let requestsTensorOps = path == .fullTensorOps2DPreferred
             || path == .fullTensorOps2DValidityV2
-        // The pinned model uses 512/16/2 only for full attention; its
-        // sliding-window layers use 256/16/8. A future model that reuses this
-        // shape for sliding attention must add a full-visibility check here.
+        // TensorOps starts its key loop at zero and ignores slidingWindow, so
+        // these are visibility guards rather than shape optimizations.
+        let windowNeverClips = effectiveParams.slidingWindow == 0
+            || effectiveParams.slidingWindow >= effectiveParams.kvValidCount
         let tensorOpsShape = requestsTensorOps
+            && layerKind == .full
             && kvRingCapacity == 0
+            && windowNeverClips
             && effectiveParams.headDim == 512
             && effectiveParams.numQHeads == 16
             && effectiveParams.numKVHeads == 2
@@ -103,7 +126,7 @@ final class PrefillAttention {
             pipeline = tensorOpsPipeline
         } else if tensorOpsShape && path == .fullTensorOps2DValidityV2 {
             preconditionFailure(
-                "TensorOps 2D prefill attention requires Apple10 MPP tensor support")
+                "TensorOps 2D prefill attention pipeline is unavailable on this Metal stack")
         } else {
             // Explicit mode also falls back for incompatible shapes. Benchmark
             // fixtures must use 512/16/2 to prove that TensorOps ran.

Sources/TurboFieldfare/Kernels/Vision/VisionResize.swift

@@ -171,7 +171,7 @@ public final class VisionResize {
 
         commandBuffer.commit()
         commandBuffer.waitUntilCompleted()
-        try checkCommandBufferError(commandBuffer.error)
+        try checkCommandBufferError(commandBuffer)
 
         let finished = sourceHeight == destinationHeight ? intermediate : output
         let rowBytes = sourceHeight == destinationHeight

Sources/TurboFieldfare/Runtime/Generation/RawCompletion.swift

@@ -307,6 +307,6 @@ private func sampleOnce(scratch: RawCompletionScratch, context: MetalContext,
                            history: history, config: config, position: position,
                            outToken: scratch.outToken)
     cb.commit(); cb.waitUntilCompleted()
-    try checkCommandBufferError(cb.error)
+    try checkCommandBufferError(cb)
     return Int32(bitPattern: scratch.outToken.contents().load(as: UInt32.self))
 }

Sources/TurboFieldfare/Runtime/Inference/RealForwardRunner.swift

@@ -861,6 +861,7 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
         guard var cb = ctx.queue.makeCommandBuffer() else {
             throw ModelError.residentBufferWrapFailed
         }
+        cb.label = "prefill start=\(startPosition) count=\(t) phase=embed"
         if let embeddingOverride {
             // The override spans the whole chunk — the planner emits image spans
             // as standalone chunks — so the INT4 gather it would cover is
@@ -1101,6 +1102,8 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
                     guard let sharedCB = ctx.queue.makeCommandBuffer() else {
                         throw ModelError.residentBufferWrapFailed
                     }
+                    sharedCB.label =
+                        "prefill start=\(startPosition) count=\(t) layer=\(L) phase=shared_expert"
                     let sharedProj = sharedExpertProjections[L]
                     try prefillSharedExpert.encodeBlock(commandBuffer: sharedCB,
                                                         x: scratch.denseX,
@@ -1241,6 +1244,8 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
                         guard let tileCB = ctx.queue.makeCommandBuffer() else {
                             throw ModelError.residentBufferWrapFailed
                         }
+                        tileCB.label =
+                            "prefill start=\(startPosition) count=\(t) layer=\(L) phase=routed_tile"
                         _ = prefillGroupedMoE.encodeStreamedBatched(
                             commandBuffer: tileCB,
                             hidden: scratch.routedX,
@@ -1300,6 +1305,8 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
                         guard let nextCB = ctx.queue.makeCommandBuffer() else {
                             throw ModelError.residentBufferWrapFailed
                         }
+                        nextCB.label =
+                            "prefill start=\(startPosition) count=\(t) layer=\(L + 1) phase=qkv_attention"
                         cb = nextCB
                     }
                     continue
@@ -1311,6 +1318,7 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
             guard let finalCB = ctx.queue.makeCommandBuffer() else {
                 throw ModelError.residentBufferWrapFailed
             }
+            finalCB.label = "prefill start=\(startPosition) count=\(t) phase=final_head"
             if outputMode == .greedyIfAvailable, useFusedGreedyHead {
                 fusionHead.encodeGreedyDecode(
                     commandBuffer: finalCB,
@@ -1395,12 +1403,12 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
                 waitUntilCompleted(pending.cb)
             }
             if let sharedCB = pending.sharedCB {
-                try checkCommandBufferError(sharedCB.error)
+                try checkCommandBufferError(sharedCB)
             }
             if let phase1HitCB = pending.phase1HitCB {
-                try checkCommandBufferError(phase1HitCB.error)
+                try checkCommandBufferError(phase1HitCB)
             }
-            try checkCommandBufferError(pending.cb.error)
+            try checkCommandBufferError(pending.cb)
             totalCb2Nanos &+= pending.encodeAndCommitNanos
         }
 
@@ -1595,7 +1603,7 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
                 try finishPendingRoutedCommand(pending, waitIfNeeded: false)
                 pendingRoutedCommand = nil
             }
-            try checkCommandBufferError(cb.error)
+            try checkCommandBufferError(cb)
             totalCb1Nanos &+= clock_gettime_nsec_np(CLOCK_UPTIME_RAW) - tCb1Start - waitNanos
 
             // CPU readback to fetch routed-expert blobs from disk.
@@ -1873,7 +1881,7 @@ public final class RealForwardRunner: ChunkedPrefillRunner, MultimodalPrefillRun
 
     private nonisolated func waitForCompletion(_ cb: MTLCommandBuffer) throws {
         waitUntilCompleted(cb)
-        try checkCommandBufferError(cb.error)
+        try checkCommandBufferError(cb)
     }
 
     private nonisolated func waitUntilCompleted(_ cb: MTLCommandBuffer) {

Sources/TurboFieldfare/Runtime/Vision/VisionRuntime.swift

@@ -244,7 +244,7 @@ public final class VisionRuntime {
     }
 
     public static var isSupportedOnDefaultDevice: Bool {
-        guard let device = MTLCreateSystemDefaultDevice() else { return false }
+        guard let device = MetalContext.makeSystemDefaultDevice() else { return false }
         return isSupported(on: device)
     }
 
@@ -519,8 +519,10 @@ public final class VisionRuntime {
             let started = clock_gettime_nsec_np(CLOCK_UPTIME_RAW)
             entry.buffer.waitUntilCompleted()
             gpuWaitNanoseconds += clock_gettime_nsec_np(CLOCK_UPTIME_RAW) - started
-            if let error = entry.buffer.error {
-                throw VisionRuntimeError.commandFailed(String(describing: error))
+            if let detail = metalCommandBufferFailureDetail(label: entry.buffer.label,
+                                                            status: entry.buffer.status,
+                                                            error: entry.buffer.error) {
+                throw VisionRuntimeError.commandFailed(detail)
             }
             gpuNanoseconds += UInt64(
                 max(0, entry.buffer.gpuEndTime - entry.buffer.gpuStartTime)
@@ -844,17 +846,21 @@ public final class VisionRuntime {
         let commandBuffer = try makeCommandBuffer()
         commandBuffer.commit()
         commandBuffer.waitUntilCompleted()
-        if let error = commandBuffer.error {
-            throw VisionRuntimeError.commandFailed(String(describing: error))
+        if let detail = metalCommandBufferFailureDetail(label: commandBuffer.label,
+                                                        status: commandBuffer.status,
+                                                        error: commandBuffer.error) {
+            throw VisionRuntimeError.commandFailed(detail)
         }
         return clock_gettime_nsec_np(CLOCK_UPTIME_RAW) - started
     }
 
     private func commitAndMeasure(_ commandBuffer: MTLCommandBuffer) throws -> UInt64 {
         commandBuffer.commit()
         commandBuffer.waitUntilCompleted()
-        if let error = commandBuffer.error {
-            throw VisionRuntimeError.commandFailed(String(describing: error))
+        if let detail = metalCommandBufferFailureDetail(label: commandBuffer.label,
+                                                        status: commandBuffer.status,
+                                                        error: commandBuffer.error) {
+            throw VisionRuntimeError.commandFailed(detail)
         }
         return UInt64(max(0, commandBuffer.gpuEndTime - commandBuffer.gpuStartTime)
                       * 1_000_000_000)

Sources/TurboFieldfareCLI/Run.swift

@@ -55,7 +55,7 @@ public func run(args: Args,
         let input = try parseInput(args: args)
         var selectedDevice: MTLDevice?
         if input.hasImages {
-            guard let device = MTLCreateSystemDefaultDevice() else {
+            guard let device = MetalContext.makeSystemDefaultDevice() else {
                 return errored(stderr, "no Metal device", 1)
             }
             try VisionRuntime.requireSupportedDevice(device)
@@ -103,7 +103,7 @@ public func run(args: Args,
         // Hoisted above the `auto` estimate, which needs a device to read image
         // geometry. Planning never touches the GPU, but building the plan does
         // need the device the run will use.
-        guard let device = selectedDevice ?? MTLCreateSystemDefaultDevice() else {
+        guard let device = selectedDevice ?? MetalContext.makeSystemDefaultDevice() else {
             return errored(stderr, "no Metal device", 1)
         }
 

Tests/TurboFieldfare/Core/Infrastructure/InteractivityWatchdogEnvironmentTests.swift

@@ -0,0 +1,52 @@
+import Darwin
+import Testing
+
+@testable import TurboFieldfare
+
+@Suite(.serialized) struct InteractivityWatchdogEnvironmentTests {
+    private static let key = "AGX_RELAX_CDM_CTXSTORE_TIMEOUT"
+
+    private static func read() -> String? {
+        guard let raw = getenv(key) else { return nil }
+        return String(cString: raw)
+    }
+
+    private static func withCleanEnvironment(_ body: () throws -> Void) rethrows {
+        let saved = read()
+        unsetenv(key)
+        defer {
+            if let saved {
+                setenv(key, saved, 1)
+            } else {
+                unsetenv(key)
+            }
+        }
+        try body()
+    }
+
+    #if os(macOS)
+    @Test func constructingAMetalContextSetsTheWatchdogRelaxationOnMacOS() throws {
+        try Self.withCleanEnvironment {
+            #expect(Self.read() == nil)
+            _ = try MetalContext()
+            #expect(Self.read() == "1")
+        }
+    }
+
+    @Test func anExplicitOperatorOverrideIsNotOverwritten() throws {
+        try Self.withCleanEnvironment {
+            setenv(Self.key, "0", 1)
+            _ = try MetalContext()
+            #expect(Self.read() == "0")
+        }
+    }
+    #else
+    @Test func constructingAMetalContextDoesNotSetTheMacOSWatchdogRelaxation() throws {
+        try Self.withCleanEnvironment {
+            #expect(Self.read() == nil)
+            _ = try MetalContext()
+            #expect(Self.read() == nil)
+        }
+    }
+    #endif
+}