Skip to content

Commit 53b36cf

Browse files
stikvesstikves
committed
Add GPU-based constrained sampling for pipelined engine
Enable grammar-constrained generation (JSON schema enforcement) on the pipelined engine by applying xgrammar bitmasks directly in the MPSGraph GPU sampler, eliminating the per-token logit transfer to CPU. Key components: - ConstrainedGenerationCapable protocol: capability signal for routing constrained generation to engines that support GPU-side bitmask application. Replaces concrete type checks with protocol-based dispatch. - ConstrainedSessionHandle: encapsulated, narrow-API handle to the grammar session. Private session storage with only the 5 operations the engine loop needs exposed. Single-writer safety guaranteed by the engine acquire/release gate. - PipelinedConstrainedDecodingStrategy: AsyncSequence-based strategy that drives the GPU constrained loop. SingleUseFlag iteration guard, proper error propagation, Task.checkCancellation, early-stop on consumer drop. - Session cache (checkout/checkin): sessions reused across calls with the same schema. Engine Task defer block owns return-to-cache, ensuring the handle is never accessed concurrently. - Bitmask expansion in MPSGraph samplers: bitwise AND + notEqual + cast graph with zero measurable overhead. Both ArgmaxSampler and CompositeSampler support constrained path via applyBitmask parameter. - ConstrainedGenerationSession.fillBitmask(into:): zero-copy write of grammar bitmask directly into GPU-visible shared MTLBuffer. - Integration tests via MockConstrainedEngine conforming to the protocol, exercising session cache, schema invalidation, stop sequences, error propagation, and maxTokens boundary without Metal hardware.
1 parent 5ed9981 commit 53b36cf

10 files changed

Lines changed: 1717 additions & 92 deletions

File tree

‎swift/Sources/CoreAILanguageModels/DecodingStrategies/PipelinedConstrainedDecodingStrategy.swift‎

