Skip to content

Commit d665c19

Browse files
sukru tikvesstikves
authored andcommitted
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.
1 parent cba2c84 commit d665c19

1 file changed

Lines changed: 5 additions & 4 deletions

File tree

swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -496,7 +496,6 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
496496
private let randomData: MPSGraphTensorData
497497
private let topPData: MPSGraphTensorData
498498
private let minPData: MPSGraphTensorData
499-
private let execDescriptor: MPSGraphExecutableExecutionDescriptor
500499

501500
/// Testing only: Override random value for deterministic tests.
502501
var testingOnlyRandomOverride: Float?
@@ -678,7 +677,6 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
678677
shape: [1 as NSNumber],
679678
dataType: .float32
680679
)
681-
self.execDescriptor = MPSGraphExecutableExecutionDescriptor()
682680
}
683681

684682
/// Encode composite sampling asynchronously (protocol conformance).
@@ -725,7 +723,10 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
725723
cachedOutputBuffer = outputBuffer
726724
}
727725

728-
execDescriptor.completionHandler = { [outputBuffer, outputOffset] (_, error) in
726+
// Per-call descriptor — reusing one across pipelined steps corrupts
727+
// intermediate scratch buffers when multiple runAsync calls overlap.
728+
let desc = MPSGraphExecutableExecutionDescriptor()
729+
desc.completionHandler = { [outputBuffer, outputOffset] (_, error) in
729730
if let error = error {
730731
print("MPSGraph composite sampler error: \(error)")
731732
completion(0)
@@ -743,7 +744,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
743744
with: queue,
744745
inputs: [logitsData, temperatureData, randomData, topPData, minPData],
745746
results: [outputData],
746-
executionDescriptor: execDescriptor
747+
executionDescriptor: desc
747748
)
748749
}
749750

0 commit comments

Comments
 (0)