From d665c19180f0d93b15aecabdf2993f282fc1d1d8 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Tue, 21 Jul 2026 20:08:22 -0700 Subject: [PATCH] Fix pipelined sampling corruption: use per-call execution descriptor The MPSGraphCompositeSampler reused a single MPSGraphExecutableExecutionDescriptor across all pipelined steps. Under pipelined execution (depth > 1), overlapping runAsync calls on the same executable corrupt intermediate scratch buffers when sharing a descriptor, producing garbled output (word repetitions, doubled punctuation) with temperature > 0. Create a fresh descriptor per encode() call, matching the pattern MPSGraphArgmaxSampler already uses. --- .../CoreAILanguageModels/Samplers/MPSGraphSamplers.swift | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift index 041c83c..a38e64e 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift @@ -496,7 +496,6 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { private let randomData: MPSGraphTensorData private let topPData: MPSGraphTensorData private let minPData: MPSGraphTensorData - private let execDescriptor: MPSGraphExecutableExecutionDescriptor /// Testing only: Override random value for deterministic tests. var testingOnlyRandomOverride: Float? @@ -678,7 +677,6 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { shape: [1 as NSNumber], dataType: .float32 ) - self.execDescriptor = MPSGraphExecutableExecutionDescriptor() } /// Encode composite sampling asynchronously (protocol conformance). @@ -725,7 +723,10 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { cachedOutputBuffer = outputBuffer } - execDescriptor.completionHandler = { [outputBuffer, outputOffset] (_, error) in + // Per-call descriptor — reusing one across pipelined steps corrupts + // intermediate scratch buffers when multiple runAsync calls overlap. + let desc = MPSGraphExecutableExecutionDescriptor() + desc.completionHandler = { [outputBuffer, outputOffset] (_, error) in if let error = error { print("MPSGraph composite sampler error: \(error)") completion(0) @@ -743,7 +744,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { with: queue, inputs: [logitsData, temperatureData, randomData, topPData, minPData], results: [outputData], - executionDescriptor: execDescriptor + executionDescriptor: desc ) }