Lines changed: 232 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,232 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import CoreAIShared
7+
import Foundation
8+
import Synchronization
9+
import Tokenizers
10+
11+
private let defaultMaxConstrainedTokens = 512
12+
13+
private final class SingleUseFlag: Sendable {
14+
private let value = Atomic<Bool>(false)
15+
func testAndSet() -> Bool {
16+
value.exchange(true, ordering: .acquiring)
17+
}
18+
}
19+
20+
/// Grammar-constrained decoding strategy using the GPU pipelined engine.
21+
///
22+
/// Unlike `ConstrainedDecodingStrategy` which forces the sequential engine path
23+
/// (CPU logits -> mask -> CPU sample), this strategy applies the xgrammar bitmask
24+
/// directly inside the MPSGraph sampler on GPU -- eliminating the ~296KB logit
25+
/// transfer and CPU-side softmax/sampling per token.
26+
///
27+
/// Only used when the engine conforms to `ConstrainedGenerationCapable`. The routing logic in
28+
/// `CoreAILanguageModel` selects this strategy based on engine type.
29+
public struct PipelinedConstrainedDecodingStrategy: DecodingStrategy {
30+
private let jsonSchema: String
31+
private let vocabSizeOverride: Int?
32+
33+
public init(jsonSchema: String, vocabSize: Int? = nil) {
34+
self.jsonSchema = jsonSchema
35+
self.vocabSizeOverride = vocabSize
36+
}
37+
38+
// MARK: - DecodingStrategy conformance
39+
40+
public func decode(
41+
from input: Input,
42+
tokenizer: any Tokenizer,
43+
inferenceEngine: any InferenceEngine,
44+
samplingConfiguration: SamplingConfiguration,
45+
options: InferenceOptions,
46+
stopSequences: StopSequences
47+
) async throws -> PipelinedConstrainedSequence {
48+
guard let constrainedEngine = inferenceEngine as? any ConstrainedGenerationCapable else {
49+
throw InferenceRuntimeError.invalidArgument(
50+
"PipelinedConstrainedDecodingStrategy requires a ConstrainedGenerationCapable engine")
51+
}
52+
53+
let vocabSize = vocabSizeOverride ?? ConstrainedDecodingStrategy.deriveVocabSize(from: tokenizer)
54+
guard let vocabSize else {
55+
throw InferenceRuntimeError.invalidArgument(
56+
"Cannot determine vocabulary size from tokenizer. "
57+
+ "Pass vocabSize explicitly via CoreAIRunner or LLMAsset metadata."
58+
)
59+
}
60+
61+
let singleTokenStops = stopSequences.sequences.filter { $0.count == 1 }.map { $0[0] }
62+
if stopSequences.sequences.contains(where: { $0.count > 1 }) {
63+
CLILogger.log(
64+
"Warning: Multi-token stop sequences not supported by xgrammar, using single-token stops only",
65+
component: "PipelinedConstrained")
66+
}
67+
let stopTokenIds: [Int32]? = singleTokenStops.isEmpty ? nil : singleTokenStops
68+
69+
let session = try constrainedEngine.getOrCreateConstrainedSession(
70+
jsonSchema: jsonSchema,
71+
tokenizer: tokenizer,
72+
vocabSize: vocabSize,
73+
stopTokenIds: stopTokenIds
74+
)
75+
76+
let inputTokens = try PromptUtils.maybeApplyTokenizerChatTemplate(input, tokenizer: tokenizer)
77+
.map(Int32.init)
78+
let maxTokens = options.maxTokens ?? defaultMaxConstrainedTokens
79+
80+
CLILogger.log(
81+
"Starting GPU pipelined constrained decoding (vocabSize=\(vocabSize), maxTokens=\(maxTokens))",
82+
component: "PipelinedConstrained")
83+
84+
return PipelinedConstrainedSequence(
85+
session: session,
86+
inputTokens: inputTokens,
87+
maxTokens: maxTokens,
88+
tokenizer: tokenizer,
89+
engine: constrainedEngine,
90+
samplingConfiguration: samplingConfiguration,
91+
stopSequences: stopSequences
92+
)
93+
}
94+
}
95+
96+
// MARK: - PipelinedConstrainedSequence
97+
98+
/// Single-use async sequence of `GenerationResult` produced by pipelined constrained decoding.
99+
///
100+
/// The engine's background Task owns session lifetime: it returns the session to
101+
/// cache in its `defer` block after generation completes (normal, error, or cancel).
102+
/// The iterator does not manage session lifecycle.
103+
public struct PipelinedConstrainedSequence: AsyncSequence {
104+
public typealias Element = GenerationResult
105+
public typealias Failure = Error
106+
107+
fileprivate let session: ConstrainedSessionHandle
108+
fileprivate let inputTokens: [Int32]
109+
fileprivate let maxTokens: Int
110+
let tokenizer: any Tokenizer
111+
let engine: any ConstrainedGenerationCapable
112+
let samplingConfiguration: SamplingConfiguration
113+
let stopSequences: StopSequences
114+
115+
/// Guards against multiple `makeAsyncIterator()` calls sharing the same session box.
116+
private let consumed = SingleUseFlag()
117+
118+
public func makeAsyncIterator() -> Iterator {
119+
precondition(!consumed.testAndSet(), "PipelinedConstrainedSequence may only be iterated once")
120+
121+
return Iterator(
122+
session: session,
123+
inputTokens: inputTokens,
124+
maxTokens: maxTokens,
125+
tokenizer: tokenizer,
126+
engine: engine,
127+
samplingConfiguration: samplingConfiguration,
128+
stopSequences: stopSequences
129+
)
130+
}
131+
}
132+
133+
extension PipelinedConstrainedSequence {
134+
public final class Iterator: AsyncIteratorProtocol {
135+
public typealias Element = GenerationResult
136+
public typealias Failure = Error
137+
138+
private let tokenizer: any Tokenizer
139+
private let engine: any ConstrainedGenerationCapable
140+
private let samplingConfiguration: SamplingConfiguration
141+
private let stopSequences: StopSequences
142+
private let session: ConstrainedSessionHandle
143+
144+
private let inputTokens: [Int32]
145+
private let maxTokens: Int
146+
147+
private var innerIterator: AsyncThrowingStream<Int32, Error>.AsyncIterator?
148+
private var generatedTokens: [Int32] = []
149+
private var previousDecodedText: String = ""
150+
private var recentTokens: [Int32] = []
151+
private var finished: Bool = false
152+
153+
fileprivate init(
154+
session: ConstrainedSessionHandle,
155+
inputTokens: [Int32],
156+
maxTokens: Int,
157+
tokenizer: any Tokenizer,
158+
engine: any ConstrainedGenerationCapable,
159+
samplingConfiguration: SamplingConfiguration,
160+
stopSequences: StopSequences
161+
) {
162+
self.session = session
163+
self.inputTokens = inputTokens
164+
self.maxTokens = maxTokens
165+
self.tokenizer = tokenizer
166+
self.engine = engine
167+
self.samplingConfiguration = samplingConfiguration
168+
self.stopSequences = stopSequences
169+
}
170+
171+
public func next() async throws -> GenerationResult? {
172+
if finished { return nil }
173+
174+
// Lazily start the engine stream on first iteration
175+
if innerIterator == nil {
176+
let stream = try engine.generateConstrained(
177+
with: inputTokens,
178+
samplingConfiguration: samplingConfiguration,
179+
maxTokens: maxTokens,
180+
session: session
181+
)
182+
self.innerIterator = stream.makeAsyncIterator()
183+
}
184+
185+
do {
186+
while true {
187+
try Task.checkCancellation()
188+
189+
guard var iterator = self.innerIterator else {
190+
finished = true
191+
return nil
192+
}
193+
194+
guard let tokenId = try await iterator.next() else {
195+
finished = true
196+
return nil
197+
}
198+
self.innerIterator = iterator
199+
200+
// Check stop sequences
201+
recentTokens.append(tokenId)
202+
if recentTokens.count > stopSequences.maxLength {
203+
recentTokens.removeFirst()
204+
}
205+
if stopSequences.matches(recentTokens: recentTokens) {
206+
finished = true
207+
return nil
208+
}
209+
210+
// Decode text incrementally
211+
generatedTokens.append(tokenId)
212+
let fullDecode = tokenizer.decode(tokens: generatedTokens.map { Int($0) })
213+
let common = fullDecode.commonPrefix(with: previousDecodedText)
214+
let delta = String(fullDecode.dropFirst(common.count))
215+
216+
if delta.unicodeScalars.contains(where: { $0 == "\u{FFFD}" }) {
217+
continue
218+
}
219+
220+
previousDecodedText = fullDecode
221+
222+
if !delta.isEmpty {
223+
return GenerationResult(text: delta, tokenId: tokenId, rawLogits: nil)
224+
}
225+
}
226+
} catch {
227+
finished = true
228+
throw error
229+
}
230+
}
231+
}
232+
}

‎swift/Sources/CoreAILanguageModels/GuidedGeneration/ConstrainedGenerationSession.swift‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,26 @@ public struct ConstrainedGenerationSession: ~Copyable {
193193
matcher.reset()
194194
allTokensBlocked = false
195195
}
196+
197+
/// Fill the bitmask directly into a caller-provided buffer (e.g., a GPU-visible MTLBuffer).
198+
///
199+
/// The caller must ensure the pointer has room for at least `(vocabularySize + 31) / 32`
200+
/// Int32 words. Returns `false` if the grammar is terminated (buffer untouched).
201+
public mutating func fillBitmask(into pointer: UnsafeMutablePointer<Int32>) -> Bool {
202+
if isTerminated { return false }
203+
204+
let hasConstraints = matcher.fillNextTokenBitmask(pointer)
205+
if !hasConstraints {
206+
allTokensBlocked = true
207+
return false
208+
}
209+
// Check for all-zeros (no tokens allowed)
210+
for i in 0..<bitmaskSize {
211+
if pointer[i] != 0 { return true }
212+
}
213+
allTokensBlocked = true
214+
return false
215+
}
196216
}
197217

198218
// MARK: - Float16 Masking

‎swift/Sources/CoreAILanguageModels/InferenceEngines/ConstrainedGenerationCapable.swift‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import Tokenizers
7+
8+
/// An inference engine that supports GPU-accelerated grammar-constrained generation.
9+
///
10+
/// Conforming engines can apply xgrammar bitmasks directly in the GPU sampling path,
11+
/// eliminating the need to transfer logits to CPU for masking. This protocol is used as
12+
/// a capability signal for routing (e.g., choosing pipelined vs. sequential constrained
13+
/// strategies) and for testability (mock conformers in unit tests).
14+
///
15+
/// The session lifecycle follows a checkout/checkin pattern:
16+
/// 1. Call `getOrCreateConstrainedSession` to obtain a handle (cached or fresh)
17+
/// 2. Pass the handle to `generateConstrained` which drives the GPU loop
18+
/// 3. The engine returns the handle to its internal cache after generation completes
19+
package protocol ConstrainedGenerationCapable: InferenceEngine {
20+
/// Obtain a constrained session handle, reusing a cached one if the schema matches.
21+
///
22+
/// The internal cache slot is emptied on checkout — concurrent calls create independent sessions.
23+
func getOrCreateConstrainedSession(
24+
jsonSchema: String,
25+
tokenizer: any Tokenizer,
26+
vocabSize: Int,
27+
stopTokenIds: [Int32]?
28+
) throws -> ConstrainedSessionHandle
29+
30+
/// Stream constrained token generation using GPU-side bitmask application.
31+
///
32+
/// The engine resets state, prefills the prompt (unconstrained), then enters
33+
/// a semi-pipelined loop: inference overlaps bitmask computation, but sampling
34+
/// waits per token (grammar state is inherently sequential).
35+
///
36+
/// The handle is returned to the engine's cache automatically when the Task completes.
37+
func generateConstrained(
38+
with input: [TokenId],
39+
samplingConfiguration: SamplingConfiguration,
40+
maxTokens: Int,
41+
session: ConstrainedSessionHandle
42+
) throws -> AsyncThrowingStream<TokenId, Error>
43+
}

0 commit comments

Comments
 (0)